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