fix(dryrun): Run authenticated dryrun queries in mozdata rather than the target project

bigquery_etl/cli/deploy.py

@@ -1208,6 +1208,7 @@ def _update_table_schema(file_path: Path, options: dict):
             tmp_dataset="tmp",  # Default dataset for temporary tables during schema updates
             tmp_tables={},
             use_cloud_function=options["use_cloud_function"],
+            billing_project=(project_id if not options["use_cloud_function"] else None),
             respect_dryrun_skip=options["respect_dryrun_skip"],
             is_init=False,
             credentials=options["credentials"],
@@ -1270,17 +1271,21 @@ def _deploy_table_artifact(file_path: Path, options: dict):
         if not options["table_force"] and str(file_path).endswith(".sql"):
             client = bigquery.Client(credentials=options["credentials"])
             try:
+                table_name = file_path.parent.name
+                dataset_name = file_path.parent.parent.name
+                project_name = file_path.parent.parent.parent.name
                 query_schema = Schema.from_query_file(
                     file_path,
                     use_cloud_function=options["use_cloud_function"],
+                    billing_project=(
+                        project_name if not options["use_cloud_function"] else None
+                    ),
                     respect_skip=options["respect_dryrun_skip"],
                     sql_dir=options["sql_dir"],
                     client=client,
                     id_token=options["id_token"],
                 )
                 if not existing_schema.equal(query_schema):
-                    dataset_name = file_path.parent.parent.name
-                    table_name = file_path.parent.name
                     raise FailedDeployException(
                         f"Query {file_path} does not match schema in {schema_path}. "
                         f"Run `./bqetl query schema update {dataset_name}.{table_name}`"

bigquery_etl/cli/query.py

@@ -2559,6 +2559,7 @@ def _update_query_schema(
     id_token=None,
     use_dataset_schema=False,
     use_global_schema=False,
+    billing_project=None,
 ):
     """
     Update the schema of a specific query file.
@@ -2651,6 +2652,7 @@ def _update_query_schema(
             query_file_path,
             content=sql_content,
             use_cloud_function=use_cloud_function,
+            billing_project=billing_project,
             respect_skip=respect_dryrun_skip,
             sql_dir=sql_dir,
             credentials=credentials,
@@ -2731,6 +2733,7 @@ def _update_query_schema(
         table_name,
         partitioned_by=partitioned_by,
         use_cloud_function=use_cloud_function,
+        billing_project=billing_project,
         respect_skip=respect_dryrun_skip,
         credentials=credentials,
         id_token=id_token,

bigquery_etl/deploy.py

@@ -104,6 +104,7 @@ def deploy_table(
         query_schema = Schema.from_query_file(
             artifact_file,
             use_cloud_function=use_cloud_function,
+            billing_project=(project_name if not use_cloud_function else None),
             respect_skip=respect_dryrun_skip,
             sql_dir=sql_dir,
             client=client,

bigquery_etl/dryrun.py

@@ -148,14 +148,18 @@ def __init__(
         self.project = project
         self.dataset = dataset
         self.table = table
+        self.billing_project = billing_project
         # if using cloud function and billing project isn't set, randomly select project to use
-        self.billing_project = (
-            billing_project
-            if billing_project or not use_cloud_function
-            else random.choice(
-                ConfigLoader.get("dry_run", "default_projects", fallback=[None])
+        if not self.billing_project:
+            self.billing_project = (
+                random.choice(
+                    ConfigLoader.get(
+                        "dry_run", "cloud_function_billing_projects", fallback=[None]
+                    )
+                )
+                if use_cloud_function
+                else ConfigLoader.get("dry_run", "default_billing_project")
             )
-        )
         try:
             self.metadata = Metadata.of_query_file(self.sqlfile)
         except FileNotFoundError:
@@ -474,17 +478,15 @@ def dry_run_result(self):
                 )
                 result = json.load(r)
             else:
-                # Prefer billing project if provided, otherwise use the project from the SQL file
-                self.client.project = (
-                    self.billing_project if self.billing_project else project
-                )
                 job_config = bigquery.QueryJobConfig(
                     dry_run=True,
                     use_query_cache=False,
                     default_dataset=f"{project}.{dataset}",
                     query_parameters=query_parameters,
                 )
-                job = self.client.query(sql, job_config=job_config)
+                job = self.client.query(
+                    sql, job_config=job_config, project=self.billing_project
+                )
                 try:
                     dataset_labels = self.client.get_dataset(job.default_dataset).labels
                 except Exception as e:

bqetl_project.yaml

@@ -296,7 +296,8 @@ dry_run:
   - sql/moz-fx-data-shared-prod/firefox_accounts_derived/event_types_history_v1/query.sql
   # Tests
   - sql/moz-fx-data-test-project/test/simple_view/view.sql
-  default_projects:  # projects to use for query dry run jobs
+  default_billing_project: mozdata  # billing project to use for authenticated query dry run jobs
+  cloud_function_billing_projects:  # billing projects to use for query dry run jobs from the cloud function
   - moz-fx-data-backfill-10
   - moz-fx-data-backfill-11
   - moz-fx-data-backfill-12