feat: add --use-cloud-function flag to validate_config

jetstream/analysis.py

@@ -1030,7 +1030,12 @@ def check_runnable(self, current_date: datetime | None = None) -> bool:
     def _app_id_to_bigquery_dataset(self, app_id: str) -> str:
         return re.sub(r"[^a-zA-Z0-9]", "_", app_id).lower()
 
-    def validate(self, use_glean_ids: bool = False, metric_slugs: list[str] | None = None) -> int:
+    def validate(
+        self,
+        use_glean_ids: bool = False,
+        metric_slugs: list[str] | None = None,
+        use_cloud_function: bool = True,
+    ) -> int:
         """Validates enrollments and metrics queries for an experiment.
 
         Returns (int) the amount of data to be processed by the metrics query.
@@ -1108,7 +1113,7 @@ def validate(self, use_glean_ids: bool = False, metric_slugs: list[str] | None =
 
         self._write_sql_output(f"enrollments_{bq_normalize_name(experiment_slug)}", enrollments_sql)
 
-        bytes_processed = dry_run_query(enrollments_sql)
+        bytes_processed = dry_run_query(enrollments_sql, use_cloud_function=use_cloud_function)
         logger.info(f"Dry running enrollments query for {experiment_slug}:")
         logger.info(enrollments_sql)
         if bytes_processed and bytes_processed > 0:
@@ -1119,7 +1124,9 @@ def validate(self, use_glean_ids: bool = False, metric_slugs: list[str] | None =
         if not metric_slugs:
             output_loc = f"metrics_{bq_normalize_name(experiment_slug)}"
             logger.info(f"Dry running metrics query for {experiment_slug}")
-            metrics_tb = self.validate_metric_query(exp, metrics, limits, output_loc, use_glean_ids)
+            metrics_tb = self.validate_metric_query(
+                exp, metrics, limits, output_loc, use_glean_ids, use_cloud_function
+            )
         else:
             selected_metrics = [m for m in metrics if m.name in metric_slugs]
             metrics_tb_async_list = []
@@ -1135,7 +1142,7 @@ def validate(self, use_glean_ids: bool = False, metric_slugs: list[str] | None =
                     metrics_tb_async_list.append(
                         pool.apply_async(
                             self.validate_metric_query,
-                            (exp, [metric], limits, output_loc, use_glean_ids),
+                            (exp, [metric], limits, output_loc, use_glean_ids, use_cloud_function),
                         )
                     )
             # ensure we have all the async results, then get the max value
@@ -1153,6 +1160,7 @@ def validate_metric_query(
         limits: TimeLimits,
         output_loc: str,
         use_glean_ids: bool = False,
+        use_cloud_function: bool = True,
     ) -> float:
         metrics_sql = experiment.build_metrics_query(
             metrics,
@@ -1189,7 +1197,7 @@ def validate_metric_query(
 
         self._write_sql_output(output_loc, metrics_sql)
 
-        bytes_processed = dry_run_query(metrics_sql)
+        bytes_processed = dry_run_query(metrics_sql, use_cloud_function=use_cloud_function)
         logger.info(metrics_sql)
 
         if bytes_processed and bytes_processed > 0:

jetstream/cli.py

@@ -1401,13 +1401,24 @@ def rerun_config_changed(
     help="Threshold above which metrics query validation returns exit code 2.",
     default=MODERATE_DATA_THRESHOLD,
 )
+@click.option(
+    "--use-cloud-function/--no-use-cloud-function",
+    "use_cloud_function",
+    help=(
+        "Dry run queries via the dry run Cloud Function (default). "
+        "Use --no-use-cloud-function to dry run with your own user credentials "
+        "instead, e.g. to validate configs that reference tables only you can read."
+    ),
+    default=True,
+)
 def validate_config(
     path: Iterable[os.PathLike],
     config_repos,
     private_config_repos,
     is_private,
     high_data_threshold,
     moderate_data_threshold,
+    use_cloud_function,
 ):
     """Validate config files.
 
@@ -1448,6 +1459,7 @@ def validate_config(
             config_getter=ConfigLoader.with_configs_from(config_repos).with_configs_from(
                 private_config_repos, is_private=True
             ),
+            use_cloud_function=use_cloud_function,
         )
         if (
             isinstance(entity, Config)
@@ -1466,6 +1478,7 @@ def validate_config(
                     private_config_repos, is_private=True
                 ),
                 experiment=experiments[0],
+                use_cloud_function=use_cloud_function,
             )
         try:
             tb_processed = call()

jetstream/config.py

@@ -225,6 +225,7 @@ def validate(
     config: Outcome | Config | DefaultConfig | DefinitionConfig,
     experiment: Experiment | None = None,
     config_getter: _ConfigLoader = ConfigLoader,
+    use_cloud_function: bool = True,
 ) -> int:
     """Validate and dry run a config.
 
@@ -292,4 +293,6 @@ def validate(
     else:
         raise Exception(f"Unable to validate config: {config}")
 
-    return Analysis("no project", "no dataset", resolved_config).validate()
+    return Analysis("no project", "no dataset", resolved_config).validate(
+        use_cloud_function=use_cloud_function
+    )

jetstream/dryrun.py

@@ -60,8 +60,17 @@ def __init__(self, error: Any, sql: str):
         super().__init__(error)
 
 
-def dry_run_query(sql: str) -> int:
-    """Dry run the provided SQL query."""
+def dry_run_query(sql: str, use_cloud_function: bool = True) -> int:
+    """Dry run the provided SQL query.
+
+    If ``use_cloud_function`` is True (default), the query is dry run via the
+    dry run Cloud Function, which executes using its own service account.
+    If False, the query is dry run directly using the current user's
+    credentials.
+    """
+    if not use_cloud_function:
+        return _dry_run_with_user_credentials(sql)
+
     try:
         # look for token created by the GitHub Actions workflow
         id_token = os.environ.get("GOOGLE_GHA_ID_TOKEN")
@@ -126,3 +135,23 @@ def dry_run_query(sql: str) -> int:
         return int(response.get("bytesProcessed", -1))
 
     raise DryRunFailedError((error and error.get("message", None)) or response["errors"], sql=sql)
+
+
+def _dry_run_with_user_credentials(sql: str) -> int:
+    """Dry run a query directly using the current user's credentials.
+
+    Unlike the Cloud Function, this executes the dry run as the authenticated
+    user, so it can access tables the user has permission to read (e.g. personal
+    tables in the ``analysis`` dataset) that the dry run service account cannot.
+    """
+    from google.cloud import bigquery
+
+    try:
+        client = bigquery.Client()
+        job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
+        query_job = client.query(sql, job_config=job_config)
+    except Exception as e:
+        raise DryRunFailedError(e, sql) from e
+
+    logger.info("Dry run OK")
+    return int(query_job.total_bytes_processed or -1)