[run-tests] Substitute UDFs

Author: scholtzan

bigquery_etl/dryrun.py

@@ -95,15 +95,59 @@ def wrap_in_view_for_dryrun(sql: str) -> str:
         if not statements or not isinstance(statements[-1], sqlglot.exp.Select):
             return sql
 
-        # Split original SQL by semicolons to preserve formatting;
-        # stripping formatting causes some query dry runs to fail
-        parts = [p for p in sql.split(";") if p.strip()]
+        # Replace CREATE TEMP FUNCTION with CREATE FUNCTION using fully qualified names
+        # CREATE VIEW doesn't support temp functions
+        test_project = ConfigLoader.get(
+            "default", "test_project", fallback="bigquery-etl-integration-test"
+        )
 
-        if len(parts) != len(statements):
-            return sql
+        def replace_temp_function(match):
+            func_name = match.group(1)
+            # If function name is already qualified, keep it; otherwise add project.dataset prefix
+            if "." not in func_name and "`" not in func_name:
+                return f"CREATE FUNCTION `{test_project}.tmp.{func_name}`"
+            else:
+                return f"CREATE FUNCTION {func_name}"
 
-        prefix_sql = ";\n".join(parts[:-1]) + ";" if len(parts) > 1 else ""
-        query_sql = parts[-1].strip()
+        sql = re.sub(
+            r"\bCREATE\s+TEMP(?:ORARY)?\s+FUNCTION\s+([^\s(]+)",
+            replace_temp_function,
+            sql,
+            flags=re.IGNORECASE,
+        )
+
+        # Single statement - just wrap it
+        if len(statements) == 1:
+            view_name = f"_dryrun_view_{random_str(8)}"
+            test_project = ConfigLoader.get(
+                "default", "test_project", fallback="bigquery-etl-integration-test"
+            )
+            query_sql = sql.strip().rstrip(";")
+            return f"CREATE VIEW `{test_project}.tmp.{view_name}` AS\n{query_sql}"
+
+        # Multiple statements: use sqlglot tokenizer to find statement boundaries
+        # This handles semicolons in strings and comments
+        tokens = list(sqlglot.tokens.Tokenizer(dialect="bigquery").tokenize(sql))
+
+        # Find semicolon tokens that separate statements (not in strings/comments)
+        semicolon_positions = []
+        for token in tokens:
+            if token.token_type == sqlglot.tokens.TokenType.SEMICOLON:
+                semicolon_positions.append(token.end)
+
+        # We need (len(statements) - 1) semicolons to separate statements
+        if len(semicolon_positions) >= len(statements) - 1:
+            # The (n-1)th semicolon separates the prefix from the last statement
+            split_pos = semicolon_positions[len(statements) - 2]
+            prefix_sql = sql[:split_pos].strip()
+            query_sql = sql[split_pos:].strip().lstrip(";").strip()
+        else:
+            # Fallback: regenerate prefix statements, use regenerated query
+            prefix_statements = statements[:-1]
+            prefix_sql = ";\n".join(
+                stmt.sql(dialect="bigquery") for stmt in prefix_statements
+            )
+            query_sql = statements[-1].sql(dialect="bigquery")
 
         # Wrap in view
         view_name = f"_dryrun_view_{random_str(8)}"
@@ -112,7 +156,7 @@ def wrap_in_view_for_dryrun(sql: str) -> str:
         )
         wrapped_query = f"CREATE VIEW `{test_project}.tmp.{view_name}` AS\n{query_sql}"
 
-        return f"{prefix_sql}\n\n{wrapped_query}" if prefix_sql else wrapped_query
+        return f"{prefix_sql};\n\n{wrapped_query}"
 
     except Exception as e:
         print(f"Warning: Failed to wrap SQL in view: {e}")
@@ -271,11 +315,6 @@ def dry_run_result(self):
         else:
             sql = self.get_sql()
 
-        # Wrap the query in a CREATE VIEW for faster dry runs
-        # Skip wrapping when strip_dml=True as it's used for special analysis modes
-        if not self.strip_dml:
-            sql = wrap_in_view_for_dryrun(sql)
-
         query_parameters = []
         scheduling_metadata = self.metadata.scheduling if self.metadata else {}
         if date_partition_parameter := scheduling_metadata.get(
@@ -299,6 +338,30 @@ def dry_run_result(self):
                 )
             )
 
+        # Wrap the query in a CREATE VIEW for faster dry runs
+        # Skip wrapping when strip_dml=True as it's used for special analysis modes
+        if not self.strip_dml:
+            # If query has parameters, replace them with literal values in the wrapped version
+            # since CREATE VIEW cannot use parameterized queries
+            sql_for_wrapping = sql
+            if query_parameters:
+                for param in query_parameters:
+                    param_name = f"@{param.name}"
+                    # Convert parameter value to SQL literal
+                    if param.type_ == "DATE":
+                        param_value = f"DATE '{param.value}'"
+                    elif param.type_ in ("STRING", "DATETIME", "TIMESTAMP"):
+                        param_value = f"'{param.value}'"
+                    elif param.type_ == "BOOL":
+                        param_value = str(param.value).upper()
+                    else:
+                        param_value = str(param.value)
+                    sql_for_wrapping = sql_for_wrapping.replace(param_name, param_value)
+
+            sql = wrap_in_view_for_dryrun(sql_for_wrapping)
+
+        # print(sql)
+
         project = basename(dirname(dirname(dirname(self.sqlfile))))
         dataset = basename(dirname(dirname(self.sqlfile)))
         try:
@@ -633,6 +696,8 @@ def validate_schema(self):
             strip_dml=self.strip_dml,
         )
 
+        # print(table_schema)
+
         # This check relies on the new schema being deployed to prod
         if not query_schema.compatible(table_schema):
             click.echo(

bigquery_etl/schema/stable_table_schema.py

@@ -60,8 +60,7 @@ def prod_schemas_uri():
     """
     dryrun = DryRun(
         "moz-fx-data-shared-prod/telemetry_derived/foo/query.sql",
-        content="SELECT 1",
-        strip_dml=True,
+        content="SELECT 1 AS field",
     )
     build_id = dryrun.get_dataset_labels()["schemas_build_id"]
     commit_hash = build_id.split("_")[-1]

sql/moz-fx-data-shared-prod/analysis/bqetl_default_task_v1/query.sql

@@ -2,6 +2,6 @@
 SELECT
   *
 FROM
-  (SELECT 1)
+  (SELECT 1 AS field)
 WHERE
   FALSE

sql_generators/derived_view_schemas/__init__.py

@@ -112,7 +112,7 @@ def _get_reference_partition_column(ref_path):
         logging.debug("No reference partition column, dry running without one.")
 
     view = View.from_file(
-        view_file, partition_column=reference_partition_column, id_token=id_token
+        view_file, partition_column=reference_partition_column, id_token=id_token, strip_dml=True
     )
 
     # `View.schema` prioritizes the configured schema over the dryrun schema, but here

tests/test_dryrun.py

@@ -15,10 +15,12 @@ def tmp_query_path(tmp_path):
 class TestDryRun:
     def test_dry_run_sql_file(self, tmp_query_path):
         query_file = tmp_query_path / "query.sql"
-        query_file.write_text("SELECT 123")
+        query_file.write_text("SELECT 123 AS field")
 
         dryrun = DryRun(str(query_file))
-        assert dryrun.is_valid()
+        print(dryrun.dry_run_result)
+        response = dryrun.dry_run_result
+        assert response["valid"]
 
     def test_dry_run_invalid_sql_file(self, tmp_query_path):
         query_file = tmp_query_path / "query.sql"
@@ -30,7 +32,7 @@ def test_dry_run_invalid_sql_file(self, tmp_query_path):
 
     def test_sql_file_valid(self, tmp_query_path):
         query_file = tmp_query_path / "query.sql"
-        query_file.write_text("SELECT 123")
+        query_file.write_text("SELECT 123 AS field")
 
         dryrun = DryRun(str(query_file))
         assert dryrun.is_valid()
@@ -60,7 +62,7 @@ def test_sql_file_invalid(self, tmp_query_path):
 
     def test_get_referenced_tables_empty(self, tmp_query_path):
         query_file = tmp_query_path / "query.sql"
-        query_file.write_text("SELECT 123")
+        query_file.write_text("SELECT 123 AS field")
 
         dryrun = DryRun(str(query_file))
         assert dryrun.get_referenced_tables() == []
@@ -69,7 +71,7 @@ def test_get_sql(self, tmp_path):
         os.makedirs(tmp_path / "telmetry_derived")
         query_file = tmp_path / "telmetry_derived" / "query.sql"
 
-        sql_content = "SELECT 123 "
+        sql_content = "SELECT 123 AS field"
         query_file.write_text(sql_content)
 
         assert DryRun(sqlfile=str(query_file)).get_sql() == sql_content