From e61baa6d0160991f7f7cd4fd1e361536cf4f80a3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:25:31 -0700 Subject: [PATCH] fix: read one quoted run as one literal, and stop before bind values --- .../check_migrations_no_data_rewrites.py | 32 ++++++++++++++--- .../test_check_migrations_no_data_rewrites.py | 34 +++++++++++++++++++ 2 files changed, 61 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 3d57691cea4..b4907425838 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -147,6 +147,8 @@ OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) +BIND_VALUES = re.compile(r"\bUSING\b", re.IGNORECASE) + GUIDANCE = """ Migrations apply at proxy boot, before it serves traffic, so a statement whose cost scales with table size is downtime. Add the column and let the application backfill @@ -286,10 +288,20 @@ def skip_block_comment(sql: str, start: int) -> int: def skip_quoted(sql: str, start: int, quote: str) -> int: - """One quoted run, up to and including its closing quote. A doubled quote needs no - special case: closing on the first and reopening on the second masks the same span.""" - stop = sql.find(quote, start + 1) - return len(sql) if stop == -1 else stop + 1 + """One quoted run, up to and including its closing quote. A doubled quote is an escaped + quote sitting inside the run rather than the end of it. Closing on the first and reopening + on the second would mask the same span, which is why this looked like it needed no special + case, but the run is also handed on whole as one literal, and splitting it there offers the + tail of a string to be read as SQL in its own right.""" + index = start + 1 + while True: + stop = sql.find(quote, index) + if stop == -1: + return len(sql) + if sql[stop + 1 : stop + 2] == quote: + index = stop + 2 + continue + return stop + 1 def strip_parens(statement: str) -> str: @@ -523,8 +535,9 @@ def scan_region( exempt = markers.exempt(start, end) if hands_off_sql(match.group(), executed) and not exempt: + commands_end = match.start() + bind_values_start(match.group()) for start, end in literals: - if match.start() <= start and end <= match.end(): + if match.start() <= start and end <= commands_end: yield from scan_region(document, region[start:end], migration, markers, offset + start) keyword = offending_keyword(match.group()) @@ -536,6 +549,15 @@ def scan_region( yield from scan_region(document, region[start:end], migration, markers, offset + start) +def bind_values_start(statement: str) -> int: + """Where a statement stops handing commands to the server and starts listing bind values. + The expressions after `USING` are values substituted into the command, never commands in + their own right, so one that merely spells out a rewrite is not running it. Read off the + masked text, so a `USING` written inside the command string is not mistaken for this one.""" + keyword = BIND_VALUES.search(statement) + return len(statement) if keyword is None else keyword.start() + + def statement_start(statement: re.Match[str]) -> int: """Where the statement's own text begins, past the whitespace and blanked comments it picked up from whatever sat between it and the statement before it, one of which can be a marker.""" 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 5413f8cdc62..21fafcd205f 100644 --- a/tests/test_litellm/test_check_migrations_no_data_rewrites.py +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -467,6 +467,40 @@ class TestDynamicSql: sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = date_trunc(''day'', \"t\")';\nEND $$;" assert _keywords(tmp_path, sql) == ("UPDATE",) + def test_a_rewrite_quoted_as_data_inside_executed_sql_is_not_run(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''UPDATE \"Foo\" SET \"a\" = 1''';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_doubled_quote_does_not_split_the_literal_it_sits_in(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('a''UPDATE \"Bar\" SET \"a\" = 1''b');" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_following_a_doubled_quote_in_the_same_payload_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''x''; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_in_a_later_command_before_bind_values_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1; DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_bind_value_naming_a_rewrite_is_not_run(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'INSERT INTO \"Audit\" (\"note\") VALUES ($1)'" + " USING 'DELETE FROM \"Foo\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_executed_with_bind_values_is_still_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_using_written_inside_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" USING \"Bar\"" + " WHERE \"Foo\".\"a\" = \"Bar\".\"a\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + def test_execute_of_ddl_passes(self, tmp_path): sql = "DO $$\nBEGIN\n EXECUTE 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nEND $$;" assert _keywords(tmp_path, sql) == ()