From a94dea6c69fab5f6a9d5c85d8ac33c49cf27aca1 Mon Sep 17 00:00:00 2001 From: Anna Scholtz Date: Thu, 9 Jul 2026 15:30:28 -0700 Subject: [PATCH] feat: add --use-cloud-function flag to validate_config Add a --use-cloud-function/--no-use-cloud-function option to validate_config (default on). With --no-use-cloud-function, queries are dry run using the current user's credentials instead of the dry run Cloud Function's service account, allowing validation of configs that reference tables only the user can read (e.g. personal analysis tables). --- jetstream/analysis.py | 18 +++++++++++++----- jetstream/cli.py | 13 +++++++++++++ jetstream/config.py | 5 ++++- jetstream/dryrun.py | 33 +++++++++++++++++++++++++++++++-- 4 files changed, 61 insertions(+), 8 deletions(-) diff --git a/jetstream/analysis.py b/jetstream/analysis.py index 4efc58aa..c8b3af39 100644 --- a/jetstream/analysis.py +++ b/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: diff --git a/jetstream/cli.py b/jetstream/cli.py index d492f375..64ea4c89 100644 --- a/jetstream/cli.py +++ b/jetstream/cli.py @@ -1401,6 +1401,16 @@ 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, @@ -1408,6 +1418,7 @@ def validate_config( 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() diff --git a/jetstream/config.py b/jetstream/config.py index 9a6b26f4..95ec8f7b 100644 --- a/jetstream/config.py +++ b/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 + ) diff --git a/jetstream/dryrun.py b/jetstream/dryrun.py index dc424209..f3a458a4 100644 --- a/jetstream/dryrun.py +++ b/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)