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 a loop body as its own statement, not as part of the header

Browse source 
A `FOR ... LOOP` header carries no semicolon of its own, so the first
statement of the loop body is written into the same semicolon-delimited
run. Reading the pair as one statement let the header's row source stand
in as the keyword for both, which hid whatever the loop repeats: a plain
`UPDATE` in a query-driven loop went unreported, and so did an `EXECUTE`
of one. That is the shape a row-by-row backfill takes, and it is the
shape this gate exists to stop.
This commit is contained in:
mateo-berri 2026-08-22 12:03:49 -07:00
parent e61baa6d01
commit 923da852fa
2 changed files with 162 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
tests/code_coverage_tests/check_migrations_no_data_rewrites.py
View file

@ -106,6 +106,7 @@ INTO_TARGETS = re.compile(
re.IGNORECASE,
)
LOOP_TARGET = re.compile(r"\bFOR(?:EACH)?\s+([A-Za-z_][A-Za-z0-9_]*)\s+IN\b", re.IGNORECASE)
LOOP_HEADER = re.compile(r"\bFOR(?:EACH)?\b.*?\bLOOP\b", re.IGNORECASE | re.DOTALL)
WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?<![<>!:=])=(?![=>])")
PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$")
EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE)
@ -530,25 +531,37 @@ def scan_region(
executed = executed_names(masked)
for match in STATEMENT.finditer(masked):
start = offset + statement_start(match)
end = offset + match.end()
exempt = markers.exempt(start, end)
exempt = markers.exempt(offset + statement_start(match), offset + match.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 <= commands_end:
yield from scan_region(document, region[start:end], migration, markers, offset + start)
for clause, base in clauses(match.group(), match.start()):
if hands_off_sql(clause, executed) and not exempt:
commands_end = base + bind_values_start(clause)
for start, end in literals:
if base <= start and end <= commands_end:
yield from scan_region(document, region[start:end], migration, markers, offset + start)
keyword = offending_keyword(match.group())
if keyword is None or exempt:
continue
yield Violation(migration, line_of(document, offset + keyword_start(match)), keyword)
keyword = offending_keyword(clause)
if keyword is None or exempt:
continue
yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword)
for start, end in bodies:
yield from scan_region(document, region[start:end], migration, markers, offset + start)
def clauses(statement: str, start: int) -> Iterator[tuple[str, int]]:
"""The statements written inside one semicolon-delimited run, each with where it begins. A
`FOR ... LOOP` header takes no semicolon of its own, so the first statement of the loop body
is written into the same run, and reading the pair as one statement lets the header's row
source stand in as the keyword for both. That hides the statement the loop repeats, which is
the shape a row-by-row backfill takes. Splitting after each header, nested ones included,
reads the header and the body as the separate statements Postgres runs them as."""
edges = (0, *(header.end() for header in LOOP_HEADER.finditer(statement)), len(statement))
for opens, closes in zip(edges, edges[1:]):
if opens < closes:
yield statement[opens:closes], start + opens
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
@ -565,9 +578,9 @@ def statement_start(statement: re.Match[str]) -> int:
return statement.start() + len(text) - len(text.lstrip())
def keyword_start(statement: re.Match[str]) -> int:
word = leading_keyword(statement.group())
return statement.start() + (0 if word is None else word.start())
def keyword_start(clause: str, base: int) -> int:
word = leading_keyword(clause)
return base + (0 if word is None else word.start())
def line_of(sql: str, offset: int) -> int:

134
tests/test_litellm/test_check_migrations_no_data_rewrites.py
View file

@ -275,6 +275,140 @@ class TestDollarQuotedBlocks:
assert _scan(tmp_path, sql)[0].line == 7
class TestLoopBodies:
def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN SELECT "id" FROM "Bar" LOOP\n'
' UPDATE "Foo" SET "a" = 1;\n'
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
assert _scan(tmp_path, sql)[0].line == 5
def test_a_delete_in_a_query_driven_loop_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN SELECT "id" FROM "Bar" LOOP\n'
' DELETE FROM "Foo" WHERE "id" = r."id";\n'
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("DELETE",)
def test_a_join_using_in_the_loop_query_does_not_hide_the_body(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN SELECT a."id" FROM "A" a JOIN "B" b USING ("id") LOOP\n'
" EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n"
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
def test_a_rewrite_executed_in_a_query_driven_loop_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN SELECT "id" FROM "Bar" LOOP\n'
" EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n"
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
def test_a_rewrite_nested_under_a_guard_inside_a_loop_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN SELECT "id" FROM "Bar" LOOP\n'
' IF r."id" > 0 THEN\n'
' UPDATE "Foo" SET "a" = 1;\n'
" END IF;\n"
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
def test_a_rewrite_inside_a_nested_loop_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE a record;\n"
"DECLARE b record;\n"
"BEGIN\n"
' FOR a IN SELECT "id" FROM "A" LOOP FOR b IN SELECT "id" FROM "B" LOOP\n'
' UPDATE "Foo" SET "a" = 1;\n'
" END LOOP; END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
def test_a_rewrite_supplying_a_nested_loop_is_flagged(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE a record;\n"
"DECLARE b record;\n"
"BEGIN\n"
' FOR a IN SELECT "id" FROM "A" LOOP\n'
' FOR b IN UPDATE "Foo" SET "x" = 1 RETURNING "id" LOOP\n'
" NULL;\n"
" END LOOP; END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
assert _scan(tmp_path, sql)[0].line == 6
def test_a_loop_running_only_ddl_passes(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN SELECT "id" FROM "Bar" LOOP\n'
' CREATE INDEX "i" ON "Foo"("a");\n'
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ()
def test_a_loop_over_a_rewrite_returning_rows_is_flagged_once(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
' FOR r IN UPDATE "Foo" SET "a" = 1 RETURNING "id" LOOP\n'
" NULL;\n"
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ("UPDATE",)
def test_a_marker_on_a_loop_exempts_the_rewrite_it_repeats(self, tmp_path):
sql = (
"DO $$\n"
"DECLARE r record;\n"
"BEGIN\n"
" -- data-migration-ok: one row\n"
' FOR r IN SELECT "id" FROM "Bar" LOOP\n'
' UPDATE "Foo" SET "a" = 1;\n'
" END LOOP;\n"
"END $$;"
)
assert _keywords(tmp_path, sql) == ()
def test_a_select_for_update_lock_is_not_read_as_a_loop(self, tmp_path):
sql = 'DO $$\nBEGIN\n PERFORM 1 FROM "Foo" FOR UPDATE;\nEND $$;'
assert _keywords(tmp_path, sql) == ()
class TestQuotingAndComments:
def test_update_inside_string_literal_passes(self, tmp_path):
sql = 'ALTER TABLE "Foo" ADD COLUMN "note" TEXT NOT NULL DEFAULT \'UPDATE nothing\';'