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)