fix: scan sql held in a variable, and bound inserts by their own row source

tests/code_coverage_tests/check_migrations_no_data_rewrites.py

@@ -11,9 +11,10 @@ Flagged, per statement, by its leading keyword:
   UPDATE      rewrites every matching row, and `WHERE` does not bound the scan
   DELETE      same scan, and the dead tuples outlive the migration
   MERGE       both of the above in one statement
-  INSERT      only when it draws rows from a `SELECT`; `INSERT ... VALUES` is bounded
-              by the literal row list and passes, scalar subqueries in that list
-              included
+  INSERT      only when it draws rows from a `SELECT`; an insert whose row source is
+              a leading `VALUES` is bounded by the rows spelled out there and passes,
+              scalar subqueries in that list included, while a `VALUES` reached
+              through a subquery or a set operation bounds nothing
   WITH        a CTE-led statement containing any of the above
 
 Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a
@@ -22,7 +23,9 @@ statement's leading keyword, so they pass.
 Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this
 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.
+same to Postgres whether it is spelled out or handed over as a string, and so is a
+literal assigned to a variable with `:=`, which is where an `EXECUTE` further down
+the body gets its statement from.
 
 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
@@ -248,11 +251,17 @@ def offending_keyword(statement: str) -> str | None:
 
 
 def draws_rows_from_a_select(statement: str) -> bool:
-    """Whether an `INSERT` takes its rows from a query rather than a literal list. A
-    top-level `VALUES` bounds the insert to the rows written out there, so the scalar
-    subqueries and helper CTEs that sit in parentheses around it do not make it a
-    rewrite."""
-    return contains(statement, "SELECT") and not contains(strip_parens(statement), "VALUES")
+    """Whether an `INSERT` takes its rows from a query rather than a literal list. Only a
+    `SELECT` the insert is built on counts, so the scalar subqueries and helper CTEs that
+    sit in parentheses around a `VALUES` list do not make it a rewrite, while one reached
+    through a set operation does."""
+    return contains(strip_parens(statement), "SELECT")
+
+
+def hands_off_sql(statement: str) -> bool:
+    """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs
+    one outright, and an assignment parks one in a variable for an `EXECUTE` further down."""
+    return leads_with(statement, "EXECUTE") or ":=" in statement
 
 
 def leads_with(statement: str, keyword: str) -> bool:
@@ -281,11 +290,10 @@ def scan_region(
     masked, bodies, literals = mask(region)
 
     for match in STATEMENT.finditer(masked):
-        if leads_with(match.group(), "EXECUTE"):
+        if hands_off_sql(match.group()):
             for start, end in literals:
                 if match.start() <= start and end <= match.end():
                     yield from scan_region(document, region[start:end], migration, exempt, offset + start)
-            continue
         keyword = offending_keyword(match.group())
         if keyword is None:
             continue

tests/test_litellm/test_check_migrations_no_data_rewrites.py

@@ -113,6 +113,18 @@ class TestInsert:
         sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM (VALUES (1), (2)) AS "v"("id");'
         assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",)
 
+    def test_values_after_a_set_operation_does_not_bound_an_insert_select(self, tmp_path):
+        sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" UNION ALL VALUES (1);'
+        assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",)
+
+    def test_values_after_an_except_does_not_bound_an_insert_select(self, tmp_path):
+        sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" EXCEPT VALUES (1);'
+        assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",)
+
+    def test_a_select_term_after_a_values_list_is_still_flagged(self, tmp_path):
+        sql = 'INSERT INTO "Foo" ("id") VALUES (1), (2) UNION ALL SELECT "id" FROM "Bar";'
+        assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",)
+
 
 class TestCommonTableExpressions:
     def test_cte_led_update_is_flagged(self, tmp_path):
@@ -272,6 +284,11 @@ class TestEscapeHatch:
         sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;'
         assert _keywords(tmp_path, sql) == ("UPDATE",)
 
+    def test_a_marker_written_below_its_statement_leaves_that_statement_flagged(self, tmp_path):
+        sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nUPDATE "Bar" SET "b" = 2;'
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+        assert _scan(tmp_path, sql)[0].line == 1
+
     def test_marker_inside_a_do_block_below_the_first_line_exempts(self, tmp_path):
         sql = (
             "-- AlterTable\n"
@@ -334,6 +351,47 @@ class TestDynamicSql:
         sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('UPDATE \"Bar\" SET \"a\" = 1');"
         assert _keywords(tmp_path, sql) == ()
 
+    def test_a_rewrite_declared_into_a_variable_is_flagged(self, tmp_path):
+        sql = (
+            "DO $$\n"
+            "DECLARE\n"
+            "    stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n"
+            "BEGIN\n"
+            "    EXECUTE stmt;\n"
+            "END $$;"
+        )
+        assert _keywords(tmp_path, sql) == ("UPDATE",)
+        assert _scan(tmp_path, sql)[0].line == 3
+
+    def test_a_rewrite_assigned_in_the_body_is_flagged(self, tmp_path):
+        sql = (
+            "DO $$\n"
+            "DECLARE\n"
+            "    stmt text;\n"
+            "BEGIN\n"
+            "    stmt := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n"
+            "    EXECUTE stmt;\n"
+            "END $$;"
+        )
+        assert _keywords(tmp_path, sql) == ("DELETE",)
+        assert _scan(tmp_path, sql)[0].line == 5
+
+    def test_ddl_assigned_to_a_variable_passes(self, tmp_path):
+        sql = "DO $$\nDECLARE\n    stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nBEGIN\n    EXECUTE stmt;\nEND $$;"
+        assert _keywords(tmp_path, sql) == ()
+
+    def test_a_marker_exempts_a_rewrite_held_in_a_variable(self, tmp_path):
+        sql = (
+            "DO $$\n"
+            "DECLARE\n"
+            "    -- data-migration-ok: one config row, keyed by its primary key\n"
+            "    stmt text := 'UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n"
+            "BEGIN\n"
+            "    EXECUTE stmt;\n"
+            "END $$;"
+        )
+        assert _keywords(tmp_path, sql) == ()
+
 
 class TestReporting:
     def test_line_number_points_at_the_statement_keyword(self, tmp_path):