feat(dryrun): Add --validate-union-schemas option (DENG-7746).

mozilla · Jan 30, 2025 · 43930ec · 43930ec
1 parent 56f30f6
commit 43930ec
Show file tree

Hide file tree

Showing 3 changed files with 388 additions and 70 deletions.
diff --git a/bigquery_etl/cli/dryrun.py b/bigquery_etl/cli/dryrun.py
@@ -5,6 +5,7 @@
 import os
 import re
 import sys
+import traceback
 from functools import partial
 from multiprocessing.pool import Pool
 from typing import List, Set, Tuple
@@ -53,6 +54,13 @@
     is_flag=True,
     default=False,
 )
+@click.option(
+    "--validate_union_schemas",
+    "--validate-union-schemas",
+    help="Require any subqueries being unioned to have exactly matching schemas.",
+    is_flag=True,
+    default=False,
+)
 @click.option(
     "--respect-skip/--ignore-skip",
     help="Respect or ignore query skip configuration. Default is --respect-skip.",
@@ -67,6 +75,7 @@ def dryrun(
     paths: List[str],
     use_cloud_function: bool,
     validate_schemas: bool,
+    validate_union_schemas: bool,
     respect_skip: bool,
     project: str,
 ):
@@ -112,9 +121,10 @@ def dryrun(
 
     sql_file_valid = partial(
         _sql_file_valid,
-        use_cloud_function,
-        respect_skip,
-        validate_schemas,
+        use_cloud_function=use_cloud_function,
+        respect_skip=respect_skip,
+        validate_schemas=validate_schemas,
+        validate_union_schemas=validate_union_schemas,
         credentials=credentials,
         id_token=id_token,
     )
@@ -133,7 +143,13 @@ def dryrun(
 
 
 def _sql_file_valid(
-    use_cloud_function, respect_skip, validate_schemas, sqlfile, credentials, id_token
+    sqlfile,
+    use_cloud_function,
+    respect_skip,
+    validate_schemas,
+    validate_union_schemas,
+    credentials,
+    id_token,
 ) -> Tuple[bool, str]:
     """Dry run the SQL file."""
     result = DryRun(
@@ -151,4 +167,15 @@ def _sql_file_valid(
             success = False
         return success, sqlfile
 
+    if validate_union_schemas:
+        try:
+            success = result.validate_union_schemas()
+        except Exception:
+            click.echo(
+                f"Failed to validate union schemas in {sqlfile}:\n{traceback.format_exc(limit=0)}",
+                err=True,
+            )
+            success = False
+        return success, sqlfile
+
     return result.is_valid(), sqlfile