From dd2e1cf7a8ab6cbbc934d78690e249cb5719ebec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:08:59 -0700 Subject: [PATCH] fix: read a single-quoted DO body for routine calls, not only dollar-quoted ones --- .../check_migrations_no_data_rewrites.py | 50 +++++++++++++++---- .../test_check_migrations_no_data_rewrites.py | 20 ++++++++ 2 files changed, 60 insertions(+), 10 deletions(-) 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 8539a258083..c8c8dd315a8 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -625,6 +625,7 @@ def scan_region( reports its real file line and lines up with the markers read from that file.""" masked, bodies, literals = mask(region) executed = executed_names(masked) + runnable = executed_literals(masked, literals, executed) for match in STATEMENT.finditer(masked): exempt = markers.exempt(offset + statement_start(match), offset + match.end()) @@ -642,14 +643,37 @@ def scan_region( yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword) for body in bodies: - if not runs_when_applied(masked, region, bodies, body): + if not runs_when_applied(masked, region, bodies, runnable, body): continue start, end = body yield from scan_region(document, region[start:end], migration, markers, offset + start) +def executed_literals( + masked: str, literals: tuple[tuple[int, int], ...], executed: frozenset[str] +) -> tuple[tuple[int, int], ...]: + """The single-quoted literals a region runs as SQL, where a call to a routine the same + migration defines is as real as one written in the open. `DO '...'` runs its body and + `EXECUTE` runs the string it is handed, so a definition named inside one of those is called, + while a name in a message string or any literal nothing executes stays text. These are the + spans the direct scan already recurses into, read here so a call written in one is found when + the migration is searched for the routine's name.""" + return tuple( + (start, end) + for match in STATEMENT.finditer(masked) + for clause, base in clauses(match.group(), match.start()) + if hands_off_sql(clause, executed) + for start, end in literals + if base <= start and end <= base + bind_values_start(clause) + ) + + def runs_when_applied( - masked: str, region: str, bodies: tuple[tuple[int, int], ...], body: tuple[int, int] + masked: str, + region: str, + bodies: tuple[tuple[int, int], ...], + runnable: tuple[tuple[int, int], ...], + body: tuple[int, int], ) -> bool: """Whether a dollar-quoted body runs while the migration is being applied. A `DO` block runs where it is written, and so does every other use of this quoting. A `CREATE FUNCTION` or a @@ -671,22 +695,28 @@ def runs_when_applied( named = ROUTINE_NAME.match(region, defined.end(), start) if named is None or named.group(1).startswith('"'): return True - return contains(outside_definition(masked, region, bodies, opens, end), re.escape(named.group(1))) + return contains(outside_definition(masked, region, bodies, runnable, opens, end), re.escape(named.group(1))) def outside_definition( - masked: str, region: str, bodies: tuple[tuple[int, int], ...], opens: int, closes: int + masked: str, + region: str, + bodies: tuple[tuple[int, int], ...], + runnable: tuple[tuple[int, int], ...], + opens: int, + closes: int, ) -> 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. 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` + """The migration's text with one routine definition blanked out and every runnable body put + back: the dollar-quoted bodies and the single-quoted literals `DO` and `EXECUTE` run as SQL. + Masking blanks all of them alike, and a `DO` block, dollar-quoted or single-quoted, is the + ordinary way a migration runs a routine it has just defined, so a call written inside one has + to stay readable. Each 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: + for start, end in (*bodies, *runnable): text[start:end] = without_comments(region[start:end]) text[opens:closes] = blank(region[opens:closes]) return "".join(text) 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 932572a0264..e047218de5f 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -452,6 +452,26 @@ class TestStoredRoutines: sql = self.DEFINITION + "DO $$ BEGIN RAISE NOTICE '--'; PERFORM backfill(); END; $$;\n" assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_call_from_inside_a_single_quoted_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_executed_literal_inside_a_single_quoted_do_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_variable_run_by_execute_in_a_single_quoted_do_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'DECLARE q text; BEGIN q := ''SELECT backfill()''; EXECUTE q; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_the_name_written_only_in_a_single_quoted_do_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN\n-- backfill() runs later\nPERFORM 1; END';\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_non_runnable_string_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "SELECT 'backfill() runs after the deploy';\n" + assert _keywords(tmp_path, sql) == () + 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"