ci: ban row-rewriting DML from prisma migrations

tests/code_coverage_tests/check_migrations_no_data_rewrites.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+"""Ban row-rewriting DML from Prisma migrations.
+
+Migrations run synchronously at proxy boot, before the process serves traffic, so
+anything whose cost scales with existing table size turns into downtime. A single
+`UPDATE` with no batching over a spend-log-sized table is minutes of unavailability
+plus a doubled heap that plain autovacuum will not give back.
+
+Flagged, per statement, by its leading keyword:
+
+  UPDATE      rewrites every matching row, and `WHERE` does not bound the scan
+  DELETE      same scan, and the dead tuples outlive the migration
+  INSERT      only when it draws rows from a `SELECT`; `INSERT ... VALUES` is bounded
+              by the literal row list and passes
+  WITH        a CTE-led statement containing any of the above
+
+Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a
+statement's leading keyword, so they pass.
+
+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.
+
+Add a column and let the application populate it, or run the rewrite as an opt-in
+batched job outside boot. When a rewrite is genuinely bounded and must ship inside
+the migration, put `-- data-migration-ok: <reason>` on the statement, naming what
+bounds it. The reason is required.
+
+`GRANDFATHERED` freezes the violations that predate this check. Prisma records a
+checksum for every applied migration and this repo treats applied files as
+immutable, so those two cannot take an inline marker. The set is closed; a new
+migration belongs nowhere in it.
+"""
+
+from __future__ import annotations
+
+import re
+import sys
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
+
+GRANDFATHERED = frozenset(
+    {
+        "20260817000000_shadow_eval_multi_key",
+        "20260818224500_add_shadow_eval_stopped_by",
+    }
+)
+
+MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTILINE)
+DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$")
+FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
+STATEMENT = re.compile(r"[^;]+")
+
+REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"})
+
+STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset(
+    {
+        "INSERT",
+        "SELECT",
+        "WITH",
+        "ALTER",
+        "CREATE",
+        "DROP",
+        "TRUNCATE",
+        "COMMENT",
+        "GRANT",
+        "REVOKE",
+        "COPY",
+        "SET",
+        "PERFORM",
+        "RAISE",
+        "RETURN",
+        "EXECUTE",
+        "CALL",
+        "REINDEX",
+        "REFRESH",
+        "VACUUM",
+        "ANALYZE",
+    }
+)
+
+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
+it, or move the rewrite to a batched job outside boot.
+
+If the rewrite is genuinely bounded and has to ship in the migration, mark the
+statement with the bound spelled out:
+
+    -- data-migration-ok: <what bounds this>
+    UPDATE ...
+"""
+
+
+@dataclass(frozen=True, slots=True)
+class Violation:
+    migration: str
+    line: int
+    keyword: str
+
+    def render(self) -> str:
+        location = f"{MIGRATIONS_DIR.relative_to(REPO_ROOT)}/{self.migration}/migration.sql"
+        return f"{location}:{self.line}: {self.keyword} rewrites existing rows at boot"
+
+
+def blank(text: str) -> str:
+    return "".join(character if character == "\n" else " " for character in text)
+
+
+def mask(sql: str) -> tuple[str, tuple[tuple[int, int], ...]]:
+    """Blank comments and quoted text, keeping offsets, and locate dollar-quoted bodies."""
+    chunks: list[str] = []
+    bodies: list[tuple[int, int]] = []
+    index = 0
+    length = len(sql)
+
+    while index < length:
+        pair = sql[index : index + 2]
+
+        if pair == "--":
+            stop = sql.find("\n", index)
+            stop = length if stop == -1 else stop
+            chunks.append(blank(sql[index:stop]))
+            index = stop
+            continue
+
+        if pair == "/*":
+            stop = skip_block_comment(sql, index)
+            chunks.append(blank(sql[index:stop]))
+            index = stop
+            continue
+
+        character = sql[index]
+
+        if character in "'\"":
+            stop = skip_quoted(sql, index, character)
+            chunks.append(blank(sql[index:stop]))
+            index = stop
+            continue
+
+        if character == "$":
+            tag = DOLLAR_TAG.match(sql, index)
+            if tag is not None:
+                closing = sql.find(tag.group(), tag.end())
+                body_end = length if closing == -1 else closing
+                stop = length if closing == -1 else closing + len(tag.group())
+                bodies.append((tag.end(), body_end))
+                chunks.append(blank(sql[index:stop]))
+                index = stop
+                continue
+
+        chunks.append(character)
+        index += 1
+
+    return "".join(chunks), tuple(bodies)
+
+
+def skip_block_comment(sql: str, start: int) -> int:
+    depth = 1
+    index = start + 2
+    while index < len(sql) and depth > 0:
+        pair = sql[index : index + 2]
+        if pair == "/*":
+            depth += 1
+            index += 2
+        elif pair == "*/":
+            depth -= 1
+            index += 2
+        else:
+            index += 1
+    return index
+
+
+def skip_quoted(sql: str, start: int, quote: str) -> int:
+    index = start + 1
+    while index < len(sql):
+        if sql[index] != quote:
+            index += 1
+        elif sql[index + 1 : index + 2] == quote:
+            index += 2
+        else:
+            return index + 1
+    return len(sql)
+
+
+def strip_parens(statement: str) -> str:
+    """Blank parenthesised groups in place, so an `IF EXISTS (SELECT ...)` guard does not
+    stand in for the statement it guards."""
+    chunks: list[str] = []
+    depth = 0
+
+    for character in statement:
+        if character == "(":
+            depth += 1
+            chunks.append(" ")
+        elif character == ")":
+            depth = max(depth - 1, 0)
+            chunks.append(" ")
+        elif depth > 0 and character != "\n":
+            chunks.append(" ")
+        else:
+            chunks.append(character)
+
+    return "".join(chunks)
+
+
+def leading_keyword(statement: str) -> re.Match[str] | None:
+    """The statement's own keyword, looking past PL/pgSQL block syntax such as
+    `BEGIN`, `IF ... THEN` and `END`."""
+    return next(
+        (word for word in FIRST_WORD.finditer(statement) if word.group().upper() in STATEMENT_KEYWORDS),
+        None,
+    )
+
+
+def offending_keyword(statement: str) -> str | None:
+    word = leading_keyword(strip_parens(statement))
+    if word is None:
+        return None
+
+    keyword = word.group().upper()
+
+    if keyword in REWRITES_ROWS:
+        return keyword
+
+    if keyword == "INSERT":
+        return "INSERT ... SELECT" if contains(statement, "SELECT") else None
+
+    if keyword == "WITH":
+        nested = next((name for name in sorted(REWRITES_ROWS) if contains(statement, name)), None)
+        if nested is not None:
+            return f"WITH ... {nested}"
+        if contains(statement, "INSERT") and contains(statement, "SELECT"):
+            return "WITH ... INSERT ... SELECT"
+
+    return None
+
+
+def contains(statement: str, keyword: str) -> bool:
+    return re.search(rf"\b{keyword}\b", statement, re.IGNORECASE) is not None
+
+
+def exempt_lines(sql: str) -> frozenset[int]:
+    return frozenset(sql.count("\n", 0, match.start()) + 1 for match in MARKER.finditer(sql))
+
+
+def scan(sql: str, migration: str, exempt: frozenset[int], offset: int = 0) -> Iterator[Violation]:
+    masked, bodies = mask(sql)
+
+    for match in STATEMENT.finditer(masked):
+        keyword = offending_keyword(match.group())
+        if keyword is None:
+            continue
+        first = line_of(sql, offset + keyword_start(match))
+        last = line_of(sql, offset + match.end())
+        if any(line in exempt for line in range(first - 1, last + 1)):
+            continue
+        yield Violation(migration, first, keyword)
+
+    for start, end in bodies:
+        yield from scan(sql[start:end], migration, exempt, offset + start)
+
+
+def keyword_start(statement: re.Match[str]) -> int:
+    word = leading_keyword(strip_parens(statement.group()))
+    return statement.start() + (0 if word is None else word.start())
+
+
+def line_of(sql: str, offset: int) -> int:
+    return sql.count("\n", 0, offset) + 1
+
+
+def scan_migration(directory: Path) -> tuple[Violation, ...]:
+    sql = (directory / "migration.sql").read_text(encoding="utf-8")
+    return tuple(scan(sql, directory.name, exempt_lines(sql)))
+
+
+def stale_grandfathers(found: Mapping[str, tuple[Violation, ...]]) -> tuple[str, ...]:
+    clean = (name for name in GRANDFATHERED & found.keys() if not found[name])
+    missing = GRANDFATHERED - found.keys()
+    return tuple(sorted((*clean, *missing)))
+
+
+def main() -> int:
+    if not MIGRATIONS_DIR.is_dir():
+        print(f"migrations directory not found: {MIGRATIONS_DIR}", file=sys.stderr)
+        return 2
+
+    directories = tuple(sorted(path for path in MIGRATIONS_DIR.iterdir() if (path / "migration.sql").is_file()))
+    found = {directory.name: scan_migration(directory) for directory in directories}
+    violations = tuple(
+        violation for name, results in found.items() if name not in GRANDFATHERED for violation in results
+    )
+
+    for violation in violations:
+        print(violation.render())
+
+    stale = stale_grandfathers(found)
+    for name in stale:
+        print(f"{name}: listed in GRANDFATHERED but no longer violates; remove it from the set")
+
+    if violations:
+        print(GUIDANCE, file=sys.stderr)
+        print(f"{len(violations)} data-rewriting statement(s) in migrations.", file=sys.stderr)
+
+    if violations or stale:
+        return 1
+
+    print(f"No data-rewriting statements in {len(directories)} migrations.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_litellm/test_check_migrations_no_data_rewrites.py

@@ -0,0 +1,234 @@
+"""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.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+
+_CHECKER_PATH = Path(__file__).resolve().parents[1] / "code_coverage_tests" / "check_migrations_no_data_rewrites.py"
+_SPEC = importlib.util.spec_from_file_location("check_migrations_no_data_rewrites", _CHECKER_PATH)
+assert _SPEC is not None and _SPEC.loader is not None
+checker = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = checker
+_SPEC.loader.exec_module(checker)
+
+
+def _scan(tmp_path: Path, sql: str) -> tuple:
+    directory = tmp_path / "20260101000000_fixture"
+    directory.mkdir(exist_ok=True)
+    (directory / "migration.sql").write_text(sql, encoding="utf-8")
+    return checker.scan_migration(directory)
+
+
+def _keywords(tmp_path: Path, sql: str) -> tuple:
+    return tuple(violation.keyword for violation in _scan(tmp_path, sql))
+
+
+class TestRowRewritesAreFlagged:
+    def test_update_is_flagged(self, tmp_path):
+        assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1;') == ("UPDATE",)
+
+    def test_delete_is_flagged(self, tmp_path):
+        assert _keywords(tmp_path, 'DELETE FROM "Foo" WHERE "a" IS NULL;') == ("DELETE",)
+
+    def test_update_without_trailing_semicolon_is_flagged(self, tmp_path):
+        assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1') == ("UPDATE",)
+
+    def test_lowercase_update_is_flagged(self, tmp_path):
+        assert _keywords(tmp_path, 'update "Foo" set "a" = 1;') == ("UPDATE",)
+
+    def test_merge_is_flagged(self, tmp_path):
+        sql = 'MERGE INTO "Foo" t USING "Bar" s ON t."id" = s."id" WHEN MATCHED THEN UPDATE SET "a" = s."a";'
+        assert _keywords(tmp_path, sql) == ("MERGE",)
+
+    def test_every_offending_statement_is_reported(self, tmp_path):
+        sql = 'UPDATE "Foo" SET "a" = 1;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nDELETE FROM "Bar";'
+        assert _keywords(tmp_path, sql) == ("UPDATE", "DELETE")
+
+    def test_the_incident_migration_is_flagged(self, tmp_path):
+        sql = (
+            'UPDATE "LiteLLM_SpendLogs"\n'
+            '   SET "created_at" = "endTime",\n'
+            '       "updated_at" = "endTime"\n'
+            ' WHERE "created_at" > "endTime" + interval \'1 hour\';\n'
+        )
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+
+
+class TestSchemaStatementsPass:
+    def test_on_delete_cascade_is_not_a_data_rewrite(self, tmp_path):
+        sql = (
+            'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") '
+            'REFERENCES "B"("id") ON DELETE CASCADE ON UPDATE CASCADE;'
+        )
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_on_delete_set_null_is_not_a_data_rewrite(self, tmp_path):
+        sql = (
+            'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") '
+            'REFERENCES "B"("id") ON DELETE SET NULL ON UPDATE CASCADE;'
+        )
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_add_column_with_default_passes(self, tmp_path):
+        sql = 'ALTER TABLE "Foo" ADD COLUMN IF NOT EXISTS "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP;'
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_drop_table_passes(self, tmp_path):
+        assert _keywords(tmp_path, 'DROP TABLE IF EXISTS "Foo";') == ()
+
+    def test_empty_file_passes(self, tmp_path):
+        assert _keywords(tmp_path, "") == ()
+
+    def test_only_comments_passes(self, tmp_path):
+        assert _keywords(tmp_path, "-- nothing to do here\n") == ()
+
+
+class TestInsert:
+    def test_insert_values_is_bounded_and_passes(self, tmp_path):
+        assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == ()
+
+    def test_insert_select_scans_and_is_flagged(self, tmp_path):
+        assert _keywords(tmp_path, 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar";') == ("INSERT ... SELECT",)
+
+
+class TestCommonTableExpressions:
+    def test_cte_led_update_is_flagged(self, tmp_path):
+        sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) UPDATE "Foo" SET "a" = 1 FROM batch;'
+        assert _keywords(tmp_path, sql) == ("WITH ... UPDATE",)
+
+    def test_cte_led_delete_is_flagged(self, tmp_path):
+        sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) DELETE FROM "Foo" USING batch;'
+        assert _keywords(tmp_path, sql) == ("WITH ... DELETE",)
+
+    def test_cte_led_insert_select_is_flagged(self, tmp_path):
+        sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") SELECT "id" FROM batch;'
+        assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",)
+
+    def test_read_only_cte_passes(self, tmp_path):
+        sql = 'WITH batch AS (SELECT "id" FROM "Foo") SELECT count(*) FROM batch;'
+        assert _keywords(tmp_path, sql) == ()
+
+
+class TestDollarQuotedBlocks:
+    def test_update_inside_do_block_is_flagged(self, tmp_path):
+        sql = 'DO $$\nBEGIN\n    UPDATE "Foo" SET "a" = 1;\nEND $$;'
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+
+    def test_conditional_ddl_do_block_passes(self, tmp_path):
+        sql = (
+            "DO $$\n"
+            "BEGIN\n"
+            "    IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'x') THEN\n"
+            '        ALTER TABLE "Foo" DROP CONSTRAINT "x";\n'
+            "    END IF;\n"
+            "END $$;"
+        )
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_tagged_dollar_quote_is_scanned(self, tmp_path):
+        sql = 'DO $body$\nBEGIN\n    DELETE FROM "Foo";\nEND $body$;'
+        assert _keywords(tmp_path, sql) == ("DELETE",)
+
+    def test_semicolons_inside_do_block_do_not_split_outer_statements(self, tmp_path):
+        sql = 'DO $$ BEGIN PERFORM 1; END $$;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;'
+        assert _keywords(tmp_path, sql) == ()
+
+
+class TestQuotingAndComments:
+    def test_update_inside_string_literal_passes(self, tmp_path):
+        sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'UPDATE nothing';"
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_escaped_quote_inside_string_does_not_leak(self, tmp_path):
+        sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'it''s fine';\n"
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_update_inside_line_comment_passes(self, tmp_path):
+        assert _keywords(tmp_path, '-- UPDATE "Foo" SET "a" = 1;\nDROP TABLE "Bar";') == ()
+
+    def test_update_inside_block_comment_passes(self, tmp_path):
+        assert _keywords(tmp_path, '/* UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";') == ()
+
+    def test_nested_block_comment_passes(self, tmp_path):
+        sql = '/* outer /* UPDATE "Foo" SET "a" = 1; */ still comment */\nDROP TABLE "Bar";'
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_update_inside_quoted_identifier_passes(self, tmp_path):
+        assert _keywords(tmp_path, 'ALTER TABLE "UPDATE Foo" ADD COLUMN "b" TEXT;') == ()
+
+    def test_positional_parameter_is_not_a_dollar_quote(self, tmp_path):
+        sql = 'ALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nUPDATE "Foo" SET "b" = $1;'
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+
+
+class TestEscapeHatch:
+    def test_marker_with_reason_exempts_the_statement(self, tmp_path):
+        sql = '-- data-migration-ok: one row per tenant, at most a few hundred\nUPDATE "Foo" SET "a" = 1;'
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_marker_without_reason_does_not_exempt(self, tmp_path):
+        assert _keywords(tmp_path, '-- data-migration-ok:\nUPDATE "Foo" SET "a" = 1;') == ("UPDATE",)
+
+    def test_marker_exempts_only_its_own_statement(self, tmp_path):
+        sql = (
+            "-- data-migration-ok: bounded to in-flight jobs\n"
+            'UPDATE "Foo" SET "a" = 1;\n'
+            'UPDATE "Bar" SET "b" = 2;\n'
+        )
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+        assert _scan(tmp_path, sql)[0].line == 3
+
+    def test_marker_works_inside_a_do_block(self, tmp_path):
+        sql = 'DO $$\nBEGIN\n    -- data-migration-ok: single row\n    UPDATE "Foo" SET "a" = 1;\nEND $$;'
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_marker_below_the_statement_does_not_exempt_the_next_one(self, tmp_path):
+        sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;'
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+
+
+class TestReporting:
+    def test_line_number_points_at_the_statement_keyword(self, tmp_path):
+        sql = '-- CreateIndex\nCREATE INDEX "i" ON "Foo"("a");\n\nUPDATE "Foo" SET "a" = 1;'
+        assert _scan(tmp_path, sql)[0].line == 4
+
+    def test_render_names_the_migration_and_line(self, tmp_path):
+        violation = _scan(tmp_path, '\n\nDELETE FROM "Foo";')[0]
+        rendered = violation.render()
+        assert "20260101000000_fixture/migration.sql:3" in rendered
+        assert "DELETE" in rendered
+
+
+class TestGrandfathering:
+    def test_every_grandfathered_migration_still_violates(self):
+        for name in sorted(checker.GRANDFATHERED):
+            directory = checker.MIGRATIONS_DIR / name
+            assert directory.is_dir(), f"{name} no longer exists; drop it from GRANDFATHERED"
+            assert checker.scan_migration(directory), f"{name} is clean; drop it from GRANDFATHERED"
+
+    def test_stale_entry_is_reported_when_a_migration_stops_violating(self):
+        found = {name: () for name in checker.GRANDFATHERED}
+        assert checker.stale_grandfathers(found) == tuple(sorted(checker.GRANDFATHERED))
+
+    def test_missing_entry_is_reported(self):
+        assert checker.stale_grandfathers({}) == tuple(sorted(checker.GRANDFATHERED))
+
+    def test_no_stale_entries_against_the_real_tree(self):
+        found = {
+            path.name: checker.scan_migration(path)
+            for path in checker.MIGRATIONS_DIR.iterdir()
+            if (path / "migration.sql").is_file()
+        }
+        assert checker.stale_grandfathers(found) == ()
+
+
+class TestShippedMigrations:
+    def test_the_repo_is_clean(self):
+        assert checker.main() == 0