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 8ebdc43fac2..1e67617cead 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/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 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 9c6871cf557..2aa61b50949 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/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):