From 5e9cd25bacdd85aacc19e4110013dcab9af4a377 Mon Sep 17 00:00:00 2001 From: Andrew Stevens Date: Mon, 17 Aug 2026 15:10:12 +0200 Subject: [PATCH] fix(security): close eight PII-denylist bypasses The denylist could be defeated eight ways. Two were reported from an adversarial review of a deployed agent; the other six surfaced while fixing them. Each was confirmed against 0.1.1 with a reproducing query before being fixed, and each has a regression test. Scope aliasing. PiiProjectionRule checked only the outermost select list via outermost_projection_names(), so any inner scope that renamed a denied column laundered it: WITH c AS (SELECT BillingCity AS city FROM t) SELECT city FROM c Derived tables, UNION arms and multi-hop alias chains worked the same way. Checking now runs over every SELECT scope and matches the underlying column names in the scope that names them, so an alias cannot launder a denied column. Value probing. The denylist gated projection only, so WHERE, GROUP BY, HAVING and ORDER BY references passed. None return the column, but each answers a yes/no question about its value, and enough queries reconstruct it. New pii_mode config ("reference" | "project", default "reference") denies any reference to a denied column. "project" is the documented loosening path and is still all-scope. Found while fixing the above: - SELECT c FROM tbl AS c returned every column of every row, PII included, and the guard auto-executed it. A bare table alias in a value position expands to the whole row; it parses as an ordinary column, so the denylist had nothing to match. Strictly worse than the SELECT * already blocked. New NoUnresolvableColumnsRule. - NATURAL JOIN joins on whichever columns the tables share, which the guard cannot enumerate without a schema. Now denied. - JOIN ... USING (email) produced zero exp.Column nodes, so reference mode never saw it - a working single-query value oracle. Same for AS g(email) column aliases and STRUCT('x' AS email) field names; all three carry names as bare exp.Identifier and are now harvested. - The star check only inspected the projection's root node, so OBJECT_CONSTRUCT(*), COLUMNS(*), * APPLY(f) and ROW(c.*) passed on non-BigQuery dialects. It is now a deep walk with COUNT(*) as the explicit carve-out. ClickHouse COLUMNS('regex'), which has no Star node at all, is matched on node type. - Qualified t.* bypassed the star rule even at top level: it parses as an exp.Column wrapping a Star. - Every exp.AggFunc counted as PII-neutralising, so MAX(email), ARRAY_AGG(email) and STRING_AGG(email) returned real values through project mode. Only aggregates reducing to a derived statistic qualify now. NoSelectStarRule replaces NoTopLevelStarRule, which stays importable as an alias. default_rules ordering semantics are unchanged. BREAKING: queries 0.1.1 allowed are now denied. Alias laundering, inner-scope and nested stars, whole-row aliases, NATURAL JOIN, identifier-only references and PII through value-preserving aggregates have no supported way back - that is the point. Denied columns in predicates can be restored with pii_mode="project", which re-opens the value-probing oracle. The bundled Q1 identity-resolution query is now denied in both modes: its CTE projects the denied columns, and COUNTIF(email_norm = 'target') is itself an oracle. Also: minimum Python is now 3.11. The package has imported enum.StrEnum since 0.1.0 while advertising >=3.10, so import failed on 3.10 and that CI job could never have passed. Metadata, classifiers, ruff target-version, mypy python_version and the CI matrix now agree. Cleared two pre-existing CI failures: a redundant int() around math.floor in format_cost (RUF046), and a strict-mypy no-untyped-call on sqlglot's unalias(), replaced with an equivalent .this read. Tests: 74 -> 204. --- .github/workflows/ci.yml | 2 +- CHANGELOG.md | 173 +++++++++- README.md | 104 +++++- pyproject.toml | 9 +- src/sql_guard/__init__.py | 32 +- src/sql_guard/sql_guard.py | 516 +++++++++++++++++++++++++++-- tests/conftest.py | 20 ++ tests/test_pii_scopes.py | 347 +++++++++++++++++++ tests/test_sql_guard.py | 107 +++++- tests/test_unresolvable_columns.py | 390 ++++++++++++++++++++++ 10 files changed, 1643 insertions(+), 57 deletions(-) create mode 100644 tests/test_pii_scopes.py create mode 100644 tests/test_unresolvable_columns.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d8f07ec..b668bd7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,7 +11,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13"] + python-version: ["3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c0cc44..f7dd955 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,176 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.2.0] — 2026-08-17 + +Security release. Closes eight PII-denylist bypasses: two found by adversarial +review of a deployed agent, and six more found while fixing those. Every one +was confirmed against 0.1.1 with a reproducing query before being fixed, and +each has a regression test. **Contains breaking behaviour changes** — queries +that 0.1.1 allowed are now denied. That is the point of the release; see +*Upgrading*. + +The most severe was not in the original report: `SELECT c FROM tbl AS c` +returned every column of every row, PII included, and the guard auto-executed +it. + +### Security +- **PII denylist is now enforced in every query scope, not just the outermost + select list.** `PiiProjectionRule` called `outermost_projection_names`, so + any inner scope that renamed a denied column laundered it past the guard. + All of these returned `confirm` on 0.1.1 and are denied as of 0.2.0: + + ```sql + WITH c AS (SELECT BillingCity AS city FROM `p.d.t`) SELECT city FROM c + SELECT x FROM (SELECT BillingCity AS x FROM `p.d.t`) + SELECT uid FROM `p.d.t` UNION ALL SELECT c FROM (SELECT email AS c FROM `p.d.t`) + WITH a AS (SELECT email AS e1 FROM `p.d.t`), b AS (SELECT e1 AS e2 FROM a) SELECT e2 FROM b + ``` + + Matching applies to the underlying column names within each scope, so an + alias never launders a denied column. +- **Denied columns can no longer be used to probe values.** The denylist gated + projection only, so `WHERE BillingCity = 'Columbus'`, `GROUP BY`, `HAVING` + and `ORDER BY` references all passed — none return the column, but each + answers a yes/no question about its value, and enough queries reconstruct it. + The new `pii_mode` config switch denies *any* reference by default. +- **`SELECT *` is rejected in every scope,** not just the outermost one. Once + PII checking is all-scope, a star inside a CTE or derived table makes that + scope's projection list unresolvable, so the guard cannot prove a denied + column is absent — same rationale as the existing top-level rule. +- **Qualified `t.*` is now caught.** It parses as an `exp.Column` wrapping a + `Star`, which the previous `isinstance(projection, exp.Star)` check missed, + so `SELECT t.* FROM tbl t` bypassed the star rule even at the top level. +- **Whole-row references by row-source alias are rejected.** `SELECT c FROM tbl AS c` + returns every column in the row as a struct — strictly more powerful than the + `SELECT *` the guard already blocked, and it parsed as an ordinary column + named `c`, so the denylist had nothing to match. Also covers + `TO_JSON_STRING(c)`, `ARRAY_AGG(c)`, `STRUCT(c)`, the unaliased + `SELECT tbl FROM tbl` form, CTE names, derived-table aliases + (`SELECT d FROM (SELECT ...) AS d`), and `VALUES` / `PIVOT` aliases. + New `NoUnresolvableColumnsRule`. + + The rule resolves ambiguity from the AST rather than denying on a bare name + collision, which would have made it unusable: a table contributes only the + name it is addressable by (an aliased `p.d.status AS s` does not reserve + `status`), and a CTE or derived table publishes its own output names, so + `WITH revenue AS (SELECT ..., SUM(x) AS revenue ...) SELECT revenue FROM revenue` + — a mainstream idiom — is correctly read as a column reference. +- **ClickHouse `COLUMNS('regex')` is rejected.** It expands to many columns but + parses with no `Star` node at all, so a star check keyed on `exp.Star` alone + never saw it. +- **`NATURAL JOIN` is rejected.** It joins on whichever columns the tables + share — a schema fact the guard does not have, so it cannot rule out a denied + column among them. `JOIN ... USING (...)` remains allowed and is now read. +- **Stars nested inside a wrapping construct are caught.** The star check only + inspected the projection's root node, so `OBJECT_CONSTRUCT(*)` (Snowflake), + `COLUMNS(*)` (DuckDB), `* APPLY(f)` (ClickHouse) and `ROW(c.*)` (Trino) all + passed. It is now a deep walk, with `COUNT(*)` / `COUNT(DISTINCT *)` as the + explicit carve-out. +- **Column names that never become an `exp.Column` are now read.** sqlglot + parses several positions as bare `exp.Identifier`, so reference mode's + `find_all(exp.Column)` sweep missed them: `JOIN ... USING (email)`, + `... AS g(email)` column aliases, and `STRUCT('x' AS email)` field names. + `JOIN ... USING (email)` was a working single-query value oracle. +- **Aggregation is no longer a blanket PII exemption.** Every `exp.AggFunc` + counted as PII-neutralising, so `MAX(email)`, `MIN(email)`, + `ARRAY_AGG(email)`, `STRING_AGG(email)` and `ANY_VALUE(email)` returned real + values through `pii_mode="project"`. Only aggregates that reduce to a derived + statistic (`COUNT`, `COUNTIF`, `SUM`, `AVG`, `APPROX_COUNT_DISTINCT`, + `STDDEV`, `VARIANCE`) qualify now. + +### Added +- `pii_mode` on `SqlGuardConfig` and `SqlGuardConfig.from_settings`, typed as + `PiiMode = Literal["reference", "project"]`. Defaults to `"reference"` + (deny any reference to a denied column). `"project"` restores 0.1.1-style + projection-only checking, still applied across all scopes. An invalid value + raises `ValueError` rather than silently falling back to a default. +- `NoSelectStarRule` — the all-scope star rule. `NoTopLevelStarRule` remains + importable as an alias of it. +- `NoUnresolvableColumnsRule` — rejects whole-row table-alias references and + `NATURAL JOIN`. Added to `default_rules` between the star rule and the PII + rule. +- AST helpers for custom rules: `all_selects`, `all_projection_names`, + `all_referenced_column_names`, `has_select_star`, `whole_row_references`. + +### Changed +- `default_rules` substitutes `NoSelectStarRule` for `NoTopLevelStarRule` + (the same object under its new name). Ordering semantics are unchanged: the + star rule still runs before `PiiProjectionRule`, which now also guarantees a + star is rejected before the PII check tries to resolve a projection list it + cannot see. +- `outermost_projection_names` and `has_top_level_select_star` are retained + and still outermost-only; their docstrings now warn against using them for + denylist enforcement. `has_top_level_select_star` does now recognise `t.*`. +- **Minimum Python is now 3.11.** The package has imported `enum.StrEnum` + (3.11+) since 0.1.0 while advertising `requires-python = ">=3.10"`, so + `import sql_guard` raised `ImportError` on 3.10 and that CI matrix job could + never have passed. `requires-python`, the classifiers, ruff's + `target-version`, mypy's `python_version` and the CI matrix now agree on + 3.11. No source change was needed — this documents what the code already + required. + +### Fixed +- `format_cost` no longer wraps `math.floor` in a redundant `int()` + (`math.floor` already returns an `int`). Behaviour is unchanged; this clears + a `ruff check` failure. +- `_projection_names` reads `.this` instead of calling sqlglot's untyped + `unalias()`. Identical result — `unalias` returns `self.this` for an `Alias` + — but it clears the only `mypy --strict` error in the package. + +### Known limits (documented, not fixed) + +The denylist matches column names in the SQL text and the guard has no schema, +so these remain out of scope and are now stated explicitly in the README: +PII inside JSON/VARIANT/STRUCT payloads addressed by string literals +(`JSON_VALUE(payload, '$.email')`), re-identification through non-PII columns, +and side channels such as row counts and dry-run byte counts. + +Also uncovered: selecting a STRUCT column whole, and `UNNEST` aliases over an +array of structs — both return every field without naming one, and neither is +distinguishable at parse time from the legitimate scalar-array form. + +Separately: several entries in `_PII_SAFE_FUNC_KEYS` (`sha256`, +`farm_fingerprint`) match no sqlglot key and never fire, so those spellings of +hashed PII are *denied* rather than treated as neutralised. `md5` does match +via the `TO_HEX(MD5(x))` folding, so that spelling is permitted under +`pii_mode="project"` — the two are inconsistent. Left as-is because the +inconsistency errs toward denial; revisit deliberately when deciding whether +hashed PII is acceptable output. + +### Known issues (unfixed, pre-existing) + +- `SelectOnlyRule` accepts `exp.Select | exp.Union`, but sqlglot 26.x derives + `Except` and `Intersect` from `SetOperation` rather than `Union`, so a + top-level `EXCEPT DISTINCT` / `INTERSECT` is rejected as "Only SELECT + statements are allowed". Fails closed, so it is an availability bug rather + than a security one, but it makes the `Except`/`Intersect` branch of + `outer_selects` dead code for top-level set operations. +- `AllowedTablesRule` runs last, so a query that is both off-allowlist and + star/PII-violating is always reported as the latter. Telemetry built on + `decision.reason` will under-count allowlist breaches. + +### Upgrading + +Queries denied by 0.2.0 that 0.1.1 allowed fall into two groups: + +1. **Bypasses** — alias laundering, inner-scope and nested stars, whole-row + aliases, `NATURAL JOIN`, identifier-only column references, and PII through + value-preserving aggregates. There is no supported way to re-permit these, + by design. Whole-row alias references are fixed by qualifying them + (`SELECT c.col` rather than `SELECT c`). +2. **Denied columns in predicates** — set `pii_mode="project"` to restore the + old behaviour. Understand that this re-opens the value-probing oracle + described above; prefer narrowing the denylist or exposing a pre-masked + view. + +The bundled `Q1` identity-resolution query (normalising `email`/`mobile` in a +CTE, projecting only `COUNTIF` aggregates) is denied in **both** modes: the CTE +scope projects the denied columns, and `COUNTIF(email_norm = 'target')` is +itself a value oracle. Aggregates over denied columns are not safe under this +threat model. Move such normalisation into a warehouse view the denylist does +not cover. + ## [0.1.1] — 2026-06-19 ### Added @@ -51,6 +221,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `SELECT * EXCEPT(...)` is rejected — the guard cannot prove the EXCEPT list enumerates every PII column. -[Unreleased]: https://github.com/sakura-sky/sql-guard/compare/v0.1.1...HEAD +[Unreleased]: https://github.com/sakura-sky/sql-guard/compare/v0.2.0...HEAD +[0.2.0]: https://github.com/sakura-sky/sql-guard/compare/v0.1.1...v0.2.0 [0.1.1]: https://github.com/sakura-sky/sql-guard/compare/v0.1.0...v0.1.1 [0.1.0]: https://github.com/sakura-sky/sql-guard/releases/tag/v0.1.0 diff --git a/README.md b/README.md index 718a9c6..e682054 100644 --- a/README.md +++ b/README.md @@ -33,19 +33,104 @@ if decision.denied: 1. **Single SELECT only.** DML, DDL, scripts, multi-statement payloads — all rejected. Even if buried in subqueries. -2. **PII column denylist.** Projections that touch denylisted columns are - rejected. Catches aliased PII (`SELECT email AS x`), PII through transforms - (`SELECT LOWER(email)`), and the right arm of a `UNION ALL`. -3. **No top-level `SELECT *`.** Bare `*`, `* EXCEPT(...)`, and `* REPLACE(...)` - are all rejected — the guard can't prove EXCEPT enumerates every PII column. -4. **Table allowlist.** Only fully-qualified tables you approved can be +2. **PII column denylist.** By default, *any reference* to a denylisted column + is rejected — in the select list, `WHERE`, `GROUP BY`, `HAVING`, `ORDER BY` + or a `JOIN` condition, at any nesting depth. Catches aliased PII + (`SELECT email AS x`), PII through transforms (`SELECT LOWER(email)`), + the right arm of a `UNION ALL`, and columns renamed inside a CTE or derived + table. See [PII modes](#pii-modes). +3. **No `SELECT *` in any scope.** Bare `*`, `* EXCEPT(...)`, `* REPLACE(...)` + and qualified `t.*` are all rejected, inside CTEs and subqueries as well as + at the top level — the guard can't prove EXCEPT enumerates every PII column, + nor what `*` expands to. The check is a deep walk, so stars wrapped in a + function (`OBJECT_CONSTRUCT(*)`, `COLUMNS(*)`, `* APPLY(f)`) are caught too; + `COUNT(*)` is the deliberate exception. +4. **Nothing whose columns the guard can't enumerate.** A bare table alias in a + value position (`SELECT c FROM tbl AS c`) returns the whole row as a struct, + and `NATURAL JOIN` joins on unknown shared columns. Both are rejected for the + same reason as `SELECT *`. +5. **Table allowlist.** Only fully-qualified tables you approved can be referenced. CTE aliases are excluded. -5. **Cost cap.** Given the bytes-processed figure from a dry-run, the guard +6. **Cost cap.** Given the bytes-processed figure from a dry-run, the guard returns `allow` below your auto threshold, `confirm` in between, and `deny` above the hard cap or bytes-billed ceiling. Every check is a `Rule` you can replace or compose with. +### PII modes + +`pii_mode` controls how far the denylist reaches. + +| Mode | Denies | Use when | +|---|---|---| +| `"reference"` (default) | Any reference to a denied column, in any clause and any scope. | The agent must not learn PII values at all. | +| `"project"` | Only projections of denied columns — checked in every scope. | Predicate access to PII is a deliberate, accepted trade-off. | + +The default is the strict one because projection-only checking leaves the +values reachable. A denied column in a `WHERE` clause never appears in the +output, but the row count still answers a yes/no question about it: + +```sql +-- Passes a projection-only guard. Returns 0 or non-zero. +SELECT COUNT(*) FROM `p.d.orders` WHERE billing_city = 'Columbus' +``` + +Repeat with `LIKE 'a%'`, `> 'm'`, and so on, and the value falls out in a +handful of queries. `GROUP BY`, `HAVING` and `ORDER BY` leak the same way. + +**The loosening path.** If your deployment genuinely needs to filter on PII — +segmenting on a hashed identifier, say, or counting non-null contact rows — +set `pii_mode="project"`: + +```python +SqlGuardConfig.from_settings(..., pii_mode="project") +``` + +That re-permits denied columns in predicates while still rejecting every +projection of them, in every scope. Prefer narrowing the denylist, or exposing +a pre-masked warehouse view the denylist doesn't cover, before reaching for it. + +Note that "aggregate" is not a safe harbour in either mode. `MAX(email)`, +`ARRAY_AGG(email)` and `STRING_AGG(email)` return real values and are rejected; +only aggregates that reduce to a statistic (`COUNT`, `SUM`, `AVG`, `STDDEV`, …) +are treated as PII-neutralising, and then only under `pii_mode="project"`. + +### What the denylist does not cover + +The guard matches **column names in the SQL text**. It has no schema, so: + +- **PII inside JSON / VARIANT / STRUCT payloads** is not covered. + `JSON_VALUE(payload, '$.email')` names only `payload`; the field name is a + string literal the engine resolves. The same applies to selecting a struct + column whole (`SELECT contact FROM t`) and to `UNNEST` aliases over an array + of structs (`SELECT s FROM t, UNNEST(t.contacts) AS s`) — both return every + field without naming one. **Denylist the containing column.** +- **Re-identification through non-PII columns** is out of scope. If `uid` maps + 1:1 to a person, blocking `email` does not prevent correlation with outside + data. +- **Side channels** — row counts, dry-run byte counts and error messages carry + bits about denied values even when every direct reference is refused. + +These are limits of a parse-level guard, not bugs. Warehouse-side column +security is the durable answer; `sql-guard` is defence in depth. + +### Scopes + +Denylist and star checks run against **every** `SELECT` scope — CTE bodies, +derived tables, scalar and `IN` subqueries, and each arm of a set operation — +matching on the underlying column names in the scope that names them. An alias +therefore cannot launder a denied column: + +```sql +-- Denied: the CTE scope still names billing_city. +WITH c AS (SELECT billing_city AS city FROM `p.d.orders`) +SELECT city FROM c +``` + +Checking only the outermost select list would see `city` and let it through. +The same applies through derived tables, `UNION` arms, and multi-hop alias +chains (`a AS (...) → b AS (...) → SELECT`). + ### The cost-cap rule in detail Three independent thresholds bound any single query: @@ -70,7 +155,8 @@ through — useful for sandboxes or onboarding a new tenant. job; pass `bytes_processed` to `evaluate_cost`. Keeps the guard testable without credentials and dialect-agnostic. - It does not introspect table schemas. If you say "this table is allowed," - the guard takes your word for it. + the guard takes your word for it. This is why `SELECT *` is rejected + everywhere: without a schema the guard cannot enumerate what `*` returns. - It does not authorise the user. Identity, IAM, row-level security: not in scope. The guard is a *policy* layer, not a *permissions* layer. @@ -159,7 +245,7 @@ pip install 'sql-guard[adk]' # + Google ADK + BigQuery client for the # FunctionTool integration ``` -Python 3.10+ supported. +Python 3.11+ supported. ## License diff --git a/pyproject.toml b/pyproject.toml index a35b6ee..b799d84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,9 +1,9 @@ [project] name = "sql-guard" -version = "0.1.1" +version = "0.2.0" description = "Deterministic policy engine for LLM-generated SQL: SELECT-only, PII column denylist, table allowlist, cost cap. Multi-dialect via sqlglot (BigQuery, Snowflake, Postgres, Trino, …)." readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" license = { text = "Apache-2.0" } authors = [{ name = "Sakura Sky Engineering" }] keywords = [ @@ -24,7 +24,6 @@ classifiers = [ "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -64,11 +63,11 @@ packages = ["src/sql_guard"] [tool.ruff] line-length = 100 -target-version = "py310" +target-version = "py311" src = ["src", "tests"] [tool.mypy] -python_version = "3.10" +python_version = "3.11" strict = true warn_return_any = true warn_unused_configs = true diff --git a/src/sql_guard/__init__.py b/src/sql_guard/__init__.py index 7be074d..6d7f273 100644 --- a/src/sql_guard/__init__.py +++ b/src/sql_guard/__init__.py @@ -4,17 +4,23 @@ from sql_guard import ( SqlGuard, SqlGuardConfig, GuardDecision, GuardOutcome, - PiiDenylist, + PiiDenylist, PiiMode, Rule, RuleContext, default_rules, # built-in rules - SelectOnlyRule, NoEmbeddedDmlRule, NoTopLevelStarRule, - PiiProjectionRule, AllowedTablesRule, + SelectOnlyRule, NoEmbeddedDmlRule, NoSelectStarRule, + NoUnresolvableColumnsRule, PiiProjectionRule, AllowedTablesRule, + SingleStatementRule, # cost models CostModel, BigQueryOnDemandCost, FlatRateCost, # AST helpers for custom rules - outer_selects, outermost_projection_names, referenced_tables, + all_selects, all_projection_names, all_referenced_column_names, + has_select_star, whole_row_references, referenced_tables, + outer_selects, outermost_projection_names, has_top_level_select_star, format_bytes, format_cost, ) + +``NoTopLevelStarRule`` remains importable as an alias of ``NoSelectStarRule``, +which now covers every query scope rather than only the outermost one. """ from __future__ import annotations @@ -28,7 +34,10 @@ GuardDecision, GuardOutcome, NoEmbeddedDmlRule, + NoSelectStarRule, NoTopLevelStarRule, + NoUnresolvableColumnsRule, + PiiMode, PiiProjectionRule, Rule, RuleContext, @@ -36,16 +45,21 @@ SingleStatementRule, SqlGuard, SqlGuardConfig, + all_projection_names, + all_referenced_column_names, + all_selects, default_rules, format_bytes, format_cost, + has_select_star, has_top_level_select_star, outer_selects, outermost_projection_names, referenced_tables, + whole_row_references, ) -__version__ = "0.1.1" +__version__ = "0.2.0" __all__ = [ "AllowedTablesRule", @@ -55,8 +69,11 @@ "GuardDecision", "GuardOutcome", "NoEmbeddedDmlRule", + "NoSelectStarRule", "NoTopLevelStarRule", + "NoUnresolvableColumnsRule", "PiiDenylist", + "PiiMode", "PiiProjectionRule", "Rule", "RuleContext", @@ -65,11 +82,16 @@ "SqlGuard", "SqlGuardConfig", "__version__", + "all_projection_names", + "all_referenced_column_names", + "all_selects", "default_rules", "format_bytes", "format_cost", + "has_select_star", "has_top_level_select_star", "outer_selects", "outermost_projection_names", "referenced_tables", + "whole_row_references", ] diff --git a/src/sql_guard/sql_guard.py b/src/sql_guard/sql_guard.py index 25d1ba9..211ca15 100644 --- a/src/sql_guard/sql_guard.py +++ b/src/sql_guard/sql_guard.py @@ -30,7 +30,7 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from enum import StrEnum -from typing import Final, Protocol, runtime_checkable +from typing import Final, Literal, Protocol, runtime_checkable import sqlglot from sqlglot import expressions as exp @@ -132,6 +132,27 @@ def bytes_to_usd(self, bytes_processed: int) -> float: # --------------------------------------------------------------------------- +PiiMode = Literal["reference", "project"] +"""How strictly the PII denylist is applied. + +``"reference"`` (default) + Deny *any* reference to a denylisted column, anywhere in the statement — + projections, ``WHERE``, ``GROUP BY``, ``HAVING``, ``ORDER BY``, ``JOIN`` + conditions. Denylisted values cannot be probed even indirectly. + +``"project"`` + Deny only projections of denylisted columns (checked in every query + scope). Predicates may reference them. + +``"project"`` is the looser setting. It permits a caller to binary-search a +denied value without ever projecting it — ``WHERE email = 'a@b.com'`` returns +zero or non-zero rows, and repeated queries reconstruct the value. Choose it +only when predicate access to PII is a deliberate, accepted trade-off. +""" + +_PII_MODES: Final[frozenset[str]] = frozenset({"reference", "project"}) + + @dataclass(frozen=True) class SqlGuardConfig: """Configuration for :class:`SqlGuard`. @@ -139,6 +160,11 @@ class SqlGuardConfig: Attributes: pii_denylist: Column names the guard refuses to project. + pii_mode: + How strictly the denylist applies — see :data:`PiiMode`. Defaults + to ``"reference"``, which denies any reference to a denylisted + column. Set to ``"project"`` to allow denylisted columns in + predicates and only block projections. allowed_tables: Fully-qualified table names the guard permits. Empty/disabled set means "no allowlist enforcement" — every table passes. @@ -166,12 +192,22 @@ class SqlGuardConfig: pii_denylist: PiiDenylist allowed_tables: frozenset[str] dialect: str = "bigquery" + pii_mode: PiiMode = "reference" cost_model: CostModel = field(default_factory=BigQueryOnDemandCost) max_cost_usd_auto: float = 0.10 max_cost_usd_hard: float = 20.00 max_bytes_billed: int = 10 * 1024**3 # 10 GiB enforce_allowed_tables: bool = True + def __post_init__(self) -> None: + # Fail loudly on a typo. Silently falling back to a default would pick + # a policy the operator did not ask for — the one thing a guard must + # never do. + if self.pii_mode not in _PII_MODES: + raise ValueError( + f"pii_mode must be one of {sorted(_PII_MODES)}; got {self.pii_mode!r}.", + ) + @classmethod def from_settings( cls, @@ -179,6 +215,7 @@ def from_settings( allowed_tables: Iterable[str], *, dialect: str = "bigquery", + pii_mode: PiiMode = "reference", cost_model: CostModel | None = None, max_cost_usd_auto: float = 0.10, max_cost_usd_hard: float = 20.00, @@ -189,6 +226,7 @@ def from_settings( pii_denylist=pii_denylist, allowed_tables=frozenset(t.lower() for t in allowed_tables), dialect=dialect, + pii_mode=pii_mode, cost_model=cost_model or BigQueryOnDemandCost(), max_cost_usd_auto=max_cost_usd_auto, max_cost_usd_hard=max_cost_usd_hard, @@ -276,37 +314,140 @@ def evaluate(self, ctx: RuleContext) -> GuardDecision | None: @dataclass(frozen=True) -class NoTopLevelStarRule: - """Reject ``SELECT *`` (and ``* EXCEPT(...)`` / ``* REPLACE(...)``). +class NoSelectStarRule: + """Reject ``SELECT *`` in *any* query scope. + + Covers bare ``*``, ``* EXCEPT(...)``, ``* REPLACE(...)`` and qualified + ``t.*``, in the outer select list and inside CTE bodies, derived tables + and subqueries alike. The guard cannot prove that EXCEPT enumerates every PII column, and new PII columns added later would silently start leaking. Callers must list columns explicitly. + + The all-scope reach is what makes :class:`PiiProjectionRule` sound. A star + anywhere makes that scope's projection list unresolvable without a schema, + so a denied column could flow out through it — + ``WITH c AS (SELECT * FROM t) SELECT city FROM c`` never names + ``BillingCity``, yet returns it. Same rationale as the docstring above: + the guard cannot prove what ``*`` contains. """ def evaluate(self, ctx: RuleContext) -> GuardDecision | None: - if has_top_level_select_star(ctx.statement): + if has_select_star(ctx.statement): return _deny( - "Top-level `SELECT *` is not allowed (including `* EXCEPT(...)` " - "and `* REPLACE(...)`). List columns explicitly so the PII " + "`SELECT *` is not allowed in any query scope — including " + "`* EXCEPT(...)`, `* REPLACE(...)`, qualified `t.*`, and stars " + "inside CTEs or subqueries. List columns explicitly so the PII " "denylist can be enforced.", referenced_tables=tuple(sorted(ctx.referenced_tables)), ) return None +# Back-compat alias: the rule outgrew its "top-level" name when it was extended +# to every scope. Consumers importing the old name keep working. +NoTopLevelStarRule = NoSelectStarRule + + +@dataclass(frozen=True) +class NoUnresolvableColumnsRule: + """Reject constructs whose column set the guard cannot enumerate. + + Sibling of :class:`NoSelectStarRule`, same argument: if the guard cannot + name the columns a construct touches, it cannot prove they are denylist- + free. Two such constructs exist beyond ``SELECT *``: + + **Whole-row references.** A bare table alias in a value position expands to + every column in the row — a ``STRUCT`` in BigQuery, a composite in Postgres, + a struct in DuckDB:: + + SELECT c FROM `p.d.customers` AS c -- every column, incl. PII + SELECT TO_JSON_STRING(c) FROM `p.d.t` AS c -- same, as JSON text + + This parses as an ordinary ``exp.Column`` named ``c``, so a denylist check + sees one unremarkable non-PII name. It is strictly more powerful than + ``SELECT *``, which the guard already rejects. + + **``NATURAL JOIN``.** Joins on whatever columns the two tables happen to + share. Which columns those are is a schema fact the guard does not have, so + it cannot rule out a denied column among them. ``JOIN ... USING (col)`` is + fine by contrast — the columns are named, and + :func:`all_referenced_column_names` reads them. + """ + + def evaluate(self, ctx: RuleContext) -> GuardDecision | None: + for join in ctx.statement.find_all(exp.Join): + method = join.args.get("method") + if isinstance(method, str) and method.upper() == "NATURAL": + return _deny( + "`NATURAL JOIN` is not allowed — it joins on whichever " + "columns the tables share, which the guard cannot " + "enumerate without a schema, so it cannot rule out a PII " + "column among them. Use an explicit `JOIN ... ON` or " + "`USING (...)` naming the join columns.", + referenced_tables=tuple(sorted(ctx.referenced_tables)), + ) + + offenders = sorted(set(whole_row_references(ctx.statement))) + if offenders: + return _deny( + f"Query references whole rows by table alias " + f"({', '.join(offenders)}). A bare table alias expands to every " + "column in the row, so the guard cannot prove the result is " + "free of PII — the same reason `SELECT *` is rejected. " + "Reference the columns you need explicitly (`alias.column`).", + referenced_tables=tuple(sorted(ctx.referenced_tables)), + ) + return None + + @dataclass(frozen=True) class PiiProjectionRule: - """Reject queries that project PII-denylisted columns.""" + """Reject queries that touch PII-denylisted columns. + + Two modes, selected by :attr:`SqlGuardConfig.pii_mode`: + + ``"reference"`` (default) + Deny any reference to a denied column anywhere in the statement. + Projection-only checking leaves the value probeable: + ``WHERE BillingCity = 'Columbus'`` never projects the column, but the + row count answers a yes/no question about its value, and enough such + questions reconstruct it. + + ``"project"`` + Deny only projections, checked in *every* scope. + + Both modes match on the underlying column names within each scope, so an + alias cannot launder a denied column: + ``WITH c AS (SELECT BillingCity AS city FROM t) SELECT city FROM c`` is + caught in the CTE scope, where ``BillingCity`` is still named. Checking + only the outermost select list would see ``city`` and pass it. + """ def evaluate(self, ctx: RuleContext) -> GuardDecision | None: - projected = outermost_projection_names(ctx.statement) - hits = ctx.config.pii_denylist.matching(projected) + if ctx.config.pii_mode == "project": + names = all_projection_names(ctx.statement) + hits = ctx.config.pii_denylist.matching(names) + if hits: + unique = tuple(sorted(set(hits))) + return _deny( + f"Query projects PII columns ({', '.join(unique)}). " + "Aggregate-only or non-PII columns are allowed.", + pii_columns=unique, + ) + return None + + names = all_referenced_column_names(ctx.statement) + hits = ctx.config.pii_denylist.matching(names) if hits: unique = tuple(sorted(set(hits))) return _deny( - f"Query projects PII columns ({', '.join(unique)}). " - "Aggregate-only or non-PII columns are allowed.", + f"Query references PII columns ({', '.join(unique)}). " + "Denylisted columns cannot be projected, filtered, grouped, or " + "sorted on — not even via an alias, CTE or subquery. Use " + 'non-PII columns, or run the guard with pii_mode="project" if ' + "predicate access to PII is an accepted trade-off.", pii_columns=unique, ) return None @@ -332,15 +473,26 @@ def evaluate(self, ctx: RuleContext) -> GuardDecision | None: def default_rules(_config: SqlGuardConfig) -> list[Rule]: """Return the built-in rules in their canonical order. - Order matters: :class:`NoTopLevelStarRule` runs before + Order matters: :class:`NoSelectStarRule` runs before :class:`PiiProjectionRule` so a ``SELECT * EXCEPT(email)`` gets the "list columns" message rather than a confusing PII-message about names that happen to appear in EXCEPT. + + That ordering also keeps the star rule's guarantee ahead of the PII check + in every scope, not just the outermost one: a star inside a CTE is + rejected before :class:`PiiProjectionRule` tries to resolve a projection + list it cannot see. + + :class:`NoUnresolvableColumnsRule` sits in the same slot and for the same + reason — whole-row aliases and ``NATURAL JOIN`` hide their column sets from + the denylist, so they are rejected before the PII check runs and reports a + misleading "no PII found". """ return [ SelectOnlyRule(), NoEmbeddedDmlRule(), - NoTopLevelStarRule(), + NoSelectStarRule(), + NoUnresolvableColumnsRule(), PiiProjectionRule(), AllowedTablesRule(), ] @@ -531,6 +683,12 @@ def outermost_projection_names(statement: exp.Expression) -> list[str]: Walks both arms of UNION/UNION ALL/EXCEPT/INTERSECT. For each projection, contributes the alias (if any) plus every internal ``exp.Column`` ref. Projections wrapped in PII-neutralizing functions return no names. + + .. warning:: + Outermost scope only. A denied column projected inside a CTE or derived + table and re-exposed under an alias is invisible here. Rules enforcing a + denylist want :func:`all_projection_names`; this helper is retained for + callers that specifically need the outer select list. """ out: list[str] = [] for select in outer_selects(statement): @@ -539,6 +697,211 @@ def outermost_projection_names(statement: exp.Expression) -> list[str]: return out +def all_selects(statement: exp.Expression) -> list[exp.Select]: + """Every ``Select`` scope in *statement*, outermost included. + + Covers CTE bodies, derived tables, scalar and ``IN`` subqueries, and every + arm of a set operation — anywhere a projection list can hide. + """ + return list(statement.find_all(exp.Select)) + + +def all_projection_names(statement: exp.Expression) -> list[str]: + """Projection names from *every* scope, not just the outermost. + + Each scope contributes its own projections' aliases plus the underlying + ``exp.Column`` refs, so a denied column is caught in the scope that names + it even if later scopes only ever see the alias. + """ + out: list[str] = [] + for select in all_selects(statement): + for projection in select.expressions: + out.extend(_projection_names(projection)) + return out + + +def all_referenced_column_names(statement: exp.Expression) -> list[str]: + """Every column name referenced anywhere in *statement*, in any clause. + + Includes projections, ``WHERE``, ``GROUP BY``, ``HAVING``, ``ORDER BY``, + ``QUALIFY`` and ``JOIN`` conditions, at every nesting depth. This is the + name set behind ``pii_mode="reference"``: a denied column used purely as a + filter still leaks its values one predicate at a time. + + Projection aliases are folded in as well, so ``SELECT other AS email`` + is caught by an ``email`` denylist entry. Star refs (``t.*``) contribute + no name — :class:`NoSelectStarRule` rejects those outright. + + Column names do not all arrive as ``exp.Column``. sqlglot parses several + positions as bare ``exp.Identifier``, and each is a way to name a denied + column without producing a single ``exp.Column`` node, so each is harvested + explicitly: + + * ``JOIN ... USING (email)`` — ``Join.args["using"]``. + * ``... AS g(email)`` column aliases — ``TableAlias.args["columns"]``. + * ``STRUCT('a@b.com' AS email)`` field names — ``exp.PropertyEQ``. + + ``NATURAL JOIN`` names no columns at all and is rejected outright by + :class:`NoUnresolvableColumnsRule`. + """ + names: list[str] = list(all_projection_names(statement)) + for column in statement.find_all(exp.Column): + if isinstance(column.this, exp.Star): + continue + if column.name: + names.append(column.name) + + for join in statement.find_all(exp.Join): + for identifier in join.args.get("using") or (): + if isinstance(identifier, exp.Identifier) and identifier.name: + names.append(identifier.name) + + for table_alias in statement.find_all(exp.TableAlias): + for identifier in table_alias.args.get("columns") or (): + if isinstance(identifier, exp.Identifier) and identifier.name: + names.append(identifier.name) + + for prop in statement.find_all(exp.PropertyEQ): + field = prop.this + if isinstance(field, exp.Identifier) and field.name: + names.append(field.name) + + return names + + +def whole_row_references(statement: exp.Expression) -> list[str]: + """Unqualified column refs that actually name a row source, not a column. + + ``SELECT c FROM tbl AS c`` parses identically to a column named ``c``, but + the engine returns the whole row. Any unqualified ``exp.Column`` whose name + matches a *range variable* — a table alias, bare table name, CTE alias, or + derived-table / ``VALUES`` / ``PIVOT`` alias — is a candidate. + + Two refinements keep the false-positive rate survivable, because a rule + that denies ordinary analytics gets switched off and protects nothing: + + * **A table contributes only the name it is actually addressable by.** For + ``FROM \\`p.d.status\\` AS s`` that is ``s``, not ``status``, so + ``SELECT order_id, status FROM orders o JOIN \\`p.d.status\\` s ...`` + keeps working — ``status`` there is a column, and no range variable of + that name exists. + * **Where the guard can see a scope's output columns, it uses them.** A CTE + or derived table names its own projections, so if the reference matches + one it is a column, not a row:: + + WITH revenue AS (SELECT uid, SUM(x) AS revenue FROM t GROUP BY uid) + SELECT uid, revenue FROM revenue -- allowed: revenue is a column + + Naming a CTE after the metric it computes is a mainstream idiom; denying + it would be untenable. For a physical table the guard has no schema, so + the reference stays denied and must be qualified (``c.col``). + + Bare ``UNNEST`` aliases are deliberately *not* treated as range variables. + ``SELECT s FROM t, UNNEST(t.tags) AS s`` is idiomatic for a scalar array, + and when the array holds structs the exposure is identical to selecting the + struct column directly (``SELECT t.tags FROM t``) — which no parse-level + rule can catch either. That whole class is a denylist-configuration + concern: denylist the containing column. See the README's coverage limits. + """ + ranges = _range_variables(statement) + if not ranges: + return [] + + out: list[str] = [] + for column in statement.find_all(exp.Column): + if isinstance(column.this, exp.Star) or column.table or not column.name: + continue + key = column.name.lower() + if key not in ranges: + continue + exposed = ranges[key] + # A scope we can read that publishes a column of this name — the + # reference resolves to that column, not to the row. + if exposed is not None and key in exposed: + continue + out.append(column.name) + return out + + +def _range_variables(statement: exp.Expression) -> dict[str, frozenset[str] | None]: + """Map each addressable row-source name to the columns it exposes. + + The value is ``None`` when the guard cannot see the column list (a physical + table, a ``PIVOT``), and a frozenset of output names when it can (a CTE + body, a derived table, a ``VALUES`` column alias list). ``None`` is the + conservative reading and always wins a collision. + """ + out: dict[str, frozenset[str] | None] = {} + + def add(name: str | None, columns: frozenset[str] | None) -> None: + if not name: + return + key = name.lower() + if key in out and out[key] is not None and columns is None: + out[key] = None + elif key not in out: + out[key] = columns + elif columns is None: + out[key] = None + + cte_aliases = {cte.alias.lower() for cte in statement.find_all(exp.CTE) if cte.alias} + + for table in statement.find_all(exp.Table): + # `FROM my_cte` parses as an exp.Table. Recording it here would mask + # the CTE's readable column list with an unknown one, and unknown wins + # collisions — which is exactly how a CTE named after the metric it + # computes ends up wrongly denied. + if not table.catalog and not table.db and table.name.lower() in cte_aliases: + continue + # Only the name the table is actually addressable by: an aliased table + # cannot be referenced by its bare name. + add(table.alias or table.name, None) + + for cte in statement.find_all(exp.CTE): + add(cte.alias, _scope_output_names(cte.this)) + + for table_alias in statement.find_all(exp.TableAlias): + parent = table_alias.parent + if isinstance(parent, exp.Table | exp.CTE): + continue # already recorded above + if isinstance(parent, exp.Unnest): + continue # see whole_row_references docstring + exposed: frozenset[str] | None = None + if isinstance(parent, exp.Subquery): + exposed = _scope_output_names(parent.this) + else: + columns = [ + identifier.name + for identifier in table_alias.args.get("columns") or () + if isinstance(identifier, exp.Identifier) and identifier.name + ] + if columns: + exposed = frozenset(c.lower() for c in columns) + add(table_alias.name, exposed) + + return out + + +def _scope_output_names(expr: exp.Expression | None) -> frozenset[str] | None: + """Output column names of a CTE body or derived table, if readable.""" + if expr is None: + return None + selects = outer_selects(expr) + if not selects: + return None + names: set[str] = set() + for select in selects: + for projection in select.expressions: + if _is_star_projection(projection): + # A star hides the real output names; NoSelectStarRule rejects + # this query anyway, but don't claim knowledge we lack. + return None + name = projection.alias_or_name + if name: + names.add(name.lower()) + return frozenset(names) + + def referenced_tables(statement: exp.Expression) -> set[str]: """Fully-qualified physical tables referenced by *statement*. @@ -559,10 +922,27 @@ def referenced_tables(statement: exp.Expression) -> set[str]: def has_top_level_select_star(statement: exp.Expression) -> bool: - """True if any outermost ``Select`` projects ``*`` in any form.""" + """True if any outermost ``Select`` projects ``*`` in any form. + + Outermost scope only — see :func:`has_select_star` for the all-scope check + that :class:`NoSelectStarRule` actually enforces. + """ for select in outer_selects(statement): for projection in select.expressions: - if isinstance(projection, exp.Star): + if _is_star_projection(projection): + return True + return False + + +def has_select_star(statement: exp.Expression) -> bool: + """True if *any* scope projects ``*`` in any form. + + Includes CTE bodies, derived tables and subqueries. ``COUNT(*)`` is not a + star projection — it is an aggregate that emits a scalar, and stays allowed. + """ + for select in all_selects(statement): + for projection in select.expressions: + if _is_star_projection(projection): return True return False @@ -596,7 +976,7 @@ def format_cost(cost_usd: float) -> str: if magnitude >= 0.01: return f"{sign}${magnitude:.2f}" # Sub-cent: enough decimals for at least two significant figures. - decimals = 1 - int(math.floor(math.log10(magnitude))) + decimals = 1 - math.floor(math.log10(magnitude)) return f"{sign}${magnitude:.{decimals}f}" @@ -607,6 +987,15 @@ def format_cost(cost_usd: float) -> str: # Scalar functions that destroy PII content (return a number or fixed-width # hash, not the original value). Names match sqlglot's lowercase ``.key``. +# +# NOTE: several entries here do not match any sqlglot key and therefore never +# fire — BigQuery's ``MD5`` parses as ``MD5Digest`` (key ``md5digest``), +# ``SHA256`` as ``SHA2`` (key ``sha2``), and ``FARM_FINGERPRINT`` as +# ``Anonymous``. The practical effect is that hashed PII is *denied*, which is +# the safe direction — a hashed email is still a stable pseudonymous +# identifier, and a short one is trivially reversible by dictionary attack — +# so the mismatch is left as-is rather than "fixed" into a loosening. Do not +# add the real keys without deciding that hashed PII is acceptable output. _PII_SAFE_FUNC_KEYS: frozenset[str] = frozenset( { "length", @@ -624,6 +1013,79 @@ def format_cost(cost_usd: float) -> str: ) +# Aggregates that reduce their input to a derived statistic, so a denied column +# inside one cannot reach the caller. Deliberately excludes MIN / MAX / +# ANY_VALUE / ARRAY_AGG / STRING_AGG / LOGICAL_OR / percentile- and mode-style +# aggregates, all of which return an input value verbatim. +_VALUE_DESTROYING_AGG_KEYS: frozenset[str] = frozenset( + { + "count", + "countif", + "sum", + "avg", + "approxdistinct", + "stddev", + "stddevpop", + "stddevsamp", + "variance", + "variancepop", + "variancesamp", + }, +) + +# Aggregates for which a ``*`` argument is a row *count*, not a row expansion. +_COUNTING_AGG_KEYS: frozenset[str] = frozenset({"count", "countif", "approxdistinct"}) + + +def _is_star_projection(projection: exp.Expression) -> bool: + """True if *projection* expands to an unknown set of columns. + + Checking the projection's root node is not enough. Every one of these is a + whole-row expansion with the ``Star`` buried one or more levels down, and a + root-node ``isinstance`` test lets all of them through:: + + SELECT t.* -- Column(Star) any dialect + SELECT OBJECT_CONSTRUCT(*) -- StarMap(Star) snowflake + SELECT COLUMNS(*) -- Columns(Star) duckdb + SELECT * APPLY(toString) -- Apply(Star) clickhouse + SELECT ROW(c.*) -- Struct(Column()) trino + + So the check is a deep walk with one carve-out: a star consumed by a + count-style aggregate is not an expansion, because the aggregate emits a + scalar rather than the row. ``COUNT(*)`` and ``COUNT(DISTINCT *)`` stay + allowed; nothing else that swallows a star does. + + ClickHouse's regex column selector ``COLUMNS('e.*')`` is also an expansion + but parses with no ``Star`` node at all — a ``Columns`` node wrapping a + string literal — so it is matched on node type. + """ + if isinstance(projection, exp.Columns) or any(projection.find_all(exp.Columns)): + return True + for star in projection.find_all(exp.Star): + if not _star_is_counted(star, projection): + return True + return False + + +def _star_is_counted(star: exp.Star, projection: exp.Expression) -> bool: + """True if *star* is an argument to a count-style aggregate. + + Walks from *star* up to *projection*. ``exp.Distinct`` is a permitted + intermediate so ``COUNT(DISTINCT *)`` resolves the same as ``COUNT(*)``. + """ + node: exp.Expression | None = star.parent + while node is not None: + key = getattr(node, "key", "") + if isinstance(node, exp.AggFunc) and isinstance(key, str): + return key.lower() in _COUNTING_AGG_KEYS + if not isinstance(node, exp.Distinct): + return False + if node is projection: + return False + node = node.parent + return False + + def _projection_names(projection: exp.Expression) -> list[str]: """Names from a single projection — every PII-relevant column ref. @@ -642,7 +1104,10 @@ def _projection_names(projection: exp.Expression) -> list[str]: * Otherwise, walk the projection but skip columns reachable only via predicate clauses (WHERE / HAVING / QUALIFY / ON / ORDER BY / GROUP BY). """ - target = projection.unalias() if isinstance(projection, exp.Alias) else projection + # `.this` rather than `.unalias()`: identical result for an Alias (that is + # all unalias does), but sqlglot leaves `unalias` untyped, which trips + # mypy's strict `no-untyped-call`. + target = projection.this if isinstance(projection, exp.Alias) else projection # Subquery projection → recurse into its SELECT list. if isinstance(target, exp.Subquery): @@ -725,10 +1190,21 @@ def _iter_args(node: exp.Expression): # type: ignore[no-untyped-def] def _is_pii_neutralizing(expr: exp.Expression) -> bool: - if isinstance(expr, exp.AggFunc): - return True + """True if *expr* cannot carry a PII value out to the caller. + + Aggregation alone does not neutralise anything. ``MAX(email)`` returns a + real address; ``ARRAY_AGG(email)`` returns all of them; ``STRING_AGG`` and + ``ANY_VALUE`` likewise. Only aggregates that reduce their input to a + derived statistic qualify, so the test is an explicit allowlist rather + than ``isinstance(expr, exp.AggFunc)``. + """ key = getattr(expr, "key", "") - return isinstance(key, str) and key.lower() in _PII_SAFE_FUNC_KEYS + if not isinstance(key, str): + return False + lowered = key.lower() + if isinstance(expr, exp.AggFunc): + return lowered in _VALUE_DESTROYING_AGG_KEYS + return lowered in _PII_SAFE_FUNC_KEYS def _table_fullname(table: exp.Table) -> str: diff --git a/tests/conftest.py b/tests/conftest.py index 72c6192..ef3fa8e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -55,6 +55,7 @@ def allowed_tables() -> frozenset[str]: @pytest.fixture def sql_guard(pii_denylist: PiiDenylist, allowed_tables: frozenset[str]) -> SqlGuard: + """Guard with default settings — i.e. ``pii_mode="reference"``.""" return SqlGuard( SqlGuardConfig.from_settings( pii_denylist=pii_denylist, @@ -64,3 +65,22 @@ def sql_guard(pii_denylist: PiiDenylist, allowed_tables: frozenset[str]) -> SqlG max_bytes_billed=10 * 1024**3, ), ) + + +@pytest.fixture +def project_mode_guard(pii_denylist: PiiDenylist, allowed_tables: frozenset[str]) -> SqlGuard: + """Guard with the looser ``pii_mode="project"`` policy. + + Denylisted columns may appear in predicates; only projections are denied. + Used to pin the behaviour of the documented loosening path. + """ + return SqlGuard( + SqlGuardConfig.from_settings( + pii_denylist=pii_denylist, + allowed_tables=allowed_tables, + pii_mode="project", + max_cost_usd_auto=0.10, + max_cost_usd_hard=20.00, + max_bytes_billed=10 * 1024**3, + ), + ) diff --git a/tests/test_pii_scopes.py b/tests/test_pii_scopes.py new file mode 100644 index 0000000..6fd07bd --- /dev/null +++ b/tests/test_pii_scopes.py @@ -0,0 +1,347 @@ +"""Regressions for the 0.2.0 PII scope fixes. + +Two classes of bypass, both found by adversarial review of a deployed agent +and both confirmed against 0.1.1 before the fix: + +**Scope aliasing.** ``PiiProjectionRule`` only inspected the outermost select +list, so any construct that renames a denied column in an inner scope laundered +it — CTEs, derived tables, UNION arms, and chains of the above. The star rule +had the same outermost-only blind spot, which matters once PII checking is +all-scope: a ``SELECT *`` inside a CTE makes that scope's projection list +unresolvable, so the guard cannot prove a denied column is absent. + +**Value probing.** The denylist gated projection only, so ``WHERE``, +``GROUP BY``, ``HAVING`` and ``ORDER BY`` references passed. None of those +return the column, but each answers questions about its value, and enough +answers reconstruct it. ``pii_mode="reference"`` (the default) closes this; +``pii_mode="project"`` is the documented loosening path. + +Every SQL string here is a *bypass* — each one returned CONFIRM on 0.1.1. +""" + +from __future__ import annotations + +import pytest + +from sql_guard import GuardOutcome, PiiDenylist, SqlGuard, SqlGuardConfig + +_CUSTOMERS = "`example-project.analytics.customers`" +_IDENTITY = "`example-project.analytics.identity`" + + +# --------------------------------------------------------------------------- +# Issue 1 — alias laundering across query scopes +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ( + "cte_alias", + f"WITH c AS (SELECT email AS city FROM {_CUSTOMERS}) SELECT city FROM c", + ), + ( + "derived_table", + f"SELECT x FROM (SELECT email AS x FROM {_CUSTOMERS})", + ), + ( + "union_arm_derived", + ( + f"SELECT uid FROM {_CUSTOMERS} " + f"UNION ALL " + f"SELECT c FROM (SELECT email AS c FROM {_CUSTOMERS})" + ), + ), + ( + "union_arm_cte", + ( + f"WITH laundered AS (SELECT email AS e FROM {_CUSTOMERS}) " + f"SELECT uid FROM {_CUSTOMERS} UNION ALL SELECT e FROM laundered" + ), + ), + ( + "multi_hop_chain", + ( + f"WITH a AS (SELECT email AS e1 FROM {_CUSTOMERS}), " + f"b AS (SELECT e1 AS e2 FROM a), " + f"c AS (SELECT e2 AS e3 FROM b) " + f"SELECT e3 FROM c" + ), + ), + ( + "nested_derived_chain", + f"SELECT z FROM (SELECT y AS z FROM (SELECT email AS y FROM {_CUSTOMERS}))", + ), + ( + "cte_feeding_join", + ( + f"WITH c AS (SELECT uid, email AS handle FROM {_CUSTOMERS}) " + f"SELECT t.uid, c.handle FROM {_IDENTITY} t JOIN c ON t.uid = c.uid" + ), + ), + ], +) +@pytest.mark.parametrize("mode", ["reference", "project"]) +def test_alias_laundering_is_denied_in_every_mode( + pii_denylist: PiiDenylist, + name: str, + sql: str, + mode: str, +) -> None: + """An alias in an inner scope must never launder a denied column. + + These are denied in *both* modes: the denied column is projected in the + scope that names it, which is a projection regardless of policy. + """ + guard = _guard(pii_denylist, mode) + decision = guard.evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"{name} passed the guard: {decision.reason}" + assert any("email" in c.lower() for c in decision.pii_columns), ( + f"{name} denied but did not name the offending column: {decision.pii_columns}" + ) + + +# --------------------------------------------------------------------------- +# Issue 1b — SELECT * hiding a scope's projection list +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ("star_in_cte", f"WITH c AS (SELECT * FROM {_CUSTOMERS}) SELECT city FROM c"), + ("star_in_derived", f"SELECT city FROM (SELECT * FROM {_CUSTOMERS})"), + ( + "star_except_in_cte", + f"WITH c AS (SELECT * EXCEPT(email) FROM {_CUSTOMERS}) SELECT city FROM c", + ), + ( + "star_in_union_arm", + f"SELECT uid FROM {_CUSTOMERS} UNION ALL SELECT * FROM {_IDENTITY}", + ), + ("qualified_star_top_level", f"SELECT t.* FROM {_CUSTOMERS} t"), + ( + "qualified_star_in_cte", + f"WITH c AS (SELECT t.* FROM {_CUSTOMERS} t) SELECT city FROM c", + ), + ( + "star_in_scalar_subquery", + f"SELECT (SELECT * FROM {_IDENTITY} LIMIT 1) AS x", + ), + ], +) +@pytest.mark.parametrize("mode", ["reference", "project"]) +def test_select_star_is_denied_in_every_scope( + pii_denylist: PiiDenylist, + name: str, + sql: str, + mode: str, +) -> None: + """A star anywhere makes that scope's projection unprovable — reject it. + + The guard cannot enumerate what ``*`` expands to without a schema, so it + cannot show the scope is free of denied columns. + """ + guard = _guard(pii_denylist, mode) + decision = guard.evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"{name} passed the guard: {decision.reason}" + assert "*" in decision.reason + + +def test_count_star_survives_the_all_scope_star_rule(pii_denylist: PiiDenylist) -> None: + """``COUNT(*)`` is an aggregate, not a star projection — still allowed. + + Guards against an over-broad star rule that walks for ``exp.Star`` nodes + instead of inspecting select-list entries. + """ + guard = _guard(pii_denylist, "reference") + sql = ( + f"WITH c AS (SELECT tier FROM {_CUSTOMERS}) SELECT tier, COUNT(*) AS n FROM c GROUP BY tier" + ) + decision = guard.evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, decision.reason + + +# --------------------------------------------------------------------------- +# Issue 2 — value probing through predicates +# --------------------------------------------------------------------------- + + +_PROBES: list[tuple[str, str]] = [ + ("where_equality", f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} WHERE email = 'a@b.com'"), + ("where_like_prefix", f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} WHERE email LIKE 'a%'"), + ("where_inequality", f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} WHERE email > 'm'"), + ("group_by_only", f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} GROUP BY email"), + ( + "having", + f"SELECT tier, COUNT(*) AS n FROM {_CUSTOMERS} GROUP BY tier HAVING MIN(email) > 'm'", + ), + ("order_by", f"SELECT tier FROM {_CUSTOMERS} ORDER BY email LIMIT 1"), + ( + "join_on", + f"SELECT a.tier FROM {_CUSTOMERS} a JOIN {_IDENTITY} b ON a.email = b.email", + ), + ( + "where_in_cte", + f"WITH c AS (SELECT uid FROM {_CUSTOMERS} WHERE email = 'a@b.com') SELECT uid FROM c", + ), + ( + "case_expression", + f"SELECT SUM(CASE WHEN email = 'a@b.com' THEN 1 ELSE 0 END) AS hit FROM {_CUSTOMERS}", + ), + ( + "aggregate_over_pii", + f"SELECT COUNTIF(email = 'a@b.com') AS hit FROM {_CUSTOMERS}", + ), +] + + +@pytest.mark.parametrize(("name", "sql"), _PROBES) +def test_value_probing_is_denied_in_reference_mode( + pii_denylist: PiiDenylist, + name: str, + sql: str, +) -> None: + """Every one of these extracts a denied value without projecting it. + + Each returns a count or an ordering that answers a yes/no question about + the column, so a caller can binary-search the value across queries. + """ + guard = _guard(pii_denylist, "reference") + decision = guard.evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"{name} passed the guard: {decision.reason}" + assert any("email" in c.lower() for c in decision.pii_columns) + + +@pytest.mark.parametrize(("name", "sql"), _PROBES) +def test_value_probing_is_permitted_in_project_mode( + pii_denylist: PiiDenylist, + name: str, + sql: str, +) -> None: + """The documented loosening path, pinned so it cannot drift silently. + + ``pii_mode="project"`` deliberately allows predicate access to denied + columns. This test records that trade-off rather than endorsing it. + """ + guard = _guard(pii_denylist, "project") + decision = guard.evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, f"{name}: {decision.reason}" + + +def test_reference_mode_deny_message_names_columns_and_the_alternative( + pii_denylist: PiiDenylist, +) -> None: + """Deny messages state the rule, name the columns, suggest the way out.""" + guard = _guard(pii_denylist, "reference") + decision = guard.evaluate_static(f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} WHERE email = 'x'") + assert decision.outcome is GuardOutcome.DENY + assert "PII" in decision.reason + assert "email" in decision.reason + assert "pii_mode" in decision.reason + + +# --------------------------------------------------------------------------- +# Legit queries that must keep working +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ( + "aggregate_over_allowed_columns", + ( + f"SELECT tier, COUNT(*) AS n, AVG(lifetime_spend) AS avg_spend " + f"FROM {_CUSTOMERS} GROUP BY tier" + ), + ), + ( + "pii_free_cte", + ( + f"WITH base AS (SELECT uid, tier, lifetime_spend FROM {_CUSTOMERS}) " + f"SELECT tier, AVG(lifetime_spend) AS avg_spend FROM base GROUP BY tier" + ), + ), + ( + "multi_hop_pii_free_chain", + ( + f"WITH a AS (SELECT uid, tier FROM {_CUSTOMERS}), " + f"b AS (SELECT uid, tier AS band FROM a) " + f"SELECT band, COUNT(*) AS n FROM b GROUP BY band" + ), + ), + ( + "pii_free_union", + f"SELECT tier FROM {_CUSTOMERS} UNION ALL SELECT tier FROM {_IDENTITY}", + ), + ( + "pii_free_derived_table", + f"SELECT band, COUNT(*) AS n FROM (SELECT tier AS band FROM {_CUSTOMERS}) GROUP BY band", + ), + ( + "where_on_allowed_column", + ( + f"SELECT tier, COUNT(*) AS n FROM {_CUSTOMERS} " + f"WHERE lifetime_spend > 5000 GROUP BY tier" + ), + ), + ( + "order_by_allowed_column", + f"SELECT tier, COUNT(*) AS n FROM {_CUSTOMERS} GROUP BY tier ORDER BY n DESC", + ), + ], +) +@pytest.mark.parametrize("mode", ["reference", "project"]) +def test_legitimate_queries_still_pass( + pii_denylist: PiiDenylist, + name: str, + sql: str, + mode: str, +) -> None: + """The fix must not cost us ordinary analytics. + + None of these name a denied column in any scope, so both modes allow them. + """ + guard = _guard(pii_denylist, mode) + decision = guard.evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, f"{name} wrongly denied: {decision.reason}" + + +# --------------------------------------------------------------------------- +# Config surface +# --------------------------------------------------------------------------- + + +def test_reference_is_the_default_mode(pii_denylist: PiiDenylist) -> None: + config = SqlGuardConfig.from_settings(pii_denylist=pii_denylist, allowed_tables=[]) + assert config.pii_mode == "reference" + + +def test_invalid_pii_mode_raises(pii_denylist: PiiDenylist) -> None: + """A typo must fail loudly, not silently pick a policy nobody chose.""" + with pytest.raises(ValueError, match="pii_mode"): + SqlGuardConfig.from_settings( + pii_denylist=pii_denylist, + allowed_tables=[], + pii_mode="referece", # type: ignore[arg-type] + ) + + +def test_no_top_level_star_rule_alias_still_importable() -> None: + """Consumers pinning the old rule name keep working.""" + from sql_guard import NoSelectStarRule, NoTopLevelStarRule + + assert NoTopLevelStarRule is NoSelectStarRule + + +def _guard(pii_denylist: PiiDenylist, mode: str) -> SqlGuard: + return SqlGuard( + SqlGuardConfig.from_settings( + pii_denylist=pii_denylist, + allowed_tables=[], + pii_mode=mode, # type: ignore[arg-type] + enforce_allowed_tables=False, + ), + ) diff --git a/tests/test_sql_guard.py b/tests/test_sql_guard.py index ca81d03..75dee9e 100644 --- a/tests/test_sql_guard.py +++ b/tests/test_sql_guard.py @@ -2,10 +2,15 @@ These tests exercise the guard's rules with synthetic ``example-project`` data: -* Positive Q1-Q4 SQL should pass the static checks. +* Positive Q2-Q4 SQL should pass the static checks. +* Q1 (identity resolution over normalised email/mobile) is DENY as of 0.2.0 — + see :func:`test_q1_identity_resolution_is_denied` for why. * Negative N1 ("emails and mobiles") must be DENY. * Negative N2 ("all transactions") must be DENY (bare ``SELECT *``). * Cost-cap rules: dry-run bytes drive ALLOW / CONFIRM / DENY. + +Scope-bypass regressions (CTE / derived-table / UNION-arm aliasing, value +probing via predicates) live in ``test_pii_scopes.py``. """ from __future__ import annotations @@ -87,7 +92,7 @@ @pytest.mark.parametrize( ("name", "sql"), - [("Q1", _Q1), ("Q2", _Q2), ("Q3", _Q3), ("Q4", _Q4)], + [("Q2", _Q2), ("Q3", _Q3), ("Q4", _Q4)], ) def test_positive_cases_pass_static(sql_guard: SqlGuard, name: str, sql: str) -> None: decision = sql_guard.evaluate_static(sql) @@ -98,6 +103,45 @@ def test_positive_cases_pass_static(sql_guard: SqlGuard, name: str, sql: str) -> assert decision.outcome is GuardOutcome.CONFIRM +@pytest.mark.parametrize("mode", ["reference", "project"]) +def test_q1_identity_resolution_is_denied( + pii_denylist: object, + allowed_tables: object, + mode: str, +) -> None: + """Q1 was a positive case until 0.2.0. It is now denied in *both* modes. + + Q1 normalises PII inside a CTE (``LOWER(TRIM(email)) AS email_norm``) and + projects only ``COUNTIF`` aggregates, so outermost-only checking saw + nothing but counts. Two independent reasons it must now fail: + + * ``pii_mode="project"``: the CTE scope projects ``email`` and ``mobile``. + Checking every scope is what stops an alias laundering a denied column, + and this query is indistinguishable from that attack at parse time. + * ``pii_mode="reference"``: the aggregates reference the normalised + columns. ``COUNTIF(email_norm = 'target@example.com')`` is precisely the + value-probing oracle reference mode exists to close — an aggregate over + a denied column still answers questions about individual values. + + Callers who need this pattern should normalise PII in a warehouse view the + guard's denylist does not cover, and point the agent at the view. + """ + from sql_guard import PiiDenylist, SqlGuardConfig + + assert isinstance(pii_denylist, PiiDenylist) + guard = SqlGuard( + SqlGuardConfig.from_settings( + pii_denylist=pii_denylist, + allowed_tables=[], + pii_mode=mode, # type: ignore[arg-type] + enforce_allowed_tables=False, + ), + ) + decision = guard.evaluate_static(_Q1) + assert decision.outcome is GuardOutcome.DENY + assert any("email" in c.lower() for c in decision.pii_columns) + + # --------------------------------------------------------------------------- # Negative cases # --------------------------------------------------------------------------- @@ -167,11 +211,12 @@ def test_select_count_star_is_allowed(sql_guard: SqlGuard) -> None: assert decision.outcome is not GuardOutcome.DENY -def test_pii_in_subquery_where_is_allowed(sql_guard: SqlGuard) -> None: - """A scalar-returning subquery that *consumes* PII in WHERE is fine. +def test_pii_in_subquery_where_is_allowed_in_project_mode(project_mode_guard: SqlGuard) -> None: + """A scalar-returning subquery that *consumes* PII in WHERE. - The subquery emits a count, not the PII value. The guard must not flag - email referenced inside the subquery's WHERE clause. + The subquery emits a count, not the PII value, so ``pii_mode="project"`` + permits it. Under the default ``"reference"`` mode this same query is + denied — see :func:`test_pii_in_subquery_where_is_denied_in_reference_mode`. """ sql = """ SELECT @@ -182,10 +227,27 @@ def test_pii_in_subquery_where_is_allowed(sql_guard: SqlGuard) -> None: WHERE email IS NOT NULL AND ARRAY_LENGTH(platform_b_emails) > 0 ) AS platform_b_count """ - decision = sql_guard.evaluate_static(sql) + decision = project_mode_guard.evaluate_static(sql) assert decision.outcome is not GuardOutcome.DENY, decision.reason +def test_pii_in_subquery_where_is_denied_in_reference_mode(sql_guard: SqlGuard) -> None: + """The same subquery-WHERE pattern is denied under the default mode. + + ``COUNT(*) ... WHERE email IS NOT NULL`` leaks one bit per query, and the + predicate can be narrowed (``WHERE email LIKE 'a%'``) to walk a value out. + """ + sql = """ + SELECT + (SELECT COUNT(*) FROM `example-project.analytics.identity` + WHERE email IS NOT NULL + ) AS c + """ + decision = sql_guard.evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY + assert any("email" in c.lower() for c in decision.pii_columns) + + def test_pii_actually_projected_by_subquery_is_denied(sql_guard: SqlGuard) -> None: """If a subquery *projects* PII, the guard must still catch it.""" sql = "SELECT (SELECT email FROM `example-project.analytics.identity` LIMIT 1) AS leaked" @@ -194,18 +256,31 @@ def test_pii_actually_projected_by_subquery_is_denied(sql_guard: SqlGuard) -> No assert any("email" in c.lower() for c in decision.pii_columns) -def test_pii_in_where_of_outer_query_is_allowed(sql_guard: SqlGuard) -> None: - """Using PII in a WHERE filter while projecting non-PII columns is fine.""" - sql = ( - "SELECT uid, COUNT(*) AS n " - "FROM `example-project.analytics.identity` " - "WHERE email IS NOT NULL " - "GROUP BY uid" - ) - decision = sql_guard.evaluate_static(sql) +_PII_IN_WHERE = ( + "SELECT uid, COUNT(*) AS n " + "FROM `example-project.analytics.identity` " + "WHERE email IS NOT NULL " + "GROUP BY uid" +) + + +def test_pii_in_where_is_allowed_in_project_mode(project_mode_guard: SqlGuard) -> None: + """Filtering on PII while projecting non-PII is permitted under "project".""" + decision = project_mode_guard.evaluate_static(_PII_IN_WHERE) assert decision.outcome is not GuardOutcome.DENY, decision.reason +def test_pii_in_where_is_denied_in_reference_mode(sql_guard: SqlGuard) -> None: + """...and denied under the default "reference" mode, which is the point. + + A denied column in a WHERE clause is a value oracle even though it is + never projected. + """ + decision = sql_guard.evaluate_static(_PII_IN_WHERE) + assert decision.outcome is GuardOutcome.DENY + assert any("email" in c.lower() for c in decision.pii_columns) + + # --------------------------------------------------------------------------- # Disallowed statements # --------------------------------------------------------------------------- diff --git a/tests/test_unresolvable_columns.py b/tests/test_unresolvable_columns.py new file mode 100644 index 0000000..f022c22 --- /dev/null +++ b/tests/test_unresolvable_columns.py @@ -0,0 +1,390 @@ +"""Regressions for the four bypasses found reviewing the 0.2.0 scope fix. + +The first pass at 0.2.0 closed alias laundering and predicate probing, then an +adversarial review broke it again four more ways. Each SQL string below was +empirically confirmed to reach ``confirm``/``allow`` against that intermediate +build; all are denied now. + +* **Whole-row alias** — ``SELECT c FROM t AS c`` returns every column in the + row as a struct. Strictly worse than ``SELECT *``, and it parses as an + ordinary column named ``c``, so the denylist saw nothing to object to. +* **Nested stars** — the star check only looked at the projection's root node, + so any construct wrapping the star (``OBJECT_CONSTRUCT(*)``, ``COLUMNS(*)``, + ``* APPLY(f)``) sailed through on non-BigQuery dialects. +* **Identifier-only column refs** — ``JOIN ... USING (email)`` and + ``AS g(email)`` carry column names as ``exp.Identifier``, never + ``exp.Column``, so reference mode's ``find_all(exp.Column)`` sweep missed + them entirely. ``NATURAL JOIN`` names no columns at all. +* **Value-preserving aggregates** — every ``exp.AggFunc`` was treated as + PII-neutralising, so ``MAX(email)`` and ``ARRAY_AGG(email)`` counted as + "aggregate only" and passed ``pii_mode="project"``. +""" + +from __future__ import annotations + +import pytest + +from sql_guard import GuardOutcome, PiiDenylist, SqlGuard, SqlGuardConfig + +_CUSTOMERS = "`example-project.analytics.customers`" + + +def _guard( + pii_denylist: PiiDenylist, + mode: str = "reference", + dialect: str = "bigquery", + *, + enforce_tables: bool = True, +) -> SqlGuard: + return SqlGuard( + SqlGuardConfig.from_settings( + pii_denylist=pii_denylist, + allowed_tables=["example-project.analytics.customers"] if enforce_tables else [], + pii_mode=mode, # type: ignore[arg-type] + dialect=dialect, + enforce_allowed_tables=enforce_tables, + ), + ) + + +# --------------------------------------------------------------------------- +# Whole-row references via a table alias +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ("bare_alias", f"SELECT c FROM {_CUSTOMERS} AS c"), + ("json_of_alias", f"SELECT TO_JSON_STRING(c) AS j FROM {_CUSTOMERS} AS c"), + ("agg_of_alias", f"SELECT ARRAY_AGG(c) AS rows FROM {_CUSTOMERS} AS c"), + ("struct_of_alias", f"SELECT STRUCT(c) AS s FROM {_CUSTOMERS} AS c"), + ("table_name_no_alias", f"SELECT customers FROM {_CUSTOMERS}"), + ( + "cte_name", + f"WITH x AS (SELECT uid, email FROM {_CUSTOMERS}) SELECT x FROM x", + ), + ("alias_in_where", f"SELECT uid FROM {_CUSTOMERS} AS c WHERE c IS NOT NULL"), + ], +) +@pytest.mark.parametrize("mode", ["reference", "project"]) +def test_whole_row_reference_is_denied( + pii_denylist: PiiDenylist, + name: str, + sql: str, + mode: str, +) -> None: + """A bare table alias returns the entire row, PII included. + + This must be denied in both modes — it is an unresolvable projection, not + a predicate question, so ``pii_mode`` has no bearing on it. + """ + decision = _guard(pii_denylist, mode).evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"{name} passed the guard: {decision.reason}" + + +def test_whole_row_deny_message_names_the_alias_and_the_fix(pii_denylist: PiiDenylist) -> None: + decision = _guard(pii_denylist).evaluate_static(f"SELECT c FROM {_CUSTOMERS} AS c") + assert decision.outcome is GuardOutcome.DENY + assert "c" in decision.reason + assert "alias.column" in decision.reason + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ( + "derived_table_alias", + f"SELECT d FROM (SELECT uid, tier FROM {_CUSTOMERS}) AS d", + ), + ( + "json_of_derived_alias", + f"SELECT TO_JSON_STRING(d) AS j FROM (SELECT uid, tier FROM {_CUSTOMERS}) AS d", + ), + ( + "cte_row_via_json", + f"WITH c AS (SELECT uid, email FROM {_CUSTOMERS}) SELECT TO_JSON_STRING(c) AS j FROM c", + ), + ( + "pivot_alias", + f"SELECT pv FROM {_CUSTOMERS} PIVOT(SUM(lifetime_spend) FOR tier IN ('a', 'b')) AS pv", + ), + ], +) +def test_non_table_row_sources_are_also_whole_row_references( + pii_denylist: PiiDenylist, + name: str, + sql: str, +) -> None: + """Derived tables, CTEs and PIVOTs are row sources too, not just tables. + + An earlier cut of this rule built its name set from ``exp.Table`` and + ``exp.CTE`` alone, so aliasing a subquery or a PIVOT re-opened the whole-row + leak. + """ + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"{name} passed the guard: {decision.reason}" + + +def test_values_alias_whole_row_is_denied(pii_denylist: PiiDenylist) -> None: + guard = _guard(pii_denylist, dialect="postgres", enforce_tables=False) + decision = guard.evaluate_static("SELECT v FROM (VALUES ('a')) AS v(x)") + assert decision.outcome is GuardOutcome.DENY + + +def test_clickhouse_regex_column_selector_is_denied(pii_denylist: PiiDenylist) -> None: + """``COLUMNS('e.*')`` expands to many columns with no ``Star`` node at all. + + A star check that only looks for ``exp.Star`` misses it entirely. + """ + guard = _guard(pii_denylist, dialect="clickhouse", enforce_tables=False) + decision = guard.evaluate_static("SELECT COLUMNS('e.*') FROM customers") + assert decision.outcome is GuardOutcome.DENY + + +# --------------------------------------------------------------------------- +# False-positive boundary — a rule that denies ordinary analytics gets disabled +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ( + "cte_named_after_its_metric", + ( + "WITH revenue AS (" + f"SELECT uid, SUM(lifetime_spend) AS revenue FROM {_CUSTOMERS} GROUP BY uid" + ") SELECT uid, revenue FROM revenue ORDER BY revenue DESC" + ), + ), + ( + "cte_named_month", + ( + "WITH month AS (" + f"SELECT DATE_TRUNC(signup_date, MONTH) AS month, lifetime_spend FROM {_CUSTOMERS}" + ") SELECT month, SUM(lifetime_spend) AS s FROM month GROUP BY month" + ), + ), + ( + "cte_passthrough_column", + ( + f"WITH sessions AS (SELECT uid, sessions FROM {_CUSTOMERS}) " + "SELECT uid, sessions FROM sessions" + ), + ), + ( + # The table is addressable only as `c`, so a bare `customers` is a + # column reference, not a row reference. + "column_sharing_name_with_aliased_table", + f"SELECT c.uid, customers FROM {_CUSTOMERS} AS c", + ), + ( + "scalar_unnest_alias", + f"SELECT s FROM {_CUSTOMERS} AS t, UNNEST(t.tags) AS s", + ), + ], +) +def test_legitimate_names_colliding_with_row_sources_still_pass( + pii_denylist: PiiDenylist, + name: str, + sql: str, +) -> None: + """Naming a CTE after the metric it computes is mainstream, not an attack. + + The rule resolves the ambiguity from the AST: a CTE or derived table + publishes its own output names, so a reference matching one is a column. + An aliased table contributes only its alias, so a column sharing a name + with some *other* table in the query is unaffected. + """ + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, f"{name} wrongly denied: {decision.reason}" + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ("qualified_column", f"SELECT c.tier FROM {_CUSTOMERS} AS c"), + ("qualified_in_agg", f"SELECT COUNT(c.uid) AS n FROM {_CUSTOMERS} AS c"), + ( + "qualified_join", + f"SELECT a.tier FROM {_CUSTOMERS} a JOIN {_CUSTOMERS} b ON a.uid = b.uid", + ), + ("unrelated_column", f"SELECT tier FROM {_CUSTOMERS} AS c"), + ], +) +def test_qualified_references_still_pass( + pii_denylist: PiiDenylist, + name: str, + sql: str, +) -> None: + """Qualifying the reference resolves the ambiguity — these must not trip. + + This is the false-positive boundary of the whole-row rule: only an + *unqualified* name matching a table or alias is treated as a whole row. + """ + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, f"{name} wrongly denied: {decision.reason}" + + +# --------------------------------------------------------------------------- +# Stars nested inside a wrapping construct +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("dialect", "sql"), + [ + ("snowflake", "SELECT OBJECT_CONSTRUCT(*) FROM customers"), + ("snowflake", "SELECT OBJECT_CONSTRUCT_KEEP_NULL(*) FROM customers"), + ("duckdb", "SELECT COLUMNS(*) FROM customers"), + ("clickhouse", "SELECT * APPLY(toString) FROM customers"), + ("trino", "SELECT ROW(c.*) FROM customers c"), + ], +) +def test_nested_star_is_denied(pii_denylist: PiiDenylist, dialect: str, sql: str) -> None: + """A star wrapped in a function still expands to the whole row. + + Checking only the projection's root node missed every one of these. + """ + guard = _guard(pii_denylist, dialect=dialect, enforce_tables=False) + decision = guard.evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"[{dialect}] {sql}: {decision.reason}" + + +@pytest.mark.parametrize( + "sql", + [ + f"SELECT COUNT(*) AS n FROM {_CUSTOMERS}", + f"SELECT COUNT(DISTINCT tier) AS n FROM {_CUSTOMERS}", + f"SELECT tier, COUNT(*) AS n FROM {_CUSTOMERS} GROUP BY tier", + ], +) +def test_counting_stars_survive_the_deep_star_check(pii_denylist: PiiDenylist, sql: str) -> None: + """``COUNT(*)`` counts rows; it does not expand them. Must stay allowed. + + Pins the one carve-out in the deep star walk, so a future tightening + cannot quietly swallow the most common aggregate in the codebase. + """ + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, decision.reason + + +# --------------------------------------------------------------------------- +# Column names that never become an exp.Column +# --------------------------------------------------------------------------- + + +def test_join_using_denied_column_is_denied(pii_denylist: PiiDenylist) -> None: + """``USING (email)`` is a value oracle and produces no ``exp.Column``. + + The join returns rows only where the denied column matches, so the row + count answers a question about its value. + """ + sql = f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} a JOIN {_CUSTOMERS} b USING (email)" + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY + assert any("email" in c.lower() for c in decision.pii_columns) + + +def test_join_using_allowed_column_still_passes(pii_denylist: PiiDenylist) -> None: + """``USING`` on a non-PII column is ordinary SQL and must keep working.""" + sql = f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} a JOIN {_CUSTOMERS} b USING (uid)" + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, decision.reason + + +def test_natural_join_is_denied(pii_denylist: PiiDenylist) -> None: + """``NATURAL JOIN`` joins on unknown shared columns — unprovable, so denied.""" + sql = f"SELECT COUNT(*) AS n FROM {_CUSTOMERS} NATURAL JOIN {_CUSTOMERS}" + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY + assert "NATURAL JOIN" in decision.reason + + +def test_values_column_alias_carrying_denied_name_is_denied(pii_denylist: PiiDenylist) -> None: + """``AS g(email)`` names a denied column as a bare Identifier.""" + guard = _guard(pii_denylist, dialect="postgres", enforce_tables=False) + sql = "SELECT COUNT(*) AS n FROM customers JOIN (VALUES ('a@b.com')) AS g(email) USING (email)" + decision = guard.evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY + assert any("email" in c.lower() for c in decision.pii_columns) + + +def test_struct_field_named_after_denied_column_is_denied(pii_denylist: PiiDenylist) -> None: + """``STRUCT('a@b.com' AS email)`` parses the field name as PropertyEQ.""" + sql = f"SELECT STRUCT('a@b.com' AS email) AS s, uid FROM {_CUSTOMERS}" + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY + assert any("email" in c.lower() for c in decision.pii_columns) + + +# --------------------------------------------------------------------------- +# Aggregates that return their input verbatim +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ("max", f"SELECT MAX(email) AS x FROM {_CUSTOMERS}"), + ("min", f"SELECT MIN(email) AS x FROM {_CUSTOMERS}"), + ("array_agg", f"SELECT ARRAY_AGG(email) AS x FROM {_CUSTOMERS}"), + ("string_agg", f"SELECT STRING_AGG(email) AS x FROM {_CUSTOMERS}"), + ("any_value", f"SELECT ANY_VALUE(email) AS x FROM {_CUSTOMERS}"), + ("max_in_cte", f"WITH c AS (SELECT MAX(email) AS m FROM {_CUSTOMERS}) SELECT m FROM c"), + ], +) +@pytest.mark.parametrize("mode", ["reference", "project"]) +def test_value_preserving_aggregates_are_denied( + pii_denylist: PiiDenylist, + name: str, + sql: str, + mode: str, +) -> None: + """These return real PII values, so "it's an aggregate" is no defence. + + ``project`` mode blocks projections of denied columns — and every one of + these *is* a projection of the value, merely routed through an aggregate. + """ + decision = _guard(pii_denylist, mode).evaluate_static(sql) + assert decision.outcome is GuardOutcome.DENY, f"{name} passed in {mode}: {decision.reason}" + assert any("email" in c.lower() for c in decision.pii_columns) + + +@pytest.mark.parametrize( + ("name", "sql"), + [ + ("count", f"SELECT COUNT(email) AS n FROM {_CUSTOMERS}"), + ("countif", f"SELECT COUNTIF(email IS NOT NULL) AS n FROM {_CUSTOMERS}"), + ], +) +def test_reducing_aggregates_over_pii_pass_only_in_project_mode( + pii_denylist: PiiDenylist, + name: str, + sql: str, +) -> None: + """Counts reduce to a statistic, so ``project`` mode permits them. + + ``reference`` mode still denies — a count is a probe. This pins the exact + line between the two modes. + """ + assert _guard(pii_denylist, "project").evaluate_static(sql).outcome is not GuardOutcome.DENY + assert _guard(pii_denylist, "reference").evaluate_static(sql).outcome is GuardOutcome.DENY + + +@pytest.mark.parametrize( + "sql", + [ + f"SELECT MAX(lifetime_spend) AS m FROM {_CUSTOMERS}", + f"SELECT ARRAY_AGG(tier) AS tiers FROM {_CUSTOMERS}", + f"SELECT tier, MIN(lifetime_spend) AS lo FROM {_CUSTOMERS} GROUP BY tier", + ], +) +def test_value_preserving_aggregates_over_allowed_columns_still_pass( + pii_denylist: PiiDenylist, + sql: str, +) -> None: + """Narrowing the aggregate exemption must not break non-PII aggregation.""" + decision = _guard(pii_denylist).evaluate_static(sql) + assert decision.outcome is not GuardOutcome.DENY, decision.reason