fix(bigquery): validate allowedDatasets before dry-run in execute-sql - #3759
fix(bigquery): validate allowedDatasets before dry-run in execute-sql#3759PratikDhanave wants to merge 1 commit into
Conversation
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
There was a problem hiding this comment.
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.
| 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 | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Efficiency Improvement: Avoid redundant allocations and parsing
- 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 inInvoke(at line 159 and line 192). Storing the result in a local variable avoids this duplicate allocation. - Redundant SQL Parsing:
bqutil.TableParseris 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 theparsedTablesslice 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
}
}
}
Summary
With
allowedDatasetsconfigured,bigquery-execute-sqlonly 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 (viaProcessGcpError) instead of the documented policy error: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
Invoke, whenallowedDatasetsis non-empty, statically parse the referenced datasets withTableParserand reject any outside the allowlist before the dry-run.ReferencedTablesenrichment), so no coverage is lost.checkDatasetAllowedhelper, used by both the pre- and post-dry-run paths.Testing
TestInvokeDatasetRestrictionscase 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.got HTTP response code 404 ... Not found: Table ...nonexistent_dataset) and passes with the fix — a genuine regression guard.TestInvokeDatasetRestrictionscases (allowed/forbidden tables, INFORMATION_SCHEMA, EXTERNAL_QUERY, etc.) still pass;go build,go vet, andgo teston the package are green.