From 528d358c0540970c8bdb3802f6c480b0c3f24a9e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:30:18 -0700 Subject: [PATCH] fix: leave a bounded insert, a bounded writable CTE and an uncalled routine alone A parenthesised VALUES list ended the search for an insert's row source only when no group followed it, so a RETURNING or an ON CONFLICT DO UPDATE carrying a subquery was read as the rows the insert copies. A writable CTE bounded by its own VALUES list was handed the query the statement ends with for the same reason: the WITH branch read the whole statement rather than the part holding the insert. A CREATE FUNCTION or CREATE PROCEDURE body was scanned as if it ran at boot, but defining a routine only stores it. The body is now read when the same migration names the routine somewhere else, so a migration that defines a backfill and then runs it is still caught, and one whose name needed quoting is read either way since quoting is blanked at the call sites too. main() had no test, so neither its exit codes nor the branch the CI gate reads were pinned; a mutant returning 0 on a violation passed the whole suite. Its four outcomes now have tests, along with both directions of each fix above. --- .../check_migrations_no_data_rewrites.py | 120 +++++++++-- .../test_check_migrations_no_data_rewrites.py | 199 +++++++++++++++++- 2 files changed, 299 insertions(+), 20 deletions(-) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 17bc90bd0d6..fa966a214b1 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -8,10 +8,12 @@ plus a doubled heap that plain autovacuum will not give back. What is banned is the row-rewriting DML behind that, not everything whose cost scales that way. A non-concurrent `CREATE INDEX`, an `ALTER COLUMN ... TYPE` that is -not binary coercible, and a volatile `DEFAULT` on a new column all read the whole -table and all pass. That is deliberate: a rule wide enough to reach them fires on -most ordinary migrations, and a marker everyone adds by reflex stops carrying -information. The outage this was written for was a backfill. +not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE ... AS +SELECT` or `SELECT ... INTO` filling a new table from an existing one, the rename +that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED VIEW` +all read the whole table and all pass. That is deliberate: a rule wide enough to +reach them fires on most ordinary migrations, and a marker everyone adds by reflex +stops carrying information. The outage this was written for was a backfill. Flagged, per statement, by its leading keyword: @@ -21,10 +23,15 @@ Flagged, per statement, by its leading keyword: INSERT only when its rows come from a query rather than a literal `VALUES` list. The query counts wherever it sits, since Postgres takes it parenthesised, and `TABLE t` is one as much as a `SELECT` is. An - insert bounded by a leading `VALUES` passes, scalar subqueries in that - list included, while a `VALUES` reached through a subquery or joined - to a query by a set operation bounds nothing - WITH a CTE-led statement containing any of the above + insert bounded by a `VALUES` list passes, written bare or in + parentheses, and so do the scalar subqueries in that list and the + `RETURNING` and `ON CONFLICT` clauses written after it, none of which + supply the rows. A `VALUES` reached through a subquery or joined to a + query by a set operation bounds nothing + WITH a CTE-led statement containing any of the above. An `INSERT` is read + against the part of the statement holding it, so a writable CTE + bounded by its own `VALUES` list is not handed the query the statement + ends with as the rows it copies Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a statement's leading keyword, so they pass. @@ -37,7 +44,12 @@ told not to run at boot, and a marker is a cheap answer if one ever does. Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise -hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the +hide. A `CREATE FUNCTION` or `CREATE PROCEDURE` body is the exception, because +defining a routine only stores it: that body is read when the same migration names +the routine somewhere else, which is what defining a backfill and then running it +looks like, and left alone when nothing calls it. A routine whose name needed +quoting is read either way, since quoting is blanked at the call sites too and a +call written there could never be found. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the same to Postgres whether it is spelled out or handed over as a string, and so is a literal parked in a variable some `EXECUTE` in the same body then runs by name, however it got there: an assignment with `:=`, the bare `=` PL/pgSQL takes as the @@ -110,6 +122,11 @@ LOOP_HEADER = re.compile(r"\bFOR(?:EACH)?\b.*?\bLOOP\b", re.IGNORECASE | re.DOTA WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?!:=])=(?![=>])") PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) +DEFINES_A_ROUTINE = re.compile( + r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE +) +QUALIFIED_NAME = r"(?:\"[^\"]*\"|[A-Za-z_][A-Za-z0-9_$]*)" +ROUTINE_NAME = re.compile(rf"\s*(?:{QUALIFIED_NAME}\s*\.\s*)?({QUALIFIED_NAME})") REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) @@ -152,6 +169,8 @@ NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) BIND_VALUES = re.compile(r"\bUSING\b", re.IGNORECASE) +WRITES_ROWS = re.compile(r"\bINSERT\b", re.IGNORECASE) + GUIDANCE = """ Migrations apply at proxy boot, before it serves traffic, so a statement whose cost scales with table size is downtime. Add the column and let the application backfill @@ -370,13 +389,30 @@ def offending_keyword(statement: str) -> str | None: if nested is not None: return f"WITH ... {nested}" if contains(statement, "INSERT"): - source = row_source_keyword(statement) + source = insert_row_source(statement) if source is not None: return f"WITH ... INSERT ... {source}" return None +def insert_row_source(statement: str) -> str | None: + """Which keyword supplies the rows to an `INSERT` written somewhere inside a `WITH` + statement. Only the parts that hold that insert are read, because a writable CTE sits + beside the query the statement ends with and reading the whole thing hands the insert + the outer `SELECT` as its row source: `WITH c AS (INSERT ... VALUES (1) RETURNING "x") + SELECT * FROM c` adds one literal row and copies nothing. A CTE keeps its insert in a + parenthesised group, and the statement's own insert, if it is the one writing, runs from + the keyword to the end, found in the text outside every parenthesis so a group's insert + is not counted twice.""" + inserts = [group for group in parenthesised_groups(statement) if contains(group, "INSERT")] + written = WRITES_ROWS.search(strip_parens(statement)) + if written is not None: + inserts.append(statement[written.start() :]) + sources = (row_source_keyword(insert) for insert in inserts) + return next((source for source in sources if source is not None), None) + + def row_source_keyword(statement: str) -> str | None: """Which keyword supplies an `INSERT` its rows, or `None` when a literal `VALUES` list does. A query outside every parenthesis is the row source outright. Failing that, a @@ -387,9 +423,12 @@ def row_source_keyword(statement: str) -> str | None: a rewrite. Failing all three, the rows come from a parenthesised group, which Postgres accepts and which reading only the unparenthesised text would let through: `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table. Each group at that level is - read on its own terms and the first to name a row source is the answer, since the ones - around it are the column list, the conflict target and the rest of the clauses an insert - is allowed to carry, and any of those can be the last group written.""" + read on its own terms until one of them supplies the rows, since the ones before it are + the column list and the ones after it are the conflict target and the rest of the clauses + an insert is allowed to carry. A wrapped `VALUES` list is the row source as much as a + wrapped query is, so it ends the search rather than being skipped over: reading past it + reaches a `RETURNING (SELECT ...)` or a `DO UPDATE SET "a" = (SELECT ...)` written after + it and calls that scalar subquery the rows the insert copies.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: @@ -402,8 +441,13 @@ def row_source_keyword(statement: str) -> str | None: groups = list(parenthesised_groups(statement)) if not groups: return row_source_in(statement) - sources = (row_source_keyword(group) for group in groups) - return next((source for source in sources if source is not None), None) + for group in groups: + if contains(strip_parens(group), "VALUES"): + return None + source = row_source_keyword(group) + if source is not None: + return source + return None def set_operation_terms(statement: str, outer: str) -> Iterator[str]: @@ -594,10 +638,54 @@ def scan_region( continue yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword) - for start, end in bodies: + for body in bodies: + if not runs_when_applied(masked, region, bodies, body): + continue + start, end = body yield from scan_region(document, region[start:end], migration, markers, offset + start) +def runs_when_applied( + masked: str, region: str, bodies: tuple[tuple[int, int], ...], body: tuple[int, int] +) -> bool: + """Whether a dollar-quoted body runs while the migration is being applied. A `DO` block runs + where it is written, and so does every other use of this quoting. A `CREATE FUNCTION` or a + `CREATE PROCEDURE` only stores its body, which runs when something calls the routine, so a + definition nothing calls rewrites no rows at boot and reporting it names a line that never + executes. Skipping every definition instead would let a migration define a backfill and then + run it unseen, which is the shape this check exists to catch, so the body is read whenever + the same migration names the routine anywhere outside the definition. The definition is + found in the masked text, where one written inside a comment has already been blanked, and + the name is read from the region at those same offsets, since masking blanks a quoted + identifier in place. A name that needed those quotes is blanked at its call sites too and + so can never be found there, which would read as uncalled however the migration runs it, + and the body is read rather than trusted.""" + start, end = body + opens = masked.rfind(";", 0, start) + 1 + defined = DEFINES_A_ROUTINE.search(masked, opens, start) + if defined is None: + return True + named = ROUTINE_NAME.match(region, defined.end(), start) + if named is None or named.group(1).startswith('"'): + return True + return contains(outside_definition(masked, region, bodies, opens, end), re.escape(named.group(1))) + + +def outside_definition( + masked: str, region: str, bodies: tuple[tuple[int, int], ...], opens: int, closes: int +) -> str: + """The migration's text with one routine definition blanked out and every dollar-quoted body + put back. Masking blanks the bodies alike, and a `DO` block is the ordinary way a migration + runs a routine it has just defined, so a call written inside one has to stay readable. The + definition is blanked after they are restored, which takes its own body with it, so a + routine that names itself recursively does not thereby count as called.""" + text = list(masked) + for start, end in bodies: + text[start:end] = region[start:end] + text[opens:closes] = blank(region[opens:closes]) + return "".join(text) + + def clauses(statement: str, start: int) -> Iterator[tuple[str, int]]: """The statements written inside one semicolon-delimited run, each with where it begins. A `FOR ... LOOP` header takes no semicolon of its own, so the first statement of the loop body diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py index b78b34e52c7..6a55a7a7c8d 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -1,10 +1,9 @@ """Tests for tests/code_coverage_tests/check_migrations_no_data_rewrites.py. The checker reads migration.sql as SQL rather than as text, so the cases that matter -are the ones a grep would get wrong: `ON DELETE CASCADE` in a foreign key (60-odd -occurrences in the shipped migrations), an `UPDATE` inside a string literal or a -comment, and an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for -conditional DDL. +are the ones a grep would get wrong: the referential actions in a foreign key, of which +the shipped migrations carry 60, an `UPDATE` inside a string literal or a comment, and +an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for conditional DDL. """ import importlib.util @@ -218,6 +217,25 @@ class TestInsert: def test_a_table_named_in_the_insert_target_does_not_flag_it(self, tmp_path): assert _keywords(tmp_path, 'INSERT INTO "audit table" ("id") VALUES (1);') == () + def test_a_returning_subquery_after_a_wrapped_values_list_is_not_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_a_conflict_update_after_a_wrapped_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES (1))' + ' ON CONFLICT ("id") DO UPDATE SET "id" = (SELECT max("id") FROM "Bar");' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_wrapped_values_list_of_several_rows_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1), (2)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_the_row_source_names_its_own_keyword_not_a_later_subquery(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (TABLE "Bar") RETURNING (SELECT count(*) FROM "Baz");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + class TestCommonTableExpressions: def test_cte_led_update_is_flagged(self, tmp_path): @@ -248,6 +266,32 @@ class TestCommonTableExpressions: sql = 'WITH latest AS (SELECT max("id") AS "id" FROM "Bar")\nINSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT "id"::text FROM latest));' assert _keywords(tmp_path, sql) == () + def test_a_writable_cte_bounded_by_values_passes(self, tmp_path): + sql = 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id") SELECT * FROM added;' + assert _keywords(tmp_path, sql) == () + + def test_a_writable_cte_copying_a_query_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_bounded_writable_cte_does_not_hide_a_copying_one_beside_it(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id"),' + ' copied AS (INSERT INTO "Baz" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added, copied;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_writable_cte_wrapping_its_row_source_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + class TestDollarQuotedBlocks: def test_update_inside_do_block_is_flagged(self, tmp_path): @@ -326,6 +370,94 @@ class TestDollarQuotedBlocks: assert _scan(tmp_path, sql)[0].line == 7 +class TestStoredRoutines: + DEFINITION = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + PROCEDURE = ( + "CREATE OR REPLACE PROCEDURE sweep() AS $$\n" + "BEGIN\n" + ' DELETE FROM "Foo";\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + + def test_a_function_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION) == () + + def test_a_procedure_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE) == () + + def test_a_function_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION + "SELECT backfill();\n") == ("UPDATE",) + + def test_a_procedure_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE + "CALL sweep();\n") == ("DELETE",) + + def test_a_call_written_above_the_definition_still_counts(self, tmp_path): + assert _keywords(tmp_path, "SELECT backfill();\n" + self.DEFINITION) == ("UPDATE",) + + def test_a_call_from_inside_a_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_trigger_wiring_the_function_up_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TRIGGER t AFTER INSERT ON "Foo" EXECUTE FUNCTION backfill();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_schema_qualified_definition_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION.replace("backfill()", "public.backfill()")) == () + + def test_a_schema_qualified_function_the_migration_calls_is_flagged(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", "public.backfill()") + "SELECT public.backfill();\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_the_name_written_only_in_a_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "-- backfill() is run by hand after the deploy\n" + assert _keywords(tmp_path, sql) == () + + def test_a_recursive_call_does_not_count_as_the_migration_calling_it(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill(n int) RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " PERFORM backfill(n - 1);\n" + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_routine_name_is_read_rather_than_trusted(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", '"back fill"()') + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_do_block_is_not_a_routine_definition(self, tmp_path): + sql = 'DO $$ BEGIN UPDATE "Foo" SET "a" = 1; END; $$;\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_definition_written_after_another_statement_is_still_recognised(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;\n' + self.DEFINITION + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_in_a_routine_the_migration_calls(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1; -- data-migration-ok: single config row\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + "SELECT backfill();\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_called_routine_reports_the_line_inside_its_body(self, tmp_path): + assert _scan(tmp_path, self.DEFINITION + "SELECT backfill();\n")[0].line == 3 + + class TestLoopBodies: def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path): sql = ( @@ -1309,3 +1441,62 @@ class TestGrandfathering: class TestShippedMigrations: def test_the_repo_is_clean(self): assert checker.main() == 0 + + +CLEAN = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;' +DIRTY = 'UPDATE "Foo" SET "a" = 1;' +FIXTURE = "20260101000000_fixture" + + +def _tree(monkeypatch, tmp_path: Path, sql: str, grandfathered: frozenset = frozenset()) -> None: + """Stand a migrations directory holding one fixture migration in for the repo's own. The + root moves with it, since a rendered violation names the migration relative to the root and + the two are read off the same checkout everywhere but here.""" + directory = tmp_path / "migrations" / FIXTURE + directory.mkdir(parents=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + monkeypatch.setattr(checker, "REPO_ROOT", tmp_path) + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "migrations") + monkeypatch.setattr(checker, "GRANDFATHERED", grandfathered) + + +class TestExitCode: + def test_a_clean_tree_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN) + assert checker.main() == 0 + + def test_a_violation_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY) + assert checker.main() == 1 + + def test_a_stale_grandfather_alone_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + assert checker.main() == 1 + + def test_a_grandfathered_violation_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY, frozenset({FIXTURE})) + assert checker.main() == 0 + + def test_a_missing_migrations_directory_is_an_error(self, tmp_path, monkeypatch): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + assert checker.main() == 2 + + def test_the_failure_names_the_migration_the_line_and_the_keyword( + self, tmp_path, monkeypatch, capsys + ): + _tree(monkeypatch, tmp_path, DIRTY) + checker.main() + printed = capsys.readouterr().out + assert f"migrations/{FIXTURE}/migration.sql:1" in printed + assert "UPDATE rewrites existing rows at boot" in printed + assert checker.GUIDANCE in printed + + def test_a_stale_grandfather_is_named(self, tmp_path, monkeypatch, capsys): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + checker.main() + assert f"{FIXTURE}: listed in GRANDFATHERED" in capsys.readouterr().out + + def test_a_missing_directory_is_reported_on_stderr(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + checker.main() + assert "migrations directory not found" in capsys.readouterr().err