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"