scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix: read dynamic SQL through the statement that hands it off

Browse source 
A marker on an EXECUTE now covers the SQL that EXECUTE runs, so it goes
where the migration reads rather than inside the string. A literal whose
first line sat below its EXECUTE was missing the marker entirely, and the
documented placement failed CI.

A literal assigned with := counts as SQL only when an EXECUTE in the same
body runs that variable by name. An error message naming a DELETE the
application handles is text, and the only way to silence it before was a
marker claiming a bounded data migration that was not there at all.
This commit is contained in:
mateo-berri 2026-08-21 17:56:02 -07:00
parent c4d9a1ac6c
commit 622d9c598d
2 changed files with 107 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
tests/code_coverage_tests/check_migrations_no_data_rewrites.py
View file

@ -24,8 +24,9 @@ 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, and so is a
literal assigned to a variable with `:=`, which is where an `EXECUTE` further down
the body gets its statement from.
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.
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
@ -36,7 +37,9 @@ batched job outside boot. When a rewrite is genuinely bounded and must ship insi
the migration, put `-- data-migration-ok: <reason>` on the statement or on the line
above it, naming what bounds it. The reason is required. A marker sharing a line
with the statement it follows exempts that statement alone, so the next statement
down is still checked rather than picking the marker up as its own.
down is still checked rather than picking the marker up as its own. A marker on an
`EXECUTE` or on the assignment feeding one covers the SQL that statement hands off,
so it goes where the migration reads rather than inside the string.
`GRANDFATHERED` freezes the violations that predate this check. Prisma records a
checksum for every applied migration and this repo treats applied files as
@ -66,6 +69,7 @@ MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTIL
DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$")
FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
STATEMENT = re.compile(r"[^;]+")
RUN_BY_NAME = re.compile(r"\bEXECUTE[ \t]+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE)
REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"})
@ -273,10 +277,27 @@ def draws_rows_from_a_select(statement: str) -> bool:
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 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)
def assigned_names(statement: str) -> frozenset[str]:
"""The candidate variable names an assignment writes to, taken as every word ahead of the
`:=`. A declaration carries its type and sometimes a leading `DECLARE` alongside the name,
and none of that is worth parsing when the only question is which name is executed."""
head, separator, _ = statement.partition(":=")
if not separator:
return frozenset()
return frozenset(word.group().lower() for word in FIRST_WORD.finditer(head))
def executed_names(masked: str) -> frozenset[str]:
"""The variables handed to an `EXECUTE` by name. Reading these off the masked text keeps
an `EXECUTE` written inside a comment or a string from counting."""
return frozenset(match.group(1).lower() for match in RUN_BY_NAME.finditer(masked))
def leads_with(statement: str, keyword: str) -> bool:
@ -312,18 +333,20 @@ def scan_region(
always counted against the whole document, so a statement nested in a dollar-quoted
body reports its real file line and lines up with the markers read from that file."""
masked, bodies, literals = mask(region)
executed = executed_names(masked)
for match in STATEMENT.finditer(masked):
if hands_off_sql(match.group()):
first = line_of(document, offset + keyword_start(match))
last = line_of(document, offset + match.end())
exempt = markers.exempt(first, last)
if hands_off_sql(match.group(), executed) and not exempt:
for start, end in literals:
if match.start() <= start and end <= match.end():
yield from scan_region(document, region[start:end], migration, markers, offset + start)
keyword = offending_keyword(match.group())
if keyword is None:
continue
first = line_of(document, offset + keyword_start(match))
last = line_of(document, offset + match.end())
if markers.exempt(first, last):
if keyword is None or exempt:
continue
yield Violation(migration, first, keyword)

71
tests/test_litellm/test_check_migrations_no_data_rewrites.py
View file

@ -415,6 +415,77 @@ class TestDynamicSql:
)
assert _keywords(tmp_path, sql) == ()
def test_a_marker_exempts_an_execute_whose_sql_starts_on_a_later_line(self, tmp_path):
sql = (
"DO $$\n"
"BEGIN\n"
" -- data-migration-ok: one config row, keyed by its primary key\n"
" EXECUTE '\n"
" UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ()
def test_a_marker_exempts_an_assignment_whose_sql_starts_on_a_later_line(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE\n"
" -- data-migration-ok: one config row, keyed by its primary key\n"
" stmt text := '\n"
" UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n"
"BEGIN\n"
" EXECUTE stmt;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ()
def test_an_unmarked_execute_whose_sql_starts_on_a_later_line_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"BEGIN\n"
" EXECUTE '\n"
" UPDATE \"Foo\" SET \"a\" = 1';\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
assert _scan(tmp_path, sql)[0].line == 4
def test_a_message_assigned_but_never_executed_passes(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE\n"
" msg text := 'UPDATE of legacy rows skipped, the application backfills them';\n"
"BEGIN\n"
" RAISE NOTICE '%', msg;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ()
def test_a_notice_naming_a_delete_it_never_runs_passes(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE\n"
" note text := 'DELETE FROM legacy rows is handled by the application';\n"
"BEGIN\n"
" RAISE NOTICE '%', note;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ()
def test_only_the_variable_that_is_executed_is_read_as_sql(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE\n"
" msg text := 'UPDATE of legacy rows skipped';\n"
" stmt text := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n"
"BEGIN\n"
" RAISE NOTICE '%', msg;\n"
" EXECUTE stmt;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("DELETE",)
assert _scan(tmp_path, sql)[0].line == 4
class TestReporting:
def test_line_number_points_at_the_statement_keyword(self, tmp_path):