Skip to content

feat(safety): SELECT * ban + column-scope guard for execute_sql - #93

Merged
ashwin-agami merged 3 commits into
mainfrom
feat/column-scope-gate
Jul 9, 2026
Merged

ashwin-agami merged 3 commits into
mainfrom
feat/column-scope-gate

Conversation

@sandeep-agami

Copy link
Copy Markdown
Collaborator

Follow-up to #91 (table-scope gate). Adds two companion refuse checks to the shared _model_safety pass so every engine entry point (CLI + MCP execute_sql) rejects a query that names an undeclared column or uses SELECT * — not just whichever path obeyed a prose rule.

What changed

runtime.py (single-source):

  • check_no_select_star(sql) — refuses a projection-level * / t.* in any select (outer, subquery, CTE body, set-operation arm). COUNT(*) and other agg(*) still pass (the star sits inside the aggregate, not the projection). A star defeats column-level scoping, so we force every projected column to be named.
  • check_column_scope(sql, org) — a column that binds to a declared physical table must be one that table declares. Strict where a column visibly binds to a physical table (per-enclosing-SELECT scoping, case-insensitive, consistent with check_table_scope); fail-open on CTE/subquery-output and select-list-alias columns we can't attribute — matching the gate's degrade-to-allow posture, so legitimate complex SQL never false-refuses. Still catches hallucinated columns, including inside CTE/subquery bodies (they bind directly to their physical table).

execute_sql.py (_model_safety): both checks wired in right after check_table_scope, before the fan/chasm + sensitive checks. New error kinds: select_star, column_out_of_scope (exit 1). Re-vendored byte-identical to plugins/agami/lib/execute_sql.py.

Set-operation bypass (also fixes #91's gate)

Surfaced while writing the adversarial tests: sqlglot.parse_one("… UNION …") returns exp.Union, not exp.Select, so the not isinstance(tree, exp.Select) → allow short-circuit made every UNION/INTERSECT/EXCEPT arm skip the guard. check_table_scope (merged in #91) had the same hole — a SELECT id FROM orders UNION SELECT id FROM secret_table bypassed table-scope. All three checks now gate on "contains a SELECT" and walk every arm. Regression test added to test_table_scope_gate.py.

Tests (full suite: 1376 passed)

  • tests/test_column_scope_gate.py — star ban, column-scope happy path, legit CTE/subquery/alias, degrade-to-allow.
  • tests/test_column_scope_adversarial.py — star evasion (subquery/CTE/UNION-arm/t.*/comment-obfuscation), column evasion (every clause, expression/window wraps, alias masquerade, UNION arm, correlated subquery, COUNT(*) not over-blocked), documented accepted fail-opens (derived-alias-qualified, CTE-shadows-table — asserted allow so a future narrowing is a conscious decision), and the upstream-owned multi-statement case.
  • tests/test_table_scope_gate.py — UNION-arm regression.

Also verified end-to-end through _model_safety: correct exit codes, error kinds, and ordering (table-scope → star → column).

🤖 Generated with Claude Code

Follow-up to #91 (table-scope gate). Two new refuse checks in the shared
_model_safety pass, so every engine entry point (CLI + MCP execute_sql)
rejects a query that names an undeclared column or uses SELECT *:

- check_no_select_star: refuses a projection-level `*` / `t.*` in ANY select
  (outer, subquery, CTE body, or set-operation arm). COUNT(*) still passes.
- check_column_scope: a column that binds to a declared physical table must
  be one that table declares. Strict where a column binds to a physical table
  (per-enclosing-SELECT scoping, case-insensitive); fail-open on CTE/subquery
  output and select-list-alias columns, matching the gate's degrade-to-allow
  posture. Catches hallucinated columns, including inside CTE/subquery bodies.

Also closes a set-operation bypass exposed while writing the adversarial
tests: sqlglot parses `... UNION ...` to exp.Union, not exp.Select, so the
`isinstance(tree, exp.Select)` short-circuit made every UNION/INTERSECT/EXCEPT
arm skip the guard. check_table_scope had the same hole; all three checks now
gate on "contains a SELECT" and walk every arm.

New error kinds: select_star, column_out_of_scope. runtime.py is single-source;
execute_sql.py re-vendored to plugins/agami/lib (byte-identical).

Tests: test_column_scope_gate.py (happy path + degrade), test_column_scope_
adversarial.py (star/column evasion, documented fail-opens, upstream-owned),
and a UNION-arm regression in test_table_scope_gate.py. Full suite: 1376 passed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings July 8, 2026 00:30

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Pull request overview

This PR strengthens the semantic-model safety layer for SQL execution by (1) banning projection-level SELECT * and (2) enforcing that referenced columns must be declared on the physical table they bind to, and it also closes a set-operation (UNION/INTERSECT/EXCEPT) bypass in the existing table-scope gate. The new checks are wired into the shared _model_safety path for both the core executor and the vendored plugin copy, with comprehensive unit + adversarial tests.

Changes:

  • Fix check_table_scope to scope set-operation arms by gating on “contains a SELECT” rather than “is a SELECT”.
  • Add check_no_select_star and check_column_scope to semantic_model.runtime, and enforce them in execute_sql._model_safety.
  • Add happy-path + adversarial test suites covering SELECT * evasion and column-scope evasion (including within CTEs/subqueries and set operations).

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
packages/agami-core/src/semantic_model/runtime.py Fixes set-op table-scope bypass; adds SELECT * ban and column-scope enforcement helpers.
packages/agami-core/src/execute_sql.py Wires new star + column-scope checks into _model_safety refusal flow.
plugins/agami/lib/execute_sql.py Mirrors the _model_safety wiring in the vendored plugin library.
tests/test_table_scope_gate.py Adds regression tests ensuring set-operation arms are table-scoped.
tests/test_column_scope_gate.py Adds baseline unit tests for star ban + column-scope behavior and degrade-to-allow cases.
tests/test_column_scope_adversarial.py Adds adversarial coverage for common evasions and documented fail-open boundaries.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +672 to +675
alias_to_table: dict[str, str] = {} # alias/name -> bare physical table (global; for qualified refs)
direct_phys: dict[int, set[str]] = {} # id(select) -> {bare physical table read directly}
has_derived: dict[int, bool] = {} # id(select) -> reads a CTE ref / derived subquery directly
for tbl in tree.find_all(exp.Table):

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.

This was addressed in 440b883 (per-SELECT alias scoping) before the current head. The alias map is no longer global: alias_by_select is keyed by id(select) and qualifiers resolve through _select_chain (innermost→outermost), so an inner alias can't validate an outer column against the wrong table. Correlated refs still see ancestor aliases.

Comment on lines +694 to +713
# select-list output names (`AS x`) are aliases, not base columns
output_aliases = {a.alias.lower() for a in tree.find_all(exp.Alias) if a.alias}

offending: set[str] = set()
for col in tree.find_all(exp.Column):
name = col.name
if not name:
continue
lname = name.lower()
if col.table:
phys = alias_to_table.get(col.table.lower())
if phys is None:
continue # qualified by a CTE/derived alias — validated at its own source
if phys in declared and lname not in declared[phys]:
offending.add(f"{phys}.{name}")
continue
# unqualified: judge against the tables its own SELECT reads directly
if lname in output_aliases:
continue # a select-list output alias, not a base column
sel = _enclosing_select(col)

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.

Also addressed in 440b883: output aliases are tracked per-SELECT (output_by_select, keyed by id(select)), so an AS x in one scope no longer masks an unrelated unqualified x in a different scope.

Comment on lines +532 to +536
# A set operation (UNION/INTERSECT/EXCEPT) parses to exp.Union, not exp.Select,
# so gate on "contains a SELECT" rather than "is a SELECT" — otherwise every
# set-operation arm would bypass the guard. A non-SELECT statement has no SELECT
# node and still degrades to allow (the upstream read-only guard owns those).
if tree is None or tree.find(exp.Select) is None:

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.

Addressed in 711ceec — see the reply on the re-raised instance of this finding. Both check_sensitive_projection and pre_flight_check now walk set-operation arms.

check_column_scope built the qualified-alias map and the output-alias set
globally across the whole AST. Two correctness holes, both confirmed and now
regression-tested:

- Alias collision across scopes: reusing an alias name (e.g. `o` for orders
  outer and customers in a correlated subquery) let the global last-write-wins
  map validate an outer qualified column against the wrong table (false refuse).
- Nested output-alias mask: an inner `SELECT ... AS bogus` polluted the global
  output-alias set, so an unrelated outer undeclared column `bogus` was skipped
  (false allow — an undeclared column slipped the guard).

Both are now resolved per enclosing SELECT: qualified refs walk their own select
chain outward (correlated refs see ancestor aliases; an inner alias shadows an
outer), and the output-alias skip applies only within the same select. direct_phys
and has_derived were already per-select; this makes alias handling consistent.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Comment on lines +532 to +536
# A set operation (UNION/INTERSECT/EXCEPT) parses to exp.Union, not exp.Select,
# so gate on "contains a SELECT" rather than "is a SELECT" — otherwise every
# set-operation arm would bypass the guard. A non-SELECT statement has no SELECT
# node and still degrades to allow (the upstream read-only guard owns those).
if tree is None or tree.find(exp.Select) is None:

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.

Fixed in 711ceec. Applied the same set-operation arm-walking to both remaining gates:

  • check_sensitive_projection now gates on tree.find(exp.Select) and scans every output-bearing arm via a new _output_selects() helper (top-level SELECT + each UNION/INTERSECT/EXCEPT arm; nested subquery/CTE selects excluded so a WHERE-subquery projecting a sensitive column isn't falsely refused). SELECT id FROM customers UNION SELECT email FROM customers now refuses.
  • pre_flight_check is split into a per-SELECT _preflight_select() + a dispatcher that analyzes each arm; a trap in any arm refuses the whole query. Arms are not auto-rewritten (allow_rewrite=False) — a rewriteable arm-trap becomes a refuse.

Regression + over-refuse-control tests added (test_sample_database.py, test_semantic_model_runtime.py), mutation-checked (reverting each guard to isinstance fails the matching test).

… gates

Copilot review (PR #93) surfaced that the UNION/INTERSECT/EXCEPT bypass this
PR fixed for check_table_scope was still open on two sibling gates: both
check_sensitive_projection and pre_flight_check short-circuited on
`not isinstance(tree, exp.Select) -> allow`. A set operation parses to
exp.SetOperation, so:

  SELECT id FROM customers UNION SELECT email FROM customers

passed table/star/column scoping (arms walked) but bypassed the sensitive-
column gate — projecting PII undetected — and skipped fan/chasm detection.

- Add _output_selects(): the top-level SELECT, or each set-operation arm.
  Nested subquery/CTE SELECTs are excluded so a WHERE-subquery that projects a
  sensitive column (not exposed) is not falsely refused.
- check_sensitive_projection: scan every output-bearing arm.
- pre_flight_check: split into a per-SELECT _preflight_select() + a dispatcher
  that analyzes each arm; a trap in any arm refuses the whole query. Arms are
  not auto-rewritten (allow_rewrite=False) — a rewriteable arm-trap refuses.

Tests (regression + over-refuse controls), mutation-checked:
- sensitive: UNION/INTERSECT/parenthesized arm projecting `email` -> refuse;
  clean UNION of non-sensitive columns -> allow.
- pre_flight: fan trap in a UNION arm -> refuse; clean UNION -> allow.

runtime.py is not in the vendored plugin closure, so no lib re-sync.
@ashwin-agami

Copy link
Copy Markdown
Contributor

Addressed Copilot's review (commit 711ceec)

Triaged all four inline comments:

Already fixed in 440b883 (per-SELECT scoping) — stale from the first review pass:

  • Global alias map → now alias_by_select keyed by id(select), resolved via _select_chain (inner shadows outer, correlated refs see ancestors).
  • Global output-aliases → now output_by_select, per-SELECT.

Fixed now in 711ceec — the live finding (set-op bypass on the sibling gates):
The UNION/INTERSECT/EXCEPT arm-walking this PR added to check_table_scope was still missing on two gates that short-circuited not isinstance(tree, exp.Select) → allow:

  • check_sensitive_projection — a … UNION SELECT email FROM customers could project a sensitive column undetected (PII leak).
  • pre_flight_check — UNION arms skipped fan/chasm detection.

New _output_selects() (top-level SELECT + each arm; nested subquery/CTE selects excluded to avoid over-refusing a WHERE-subquery). pre_flight_check split into a per-SELECT _preflight_select() + a dispatcher; a trap in any arm refuses (arms aren't auto-rewritten).

Regression + over-refuse-control tests added; mutation-checked (reverting each guard to isinstance fails the matching test). Full gate green (1382 passed). runtime.py isn't vendored, so no lib re-sync.

@ashwin-agami
ashwin-agami merged commit 897a0cc into main Jul 9, 2026
6 checks passed
@ashwin-agami
ashwin-agami deleted the feat/column-scope-gate branch July 9, 2026 08:40
@github-actions github-actions Bot locked and limited conversation to collaborators Jul 9, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants