From 3cca3f540286fd45514856f43eac95432c1dbd5a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 21 Aug 2026 18:04:16 -0700
Subject: [PATCH] fix: scan a DO body written in single quotes

DO takes its body as a string literal, and dollar quoting is a convenience
rather than a requirement. A migration spelling the body in single quotes
got its rewrite through untouched, since nothing was reading that literal
as SQL. It is ordinary syntax rather than an attempt to hide anything, so
the miss was reachable by accident.

The module docstring now also records where concatenated dynamic SQL stops
being readable, which is a keyword split across fragments that do not hold
it. Every fragment is scanned, so the shapes people actually write are all
still caught.
---
 .../check_migrations_no_data_rewrites.py      | 23 +++++++++++----
 .../test_check_migrations_no_data_rewrites.py | 29 +++++++++++++++++++
 2 files changed, 47 insertions(+), 5 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 2d5bb6ec726..bf405a8aa31 100644
--- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py
+++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py
@@ -25,8 +25,17 @@ 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. 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.
+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
+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
+do not contain it (`'UPD' || 'ATE ...'`) is not caught. Every fragment is scanned,
+so a concatenation is caught wherever the keyword survives whole in one of them,
+which covers `'UPDATE ' || quote_ident(t)` and the rest of the readable shapes. The
+gap needs a keyword deliberately split down the middle, and this check is a guard
+against a rewrite reaching a boot unnoticed, not a defence against someone hiding
+one on purpose.
 
 Line numbers always count against the whole migration file, however deeply the
 statement is nested, so a reported line points at the statement and the markers
@@ -91,6 +100,7 @@ STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset(
         "RAISE",
         "RETURN",
         "EXECUTE",
+        "DO",
         "CALL",
         "REINDEX",
         "REFRESH",
@@ -279,9 +289,12 @@ def draws_rows_from_a_select(statement: str) -> bool:
 
 def hands_off_sql(statement: str, executed: frozenset[str]) -> bool:
     """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs one
-    outright. An assignment parks one in a variable, which counts only when something further
-    down runs that variable by name, since a string the migration never executes is text."""
-    return leads_with(statement, "EXECUTE") or bool(assigned_names(statement) & executed)
+    outright, and so does `DO`, whose body is a string wherever it is not dollar-quoted. An
+    assignment parks one in a variable, which counts only when something further down runs
+    that variable by name, since a string the migration never executes is text."""
+    if leads_with(statement, "EXECUTE") or leads_with(statement, "DO"):
+        return True
+    return bool(assigned_names(statement) & executed)
 
 
 def assigned_names(statement: 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 926202f1d0e..c394eb30c37 100644
--- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py
+++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py
@@ -472,6 +472,35 @@ class TestDynamicSql:
         )
         assert _keywords(tmp_path, sql) == ()
 
+    def test_a_do_body_in_single_quotes_is_scanned(self, tmp_path):
+        sql = "DO 'BEGIN UPDATE \"Foo\" SET \"a\" = 1; END';"
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+        assert _scan(tmp_path, sql)[0].line == 1
+
+    def test_a_quoted_do_body_with_a_language_clause_is_scanned(self, tmp_path):
+        sql = "DO LANGUAGE plpgsql 'BEGIN DELETE FROM \"Foo\"; END';"
+        assert _keywords(tmp_path, sql) == ("DELETE",)
+
+    def test_a_quoted_do_body_holding_only_ddl_passes(self, tmp_path):
+        sql = "DO 'BEGIN ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT; END';"
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_a_marker_exempts_a_quoted_do_body(self, tmp_path):
+        sql = (
+            "-- data-migration-ok: one config row, keyed by its primary key\n"
+            "DO 'BEGIN UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''; END';"
+        )
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_fragment(self, tmp_path):
+        sql = "DO $$\nBEGIN\n    EXECUTE 'UPDATE ' || quote_ident('Foo') || ' SET \"a\" = 1';\nEND $$;"
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+        assert _scan(tmp_path, sql)[0].line == 3
+
+    def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_later_fragment(self, tmp_path):
+        sql = "DO $$\nBEGIN\n    EXECUTE 'WITH x AS (SELECT 1) ' || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;"
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+
     def test_only_the_variable_that_is_executed_is_read_as_sql(self, tmp_path):
         sql = (
             "DO $$\n"