Transform UDFs to temporary UDFs in tests

bigquery_etl/parse_udf.py

@@ -133,6 +133,7 @@ def parse_udf_dirs(*udf_dirs):
 def accumulate_dependencies(deps, raw_udfs, udf_name):
     """
     Accumulate a list of dependent UDF names.
+
     Given a dict of raw_udfs and a udf_name string, recurse into the
     UDF's dependencies, adding the names to deps in depth-first order.
     """
@@ -178,14 +179,15 @@ def udf_usage_definitions(text, raw_udfs=None):
 
 def udf_tests_sql(raw_udf, raw_udfs):
     """
-    Create tests for testing persistent UDFs. 
+    Create tests for testing persistent UDFs.
+
     Persistent UDFs need to be rewritten as temporary UDF so that changes
     can be tested.
     """
     tests_full_sql = []
     for test in raw_udf.tests:
         test_sql = prepend_udf_usage_definitions(test, raw_udfs)
-        test_sql = sub_persisent_udfs_as_temp(test_sql)
+        test_sql = sub_persistent_udfs_as_temp(test_sql)
         test_sql = test_sql.replace(
             "CREATE OR REPLACE FUNCTION", "CREATE TEMP FUNCTION"
         )
@@ -201,6 +203,6 @@ def prepend_udf_usage_definitions(text, raw_udfs=None):
     return "\n\n".join(statements + [text])
 
 
-def sub_persisent_udfs_as_temp(text):
+def sub_persistent_udfs_as_temp(text):
     """Substitute persistent UDF references with temporary UDF references."""
     return PERSISTENT_UDF_RE.sub(r"\1_\2", text)

bigquery_etl/pytest_plugin/sql.py

@@ -107,11 +107,12 @@ class SqlTest(pytest.Item, pytest.File):
                 views[table_name] = read(self.fspath.strpath, resource)
 
         # rewrite all udfs as temporary
-        temp_udfs = parse_udf.sub_persisent_udfs_as_temp(query)
+        temp_udfs = parse_udf.sub_persistent_udfs_as_temp(query)
         if temp_udfs != query:
-            query = temp_udfs
             # prepend udf definitions
             query = parse_udf.prepend_udf_usage_definitions(query)
+            query = parse_udf.sub_persistent_udfs_as_temp(query)
+            query = query.replace("CREATE OR REPLACE FUNCTION", "CREATE TEMP FUNCTION")
 
         dataset_id = "_".join(self.fspath.strpath.split(os.path.sep)[-3:])
         if "CIRCLE_BUILD_NUM" in os.environ:

script/publish_persistent_udfs

@@ -34,7 +34,7 @@ parser.add_argument(
 )
 parser.add_argument(
     "--udf-dirs",
-    nargs='+',
+    nargs="+",
     default=["udf/", "udf_js/"],
     help="The directories where declarations of UDFs are stored.",
 )
@@ -63,7 +63,9 @@ def main():
     client = bigquery.Client(args.project_id)
 
     if args.dep_dir:
-        push_dependencies_to_gcs(args.gcs_bucket, args.gcs_path, args.dep_dir, args.project_id)
+        push_dependencies_to_gcs(
+            args.gcs_bucket, args.gcs_path, args.dep_dir, args.project_id
+        )
 
     for raw_udf in raw_udfs:
         # get all dependencies for UDF and publish as persistent UDF