diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..fc0eafb3b16 --- /dev/null +++ b/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: ` 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: + 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()) diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..40b55d741bc --- /dev/null +++ b/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