diff --git a/CHANGELOG.md b/CHANGELOG.md index 6df490b..ce5a3de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -227,6 +227,15 @@ what shipped rather than what was written down at the time. ### Internal +- `tools/fuzz` generates random `SELECT`s and compares sqlglot's executor against a + real PostgreSQL, minimising every divergence to the smallest query that still + shows it. The tripwires above were written by hand from bugs that happened to be + noticed; this looks for the ones nobody tripped over, and found fourteen on its + first run — six of them wrong answers rather than refusals, including + `count(DISTINCT x)` ignoring the DISTINCT and an outer join's NULL-padded row + surviving a `WHERE` that tests the padded column. Not part of any CI job: it needs + a live server, and a suite that fails on a new random seed is one people delete. + (#122) - `SessionState` is extracted and shared by the middleware and the session, so there is one place where per-connection state lives. (#61) - Strict-xfail tripwires assert what real Postgres answers for each sqlglot bug diff --git a/docs/limitations.md b/docs/limitations.md index 0e1abd0..63bcd73 100644 --- a/docs/limitations.md +++ b/docs/limitations.md @@ -42,6 +42,15 @@ Postgres makes it, so `select 3000000000` describes as `int8` rather than telling a binary client `int4` and crashing its decoder. + **Text is compared by code point, not by collation.** The executor compares strings the way Python + does, so `'abc' < 'ABC'` is false and `ORDER BY name` sorts every capital ahead of every lower-case + letter. A real PostgreSQL compares by its database collation, and on the usual locale-derived one + (`en_US.UTF-8`) both of those come out the other way round. This is a modelling gap rather than a + bug with a fix: the executor has no collation to consult, and Python's `str` ordering is the only + total order available to it. It matches a database created with `--locale=C` exactly, and differs + from any other on mixed-case or non-ASCII text. If your test asserts on the *order* of text rows, + that is a reason to want a real Postgres. + `TABLESAMPLE` is refused outright, because the executor ignores it and nothing here can repair that. An `OFFSET` or `DISTINCT ON` nested inside a subquery is refused for the same reason: the repair reaches the query's own rows, and one buried in a subquery would be silently ignored as before. diff --git a/tests/test_sqlglot_workarounds.py b/tests/test_sqlglot_workarounds.py index b78f82b..a1554ef 100644 --- a/tests/test_sqlglot_workarounds.py +++ b/tests/test_sqlglot_workarounds.py @@ -23,6 +23,13 @@ https://github.com/jbylund/pg_mimic/issues/49, and the ones reaching users untreated are https://github.com/jbylund/pg_mimic/issues/58 +The tests from `test_distinct_keeps_columns_that_share_a_name` down were found by +`tools/fuzz`, which generates random SELECTs and compares the executor against a +real PostgreSQL. Reading them as a group is the honest summary of how far the +executor is from Postgres: it is a good query engine for the shapes sqlglot's own +tests cover, and every one of these is a wrong answer or a refusal for SQL that a +client has every reason to send. + A failure here is never a pg_mimic regression: it means good news upstream. """ @@ -84,6 +91,10 @@ def test_descending_order_places_nulls_instead_of_raising(): Workaround: `_rewrite_null_ordering` in tables.py. Ascending coincidentally matches Postgres; descending raises TypeError. + + Fixed upstream but in no release, so the mark comes off when the floor moves: + https://github.com/jbylund/pg_mimic/issues/109 is the checklist. The daily + `Upstream sqlglot` job is already red on its `main` leg for this. """ tables, schema = {"n": [{"a": 1}, {"a": None}]}, {"n": {"a": "INT"}} assert _rows("SELECT a FROM n ORDER BY a DESC", tables, schema) == [(None,), (1,)] @@ -334,3 +345,227 @@ def test_full_outer_join_preserves_unmatched_rows(): schema = {"x": {"a": "INT"}, "y": {"b": "INT"}} rows = _rows("SELECT x.a, y.b FROM x FULL OUTER JOIN y ON x.a = y.b", tables, schema) assert sorted(rows, key=str) == [(1, None), (2, 2), (None, 3)] + + +# Everything below was found by tools/fuzz. See the module docstring. + +# A NULL in here on purpose: without one, count(DISTINCT a) and count(*) are the +# same number, and the test cannot tell which of the two the executor answered. +_DUPLICATES = ({"t": [{"a": 1}, {"a": 1}, {"a": 2}, {"a": None}]}, {"t": {"a": "INT"}}) +_MIXED = ({"m": [{"n": Decimal("2.5"), "v": 2.0}]}, {"m": {"n": "DECIMAL", "v": "DOUBLE"}}) + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_distinct_keeps_columns_that_share_a_name(): + """https://github.com/jbylund/pg_mimic/issues/58 + + Two output columns with the same name are silently merged into one, so the + query answers a *narrower* result than it was asked for -- and with DISTINCT + over a join, a shorter one too: nine rows of two columns become three of one. + Dropping a column the client's Describe already promised is worse than an + error, and no workaround in tables.py would catch it, since nothing peeks at + result rows. + + `SELECT a, a` without DISTINCT is fine, which is what makes this a DISTINCT bug + rather than a projection one. + """ + assert _rows("SELECT DISTINCT a, a FROM t", *_NUMBERS) == [(1, 1), (2, 2), (3, 3)] + rows = _rows("SELECT DISTINCT x.a, y.a FROM t AS x CROSS JOIN t AS y", *_NUMBERS) + assert len(rows) == 9 + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_star_expands_columns_that_share_a_name(): + """https://github.com/jbylund/pg_mimic/issues/58 + + The same merge-by-name as above, reached through `*` instead of DISTINCT, and + the reason `WITH c AS (SELECT 1, 1 FROM t) SELECT * FROM c` answers one column: + both are named `_col_0`. A CTE or derived table whose columns are expressions + rather than plain column references is the common way to hit it. + """ + assert _rows("WITH c AS (SELECT a, a FROM t) SELECT * FROM c", *_NUMBERS) == [(1, 1), (2, 2), (3, 3)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_distinct_inside_an_aggregate_is_honored(): + """https://github.com/jbylund/pg_mimic/issues/58 + + Every `agg(DISTINCT x)` is broken, not just `count`, and the cause is one line + of codegen: the Python generator renders `exp.Distinct` as `set()`, discarding + the expression inside it, so the aggregate is handed an empty set literal + instead of the column. `COUNT(DISTINCT a)` compiles to `COUNT(set())`. + + `count` then answers `count(*)` -- 4 here, so both the DISTINCT *and* the + NULL-skipping that plain `count(a)` gets right are lost -- while `sum` and + `avg` raise on the empty set. `avg(DISTINCT a)` is 1.5 in Postgres and is left + out of the assertions only to keep them integer-exact. + + The parser is fine: the `Distinct` node survives in the AST, in every dialect. + And the executor is wrong in every dialect, not only `postgres`, while + single-argument `count(DISTINCT x)` is SQL-92 and agrees exactly across + PostgreSQL 18.4, SQLite 3.51 and MySQL 8.0.46 -- so a fix upstream is not + dialect-gated and cannot break a dialect that wanted the old behaviour. + """ + assert _rows("SELECT COUNT(a) FROM t", *_DUPLICATES) == [(3,)] # correct today; the contrast + assert _rows("SELECT COUNT(DISTINCT a) FROM t", *_DUPLICATES) == [(2,)] + assert _rows("SELECT COUNT(DISTINCT 1) FROM t", *_DUPLICATES) == [(1,)] + assert _rows("SELECT SUM(DISTINCT a) FROM t", *_DUPLICATES) == [(3,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_casting_to_an_integer_rounds_rather_than_truncates(): + """https://github.com/jbylund/pg_mimic/issues/58 + + `env.cast` sends the value through `int()`, which truncates toward zero, where + Postgres rounds -- and by the same two rules as `round()`: an exact type goes + half away from zero, a binary float goes through rint() and lands half to even. + Verified against PostgreSQL 18. + + So `CAST(0.5 AS INT)` is 1 rather than 0, and `CAST(-1.5::double AS INT)` is -2 + rather than -1. Off-by-one in a value, silently. + """ + exact = {"d": [{"n": Decimal(text)} for text in ("0.5", "-0.5", "2.5")]} + assert _rows("SELECT CAST(n AS INT) FROM d", exact, {"d": {"n": "DECIMAL"}}) == [(1,), (-1,), (3,)] + binary = {"d": [{"v": 0.5}, {"v": 2.5}, {"v": -1.5}]} + assert _rows("SELECT CAST(v AS INT) FROM d", binary, {"d": {"v": "DOUBLE"}}) == [(0,), (2,), (-2,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_round_to_a_scale_goes_half_away_from_zero(): + """https://github.com/jbylund/pg_mimic/issues/58 + + The two-argument `round()`, split from + `test_round_on_numeric_goes_half_away_from_zero` above because it is a separate + code path with a separate symptom: ENV["ROUND"] passes the scale to Python's + builtin, so 0.25 to one place is 0.2 where Postgres says 0.3. + """ + assert _rows("SELECT ROUND(n, 1) FROM d", {"d": [{"n": Decimal("0.25")}]}, {"d": {"n": "DECIMAL"}}) == [(Decimal("0.3"),)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_like_honors_a_backslash_escape(): + """https://github.com/jbylund/pg_mimic/issues/58 + + Postgres' LIKE takes backslash as its default escape character, so + `LIKE 'a\\%b'` matches the literal text `a%b` and nothing else. The executor + treats the backslash as a literal, matches nothing, and answers no rows. + + Distinct from the anchoring and metacharacter bugs fixed in v30.17.0 and + covered by test_like_is_anchored_and_takes_metacharacters_literally: escaping + is the part that still differs. + """ + tables, schema = {"q": [{"s": "a%b"}, {"s": "axb"}]}, {"q": {"s": "TEXT"}} + assert _rows("SELECT s FROM q WHERE s LIKE 'a\\%b'", tables, schema) == [("a%b",)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_exists_subquery_runs(): + """https://github.com/jbylund/pg_mimic/issues/58 + + `EXISTS (subquery)` emits invalid Python, in a WHERE clause or a select list, + correlated or not -- the client gets a SyntaxError out of generated code it + never wrote. The same failure mode as the scalar subquery above and probably + the same cause, but its own tripwire because EXISTS is the more common of the + two by far: it is how an ORM asks whether a related row exists, and how + `NOT IN (subquery)` is rewritten before it runs. + """ + tables, schema = {"t": [{"a": 1}], "u": [{"b": 1}]}, {"t": {"a": "INT"}, "u": {"b": "INT"}} + assert _rows("SELECT a FROM t WHERE EXISTS (SELECT 1 FROM u)", tables, schema) == [(1,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_arithmetic_mixes_numeric_and_double(): + """https://github.com/jbylund/pg_mimic/issues/58 + + A NUMERIC column against a DOUBLE one -- or against a float literal -- raises + `unsupported operand type(s) for *: 'decimal.Decimal' and 'float'`, because the + executor does Python arithmetic on the storage types and Python refuses that + pair. Postgres resolves both to double precision and answers 5. + + The one on this list most likely to be the first thing a user hits, since it + needs no subquery, no join and no unusual function: two numeric columns of + different types multiplied together is enough. + """ + assert _rows("SELECT n * v FROM m", *_MIXED) == [(5.0,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_in_a_subquery_is_null_for_a_null_left_operand(): + """https://github.com/jbylund/pg_mimic/issues/58 + + `NULL IN (SELECT ...)` is NULL in Postgres, not FALSE -- three-valued logic, + and the reason `NOT IN (subquery)` excludes a NULL row rather than keeping it. + The executor answers FALSE, so the negation keeps it, and a `NOT IN` against a + nullable column returns rows Postgres would not. + """ + tables, schema = {"t": [{"a": 1}], "u": [{"b": 1}]}, {"t": {"a": "INT"}, "u": {"b": "INT"}} + assert _rows("SELECT NULL IN (SELECT b FROM u) FROM t", tables, schema) == [(None,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_an_outer_join_pads_rows_a_where_clause_then_removes(): + """https://github.com/jbylund/pg_mimic/issues/58 + + The NULL-padded row an outer join adds survives a WHERE clause that tests the + padded column, so `WHERE x.a = 1` keeps a row whose `x.a` is NULL. Extra rows + out of a filter that should have removed them, which is the worst shape a + wrong answer can take: nothing about the result looks suspicious. + """ + tables = {"x": [{"a": 1}, {"a": 2}], "y": [{"b": 9}]} + schema = {"x": {"a": "INT"}, "y": {"b": "INT"}} + assert _rows("SELECT x.a FROM x FULL JOIN y ON NULL WHERE x.a = 1", tables, schema) == [(1,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_the_trim_and_reverse_functions_exist(): + """https://github.com/jbylund/pg_mimic/issues/58 + + Missing from ENV, like `length()` was before v30.17.0, and fixable the same + way. Grouped into one test because they are one gap: whichever of them is + added first, the rest are a line each. + + Worth a tripwire despite being a plain omission because these are not exotic -- + `btrim` is what Postgres compiles a bare `trim(x)` to. + """ + tables, schema = {"q": [{"s": " ab "}]}, {"q": {"s": "TEXT"}} + assert _rows("SELECT BTRIM(s), LTRIM(s), RTRIM(s), REVERSE(BTRIM(s)) FROM q", tables, schema) == [("ab", "ab ", " ab", "ba")] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_extract_takes_a_day_of_week(): + """https://github.com/jbylund/pg_mimic/issues/58 + + `EXTRACT(DOW FROM ts)` reaches `datetime.dow`, an attribute datetime does not + have, so it raises rather than answering. YEAR, MONTH, DAY and HOUR all work, + which is what makes this a per-field gap rather than a missing EXTRACT. + + 2024-03-15 is a Friday, and Postgres numbers Sunday 0, so 5. + """ + tables, schema = {"d": [{"ts": datetime.datetime(2024, 3, 15)}]}, {"d": {"ts": "TIMESTAMP"}} + assert _rows("SELECT EXTRACT(DOW FROM ts) FROM d", tables, schema) == [(5,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_dividing_null_is_null(): + """https://github.com/jbylund/pg_mimic/issues/58 + + `NULL / 1` raises `int() argument must be a string, a bytes-like object or a + real number, not 'NoneType'`: integer division coerces with `int()` before + checking for NULL. Any nullable integer column divided by anything is enough, + and the other operators handle NULL correctly, so it is division specifically. + """ + assert _rows("SELECT a / 1 FROM n", {"n": [{"a": None}]}, {"n": {"a": "INT"}}) == [(None,)] + + +@pytest.mark.xfail(strict=True, reason=_UPSTREAM_FIXED) +def test_is_not_null_answers_a_boolean(): + """https://github.com/jbylund/pg_mimic/issues/58 + + `x IS NOT NULL` answers `x` itself when x is not null, rather than TRUE -- the + generated code is `x` where it should be `x is not None`. It reads as correct + for as long as the value happens to be truthy, and a column of strings in a + boolean-typed output column is a wire encoding error rather than a wrong row. + + `IS NULL` is fine; only the negation is wrong. + """ + assert _rows("SELECT (CASE WHEN a = 1 THEN 'x' END) IS NOT NULL FROM t", *_NUMBERS) == [(True,), (False,), (False,)] diff --git a/tools/fuzz/README.md b/tools/fuzz/README.md new file mode 100644 index 0000000..0da8c78 --- /dev/null +++ b/tools/fuzz/README.md @@ -0,0 +1,133 @@ +# Differential fuzzing against a real PostgreSQL + +A bug-finding tool, not a test. It generates random `SELECT`s, runs each against a real +PostgreSQL and against sqlglot's executor, and reports where they disagree — minimised to the +smallest query that still disagrees. + +Its *output* is what becomes a test: a finding here turns into a strict-xfail tripwire in +[`tests/test_sqlglot_workarounds.py`](../../tests/test_sqlglot_workarounds.py), which is what runs in +CI and tells us when upstream fixes it. Nothing here runs in CI — it needs a live server, and a +fuzzer that fails the build on a new random seed is a fuzzer people delete. + +## Running it + +```bash +createdb pg_mimic_fuzz +python -m tools.fuzz --count 4000 +``` + +Roughly 150 queries a second, so a few thousand takes seconds to sweep; shrinking is what takes the +time. Exit status is 1 when anything was found. + +| flag | | +|---|---| +| `--dsn` | the oracle. Default `postgres://localhost/pg_mimic_fuzz`. The dataset is dropped and recreated on every run | +| `--target` | `mimic` (default) goes through `TableSession`, so through every rewrite in `tables.py`. `raw` calls `sqlglot.executor.execute` directly | +| `--count`, `--seed` | how many queries, and reproducibly which ones | +| `--ordered` | add a total `ORDER BY` and compare row *order* and `LIMIT`/`OFFSET` too | +| `--only` | severities to report, e.g. `--only VALUE,COUNT` for wrong answers only | +| `--without` | features to switch off, e.g. `--without concat,string_functions` | +| `--json` | write the findings to a file as well | + +## Reading the report + +Findings are graded, worst first, because Postgres and the executor disagree constantly in ways that +are not bugs and a report that ranks them all equally is a report nobody reads. + +| | | +|---|---| +| `COUNT` | different number of rows or columns. Always a bug | +| `VALUE` | genuinely different values. Always a bug | +| `REFUSED` | the executor could not run it. A missing function, or invalid generated Python | +| `ORDER` | same rows, different order. Only under `--ordered` | +| `EXACT` | numerically equal but one side lost exactness — a float where Postgres gave a `Decimal` | +| `TYPE` | equal and equally exact, different Python type. Mostly invisible to a client, since pg_mimic encodes from the *declared* type | + +Each finding also says whether it reproduces in raw sqlglot, which is what says whose bug it is: + +- **also fails raw sqlglot** — an upstream bug reaching our users. +- **raw sqlglot gets this right** — a workaround in `tables.py` is the thing that broke it. Ours. + +## Digging past the top finding + +One unimplemented function that the generator reaches often will account for most of the failures +and hide everything underneath it — `BTRIM` was 206 of 217 on the first real run. That is what +`--without` is for: switch off what you have already recorded and run again. + +```bash +python -m tools.fuzz --count 4000 --without string_functions,concat,extract --only VALUE,COUNT +``` + +## What the first run found + +Fourteen new upstream bugs, now tripwires in +[`tests/test_sqlglot_workarounds.py`](../../tests/test_sqlglot_workarounds.py) — read that file from +`test_distinct_keeps_columns_that_share_a_name` down. Six are wrong answers rather than refusals, +including three that a client cannot detect: `count(DISTINCT x)` ignoring the DISTINCT, +`SELECT DISTINCT a, a` answering one column instead of two, and an outer join's NULL-padded row +surviving a `WHERE` that tests the padded column. + +Four are pg_mimic's own — the executor gets them right and a rewrite in `tables.py` breaks them, so +they are ours to fix and are **not** covered by any test yet: + +```sql +-- _rewrite_not_in turns a correct answer into a wrong one. Postgres 7, raw sqlglot 7, TableSession 8. +SELECT count(*) FROM t WHERE NOT t.f IN (SELECT 1 FROM t) +-- ...and makes a working select-list NOT IN raise. Raw sqlglot answers true. +SELECT NOT 0 IN (SELECT uid FROM u) FROM u +-- The exact-decimal CAST rewrite creates a Decimal/float pair Python won't divide. Raw sqlglot answers. +SELECT 1 FROM u WHERE CAST(1 AS DECIMAL) / 0.5 < 1 +``` + +One thing to know before reading a report: **text ordering is locale-dependent and dominates**. Every +`ORDER` finding on the first `--ordered` run was a text sort, because the executor compares strings +by code point and the oracle compared them by its `en_US.UTF-8` collation — a modelling gap, recorded +in [Known limitations](../../docs/limitations.md), not a bug anyone can fix. Create the oracle with +`createdb --locale=C --template=template0 pg_mimic_fuzz` to get it out of the way and see what else is +there. + +## Checking a finding against upstream main + +Do this before reporting anything: a bug already fixed on main and waiting for a release is not a +finding, and two of the tripwires in this repo are in exactly that state. + +Usually you do not have to. [`.github/workflows/upstream-sqlglot.yml`](../../.github/workflows/upstream-sqlglot.yml) +already runs the whole suite against sqlglot's main branch daily — including how to install it, which +has a trap in it — so an `XPASS(strict)` there is the answer for anything with a tripwire. Read that +job's last run first. + +Reach for a local main build only for a finding that has no tripwire yet, such as one you are still +minimising. `--target raw` needs nothing but sqlglot and psycopg, and the test file needs nothing but +sqlglot and pytest, so add `--noconftest -o addopts=` to skip this repo's conftest and its dev +dependencies: + +```bash +/tmp/mainvenv/bin/python -m pytest tests/test_sqlglot_workarounds.py --noconftest -o addopts= -rX +/tmp/mainvenv/bin/python -m tools.fuzz --target raw --count 2000 +``` + +## How it works + +Four pieces, each answering one of the four questions a differential fuzzer has to get right. + +[`dataset.py`](dataset.py) — *what to query.* Two small tables where every row carries an edge case: +NULLs, duplicates, ties, zero, exact halves for rounding, `'a%b'` for LIKE, and foreign keys that +match nothing, match twice, or are NULL. + +[`generate.py`](generate.py) — *what to ask.* A typed generator, because the hard part of fuzzing SQL +is validity rather than randomness. It knows which of its productions Postgres rejects at runtime +(division by zero, `mod()` on a double, a `FULL JOIN` on a non-hashable condition) and avoids them; +about 1% of what it emits is still invalid, and those samples are discarded. + +[`compare.py`](compare.py) — *whether the answers match.* The graded comparison above, with a +relative tolerance on floats, since both engines do the same arithmetic in a different order and are +allowed to land a few ULPs apart. + +[`shrink.py`](shrink.py) — *what to report.* Without this the tool is useless: hundreds of failures +that are a dozen bugs, each buried under forty lines of generated arithmetic. The reduction moves are +deliberately reckless — hoist any node to its own child, replace any node with a literal, drop any +clause — because **Postgres is the validity filter**: a variant it refuses is rejected exactly like a +variant that stops reproducing. A variant is kept when it is smaller and fails *the same way*, which +is what stops the shrinker sliding off a subtle wrong-answer bug onto some unrelated missing +function. Minimised queries then de-duplicate themselves, since two seeds that found the same bug +minimise to the same text. diff --git a/tools/fuzz/__init__.py b/tools/fuzz/__init__.py new file mode 100644 index 0000000..7bd845d --- /dev/null +++ b/tools/fuzz/__init__.py @@ -0,0 +1,8 @@ +"""Differential fuzzing of sqlglot's executor against a real PostgreSQL. + +Not a test suite: a bug-finding tool whose *output* becomes tests. A finding here +is turned into a strict-xfail tripwire in tests/test_sqlglot_workarounds.py, which +is the thing that runs in CI and tells us when upstream fixes it. + +See tools/fuzz/README.md for how to run it and what the report means. +""" diff --git a/tools/fuzz/__main__.py b/tools/fuzz/__main__.py new file mode 100644 index 0000000..b993948 --- /dev/null +++ b/tools/fuzz/__main__.py @@ -0,0 +1,296 @@ +"""The fuzzer's entry point. + + python -m tools.fuzz --dsn "postgres://localhost/pg_mimic_fuzz" --count 5000 + +Three phases, because shrinking is expensive and most failures are duplicates: + +1. *Sweep.* Generate `count` queries, run each against Postgres and each target, + and record a cheap pre-bucket key per failure -- the normalised message for a + refusal, or the severity plus the set of SQL constructs in the query for a + wrong answer. Queries Postgres itself rejects are discarded. +2. *Shrink.* One representative per pre-bucket, reduced until it stops getting + smaller. This is where the budget goes: a few hundred evaluations per bucket + rather than per failure. +3. *Report.* Re-bucket by the minimised SQL, which merges pre-buckets that turned + out to be the same bug, and print each with the count of raw failures behind + it and whether it also reproduces in raw sqlglot. + +Exit status is 1 when anything was found, so this can gate a nightly job. +""" + +from __future__ import annotations + +import argparse +import collections +import json +import re +import sys +import time + +import sqlglot +from sqlglot import exp + +from . import compare, engines, generate +from .engines import Failed + + +class Finding: + def __init__(self, kind: str, detail: str, sql: str, engine: str): + self.kind = kind + self.detail = detail + self.sql = sql + self.engine = engine + self.count = 1 + self.minimised = sql + self.also_raw: bool | None = None + + @property + def verdict(self) -> tuple[str, str]: + return (self.kind, engines.signature(self.detail) if self.kind == "REFUSED" else "") + + +def canonical_aliases(sql: str) -> str: + """The same query with generated table aliases renumbered from one. + + Purely so identical bugs de-duplicate. Two seeds that both minimise to + `SELECT AVG(x.n) FROM t AS x` differ only in whether the generator had reached + a298 or a346 by then, and without this they are reported as two findings. + """ + try: + tree = sqlglot.parse_one(sql, dialect="postgres") + except Exception: + return sql + mapping: dict[str, str] = {} + for node in tree.walk(): + if isinstance(node, exp.Identifier) and re.fullmatch(r"a\d+", node.name): + mapping.setdefault(node.name, f"x{len(mapping) + 1}") + node.set("this", mapping[node.name]) + try: + return tree.sql(dialect="postgres") + except Exception: + return sql + + +def _constructs(sql: str) -> str: + try: + tree = sqlglot.parse_one(sql, dialect="postgres") + except Exception: + return "unparsed" + names = {type(node).__name__ for node in tree.walk() if not isinstance(node, (exp.Identifier, exp.Literal, exp.Column))} + return ",".join(sorted(names)) + + +def covers_every_column(sql: str, width: int) -> bool: + """Whether `sql` ends in an ORDER BY over ordinals 1..width, so that comparing + Postgres' row *sequence* against the executor's is legitimate. + + Two rows that tie on every output column are equal tuples, so any tie-break + gives the same sequence -- but only if the sort keys cover every column. The + generator emits exactly that, and the shrinker is supposed to preserve it; + checking here rather than trusting either of them is what keeps a dropped sort + key from being reported as a wrong row order. + """ + if width == 0: + return True + try: + tree = sqlglot.parse_one(sql, dialect="postgres") + except Exception: + return False + order = tree.args.get("order") + if order is None: + return False + ordinals = set() + for item in order.expressions: + key = item.this + if not isinstance(key, exp.Literal) or key.is_string: + return False + ordinals.add(int(key.name)) + return ordinals == set(range(1, width + 1)) + + +def _evaluate(sql: str, oracle: engines.Postgres, target, ordered: bool) -> tuple[str, str] | None: + """None if the target agrees with Postgres, else (kind, detail). + + A query Postgres refuses raises Rejected, which the caller counts as a + discarded sample rather than a finding -- it means the generator or the + shrinker built something invalid, and only Postgres gets to say so. + """ + try: + expected = oracle.run(sql) + except Failed as exc: + raise Rejected(exc.message) from None + try: + actual = target.run(sql) + except Failed as exc: + return ("REFUSED", exc.message) + total_order = ordered and covers_every_column(sql, len(expected[0]) if expected else 0) + difference = compare.compare(expected, actual, ordered=total_order) + return None if difference is None else difference + + +class Rejected(Exception): + pass + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="python -m tools.fuzz", description=__doc__) + parser.add_argument("--dsn", default="postgres://localhost/pg_mimic_fuzz", help="a Postgres to use as the oracle") + parser.add_argument("--count", type=int, default=2000, help="queries to generate") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--target", default="mimic", choices=["mimic", "raw"], help="which side of tables.py to fuzz") + parser.add_argument("--ordered", action="store_true", help="add a total ORDER BY and compare row order and LIMIT too") + parser.add_argument("--without", default="", help=f"comma-separated features to disable, from: {','.join(generate.FEATURES)}") + parser.add_argument("--only", default="", help="comma-separated severities to report, e.g. VALUE,COUNT,REFUSED") + parser.add_argument("--depth", type=int, default=3, help="maximum expression nesting") + parser.add_argument("--shrink-budget", type=int, default=600, help="evaluations per finding") + parser.add_argument("--no-shrink", action="store_true") + parser.add_argument("--json", dest="json_path", default=None, help="also write the findings here") + arguments = parser.parse_args(argv) + + without = {feature for feature in arguments.without.split(",") if feature} + unknown = without - set(generate.FEATURES) + if unknown: + parser.error(f"unknown feature(s) {sorted(unknown)}") + wanted = {severity for severity in arguments.only.split(",") if severity} + + oracle = engines.Postgres(arguments.dsn) + oracle.load() + target = engines.build([arguments.target])[0] + + generator = generate.Generator(seed=arguments.seed, without=without, ordered=arguments.ordered, max_depth=arguments.depth) + + buckets: dict[tuple, Finding] = {} + rejected: collections.Counter = collections.Counter() + ran = 0 + started = time.time() + + for index in range(arguments.count): + sql = generator.select() + try: + difference = _evaluate(sql, oracle, target, arguments.ordered) + except Rejected as exc: + rejected[engines.signature(str(exc))] += 1 + continue + ran += 1 + if difference is None: + continue + kind, detail = difference + if wanted and kind not in wanted: + continue + key = (kind, engines.signature(detail) if kind == "REFUSED" else _constructs(sql)) + if key in buckets: + buckets[key].count += 1 + else: + buckets[key] = Finding(kind, detail, sql, arguments.target) + if (index + 1) % 250 == 0: + print( + f" {index + 1}/{arguments.count} generated, {ran} valid, {sum(f.count for f in buckets.values())} failures," + f" {len(buckets)} buckets", + file=sys.stderr, + ) + + print( + f"\nswept {arguments.count} queries in {time.time() - started:.1f}s: {ran} ran on postgres," + f" {arguments.count - ran} rejected as invalid", + file=sys.stderr, + ) + + findings = list(buckets.values()) + raw = target if arguments.target == "raw" else engines.RawSqlglot() + if findings and not arguments.no_shrink: + print(f"shrinking {len(findings)} findings...", file=sys.stderr) + for finding in findings: + + def check(candidate: str, finding: Finding = finding) -> tuple[str, str] | None: + try: + outcome = _evaluate(candidate, oracle, target, arguments.ordered) + except Rejected: + return None + if outcome is None: + return None + kind, detail = outcome + return (kind, engines.signature(detail) if kind == "REFUSED" else "") + + if not arguments.no_shrink: + minimised, spent = shrink_finding(finding, check, arguments) + finding.minimised = canonical_aliases(minimised) + # The detail describes which row and column diverged and what the two + # values were, and after shrinking it describes a query that is no + # longer being shown. Re-running the minimised query is what keeps the + # report's two halves talking about the same thing. + try: + refreshed = _evaluate(finding.minimised, oracle, target, arguments.ordered) + except Rejected: + refreshed = None + if refreshed is not None and refreshed[0] == finding.kind: + finding.detail = refreshed[1] + else: + finding.minimised = finding.sql + print(f" {spent:4d} evaluations: {finding.minimised}", file=sys.stderr) + # Whether the same query also diverges without pg_mimic's rewrites in the + # way, which is what says who owns the bug. Only meaningful when the + # target is `mimic`; fuzzing raw sqlglot answers it by construction. + if arguments.target == "raw": + finding.also_raw = True + else: + try: + finding.also_raw = _evaluate(finding.minimised, oracle, raw, arguments.ordered) is not None + except Rejected: + finding.also_raw = None + + merged: dict[tuple, Finding] = {} + for finding in findings: + key = (finding.kind, finding.minimised) + if key in merged: + merged[key].count += finding.count + else: + merged[key] = finding + report(sorted(merged.values(), key=lambda f: (compare.SEVERITY.index(f.kind) if f.kind in compare.SEVERITY else -1, -f.count))) + + if rejected: + print("\nmost common reasons postgres rejected a generated query (generator quality, not findings):") + for message, count in rejected.most_common(5): + print(f" {count:5d} {message}") + + if arguments.json_path: + with open(arguments.json_path, "w") as handle: + json.dump( + [ + { + "kind": finding.kind, + "sql": finding.minimised, + "original_sql": finding.sql, + "detail": finding.detail, + "occurrences": finding.count, + "also_fails_raw_sqlglot": finding.also_raw, + } + for finding in merged.values() + ], + handle, + indent=2, + ) + return 1 if merged else 0 + + +def shrink_finding(finding: Finding, check, arguments) -> tuple[str, int]: + from .shrink import shrink + + return shrink(finding.sql, finding.verdict, check, keep_order=arguments.ordered, budget=arguments.shrink_budget) + + +def report(findings: list[Finding]) -> None: + if not findings: + print("\nno divergences found.") + return + print(f"\n{len(findings)} distinct divergence(s), worst first:\n") + for number, finding in enumerate(findings, start=1): + origin = {True: "also fails raw sqlglot", False: "raw sqlglot gets this right", None: "not reproducible standalone"}[ + finding.also_raw + ] + print(f"{number}. {finding.kind} ({finding.count} occurrence(s), {origin})") + print(f" {finding.minimised}") + print(f" {finding.detail}\n") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/fuzz/compare.py b/tools/fuzz/compare.py new file mode 100644 index 0000000..f7a0f1e --- /dev/null +++ b/tools/fuzz/compare.py @@ -0,0 +1,165 @@ +"""Deciding whether two result sets are the same answer. + +This is where a differential fuzzer is won or lost. Postgres and sqlglot's +executor disagree constantly in ways that are not bugs -- `sum(int)` is a Python +int on one side and a bigint-shaped int on the other, `avg(int)` is a Decimal in +Postgres and a float in the executor -- and a comparison that calls those +failures buries the real findings. + +So a difference is graded rather than merely detected: + +VALUE the numbers, strings or timestamps genuinely differ. Always a bug. +COUNT different number of rows. Always a bug. +EXACT numerically equal, but one side lost exactness -- a float where Postgres + returned a Decimal, so 0.1 + 0.2 is 0.30000000000000004 rather than + 0.3. Reported separately because it is real (issue #33) but of a + different kind, and because it fires often. +TYPE numerically equal, same exactness, different Python type -- an int where + Postgres said bool, say. Lowest severity; a client mostly cannot tell, + since pg_mimic encodes from the *declared* type rather than the value. +ORDER same multiset, different sequence. Only meaningful under --ordered. + +Floats are compared with a relative tolerance because both engines do their own +arithmetic in a different order, and `(a + b) * c` is allowed to land a few ULPs +apart. The tolerance is tight enough that a wrong answer -- integer division +giving 0 where Postgres gives 0.5, a rounding mode giving 2 where Postgres gives +3 -- is never inside it. +""" + +from __future__ import annotations + +import datetime +import math +from decimal import Decimal + +_TOLERANCE = 1e-9 + +# Ordered worst-first: a single row can produce several kinds of difference and +# the run should be bucketed by the most serious one. +SEVERITY = ["COUNT", "VALUE", "EXACT", "TYPE"] + + +def _is_number(value: object) -> bool: + return isinstance(value, (int, float, Decimal)) and not isinstance(value, bool) + + +def _numerically_equal(left: object, right: object) -> bool: + left_float, right_float = float(left), float(right) # type: ignore[arg-type] + if math.isnan(left_float) and math.isnan(right_float): + return True + if left_float == right_float: + return True + scale = max(abs(left_float), abs(right_float)) + return abs(left_float - right_float) <= _TOLERANCE * scale + + +def compare_value(expected: object, actual: object) -> str | None: + """None if the two are the same answer, else the severity of the difference. + + `expected` is Postgres. + """ + if expected is None and actual is None: + return None + if expected is None or actual is None: + return "VALUE" + + expected_number, actual_number = _is_number(expected), _is_number(actual) + + # A bool is not a number here, but Postgres' bool against the executor's 1/0 + # is a type difference and not a wrong answer. + if isinstance(expected, bool) or isinstance(actual, bool): + if isinstance(expected, bool) and isinstance(actual, bool): + return None if expected == actual else "VALUE" + if actual_number or expected_number: + return "TYPE" if int(expected) == int(actual) else "VALUE" # type: ignore[call-overload] + return "VALUE" + + if expected_number and actual_number: + if not _numerically_equal(expected, actual): + return "VALUE" + if isinstance(expected, Decimal) and isinstance(actual, float): + return "EXACT" + if type(expected) is not type(actual): + return "TYPE" + return None + + if isinstance(expected, str) and isinstance(actual, str): + return None if expected == actual else "VALUE" + + if isinstance(expected, datetime.datetime) and isinstance(actual, datetime.datetime): + if expected.replace(tzinfo=None) != actual.replace(tzinfo=None): + return "VALUE" + return None if (expected.tzinfo is None) == (actual.tzinfo is None) else "TYPE" + + if isinstance(expected, datetime.date) and isinstance(actual, datetime.datetime): + return None if actual.date() == expected and actual.time() == datetime.time() else "VALUE" + + return None if expected == actual else "VALUE" + + +def _sort_key(row: tuple) -> tuple: + """A canonical order that puts equal-valued rows of unequal Python type together. + + Sorting each side independently is only sound if the key ignores exactly the + differences compare_value forgives -- otherwise a Decimal row and its float + twin sort to different positions and every comparison downstream is offset by + one. + """ + key = [] + for value in row: + if value is None: + key.append((0, 0.0, "")) + elif isinstance(value, bool): + key.append((1, float(value), "")) + elif _is_number(value): + number = float(value) # type: ignore[arg-type] + key.append((1, 0.0 if math.isnan(number) else number, "")) + elif isinstance(value, str): + key.append((2, 0.0, value)) + else: + key.append((3, 0.0, str(value))) + return tuple(key) + + +def compare(expected: list[tuple], actual: list[tuple], ordered: bool) -> tuple[str, str] | None: + """None if the two result sets are the same answer, else (severity, detail).""" + if len(expected) != len(actual): + return ("COUNT", f"{len(expected)} rows from postgres, {len(actual)} from the executor") + + if ordered: + sequence_difference = _first_difference(expected, actual) + + left = sorted(expected, key=_sort_key) + right = sorted(actual, key=_sort_key) + multiset_difference = _first_difference(left, right) + if multiset_difference: + return multiset_difference + + if ordered and sequence_difference: + index = next( + (position for position, (a, b) in enumerate(zip(expected, actual)) if _sort_key(a) != _sort_key(b)), + 0, + ) + return ( + "ORDER", + f"same rows, different order: at row {index} postgres has {expected[index]!r}, the executor {actual[index]!r}", + ) + return None + + +def _first_difference(expected: list[tuple], actual: list[tuple]) -> tuple[str, str] | None: + worst: tuple[str, str] | None = None + for index, (expected_row, actual_row) in enumerate(zip(expected, actual)): + if len(expected_row) != len(actual_row): + return ("COUNT", f"row {index} has {len(expected_row)} columns in postgres, {len(actual_row)} in the executor") + for position, (expected_value, actual_value) in enumerate(zip(expected_row, actual_row)): + severity = compare_value(expected_value, actual_value) + if severity is None: + continue + detail = ( + f"row {index} column {position}: postgres {expected_value!r} ({type(expected_value).__name__}), " + f"executor {actual_value!r} ({type(actual_value).__name__})" + ) + if worst is None or SEVERITY.index(severity) < SEVERITY.index(worst[0]): + worst = (severity, detail) + return worst diff --git a/tools/fuzz/dataset.py b/tools/fuzz/dataset.py new file mode 100644 index 0000000..839a119 --- /dev/null +++ b/tools/fuzz/dataset.py @@ -0,0 +1,151 @@ +"""The fixed dataset every generated query runs against. + +Small on purpose -- eight rows and six -- because a differential fuzzer finds bugs +with *shapes*, not with volume: a wrong join, a mishandled NULL, a tie broken the +wrong way. Every row here is carrying at least one edge case, and a bigger table +would only make the shrinker slower and the repros harder to read. + +The values are chosen so that a plausible-but-wrong implementation produces a +*different* answer rather than an accidentally equal one: + +- NULLs in every nullable column, including one all-NULL-ish float and a NULL + timestamp, so outer joins and aggregates have something to get wrong. +- Duplicates (``n`` is 5 twice, ``s`` is 'abc' twice) so DISTINCT and GROUP BY + have work to do and ORDER BY has ties. +- Zero and negative zero, so division and sign() are interesting. +- ``d`` holds exact halves (2.5, -2.5, 0.5) because rounding mode is a real + divergence: Postgres' round(numeric) goes half away from zero and Python's + builtin goes half to even. +- 'ABC' next to 'abc' for ILIKE and case-sensitive ordering, '' for the + empty-string-is-not-NULL rule, and 'a%b' so a LIKE pattern that is taken as a + regex instead of a pattern gives itself away. +- ``u.t_id`` has a value matching nothing (99), a NULL, and two rows pointing at + the same ``t`` row, so INNER/LEFT/RIGHT/FULL joins all differ from each other. +""" + +from __future__ import annotations + +import datetime +from decimal import Decimal +from typing import Any + +Kind = str + +NUM: Kind = "num" +TEXT: Kind = "text" +BOOL: Kind = "bool" +TS: Kind = "ts" + + +class Column: + def __init__(self, name: str, kind: Kind, pg_type: str, py_type: type, exact: bool = False): + self.name = name + self.kind = kind + self.pg_type = pg_type + self.py_type = py_type + # An exact column is one Postgres stores without a binary float: only + # these are safe to compare for equality rather than within a tolerance. + self.exact = exact + + def __repr__(self) -> str: + return f"Column({self.name!r}, {self.kind!r})" + + +class Table: + def __init__(self, name: str, columns: list[Column], rows: list[tuple]): + self.name = name + self.columns = columns + self.rows = [dict(zip([column.name for column in columns], row)) for row in rows] + + def of_kind(self, kind: Kind) -> list[Column]: + return [column for column in self.columns if column.kind == kind] + + @property + def ddl(self) -> str: + body = ", ".join(f"{column.name} {column.pg_type}" for column in self.columns) + return f"CREATE TABLE {self.name} ({body})" + + @property + def py_types(self) -> dict[str, type]: + return {column.name: column.py_type for column in self.columns} + + +T = Table( + "t", + [ + Column("id", NUM, "INTEGER", int, exact=True), + Column("n", NUM, "INTEGER", int, exact=True), + Column("s", TEXT, "TEXT", str), + Column("f", NUM, "DOUBLE PRECISION", float), + Column("d", NUM, "NUMERIC(12,4)", Decimal, exact=True), + Column("b", BOOL, "BOOLEAN", bool), + Column("ts", TS, "TIMESTAMP", datetime.datetime), + ], + [ + (1, 1, "abc", 1.5, Decimal("2.5"), True, datetime.datetime(2024, 1, 15, 10, 30)), + (2, None, "ABC", -1.5, Decimal("-2.5"), False, datetime.datetime(2024, 3, 15)), + (3, 0, "", 0.0, Decimal("0.5"), None, datetime.datetime(2024, 3, 15)), + (4, 5, None, 2.5, Decimal("9.99"), True, datetime.datetime(2023, 12, 31, 23, 59, 59)), + (5, 5, "a%b", 0.25, Decimal("0"), False, datetime.datetime(2024, 2, 29, 12, 0)), + (6, -3, "Bump version", -0.0, Decimal("-0.0050"), True, None), + (7, None, "zzz", None, None, None, datetime.datetime(2025, 6, 1)), + (8, 2, "abc", 3.5, Decimal("3.5"), False, datetime.datetime(2024, 1, 15, 10, 30)), + ], +) + +U = Table( + "u", + [ + Column("uid", NUM, "INTEGER", int, exact=True), + Column("t_id", NUM, "INTEGER", int, exact=True), + Column("us", TEXT, "TEXT", str), + Column("un", NUM, "NUMERIC(12,4)", Decimal, exact=True), + ], + [ + (10, 1, "x", Decimal("100")), + (11, 1, "y", Decimal("-1")), + (12, 3, None, Decimal("0")), + (13, None, "x", Decimal("7")), + (14, 99, "z", None), + (15, 8, "x", Decimal("100")), + ], +) + +TABLES = [T, U] + + +def sqlglot_schema() -> dict[str, dict[str, str]]: + """What the raw executor wants: a type name per column, as sqlglot spells it.""" + names = { + "INTEGER": "INT", + "DOUBLE PRECISION": "DOUBLE", + "NUMERIC(12,4)": "DECIMAL", + "TEXT": "TEXT", + "BOOLEAN": "BOOLEAN", + "TIMESTAMP": "TIMESTAMP", + } + return {table.name: {column.name: names[column.pg_type] for column in table.columns} for table in TABLES} + + +def rows() -> dict[str, list[dict[str, Any]]]: + return {table.name: table.rows for table in TABLES} + + +def declared_columns() -> dict[str, dict[str, type]]: + """Types handed to TableSession rather than left to inference. + + Inference would get most of these right, but not all -- ``t.f`` is NULL in one + row and ``t.b`` is NULL in three -- and a fuzzer whose schema drifts with its + data is a fuzzer whose failures are its own fault. + """ + return {table.name: table.py_types for table in TABLES} + + +def load(connection) -> None: + with connection.cursor() as cursor: + for table in TABLES: + cursor.execute(f"DROP TABLE IF EXISTS {table.name}") + cursor.execute(table.ddl) + placeholders = ", ".join(["%s"] * len(table.columns)) + values = [tuple(row[column.name] for column in table.columns) for row in table.rows] + cursor.executemany(f"INSERT INTO {table.name} VALUES ({placeholders})", values) diff --git a/tools/fuzz/engines.py b/tools/fuzz/engines.py new file mode 100644 index 0000000..a0ba43f --- /dev/null +++ b/tools/fuzz/engines.py @@ -0,0 +1,140 @@ +"""The three things a generated query can be run against. + +`postgres` is the oracle. The other two are the same executor seen from either +side of pg_mimic's workaround layer, and running both is what makes a finding +actionable: + +- fails in `raw` and in `mimic` -> an upstream sqlglot bug reaching our users +- fails in `raw` only -> a workaround in tables.py is doing its job +- fails in `mimic` only -> the workaround itself is wrong. Ours to fix. + +`mimic` deliberately stops short of the wire protocol. TableSession.query() is +where every rewrite in tables.py has already been applied, and going through a +socket per query would cost a hundredfold for no extra coverage -- the encoding +path has its own tests. +""" + +from __future__ import annotations + +import asyncio +import re +import warnings +from typing import Any + +from . import dataset + +# The executor compiles a query to Python source and exec()s it, so a query whose +# generated code is questionable -- `NOT(-1 is None)` -- emits a SyntaxWarning +# from a file the user never wrote. Interesting once; at fuzzing volume it is +# thousands of lines of stderr on top of the report. +warnings.filterwarnings("ignore", category=SyntaxWarning) + + +class Failed(Exception): + """The engine refused the query. Carries the message, not a traceback.""" + + def __init__(self, message: str, sqlstate: str | None = None): + super().__init__(message) + self.message = message + self.sqlstate = sqlstate + + +_NOISE = [ + (re.compile(r"0x[0-9a-f]+", re.I), "ADDR"), + (re.compile(r"\(\d+\)"), "(N)"), + (re.compile(r"'[^']*'"), "'S'"), + (re.compile(r'"[^"]*"'), '"S"'), + (re.compile(r"\b\d+\b"), "N"), + (re.compile(r"\s+"), " "), +] + + +def signature(message: str) -> str: + """A message with the query-specific parts taken out, so two failures with the + same cause land in the same bucket. + + Without this, `name 'DPIPE' is not defined` at step `Sort: _0 (4400334928)` + and the same failure at step `Sort: _0 (4400339104)` are two findings. + """ + text = message.split("sqlglot's executor is not a complete SQL engine")[0] + for pattern, replacement in _NOISE: + text = pattern.sub(replacement, text) + return text.strip()[:200] + + +class Postgres: + name = "postgres" + + def __init__(self, dsn: str, statement_timeout_ms: int = 5000): + import psycopg + + self._connection = psycopg.Connection.connect(dsn, autocommit=True) + with self._connection.cursor() as cursor: + cursor.execute("SET TIME ZONE 'UTC'") + cursor.execute(f"SET statement_timeout = {statement_timeout_ms}") + # Not cosmetic: a server whose default_table_access_method is an + # extension's -- a developer machine with a table AM under test, which + # is exactly the kind of machine this gets run on -- fails CREATE TABLE + # outright, and a non-heap AM could answer a query differently. + cursor.execute("SET default_table_access_method = 'heap'") + self._error = psycopg.Error + + def load(self) -> None: + dataset.load(self._connection) + + def run(self, sql: str) -> list[tuple]: + try: + with self._connection.cursor() as cursor: + cursor.execute(sql) # type: ignore[arg-type] + return [tuple(row) for row in cursor.fetchall()] + except self._error as exc: + sqlstate = getattr(getattr(exc, "diag", None), "sqlstate", None) + raise Failed(str(exc).strip().splitlines()[0], sqlstate) from None + + def close(self) -> None: + self._connection.close() + + +class RawSqlglot: + name = "raw" + + def __init__(self) -> None: + self._schema = dataset.sqlglot_schema() + self._tables = dataset.rows() + + def run(self, sql: str) -> list[tuple]: + from sqlglot.executor import execute + + try: + result = execute(sql, schema=self._schema, tables=self._tables, dialect="postgres") + except Exception as exc: + raise Failed(f"{type(exc).__name__}: {exc}") from None + return [tuple(row) for row in result.rows] + + +class Mimic: + name = "mimic" + + def __init__(self) -> None: + from pg_mimic import TableSession + + self._session = TableSession(dataset.rows(), columns=dataset.declared_columns()) + self._loop = asyncio.new_event_loop() + + def run(self, sql: str) -> list[tuple]: + from pg_mimic.errors import PgError + + try: + rows = self._loop.run_until_complete(self._session.query(sql, [])) + except PgError as exc: + raise Failed(f"{exc.sqlstate}: {exc}", exc.sqlstate) from None + except Exception as exc: + raise Failed(f"{type(exc).__name__}: {exc}") from None + return [tuple(row) for row in rows] + + def close(self) -> None: + self._loop.close() + + +def build(names: list[str]) -> list[Any]: + return [RawSqlglot() if name == "raw" else Mimic() for name in names] diff --git a/tools/fuzz/generate.py b/tools/fuzz/generate.py new file mode 100644 index 0000000..152176d --- /dev/null +++ b/tools/fuzz/generate.py @@ -0,0 +1,499 @@ +"""A random query generator that emits SQL Postgres will actually accept. + +The hard part of fuzzing SQL is not randomness, it is *validity*: a generator that +produces garbage 90% of the time spends its budget on Postgres syntax errors and +finds nothing. So this one is typed -- every production knows whether it is +building a number, a string, a boolean or a timestamp -- and it is careful about +the handful of things Postgres rejects at runtime rather than at parse time: + +- division and mod by zero, avoided by only ever dividing by a non-zero literal or + by ``NULLIF(x, 0)``; +- ``mod()`` and two-argument ``round()``, which Postgres has for numeric but not + for double precision, so both operands get a NUMERIC cast; +- negative lengths to ``substring``. + +The remaining invalid queries are simply discarded when Postgres refuses them, +which also means the shrinker in shrink.py can propose reckless variants and let +Postgres filter them. + +Features are named and weighted so a run can drop one with ``--without``. That +matters more than it sounds: a single broken construct that the generator reaches +often -- ``||``, say, before sqlglot implemented it -- otherwise accounts for most +of the failures and hides everything underneath it. +""" + +from __future__ import annotations + +import random + +from . import dataset +from .dataset import BOOL, NUM, TEXT, TS, Kind + +FEATURES = [ + "arithmetic", + "division", + "mod", + "round", + "numeric_functions", + "case", + "cast", + "coalesce", + "nullif", + "least_greatest", + "concat", + "string_functions", + "substring", + "like", + "in_list", + "between", + "is_distinct_from", + "logical", + "timestamp", + "date_trunc", + "extract", + "aggregate", + "count_distinct", + "group_by", + "having", + "distinct", + "join", + "outer_join", + "cross_join", + "self_join", + "set_operation", + "cte", + "derived_table", + "exists_subquery", + "in_subquery", + "scalar_subquery", + "limit", +] + +_TEXT_LITERALS = ["'abc'", "'ABC'", "''", "'a%b'", "'zzz'", "'x'", "'Bump version'"] +_TEXT_PATTERNS = ["'abc'", "'a%'", "'%b%'", "'_bc'", "'a\\%b'", "'%'", "'A%'"] +_NUM_LITERALS = ["0", "1", "2", "-1", "5", "2.5", "-0.5", "100", "0.0"] +_NONZERO_LITERALS = ["1", "2", "-1", "5", "2.5", "-0.5"] +_TS_LITERALS = ["TIMESTAMP '2024-03-15 00:00:00'", "TIMESTAMP '2024-01-15 10:30:00'"] +_DATE_PARTS = ["'year'", "'month'", "'day'", "'hour'"] +_EXTRACT_PARTS = ["YEAR", "MONTH", "DAY", "HOUR", "DOW"] + + +class Source: + """One entry in a FROM clause: a relation, the alias it is under, and its columns.""" + + def __init__(self, alias: str, table: dataset.Table, sql: str | None = None): + self.alias = alias + self.table = table + self.sql = sql or table.name + + def columns_of(self, kind: Kind) -> list[str]: + return [f"{self.alias}.{column.name}" for column in self.table.of_kind(kind)] + + @property + def key(self) -> str: + return f"{self.alias}.{self.table.columns[0].name}" + + @property + def from_sql(self) -> str: + return self.sql if self.sql == self.alias else f"{self.sql} AS {self.alias}" + + +class Generator: + def __init__(self, seed: int, without: set[str] | None = None, ordered: bool = False, max_depth: int = 3): + self.random = random.Random(seed) + self.disabled = without or set() + self.ordered = ordered + self.max_depth = max_depth + self._alias_counter = 0 + + # -- plumbing ------------------------------------------------------------ + + def on(self, feature: str) -> bool: + return feature not in self.disabled + + def _pick(self, options: list) -> object: + return self.random.choice(options) + + def _chance(self, probability: float) -> bool: + return self.random.random() < probability + + def _alias(self) -> str: + self._alias_counter += 1 + return f"a{self._alias_counter}" + + def _choices(self, options: list[tuple[str, float]]) -> str: + """Weighted choice over (feature, weight), skipping disabled features.""" + live = [(name, weight) for name, weight in options if self.on(name) or name == "column"] + total = sum(weight for _, weight in live) + threshold = self.random.random() * total + for name, weight in live: + threshold -= weight + if threshold <= 0: + return name + return live[-1][0] + + # -- expressions --------------------------------------------------------- + + def expr(self, kind: Kind, sources: list[Source], depth: int) -> str: + if kind == NUM: + return self.num(sources, depth) + if kind == TEXT: + return self.text(sources, depth) + if kind == BOOL: + return self.boolean(sources, depth) + return self.timestamp(sources, depth) + + def _leaf(self, kind: Kind, sources: list[Source]) -> str: + available = [column for source in sources for column in source.columns_of(kind)] + literals = {NUM: _NUM_LITERALS, TEXT: _TEXT_LITERALS, BOOL: ["TRUE", "FALSE", "NULL"], TS: _TS_LITERALS}[kind] + if not available or self._chance(0.25): + return str(self._pick(literals)) + return str(self._pick(available)) + + def num(self, sources: list[Source], depth: int) -> str: + if depth <= 0: + return self._leaf(NUM, sources) + choice = self._choices( + [ + ("column", 34), + ("arithmetic", 14), + ("division", 5), + ("mod", 3), + ("round", 4), + ("numeric_functions", 8), + ("case", 5), + ("cast", 4), + ("coalesce", 5), + ("nullif", 3), + ("least_greatest", 4), + ("string_functions", 4), + ("extract", 3), + ("scalar_subquery", 2), + ] + ) + if choice == "column": + return self._leaf(NUM, sources) + if choice == "arithmetic": + operator = self._pick(["+", "-", "*"]) + return f"({self.num(sources, depth - 1)} {operator} {self.num(sources, depth - 1)})" + if choice == "division": + return f"({self.num(sources, depth - 1)} / {self._divisor(sources, depth - 1)})" + if choice == "mod": + left = f"CAST({self.num(sources, depth - 1)} AS NUMERIC)" + return f"MOD({left}, CAST({self._divisor(sources, depth - 1)} AS NUMERIC))" + if choice == "round": + inner = f"CAST({self.num(sources, depth - 1)} AS NUMERIC)" + return f"ROUND({inner})" if self._chance(0.5) else f"ROUND({inner}, {self._pick(['0', '1', '2'])})" + if choice == "numeric_functions": + function = self._pick(["ABS", "CEIL", "FLOOR", "SIGN", "TRUNC"]) + return f"{function}({self.num(sources, depth - 1)})" + if choice == "case": + return ( + f"CASE WHEN {self.boolean(sources, depth - 1)} THEN {self.num(sources, depth - 1)} " + f"ELSE {self.num(sources, depth - 1)} END" + ) + if choice == "cast": + return f"CAST({self.num(sources, depth - 1)} AS {self._pick(['NUMERIC', 'INTEGER', 'DOUBLE PRECISION'])})" + if choice == "coalesce": + return f"COALESCE({self.num(sources, depth - 1)}, {self._leaf(NUM, [])})" + if choice == "nullif": + return f"NULLIF({self.num(sources, depth - 1)}, {self._leaf(NUM, [])})" + if choice == "least_greatest": + function = self._pick(["LEAST", "GREATEST"]) + return f"{function}({self.num(sources, depth - 1)}, {self.num(sources, depth - 1)})" + if choice == "string_functions": + return f"LENGTH({self.text(sources, depth - 1)})" + if choice == "extract": + return f"EXTRACT({self._pick(_EXTRACT_PARTS)} FROM {self.timestamp(sources, depth - 1)})" + return self._scalar_subquery(NUM) + + def _divisor(self, sources: list[Source], depth: int) -> str: + if self._chance(0.6): + return str(self._pick(_NONZERO_LITERALS)) + return f"NULLIF({self.num(sources, max(depth, 0))}, 0)" + + def text(self, sources: list[Source], depth: int) -> str: + if depth <= 0: + return self._leaf(TEXT, sources) + choice = self._choices( + [ + ("column", 40), + ("concat", 14), + ("string_functions", 14), + ("substring", 8), + ("coalesce", 8), + ("nullif", 4), + ("case", 6), + ("least_greatest", 4), + ] + ) + if choice == "column": + return self._leaf(TEXT, sources) + if choice == "concat": + return f"({self.text(sources, depth - 1)} || {self.text(sources, depth - 1)})" + if choice == "string_functions": + function = self._pick(["LOWER", "UPPER", "BTRIM", "LTRIM", "RTRIM", "REVERSE"]) + return f"{function}({self.text(sources, depth - 1)})" + if choice == "substring": + start, length = self._pick(["1", "2", "3"]), self._pick(["1", "2", "5"]) + return f"SUBSTRING({self.text(sources, depth - 1)}, {start}, {length})" + if choice == "coalesce": + return f"COALESCE({self.text(sources, depth - 1)}, {self._leaf(TEXT, [])})" + if choice == "nullif": + return f"NULLIF({self.text(sources, depth - 1)}, {self._leaf(TEXT, [])})" + if choice == "case": + return ( + f"CASE WHEN {self.boolean(sources, depth - 1)} THEN {self.text(sources, depth - 1)} " + f"ELSE {self.text(sources, depth - 1)} END" + ) + function = self._pick(["LEAST", "GREATEST"]) + return f"{function}({self.text(sources, depth - 1)}, {self.text(sources, depth - 1)})" + + def timestamp(self, sources: list[Source], depth: int) -> str: + if depth <= 0 or not self.on("timestamp"): + return self._leaf(TS, sources) + choice = self._choices([("column", 60), ("date_trunc", 20), ("coalesce", 10), ("case", 10)]) + if choice == "date_trunc": + return f"DATE_TRUNC({self._pick(_DATE_PARTS)}, {self.timestamp(sources, depth - 1)})" + if choice == "coalesce": + return f"COALESCE({self.timestamp(sources, depth - 1)}, {self._pick(_TS_LITERALS)})" + if choice == "case": + return ( + f"CASE WHEN {self.boolean(sources, depth - 1)} THEN {self.timestamp(sources, depth - 1)} " + f"ELSE {self.timestamp(sources, depth - 1)} END" + ) + return self._leaf(TS, sources) + + def boolean(self, sources: list[Source], depth: int) -> str: + if depth <= 0: + return self._comparison(sources, 0) + choice = self._choices( + [ + ("column", 32), + ("logical", 16), + ("like", 12), + ("in_list", 8), + ("between", 6), + ("is_distinct_from", 5), + ("exists_subquery", 5), + ("in_subquery", 6), + ] + ) + if choice == "logical": + if self._chance(0.25): + return f"(NOT {self.boolean(sources, depth - 1)})" + operator = self._pick(["AND", "OR"]) + return f"({self.boolean(sources, depth - 1)} {operator} {self.boolean(sources, depth - 1)})" + if choice == "like": + operator = self._pick(["LIKE", "NOT LIKE", "ILIKE", "NOT ILIKE"]) + return f"({self.text(sources, depth - 1)} {operator} {self._pick(_TEXT_PATTERNS)})" + if choice == "in_list": + if self._chance(0.5): + values = ", ".join(str(self._pick(_NUM_LITERALS)) for _ in range(2)) + subject = self.num(sources, depth - 1) + else: + values = ", ".join(str(self._pick(_TEXT_LITERALS)) for _ in range(2)) + subject = self.text(sources, depth - 1) + negation = "NOT " if self._chance(0.4) else "" + return f"({subject} {negation}IN ({values}))" + if choice == "between": + low, high = "0", self._pick(["3", "5", "100"]) + return f"({self.num(sources, depth - 1)} BETWEEN {low} AND {high})" + if choice == "is_distinct_from": + negation = "NOT " if self._chance(0.5) else "" + return f"({self.num(sources, depth - 1)} IS {negation}DISTINCT FROM {self._leaf(NUM, sources)})" + if choice == "exists_subquery": + negation = "NOT " if self._chance(0.4) else "" + inner = self._pick([dataset.T, dataset.U]) + correlation = f" WHERE {self._correlated_predicate(sources, inner)}" if sources and self._chance(0.7) else "" + return f"({negation}EXISTS (SELECT 1 FROM {inner.name}{correlation}))" + if choice == "in_subquery": + negation = "NOT " if self._chance(0.4) else "" + inner = self._pick([dataset.T, dataset.U]) + column = self._pick(inner.of_kind(NUM)).name + return f"({self.num(sources, depth - 1)} {negation}IN (SELECT {column} FROM {inner.name}))" + return self._comparison(sources, depth - 1) + + def _correlated_predicate(self, sources: list[Source], inner: dataset.Table) -> str: + outer = self._pick([column for source in sources for column in source.columns_of(NUM)] or ["1"]) + return f"{inner.name}.{self._pick(inner.of_kind(NUM)).name} = {outer}" + + def _comparison(self, sources: list[Source], depth: int) -> str: + kind = self._pick([NUM, NUM, NUM, TEXT, TEXT, BOOL, TS] if self.on("timestamp") else [NUM, NUM, TEXT, BOOL]) + if self._chance(0.2): + negation = "NOT " if self._chance(0.5) else "" + return f"({self.expr(kind, sources, depth)} IS {negation}NULL)" + operator = self._pick(["=", "<>", "<", "<=", ">", ">="]) + return f"({self.expr(kind, sources, depth)} {operator} {self.expr(kind, sources, depth)})" + + def _scalar_subquery(self, kind: Kind) -> str: + inner = self._pick([dataset.T, dataset.U]) + function = self._pick(["MAX", "MIN", "COUNT"]) + return f"(SELECT {function}({self._pick(inner.of_kind(kind)).name}) FROM {inner.name})" + + # -- aggregates ---------------------------------------------------------- + + def aggregate(self, sources: list[Source], depth: int) -> str: + function = self._pick(["COUNT", "SUM", "AVG", "MIN", "MAX", "BOOL_AND", "BOOL_OR", "STDDEV"]) + if function in ("BOOL_AND", "BOOL_OR"): + return f"{function}({self.boolean(sources, depth)})" + if function in ("MIN", "MAX"): + return f"{function}({self.expr(self._pick([NUM, TEXT]), sources, depth)})" + return self.numeric_aggregate(sources, depth, function) + + def numeric_aggregate(self, sources: list[Source], depth: int, function: str | None = None) -> str: + function = function or self._pick(["COUNT", "SUM", "AVG", "MIN", "MAX", "STDDEV"]) + if function == "COUNT": + if self._chance(0.3): + return "COUNT(*)" + distinct = "DISTINCT " if self.on("count_distinct") and self._chance(0.3) else "" + return f"COUNT({distinct}{self.expr(self._pick([NUM, TEXT]), sources, depth)})" + if function in ("MIN", "MAX"): + return f"{function}({self.num(sources, depth)})" + return f"{function}({self.num(sources, depth)})" + + # -- queries ------------------------------------------------------------- + + def _from_clause(self) -> tuple[str, list[Source]]: + base = self._pick(dataset.TABLES) + sources = [Source(self._alias(), base)] + clause = sources[0].from_sql + + if self.on("derived_table") and self._chance(0.1): + alias = self._alias() + columns = ", ".join(f"{column.name}" for column in base.columns) + sources = [Source(alias, base, sql=f"(SELECT {columns} FROM {base.name})")] + clause = sources[0].from_sql + + if self.on("join") and self._chance(0.45): + partner = base if self.on("self_join") and self._chance(0.3) else self._pick(dataset.TABLES) + joined = Source(self._alias(), partner) + kinds = ["INNER"] + if self.on("outer_join"): + kinds += ["LEFT", "RIGHT", "FULL"] + if self.on("cross_join"): + kinds += ["CROSS"] + join_kind = self._pick(kinds) + if join_kind == "CROSS": + clause += f" CROSS JOIN {joined.from_sql}" + else: + condition = self._join_condition(sources[0], joined, join_kind) + clause += f" {join_kind} JOIN {joined.from_sql} ON {condition}" + sources.append(joined) + return clause, sources + + def _join_condition(self, left: Source, right: Source, join_kind: str) -> str: + """An equality between the two sides, sometimes narrowed by a predicate. + + Not merely realistic -- necessary. Postgres refuses a FULL JOIN whose + condition is not merge- or hash-joinable, so a generator that hands it an + arbitrary boolean throws away every FULL JOIN sample it produces, which is + the join type most worth fuzzing. An equality on two columns qualifies; + FULL gets nothing else bolted on, the rest may be narrowed. + """ + kind = self._pick([NUM, NUM, TEXT]) + pairs = [(a, b) for a in left.columns_of(kind) for b in right.columns_of(kind)] + equality = "{} = {}".format(*self._pick(pairs)) if pairs else f"{left.key} = {right.key}" + if join_kind == "FULL" or not self._chance(0.35): + return equality + return f"{equality} AND {self.boolean([left, right], 1)}" + + def _order_and_limit(self, column_count: int) -> str: + """A total order over the whole output, or nothing at all. + + Ordering by *every* output column is what makes a sequence comparison + against Postgres legitimate: two rows that tie on all of them are equal + tuples, so any tie-break gives the same sequence. Ordering by a subset + would make half the reported "wrong order" findings artifacts of an + unspecified order, which is the classic way a SQL fuzzer wastes its + owner's afternoon. + + Without --ordered there is no ORDER BY and rows are compared as a + multiset -- and then LIMIT is suppressed too, since `LIMIT` without + `ORDER BY` picks an arbitrary subset in Postgres and would report a + divergence on every run. + """ + if not self.ordered: + return "" + keys = [] + for position in range(1, column_count + 1): + direction = " DESC" if self._chance(0.35) else "" + nulls = "" + if self._chance(0.3): + nulls = f" NULLS {self._pick(['FIRST', 'LAST'])}" + keys.append(f"{position}{direction}{nulls}") + clause = " ORDER BY " + ", ".join(keys) + if self.on("limit") and self._chance(0.25): + clause += f" LIMIT {self._pick(['1', '2', '3'])}" + if self._chance(0.4): + clause += f" OFFSET {self._pick(['1', '2'])}" + return clause + + def select(self, depth: int | None = None, top_level: bool = True) -> str: + depth = self.max_depth if depth is None else depth + + if top_level and self.on("cte") and self._chance(0.07): + inner = self.select(depth=depth - 1, top_level=False) + return f"WITH c AS ({inner}) SELECT * FROM c" + self._order_and_limit(self._arity(inner)) + + if top_level and self.on("set_operation") and self._chance(0.08): + # Both branches from one list of kinds: Postgres refuses to match a + # numeric column against a text one, and a set operation whose arms + # disagree is a sample spent on the generator's own bug. + kinds = [self._pick([NUM, NUM, TEXT, BOOL]) for _ in range(self.random.randint(1, 2))] + left = self._plain_select(depth, len(kinds), with_tail=False, kinds=kinds) + right = self._plain_select(depth, len(kinds), with_tail=False, kinds=kinds) + operator = self._pick(["UNION", "UNION ALL", "INTERSECT", "EXCEPT"]) + return f"({left}) {operator} ({right})" + self._order_and_limit(len(kinds)) + + return self._plain_select(depth, self.random.randint(1, 3), with_tail=top_level) + + def _plain_select(self, depth: int, arity: int, with_tail: bool = True, kinds: list[Kind] | None = None) -> str: + clause, sources = self._from_clause() + # A set-operation branch has its column types dictated to it, so it takes + # the plain projection path and nothing else. + grouped = kinds is None and self.on("group_by") and self._chance(0.22) + aggregated = kinds is None and not grouped and self.on("aggregate") and self._chance(0.12) + distinct = not grouped and not aggregated and self.on("distinct") and self._chance(0.12) + + if grouped: + group_count = min(arity, self.random.randint(1, 2)) + group_exprs = [self.expr(self._pick([NUM, NUM, TEXT, BOOL]), sources, depth - 1) for _ in range(group_count)] + projections = list(group_exprs) + while len(projections) < arity: + projections.append(self.aggregate(sources, depth - 1)) + select_list = ", ".join(projections) + tail = " GROUP BY " + ", ".join(str(index + 1) for index in range(group_count)) + if self.on("having") and self._chance(0.3): + # A numeric aggregate specifically: BOOL_AND(x) = 0 is a type error + # in Postgres, and one that costs a whole sample every time. + comparison = f"{self._pick(['>', '<', '='])} {self._pick(_NUM_LITERALS)}" + tail += f" HAVING {self.numeric_aggregate(sources, depth - 1)} {comparison}" + elif aggregated: + select_list = ", ".join(self.aggregate(sources, depth - 1) for _ in range(arity)) + tail = "" + else: + kinds = kinds or [self._pick([NUM, NUM, NUM, TEXT, TEXT, BOOL, TS]) for _ in range(arity)] + select_list = ", ".join(self.expr(kind, sources, depth) for kind in kinds) + select_list = ("DISTINCT " if distinct else "") + select_list + tail = "" + + where = f" WHERE {self.boolean(sources, depth)}" if self._chance(0.55) else "" + query = f"SELECT {select_list} FROM {clause}{where}{tail}" + return query + self._order_and_limit(arity) if with_tail else query + + def _arity(self, inner_sql: str) -> int: + # The CTE wrapper selects *, so its width is the inner query's -- and the + # inner query was built with a known arity, but only its text came back. + # Counting top-level commas in the select list is enough here because no + # production emits a bare comma outside of parentheses. + head = inner_sql[len("SELECT ") : inner_sql.find(" FROM ")] + depth_counter, count = 0, 1 + for character in head: + if character == "(": + depth_counter += 1 + elif character == ")": + depth_counter -= 1 + elif character == "," and depth_counter == 0: + count += 1 + return count diff --git a/tools/fuzz/shrink.py b/tools/fuzz/shrink.py new file mode 100644 index 0000000..0c8d39f --- /dev/null +++ b/tools/fuzz/shrink.py @@ -0,0 +1,216 @@ +"""Reducing a failing query to the smallest one that still fails the same way. + +Without this the tool is useless. A run of five thousand queries produces +hundreds of failures which are perhaps a dozen distinct bugs, each buried under +forty lines of generated arithmetic. Shrinking turns +``SELECT (ABS(a1.n) * CASE WHEN ... END) FROM t AS a1 LEFT JOIN u ...`` into +``SELECT a1.s || 'x' FROM t AS a1`` -- which is a bug report, and small enough +that identical bugs from unrelated seeds collapse onto the same text and +de-duplicate themselves. + +The reduction moves are deliberately reckless: hoist any node to its own child, +replace any node with a literal, drop any clause. Most of what that proposes is +invalid SQL, and that is fine, because *Postgres is the validity filter* -- a +variant it refuses is simply rejected, exactly like a variant that stops +reproducing. Being allowed to propose nonsense is what keeps the move set to +thirty lines instead of a typed rewriter as large as the generator. + +A variant is accepted when it is smaller and fails the same way, where "the same +way" is the severity for a wrong answer and the normalised message for a refusal. +Requiring the *same* failure rather than merely *a* failure is what stops the +shrinker from sliding off a subtle wrong-answer bug onto some trivial unsupported +function it happens to pass through. +""" + +from __future__ import annotations + +from typing import Callable + +import sqlglot +from sqlglot import exp + +_LITERALS = ["1", "'x'", "TRUE", "NULL"] +_DROPPABLE = ["where", "having", "group", "limit", "offset", "distinct", "qualify", "with", "joins", "order"] + +# Arg positions that name something rather than compute something. Postgres will +# happily accept `SELECT btrim('a%b') TRUE` -- TRUE is a legal bare column alias -- +# so it is no use as a filter here, and a shrinker left to rewrite aliases and +# type names spends its budget renaming columns instead of deleting them. +_STRUCTURAL = {"alias", "alias_column_names", "table", "db", "catalog", "to", "kind", "unit", "format", "collate"} + +# Nodes that hold the query together rather than compute a value, and so must not +# be replaced by a literal. Postgres is the validity filter for most reckless +# moves, but not for these: replacing the FROM clause of `SELECT avg(a.n) FROM t +# AS a` with TRUE renders `SELECT AVG(a.n)TRUE`, which Postgres accepts -- as a +# FROM-less select with a column aliased "true". Valid, reproducing, and about a +# different query than the one being shrunk. Nothing is lost by refusing them: +# whole clauses are already removable by arg key through _DROPPABLE. +_SCAFFOLDING = ( + exp.Identifier, + exp.Literal, + exp.Boolean, + exp.Null, + exp.Star, + exp.DataType, + exp.From, + exp.Join, + exp.Table, + exp.TableAlias, + exp.Group, + exp.Order, + exp.Ordered, + exp.Limit, + exp.Offset, + exp.Where, + exp.Having, + exp.Distinct, + exp.With, + exp.CTE, + exp.Subquery, + exp.Lateral, +) + +Verdict = tuple[str, str] +Check = Callable[[str], Verdict | None] + + +def _variants(expression: exp.Expression, keep_order: bool) -> list[exp.Expression]: + candidates: list[exp.Expression] = [] + + droppable = [key for key in _DROPPABLE if not (keep_order and key == "order")] + for key in droppable: + if expression.args.get(key): + variant = expression.copy() + variant.set(key, None) + candidates.append(variant) + + selects = expression.args.get("expressions") or [] + if isinstance(expression, exp.Select) and len(selects) > 1: + for index in range(len(selects)): + variant = expression.copy() + kept = [item for position, item in enumerate(variant.args["expressions"]) if position != index] + variant.set("expressions", kept) + if keep_order and variant.args.get("order"): + # ORDER BY is by ordinal and must keep covering every output + # column, or the sequence comparison stops being legitimate. + variant.set("order", sqlglot.parse_one(_ordinal_order(len(kept)), dialect="postgres").args["order"]) + if variant.args.get("group"): + continue + candidates.append(variant) + + # A set operation or a CTE reduced to one of its parts. + if isinstance(expression, exp.SetOperation): + candidates.extend([expression.left.copy(), expression.right.copy()]) + + for node in list(expression.walk())[:400]: + if isinstance(node, exp.Subquery) and node.arg_key in ("this", "from") and node.this.args.get("from"): + # `FROM (SELECT ... FROM t) AS a` reduced to `FROM t AS a`. The generic + # hoist cannot do this one -- it would leave a bare SELECT in a FROM + # clause -- and a derived table that survives shrinking is the + # difference between two reports and one. + inner = node.this.args["from"].this + if isinstance(inner, exp.Table): + relation = inner.copy() + relation.set("alias", node.args.get("alias")) + candidates.append(_replaced(expression, node, relation)) + if node is expression or isinstance(node, _SCAFFOLDING) or node.arg_key in _STRUCTURAL: + continue + for child in node.iter_expressions(): + if isinstance(child, (exp.Identifier, exp.TableAlias, exp.DataType)) or child.arg_key in _STRUCTURAL: + continue + candidates.append(_replaced(expression, node, child.copy())) + for literal in _LITERALS: + candidates.append(_replaced(expression, node, sqlglot.parse_one(literal, dialect="postgres"))) + + return [candidate for candidate in candidates if candidate is not None] + + +def _rendered(tree: exp.Expression) -> str | None: + """The candidate as SQL, or None if it is not renderable or not reparseable. + + A reckless move set produces trees sqlglot's own generator refuses -- hoisting + a child into a position whose parent expects something else -- and one of those + raising ValueError halfway through a run would lose every finding after it. The + reparse is the same guard one step further out: a tree that renders to SQL + which does not read back as the same shape would have the shrinker chasing a + query nobody can reproduce. + """ + try: + sql = tree.sql(dialect="postgres") + sqlglot.parse_one(sql, dialect="postgres") + except Exception: + return None + return sql + + +def _ordinal_order(count: int) -> str: + return "SELECT 1 ORDER BY " + ", ".join(str(position) for position in range(1, count + 1)) + + +def _replaced(root: exp.Expression, node: exp.Expression, replacement: exp.Expression) -> exp.Expression | None: + path = _path_to(root, node) + if path is None: + return None + variant = root.copy() + target = variant + for key, index in path: + value = target.args[key] + target = value[index] if index is not None else value + parent, (key, index) = target.parent, path[-1] + if parent is None: + return None + if index is None: + parent.set(key, replacement) + else: + items = list(parent.args[key]) + items[index] = replacement + parent.set(key, items) + return variant + + +def _path_to(root: exp.Expression, node: exp.Expression) -> list[tuple[str, int | None]] | None: + """The arg keys leading from `root` to `node`, so the same position can be + addressed inside a fresh copy of the tree.""" + for key, value in root.args.items(): + if isinstance(value, exp.Expression): + if value is node: + return [(key, None)] + deeper = _path_to(value, node) + if deeper is not None: + return [(key, None), *deeper] + elif isinstance(value, list): + for index, item in enumerate(value): + if not isinstance(item, exp.Expression): + continue + if item is node: + return [(key, index)] + deeper = _path_to(item, node) + if deeper is not None: + return [(key, index), *deeper] + return None + + +def shrink(sql: str, verdict: Verdict, check: Check, keep_order: bool, budget: int = 600) -> tuple[str, int]: + """The smallest query found that still produces `verdict`, and the evaluations spent.""" + try: + best = sqlglot.parse_one(sql, dialect="postgres") + except Exception: + return sql, 0 + best_sql, spent = sql, 0 + + improved = True + while improved and spent < budget: + improved = False + rendered = [(tree, _rendered(tree)) for tree in _variants(best, keep_order)] + smaller = sorted( + ((tree, sql) for tree, sql in rendered if sql is not None and len(sql) < len(best_sql)), + key=lambda pair: len(pair[1]), + ) + for candidate, candidate_sql in smaller: + if spent >= budget: + break + spent += 1 + if check(candidate_sql) == verdict: + best, best_sql, improved = candidate, candidate_sql, True + break + return best_sql, spent