chore: ensure udfs are actively deleted, make pre-commit mypy dep consistent (#1522)

Author: shobsi

.pre-commit-config.yaml

@@ -38,6 +38,6 @@ repos:
     rev: v1.10.0
     hooks:
     -   id: mypy
-        additional_dependencies: [types-requests, types-tabulate, pandas-stubs]
+        additional_dependencies: [types-requests, types-tabulate, pandas-stubs<=2.2.3.241126]
         exclude: "^third_party"
         args: ["--check-untyped-defs", "--explicit-package-bases", "--ignore-missing-imports"]

tests/system/large/functions/test_managed_function.py

@@ -72,7 +72,7 @@ def multiply(x, y):
         )
     finally:
         # clean up the gcp assets created for the managed function.
-        cleanup_function_assets(multiply, bigquery_client)
+        cleanup_function_assets(multiply, bigquery_client, ignore_failures=False)
 
 
 def test_managed_function_stringify_with_ibis(
@@ -118,7 +118,7 @@ def stringify(x):
         )
     finally:
         # clean up the gcp assets created for the managed function.
-        cleanup_function_assets(stringify, bigquery_client)
+        cleanup_function_assets(stringify, bigquery_client, ignore_failures=False)
 
 
 @pytest.mark.parametrize(
@@ -167,7 +167,7 @@ def featurize(x: int) -> list[array_dtype]:  # type: ignore
 
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(featurize, session.bqclient)
+        cleanup_function_assets(featurize, session.bqclient, ignore_failures=False)
 
 
 @pytest.mark.parametrize(
@@ -234,7 +234,7 @@ def foo(x: int) -> typ:  # type:ignore
         pandas.testing.assert_frame_equal(bf_result_gbq, pd_result, check_dtype=False)
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(foo, session.bqclient)
+        cleanup_function_assets(foo, session.bqclient, ignore_failures=False)
 
 
 @pytest.mark.parametrize(
@@ -274,7 +274,7 @@ def foo_list(x: int) -> list[typ]:  # type:ignore
         pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False)
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(foo_list, session.bqclient)
+        cleanup_function_assets(foo_list, session.bqclient, ignore_failures=False)
 
 
 def test_managed_function_series_combine(session, scalars_dfs):
@@ -330,7 +330,9 @@ def add(x: int, y: int) -> int:
         pandas.testing.assert_series_equal(bf_result, pd_result, check_dtype=False)
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(add_managed_func, session.bqclient)
+        cleanup_function_assets(
+            add_managed_func, session.bqclient, ignore_failures=False
+        )
 
 
 def test_managed_function_series_combine_array_output(session, scalars_dfs):
@@ -391,7 +393,9 @@ def add_list(x: int, y: int) -> list[int]:
         pandas.testing.assert_series_equal(bf_result_gbq, pd_result, check_dtype=False)
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(add_list_managed_func, session.bqclient)
+        cleanup_function_assets(
+            add_list_managed_func, session.bqclient, ignore_failures=False
+        )
 
 
 def test_managed_function_dataframe_map(session, scalars_dfs):
@@ -425,7 +429,7 @@ def add_one(x):
         pandas.testing.assert_frame_equal(bf_result, pd_result)
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(mf_add_one, session.bqclient)
+        cleanup_function_assets(mf_add_one, session.bqclient, ignore_failures=False)
 
 
 def test_managed_function_dataframe_map_array_output(
@@ -464,7 +468,9 @@ def add_one_list(x):
         pandas.testing.assert_frame_equal(bf_result_gbq, pd_result, check_dtype=False)
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(mf_add_one_list, session.bqclient)
+        cleanup_function_assets(
+            mf_add_one_list, session.bqclient, ignore_failures=False
+        )
 
 
 def test_managed_function_dataframe_apply_axis_1(session, scalars_dfs):
@@ -500,7 +506,7 @@ def add_ints(x, y):
         )
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(add_ints_mf, session.bqclient)
+        cleanup_function_assets(add_ints_mf, session.bqclient, ignore_failures=False)
 
 
 def test_managed_function_dataframe_apply_axis_1_array_output(session):
@@ -605,4 +611,4 @@ def foo(x, y, z):
 
     finally:
         # Clean up the gcp assets created for the managed function.
-        cleanup_function_assets(foo, session.bqclient)
+        cleanup_function_assets(foo, session.bqclient, ignore_failures=False)

tests/system/utils.py

@@ -401,14 +401,20 @@ def cleanup_function_assets(
 ) -> None:
     """Clean up the GCP assets behind a bigframess function."""
 
-    # Clean up bigframes function.
+    # Clean up bigframes bigquery function.
     try:
         bigquery_client.delete_routine(bigframes_func.bigframes_bigquery_function)
     except Exception:
         # By default don't raise exception in cleanup.
         if not ignore_failures:
             raise
 
+    if not ignore_failures:
+        # Make sure that the BQ routins is actually deleted
+        with pytest.raises(google.api_core.exceptions.NotFound):
+            bigquery_client.get_routine(bigframes_func.bigframes_bigquery_function)
+
+    # Clean up bigframes cloud run function
     if cloudfunctions_client:
         # Clean up cloud function
         try:
@@ -420,6 +426,16 @@ def cleanup_function_assets(
             if not ignore_failures:
                 raise
 
+        if not ignore_failures:
+            # Make sure the cloud run function is actually deleted
+            try:
+                gcf = cloudfunctions_client.get_function(
+                    name=bigframes_func.bigframes_cloud_function
+                )
+                assert gcf.state is functions_v2.Function.State.DELETING
+            except google.cloud.exceptions.NotFound:
+                pass
+
 
 def get_function_name(func, package_requirements=None, is_row_processor=False):
     """Get a bigframes function name for testing given a udf."""