Skip to content

fix(bigquery): validate allowedDatasets before dry-run in execute-sql - #3759

Open
PratikDhanave wants to merge 1 commit into
googleapis:mainfrom
PratikDhanave:fix/bigquery-allowlist-before-dryrun
Open

fix(bigquery): validate allowedDatasets before dry-run in execute-sql#3759
PratikDhanave wants to merge 1 commit into
googleapis:mainfrom
PratikDhanave:fix/bigquery-allowlist-before-dryrun

Conversation

@PratikDhanave

Copy link
Copy Markdown

Summary

With allowedDatasets configured, bigquery-execute-sql only enforced the allowlist after a successful dry-run. A query referencing a non-existent dataset outside the allowlist therefore failed the dry-run with an opaque BigQuery 404 Not Found (via ProcessGcpError) instead of the documented policy error:

query accesses dataset 'PROJECT.DATASET', which is not in the allowed list

Security stayed fail-safe (no data returned), but the signal was wrong — agents and evals couldn't distinguish "table missing" from "dataset not allowlisted".

Fixes #3717

Change

  • In Invoke, when allowedDatasets is non-empty, statically parse the referenced datasets with TableParser and reject any outside the allowlist before the dry-run.
  • The dry-run and the existing post-dry-run check still run afterward to catch views/wildcards the static parser can't see (statement-type restrictions, ReferencedTables enrichment), so no coverage is lost.
  • Extracted the shared "is this dataset allowed?" logic into a checkDatasetAllowed helper, used by both the pre- and post-dry-run paths.

Testing

  • Added a TestInvokeDatasetRestrictions case where the mocked dry-run returns a 404 for a non-existent dataset; it asserts the tool now returns the allowlist policy error rather than the 404.
  • Verified the new case fails on the pre-fix code (got HTTP response code 404 ... Not found: Table ...nonexistent_dataset) and passes with the fix — a genuine regression guard.
  • All existing TestInvokeDatasetRestrictions cases (allowed/forbidden tables, INFORMATION_SCHEMA, EXTERNAL_QUERY, etc.) still pass; go build, go vet, and go test on the package are green.

With allowedDatasets configured, the dataset allowlist was only enforced
after a successful dry-run. A query referencing a non-existent dataset
outside the allowlist therefore failed the dry-run with an opaque BigQuery
404 (via ProcessGcpError) instead of the documented policy error:

  query accesses dataset 'PROJECT.DATASET', which is not in the allowed list

Security stayed fail-safe (no data returned), but the signal was wrong:
agents and evals could not distinguish "table missing" from "dataset not
allowlisted".

Statically parse the referenced datasets with TableParser and reject any
outside the allowlist before the dry-run. The dry-run and the existing
post-dry-run check still run to catch views/wildcards the parser cannot
see. Extracts the shared allowlist check into checkDatasetAllowed.

Fixes googleapis#3717

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a pre-dry-run static validation step for BigQuery dataset restrictions to ensure queries referencing non-existent datasets outside the allowlist fail with a clear policy error rather than an opaque BigQuery 404. It also refactors the dataset validation logic into a reusable helper function and adds corresponding regression tests. The reviewer suggested optimizing performance by caching the allowed datasets slice and the parsed SQL tables to avoid redundant allocations and parsing overhead.

Comment on lines +159 to +169
if len(source.BigQueryAllowedDatasets()) > 0 {
parsedTables, parseErr := bqutil.TableParser(sql, bqClient.Project())
if parseErr != nil {
return nil, util.NewAgentError("could not parse tables from query to validate against allowed datasets", parseErr)
}
for _, tableID := range parsedTables {
if agentErr := checkDatasetAllowed(source, tableID); agentErr != nil {
return nil, agentErr
}
}
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Efficiency Improvement: Avoid redundant allocations and parsing

  1. Redundant Map-to-Slice Allocations: Calling source.BigQueryAllowedDatasets() allocates a new slice and copies the map keys every time it is called. Currently, it is called twice in Invoke (at line 159 and line 192). Storing the result in a local variable avoids this duplicate allocation.
  2. Redundant SQL Parsing: bqutil.TableParser is called once before the dry-run (line 160) and once after the dry-run (line 220). Since parsing SQL is relatively expensive, we can parse it once and reuse the parsedTables slice in both places.

We can optimize this by declaring allowedDatasets and parsedTables in the outer scope of Invoke so they can be reused in the post-dry-run validation block.

	var parsedTables []string
	allowedDatasets := source.BigQueryAllowedDatasets()
	if len(allowedDatasets) > 0 {
		var parseErr error
		parsedTables, parseErr = bqutil.TableParser(sql, bqClient.Project())
		if parseErr != nil {
			return nil, util.NewAgentError("could not parse tables from query to validate against allowed datasets", parseErr)
		}
		for _, tableID := range parsedTables {
			if agentErr := checkDatasetAllowed(source, tableID); agentErr != nil {
				return nil, agentErr
			}
		}
	}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bigquery-execute-sql: run allowedDatasets check before dry-run (non-existent foreign datasets return BQ 404)

2 participants