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 a8975e2529e..912943265bd 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -32,9 +32,10 @@ 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 same to Postgres whether it is spelled out or handed over as a string, and so is a -literal assigned with `:=` to a variable some `EXECUTE` in the same body then runs -by name, and so is the body of a `DO` written in single quotes rather than dollar -quotes. A literal nothing runs is text, however much it reads like a statement, so +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 +same operator, or a query returning it through `INTO`. So is the body of a `DO` +written in single quotes rather than dollar quotes. A literal nothing runs is text, however much it reads like a statement, so an error message naming a `DELETE` the application handles stays a message. Each literal is read on its own, so a keyword built by concatenating fragments that @@ -91,10 +92,17 @@ 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"[^;]+") RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) +INTO_TARGETS = re.compile( + r"\bINTO[ \t]+(?:STRICT[ \t]+)?" + r"([A-Za-z_][A-Za-z0-9_]*(?:[ \t]*,[ \t]*[A-Za-z_][A-Za-z0-9_]*)*)", + re.IGNORECASE, +) EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) +JOINS_QUERIES = ("UNION", "INTERSECT", "EXCEPT") + STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( { "INSERT", @@ -122,6 +130,10 @@ STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( } ) +GUARDS_A_CONDITION = frozenset({"IF", "ELSIF", "ELSEIF", "CASE", "WHEN", "WHILE", "EXIT", "ASSERT"}) + +OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) + 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 @@ -311,17 +323,21 @@ def offending_keyword(statement: str) -> str | 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, including one a set - operation joins to a `VALUES` list. Failing that, a `VALUES` outside every parenthesis - is itself the row source, so the scalar subqueries and helper CTEs nested within that - list do not make the insert a rewrite. Failing both, the rows come from a parenthesised - query, which Postgres accepts and which reading only the unparenthesised text would let - through: `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" + does. A query outside every parenthesis is the row source outright. Failing that, a + `VALUES` outside every parenthesis is itself the row source, so the scalar subqueries + and helper CTEs nested within that list do not make the insert a rewrite, though only + while no set operation sits beside it at that same level: one that does joins the list + to a second query term, and that term is the row source however deeply it is + parenthesised. Failing both, the rows come from a parenthesised query, which Postgres + accepts and which reading only the unparenthesised text would let through: + `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: return joined - return None if contains(outer, "VALUES") else row_source_in(statement) + if contains(outer, "VALUES") and not any(contains(outer, word) for word in JOINS_QUERIES): + return None + return row_source_in(statement) def row_source_in(text: str) -> str | None: @@ -339,13 +355,40 @@ def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: def assigned_names(statement: str) -> frozenset[str]: - """The candidate variable names an assignment writes to, taken as every word ahead of the - `:=`. A declaration carries its type and sometimes a leading `DECLARE` alongside the name, - and none of that is worth parsing when the only question is which name is executed.""" - head, separator, _ = statement.partition(":=") - if not separator: - return frozenset() - return frozenset(word.group().lower() for word in FIRST_WORD.finditer(head)) + """The candidate variable names a statement writes to, taken as every word ahead of the + assignment operator. A declaration carries its type and sometimes a leading `DECLARE` + alongside the name, and none of that is worth parsing when the only question is which + name is executed. PL/pgSQL spells that operator `:=` and takes a bare `=` as the same + thing, so both count, the second only where `assigns_rather_than_compares` reads it as + an assignment. A query assigns through the target list after its `INTO` instead, which + is how a rewrite reaches a variable with neither operator appearing at all.""" + names: set[str] = set() + + head, operator, _ = statement.partition(":=") + assigns = bool(operator) + if not assigns: + head, operator, _ = statement.partition("=") + assigns = bool(operator) and assigns_rather_than_compares(head) + if assigns: + names.update(word.group().lower() for word in FIRST_WORD.finditer(head)) + + for targets in INTO_TARGETS.finditer(statement): + names.update(word.group().lower() for word in FIRST_WORD.finditer(targets.group(1))) + + return frozenset(names) + + +def assigns_rather_than_compares(head: str) -> bool: + """Whether the bare `=` this text runs up to writes a variable or tests one. Only the + words ahead of it tell the two apart: an assignment is reached with a name and perhaps a + type, while a comparison is reached either through a statement carrying its own keyword + or through a word that guards a condition. Those words stop counting once something + opens a block after them, since a `THEN` ends the condition its `IF` began and the + assignment that follows on the same line is an assignment like any other.""" + words = [word.group().upper() for word in FIRST_WORD.finditer(head)] + opened = max((index for index, word in enumerate(words) if word in OPENS_A_BLOCK), default=-1) + reached = set(words[opened + 1 :]) + return not (reached & STATEMENT_KEYWORDS) and not (reached & GUARDS_A_CONDITION) def executed_names(masked: str) -> frozenset[str]: 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 15603218a27..f70d49ddb83 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -145,6 +145,22 @@ class TestInsert: sql = 'INSERT INTO "Foo" ("id") (SELECT 1) UNION (SELECT 2);' assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + def test_a_values_list_joined_to_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_excepting_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) EXCEPT (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_joined_to_a_parenthesised_table_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_set_operation_inside_a_values_list_does_not_flag_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES ((SELECT 1 UNION SELECT 2 LIMIT 1));' + assert _keywords(tmp_path, sql) == () + def test_a_table_row_source_is_flagged(self, tmp_path): assert _keywords(tmp_path, 'INSERT INTO "Foo" TABLE "Bar";') == ("INSERT ... TABLE",) @@ -469,6 +485,92 @@ class TestDynamicSql: assert _keywords(tmp_path, sql) == ("DELETE",) assert _scan(tmp_path, sql)[0].line == 5 + def test_a_rewrite_selected_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1' INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_selected_into_a_strict_target_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_assigned_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_assigned_with_a_bare_equals_after_then_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " IF true THEN stmt = 'DELETE FROM \"Foo\"'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_declared_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text;\n" + "BEGIN\n" + " SELECT 'UPDATE of legacy rows is skipped' INTO msg;\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_comparing_an_executed_variable_does_not_flag_the_comparison(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'DELETE FROM \"Foo\"' THEN RAISE NOTICE 'never'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + def test_ddl_assigned_to_a_variable_passes(self, tmp_path): sql = "DO $$\nDECLARE\n stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nBEGIN\n EXECUTE stmt;\nEND $$;" assert _keywords(tmp_path, sql) == ()