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 fa966a214b1..8539a258083 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -428,7 +428,10 @@ def row_source_keyword(statement: str) -> str | None: an insert is allowed to carry. A wrapped `VALUES` list is the row source as much as a wrapped query is, so it ends the search rather than being skipped over: reading past it reaches a `RETURNING (SELECT ...)` or a `DO UPDATE SET "a" = (SELECT ...)` written after - it and calls that scalar subquery the rows the insert copies.""" + it and calls that scalar subquery the rows the insert copies. The group is read on its + own terms before it is allowed to end the search, because a `VALUES` list joined to a + query by a set operation inside the group supplies every row the query does, and + stopping on the word `VALUES` alone would pass the whole copy.""" outer = strip_parens(statement) joined = row_source_in(outer) if joined is not None: @@ -442,11 +445,11 @@ def row_source_keyword(statement: str) -> str | None: if not groups: return row_source_in(statement) for group in groups: - if contains(strip_parens(group), "VALUES"): - return None source = row_source_keyword(group) if source is not None: return source + if contains(strip_parens(group), "VALUES"): + return None return None @@ -676,16 +679,70 @@ def outside_definition( ) -> str: """The migration's text with one routine definition blanked out and every dollar-quoted body put back. Masking blanks the bodies alike, and a `DO` block is the ordinary way a migration - runs a routine it has just defined, so a call written inside one has to stay readable. The - definition is blanked after they are restored, which takes its own body with it, so a - routine that names itself recursively does not thereby count as called.""" + runs a routine it has just defined, so a call written inside one has to stay readable. Each + body comes back with its comments blanked, since a name written in a comment is + documentation rather than a call, while its string literals stay readable because `EXECUTE` + runs one as SQL and the call can be written inside it. The definition is blanked after they + are restored, which takes its own body with it, so a routine that names itself recursively + does not thereby count as called.""" text = list(masked) for start, end in bodies: - text[start:end] = region[start:end] + text[start:end] = without_comments(region[start:end]) text[opens:closes] = blank(region[opens:closes]) return "".join(text) +def without_comments(sql: str) -> str: + """The text with its comments blanked in place and everything else kept, read with the same + lexing as `mask` so a `--` inside a string literal blanks nothing. A dollar-quoted body + nested within is read the same way on its own, which keeps a stray quote inside it from + reaching past its closing tag.""" + chunks: list[str] = [] + 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(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()) + chunks.append(sql[index : tag.end()]) + chunks.append(without_comments(sql[tag.end() : body_end])) + chunks.append(sql[body_end:stop]) + index = stop + continue + + chunks.append(character) + index += 1 + + return "".join(chunks) + + 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 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 6a55a7a7c8d..932572a0264 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -236,6 +236,18 @@ class TestInsert: sql = 'INSERT INTO "Foo" ("id") (TABLE "Bar") RETURNING (SELECT count(*) FROM "Baz");' assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + def test_a_select_term_wrapped_beside_a_values_term_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_table_term_wrapped_beside_a_values_term_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_wrapped_set_operation_of_values_lists_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL VALUES (2));' + assert _keywords(tmp_path, sql) == () + class TestCommonTableExpressions: def test_cte_led_update_is_flagged(self, tmp_path): @@ -420,6 +432,26 @@ class TestStoredRoutines: sql = self.DEFINITION + "-- backfill() is run by hand after the deploy\n" assert _keywords(tmp_path, sql) == () + def test_the_name_written_only_in_a_do_body_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN\n-- backfill() is run by hand after the deploy\nPERFORM 1;\nEND; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_do_body_block_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN /* backfill() runs later */ PERFORM 1; END; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_nested_body_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN EXECUTE $q$SELECT 1 -- backfill() runs later\n$q$; END; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_in_an_executed_literal_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN EXECUTE 'SELECT backfill()'; END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_literal_holding_comment_dashes_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN RAISE NOTICE '--'; PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_recursive_call_does_not_count_as_the_migration_calling_it(self, tmp_path): sql = ( "CREATE FUNCTION backfill(n int) RETURNS void AS $$\n"