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 9e6df46899b..dcb81593e3e 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -41,8 +41,9 @@ hide. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads t 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 -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 +same operator, a query returning it through `INTO`, or a loop walking the query it +came out of. 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 @@ -104,7 +105,8 @@ INTO_TARGETS = re.compile( r"([A-Za-z_][A-Za-z0-9_]*(?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)", re.IGNORECASE, ) -ASSIGNS = re.compile(r":=|(?!:=])=(?!=)") +LOOP_TARGET = re.compile(r"\bFOR(?:EACH)?\s+([A-Za-z_][A-Za-z0-9_]*)\s+IN\b", re.IGNORECASE) +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) @@ -143,6 +145,8 @@ GUARDS_A_CONDITION = frozenset({"IF", "ELSIF", "ELSEIF", "CASE", "WHEN", "WHILE" OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) +NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) + 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 @@ -364,28 +368,21 @@ def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: def assigned_names(statement: str) -> frozenset[str]: - """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. Every operator in the statement is read rather than only the first, since - a comparison earlier on the line would otherwise claim the one slot and hide the - assignment after it: `IF n = 1 THEN stmt = '...'` writes `stmt` at its second `=`. 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() - - for operator in ASSIGNS.finditer(statement): - reached = reached_words(statement[: operator.start()]) - if operator.group() == ":=" or assigns_rather_than_compares(reached): - names.update(word.lower() for word in reached) + """The candidate variable names a statement writes to. An assignment is read as every + word ahead of its operator, since 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. A query assigns through the target list after its `INTO` + instead, and a loop through the variable it walks its query with, which is how a rewrite + reaches a variable with no operator appearing at all.""" + names = {word.lower() for word in assignment_reach(statement)} for targets in INTO_TARGETS.finditer(statement): if names_a_table(statement[: targets.start()]): continue names.update(word.group().lower() for word in FIRST_WORD.finditer(targets.group(1))) + names.update(loop.group(1).lower() for loop in LOOP_TARGET.finditer(statement)) + return frozenset(names) @@ -399,33 +396,61 @@ def names_a_table(before: str) -> bool: return word is not None and word.group(1).upper() == "INSERT" -def reached_words(head: str) -> tuple[str, ...]: - """The words an assignment operator is reached through, which is everything since the last - word to open a block. A `THEN` ends the condition its `IF` began, so nothing ahead of it - describes what follows, and neither the name being written nor the keywords that would - mark a comparison ever sit further back than that.""" - 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) - return tuple(words[opened + 1 :]) +def assignment_reach(statement: str) -> tuple[str, ...]: + """The words the statement's assignment is reached through, empty where it holds none. + PL/pgSQL spells the operator `:=` and takes a bare `=` as the same thing, so both count, + the second only where none of the words reached so far `marks_a_comparison`. The search + stops at the first operator that reads as an assignment, because a statement holds one + at most and everything after it is the expression being assigned, where an `=` only ever + compares: that is what keeps `ok := stmt = ''` from reading as a write to `stmt`. + What comes before can still be a comparison the assignment sits behind, as in + `IF n = 1 THEN stmt = ''`, and a word opening a block ends what it is reached + through, since nothing ahead of the `THEN` describes what follows it.""" + reached: list[str] = [] + compares = False + + for token in WORD_OR_ASSIGN.finditer(statement): + word = token.group().upper() + + if word == ":=": + return tuple(reached) + + if word == "=": + if not compares: + return tuple(reached) + continue + + if word in OPENS_A_BLOCK: + reached.clear() + compares = False + continue + + reached.append(word) + compares = compares or marks_a_comparison(word) + + return () -def assigns_rather_than_compares(reached: tuple[str, ...]) -> bool: - """Whether a bare `=` reached through these words writes a variable or tests one. They are - all that tells 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.""" - words = set(reached) - return not (words & STATEMENT_KEYWORDS) and not (words & GUARDS_A_CONDITION) +def marks_a_comparison(word: str) -> bool: + """Whether reaching a bare `=` through this word means the operator tests a variable + rather than writing one. These are all that 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.""" + return word in STATEMENT_KEYWORDS or word in GUARDS_A_CONDITION def executed_names(masked: str) -> frozenset[str]: """The variables handed to an `EXECUTE` by name. Reading these off the masked text keeps an `EXECUTE` written inside a comment or a string from counting. Masking blanks a literal - in place rather than removing it, so `EXECUTE '...'` can leave the word after it looking - like the name being run. Reaching that word means crossing no semicolon, which leaves only - the syntax `INTO`, `USING` and `END`, and no assignment ever writes one of those, so the - stray name has nothing to match.""" - return frozenset(match.group(1).lower() for match in RUN_BY_NAME.finditer(masked)) + in place rather than removing it, so `EXECUTE '...'` leaves whatever follows the literal + looking like the name being run. Only `INTO` and `USING` can sit there, since the syntax + allows nothing else between an `EXECUTE` and the semicolon ending it, and neither is ever + a variable, so both are dropped rather than left to collide with a query reaching one.""" + return frozenset( + match.group(1).lower() + for match in RUN_BY_NAME.finditer(masked) + if match.group(1).upper() not in NEVER_A_VARIABLE + ) def leads_with(statement: str, keyword: str) -> bool: 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 fe53e1ac823..b93e8422933 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -663,6 +663,140 @@ class TestDynamicSql: ) assert _keywords(tmp_path, sql) == () + def test_a_rewrite_compared_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total = 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " RAISE NOTICE 'purge script? %', ok;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_named_as_an_argument_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := probe_match(subject => stmt, wanted => 'DELETE FROM \"Foo\"');\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_compared_after_a_wider_comparison_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total >= 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_assigned_through_a_case_expression_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " stmt := CASE WHEN total = 1 THEN 'DELETE FROM \"Foo\"' ELSE 'SELECT 1' END;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_query_reaching_into_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\"'\n" + " INTO total;\n" + " SELECT (CASE WHEN total > 0 THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\"\n" + " WHERE \"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_reaching_using_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\" WHERE \"a\" = $1'\n" + " USING 'k1';\n" + " SELECT (CASE WHEN true THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\" x JOIN \"Foo\" y USING (\"a\")\n" + " WHERE x.\"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_walked_by_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOR stmt IN SELECT 'DELETE FROM \"Foo\"' LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_walked_by_a_foreach_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOREACH stmt IN ARRAY ARRAY['UPDATE \"Foo\" SET \"a\" = 1'] LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_loop_over_a_query_running_nothing_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " rec record;\n" + "BEGIN\n" + " FOR rec IN SELECT \"a\" FROM \"Foo\" LOOP\n" + " RAISE NOTICE 'the DELETE FROM \"Foo\" path is the application''s: %', rec;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): sql = ( "DO $$\n"