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 b4907425838..c82d52167b4 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -106,6 +106,7 @@ INTO_TARGETS = re.compile( re.IGNORECASE, ) LOOP_TARGET = re.compile(r"\bFOR(?:EACH)?\s+([A-Za-z_][A-Za-z0-9_]*)\s+IN\b", re.IGNORECASE) +LOOP_HEADER = re.compile(r"\bFOR(?:EACH)?\b.*?\bLOOP\b", re.IGNORECASE | re.DOTALL) 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) @@ -530,25 +531,37 @@ def scan_region( executed = executed_names(masked) for match in STATEMENT.finditer(masked): - start = offset + statement_start(match) - end = offset + match.end() - exempt = markers.exempt(start, end) + exempt = markers.exempt(offset + statement_start(match), offset + match.end()) - if hands_off_sql(match.group(), executed) and not exempt: - commands_end = match.start() + bind_values_start(match.group()) - for start, end in literals: - if match.start() <= start and end <= commands_end: - yield from scan_region(document, region[start:end], migration, markers, offset + start) + for clause, base in clauses(match.group(), match.start()): + if hands_off_sql(clause, executed) and not exempt: + commands_end = base + bind_values_start(clause) + for start, end in literals: + if base <= start and end <= commands_end: + yield from scan_region(document, region[start:end], migration, markers, offset + start) - keyword = offending_keyword(match.group()) - if keyword is None or exempt: - continue - yield Violation(migration, line_of(document, offset + keyword_start(match)), keyword) + keyword = offending_keyword(clause) + if keyword is None or exempt: + continue + yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword) for start, end in bodies: yield from scan_region(document, region[start:end], migration, markers, offset + start) +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 + is written into the same run, and reading the pair as one statement lets the header's row + source stand in as the keyword for both. That hides the statement the loop repeats, which is + the shape a row-by-row backfill takes. Splitting after each header, nested ones included, + reads the header and the body as the separate statements Postgres runs them as.""" + edges = (0, *(header.end() for header in LOOP_HEADER.finditer(statement)), len(statement)) + for opens, closes in zip(edges, edges[1:]): + if opens < closes: + yield statement[opens:closes], start + opens + + def bind_values_start(statement: str) -> int: """Where a statement stops handing commands to the server and starts listing bind values. The expressions after `USING` are values substituted into the command, never commands in @@ -565,9 +578,9 @@ def statement_start(statement: re.Match[str]) -> int: return statement.start() + len(text) - len(text.lstrip()) -def keyword_start(statement: re.Match[str]) -> int: - word = leading_keyword(statement.group()) - return statement.start() + (0 if word is None else word.start()) +def keyword_start(clause: str, base: int) -> int: + word = leading_keyword(clause) + return base + (0 if word is None else word.start()) def line_of(sql: str, offset: int) -> int: 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 21fafcd205f..970bc9f5264 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -275,6 +275,140 @@ class TestDollarQuotedBlocks: assert _scan(tmp_path, sql)[0].line == 7 +class TestLoopBodies: + def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_delete_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' DELETE FROM "Foo" WHERE "id" = r."id";\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_join_using_in_the_loop_query_does_not_hide_the_body(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT a."id" FROM "A" a JOIN "B" b USING ("id") LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_executed_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_nested_under_a_guard_inside_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' IF r."id" > 0 THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_inside_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP FOR b IN SELECT "id" FROM "B" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_supplying_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP\n' + ' FOR b IN UPDATE "Foo" SET "x" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_loop_running_only_ddl_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' CREATE INDEX "i" ON "Foo"("a");\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_loop_over_a_rewrite_returning_rows_is_flagged_once(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN UPDATE "Foo" SET "a" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_marker_on_a_loop_exempts_the_rewrite_it_repeats(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + " -- data-migration-ok: one row\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_select_for_update_lock_is_not_read_as_a_loop(self, tmp_path): + sql = 'DO $$\nBEGIN\n PERFORM 1 FROM "Foo" FOR UPDATE;\nEND $$;' + 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\';'