mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge pull request #41460 from BerriAI/litellm_pg10_migration_default_guard
ci(migrations): flag defaulted ADD COLUMN on request-log tables
This commit is contained in:
Yassin Kortam
GitHub
commit
8cdb275ec0
2 changed files with 189 additions and 3 deletions
|
|
@ -15,6 +15,14 @@ all read the whole table and all pass. That is deliberate: a rule wide enough to
|
|||
reach them fires on most ordinary migrations, and a marker everyone adds by reflex
|
||||
stops carrying information. The outage this was written for was a backfill.
|
||||
|
||||
The one schema change banned outright is `ADD COLUMN ... DEFAULT` on a table in
|
||||
`REQUEST_LOG_TABLES`, the tables that hold a row per request. Postgres 11 stores such
|
||||
a default as metadata and touches no rows, but Postgres 10, which is supported,
|
||||
rewrites the whole heap and rebuilds every index under an `ACCESS EXCLUSIVE` lock,
|
||||
which on a spend-log-sized table is the same outage as a backfill. Every other table
|
||||
is small enough that the rewrite is not worth a rule, and a column added to a log
|
||||
table without a default is still free on every version.
|
||||
|
||||
Flagged, per statement, by its leading keyword:
|
||||
|
||||
UPDATE rewrites every matching row, and `WHERE` does not bound the scan
|
||||
|
|
@ -32,6 +40,10 @@ Flagged, per statement, by its leading keyword:
|
|||
against the part of the statement holding it, so a writable CTE
|
||||
bounded by its own `VALUES` list is not handed the query the statement
|
||||
ends with as the rows it copies
|
||||
ALTER only `ALTER TABLE` on a request-log table, and only when one of its
|
||||
actions adds a column with a `DEFAULT`. An `ALTER COLUMN ... SET
|
||||
DEFAULT` written after the column exists changes metadata alone, so it
|
||||
passes, as does an `ADD CONSTRAINT`
|
||||
|
||||
Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a
|
||||
statement's leading keyword, so they pass.
|
||||
|
|
@ -85,7 +97,7 @@ would let one written for a `DO` block silence a rewrite added to that block lat
|
|||
|
||||
`GRANDFATHERED` freezes the violations that predate this check. Prisma records a
|
||||
checksum for every applied migration and this repo treats applied files as
|
||||
immutable, so those two cannot take an inline marker. The set is closed; a new
|
||||
immutable, so those files cannot take an inline marker. The set is closed; a new
|
||||
migration belongs nowhere in it.
|
||||
"""
|
||||
|
||||
|
|
@ -102,11 +114,15 @@ MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "
|
|||
|
||||
GRANDFATHERED = frozenset(
|
||||
{
|
||||
"20250425182129_add_session_id",
|
||||
"20260817000000_shadow_eval_multi_key",
|
||||
"20260818000000_add_spend_log_timestamps",
|
||||
"20260818224500_add_shadow_eval_stopped_by",
|
||||
}
|
||||
)
|
||||
|
||||
REQUEST_LOG_TABLES = frozenset({"LiteLLM_SpendLogs", "LiteLLM_ErrorLogs"})
|
||||
|
||||
MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTILINE)
|
||||
DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$")
|
||||
FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
|
||||
|
|
@ -128,6 +144,8 @@ DEFINES_A_ROUTINE = re.compile(
|
|||
)
|
||||
QUALIFIED_NAME = r"(?:\"[^\"]*\"|[A-Za-z_][A-Za-z0-9_$]*)"
|
||||
ROUTINE_NAME = re.compile(rf"\s*(?:{QUALIFIED_NAME}\s*\.\s*)?({QUALIFIED_NAME})")
|
||||
TABLE_NAME = ROUTINE_NAME
|
||||
ALTERS_A_TABLE = re.compile(r"\bALTER\s+TABLE\b(?:\s+IF\s+EXISTS)?(?:\s+ONLY)?", re.IGNORECASE)
|
||||
OPENS_A_CALL = re.compile(r"\s*\(")
|
||||
NAMES_AN_INDEX = re.compile(r"\bCREATE\b.+\bINDEX\b", re.IGNORECASE | re.DOTALL)
|
||||
INTRODUCES_A_RELATION = frozenset({"TABLE", "INTO", "REFERENCES", "EXISTS", "COPY"})
|
||||
|
|
@ -185,6 +203,10 @@ statement with the bound spelled out:
|
|||
|
||||
-- data-migration-ok: <what bounds this>
|
||||
UPDATE ...
|
||||
|
||||
On Postgres 10 an `ADD COLUMN ... DEFAULT` on a request-log table rewrites the table
|
||||
too. Add the column nullable with no default, then set the default in a separate
|
||||
`ALTER COLUMN ... SET DEFAULT`, which never touches existing rows.
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -537,6 +559,51 @@ def row_source_in(text: str) -> str | None:
|
|||
return next((word for word in ("SELECT", "TABLE") if contains(text, word)), None)
|
||||
|
||||
|
||||
def rewrites_a_log_table(clause: str, region: str, base: int) -> str | None:
|
||||
"""The keyword to report when an `ALTER TABLE` adds a defaulted column to a request-log
|
||||
table, which Postgres 10 answers by rewriting the whole table. The table is read from the
|
||||
region rather than the masked clause, since masking blanks the quoted name in place, after
|
||||
stepping over any comment sitting between `TABLE` and the name, which masking blanked as
|
||||
well. Each action of the statement is read on its own so that a `SET DEFAULT` on one column
|
||||
does not stand in for a default on a column another action adds."""
|
||||
altered = ALTERS_A_TABLE.search(clause)
|
||||
if altered is None:
|
||||
return None
|
||||
named = TABLE_NAME.match(region, skip_comments(region, base + altered.end()))
|
||||
if named is None or named.group(1).strip('"') not in REQUEST_LOG_TABLES:
|
||||
return None
|
||||
actions = strip_parens(clause[named.end() - base :]).split(",")
|
||||
if not any(adds_a_defaulted_column(action) for action in actions):
|
||||
return None
|
||||
return f"ADD COLUMN ... DEFAULT on {named.group(1)}"
|
||||
|
||||
|
||||
def skip_comments(sql: str, start: int) -> int:
|
||||
index = start
|
||||
while index < len(sql):
|
||||
pair = sql[index : index + 2]
|
||||
if pair == "--":
|
||||
stop = sql.find("\n", index)
|
||||
index = len(sql) if stop == -1 else stop
|
||||
elif pair == "/*":
|
||||
index = skip_block_comment(sql, index)
|
||||
elif sql[index].isspace():
|
||||
index += 1
|
||||
else:
|
||||
return index
|
||||
return index
|
||||
|
||||
|
||||
def adds_a_defaulted_column(action: str) -> bool:
|
||||
"""Whether an `ALTER TABLE` action is an `ADD COLUMN` carrying a column default. A `DEFAULT`
|
||||
right after `SET` is the referential action of an inline foreign key, which fills nothing
|
||||
in, so it does not count."""
|
||||
words = tuple(word.group().upper() for word in FIRST_WORD.finditer(action))
|
||||
if words[:1] != ("ADD",) or words[1:2] == ("CONSTRAINT",):
|
||||
return False
|
||||
return any(word == "DEFAULT" and previous != "SET" for previous, word in zip(words, words[1:]))
|
||||
|
||||
|
||||
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, and so does `DO`, whose body is a string wherever it is not dollar-quoted. An
|
||||
|
|
@ -724,9 +791,12 @@ def scan_region(
|
|||
)
|
||||
|
||||
keyword = offending_keyword(clause)
|
||||
if keyword is None or exempt:
|
||||
if exempt:
|
||||
continue
|
||||
yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword)
|
||||
found = keyword or rewrites_a_log_table(clause, region, base)
|
||||
if found is None:
|
||||
continue
|
||||
yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), found)
|
||||
|
||||
for body in bodies:
|
||||
if not runs_when_applied(masked, region, bodies, runnable, identifiers, body):
|
||||
|
|
|
|||
|
|
@ -89,6 +89,122 @@ class TestSchemaStatementsPass:
|
|||
assert _keywords(tmp_path, "-- nothing to do here\n") == ()
|
||||
|
||||
|
||||
SPEND_LOGS_DEFAULT = 'ADD COLUMN ... DEFAULT on "LiteLLM_SpendLogs"'
|
||||
|
||||
|
||||
class TestDefaultedColumnsOnRequestLogTables:
|
||||
def test_the_shipped_timestamp_migration_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
'ALTER TABLE "LiteLLM_SpendLogs"\n'
|
||||
'ADD COLUMN IF NOT EXISTS "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,\n'
|
||||
'ADD COLUMN IF NOT EXISTS "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP;\n'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_a_nullable_column_with_a_default_is_flagged(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "proxy_server_request" JSONB DEFAULT \'{}\';'
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_error_logs_is_a_request_log_table(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_ErrorLogs" ADD COLUMN "status" TEXT DEFAULT \'failure\';'
|
||||
assert _keywords(tmp_path, sql) == ('ADD COLUMN ... DEFAULT on "LiteLLM_ErrorLogs"',)
|
||||
|
||||
def test_a_column_without_a_default_passes(self, tmp_path):
|
||||
assert _keywords(tmp_path, 'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "status" TEXT;') == ()
|
||||
|
||||
def test_set_default_on_an_existing_column_passes(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_SpendLogs" ALTER COLUMN "status" SET DEFAULT \'success\';'
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_adding_a_column_and_defaulting_another_passes(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" TEXT, ALTER COLUMN "b" SET DEFAULT 1;'
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_a_referential_set_default_on_the_new_column_passes(self, tmp_path):
|
||||
sql = (
|
||||
'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "team_id" TEXT '
|
||||
'REFERENCES "LiteLLM_TeamTable"("team_id") ON DELETE SET DEFAULT;'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_a_column_default_beside_a_referential_set_default_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "team_id" TEXT DEFAULT \'t\' '
|
||||
'REFERENCES "LiteLLM_TeamTable"("team_id") ON DELETE SET DEFAULT;'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_a_block_comment_before_the_table_name_is_flagged(self, tmp_path):
|
||||
sql = 'ALTER TABLE /* audit */ "LiteLLM_SpendLogs" ADD COLUMN "a" TEXT DEFAULT \'x\';'
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_a_line_comment_before_the_table_name_is_flagged(self, tmp_path):
|
||||
sql = 'ALTER TABLE IF EXISTS -- audit\n"LiteLLM_SpendLogs" ADD COLUMN "a" TEXT DEFAULT \'x\';'
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_a_defaulted_column_among_other_actions_is_flagged(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" TEXT, ADD COLUMN "b" INTEGER DEFAULT 0;'
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_a_comma_inside_the_type_does_not_split_the_action(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" NUMERIC(10, 2) DEFAULT 0;'
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_a_foreign_key_set_default_action_passes(self, tmp_path):
|
||||
sql = (
|
||||
'ALTER TABLE "LiteLLM_SpendLogs" ADD CONSTRAINT "fk" FOREIGN KEY ("team_id") '
|
||||
'REFERENCES "LiteLLM_TeamTable"("team_id") ON DELETE SET DEFAULT;'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_a_default_inside_a_check_constraint_passes(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_SpendLogs" ADD CONSTRAINT "c" CHECK ("status" IS DISTINCT FROM DEFAULT);'
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_other_tables_pass(self, tmp_path):
|
||||
sql = 'ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "a" INTEGER NOT NULL DEFAULT 0;'
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_schema_qualified_and_if_exists_forms_are_flagged(self, tmp_path):
|
||||
sql = (
|
||||
'ALTER TABLE "public"."LiteLLM_SpendLogs" ADD COLUMN "a" INTEGER DEFAULT 0;\n'
|
||||
'ALTER TABLE IF EXISTS ONLY "LiteLLM_SpendLogs" ADD COLUMN "b" INTEGER DEFAULT 0;\n'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT, SPEND_LOGS_DEFAULT)
|
||||
|
||||
def test_inside_a_do_block_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
"DO $$\nBEGIN\n"
|
||||
" IF NOT EXISTS (SELECT 1 FROM information_schema.columns\n"
|
||||
" WHERE table_name = 'LiteLLM_SpendLogs' AND column_name = 'a') THEN\n"
|
||||
' ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" INTEGER DEFAULT 0;\n'
|
||||
" END IF;\nEND $$;\n"
|
||||
)
|
||||
violations = _scan(tmp_path, sql)
|
||||
assert [(violation.line, violation.keyword) for violation in violations] == [(5, SPEND_LOGS_DEFAULT)]
|
||||
|
||||
def test_handed_to_execute_is_flagged(self, tmp_path):
|
||||
sql = 'DO $$ BEGIN EXECUTE \'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" INTEGER DEFAULT 0\'; END $$;'
|
||||
assert _keywords(tmp_path, sql) == (SPEND_LOGS_DEFAULT,)
|
||||
|
||||
def test_in_a_comment_passes(self, tmp_path):
|
||||
sql = '-- ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" INTEGER DEFAULT 0;\nSELECT 1;'
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_a_marker_exempts_it(self, tmp_path):
|
||||
sql = (
|
||||
"-- data-migration-ok: table is created empty two statements up\n"
|
||||
'ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" INTEGER DEFAULT 0;'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_the_report_names_the_table(self, tmp_path):
|
||||
sql = '\nALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "a" INTEGER DEFAULT 0;'
|
||||
rendered = _scan(tmp_path, sql)[0].render()
|
||||
assert "20260101000000_fixture/migration.sql:2" in rendered
|
||||
assert 'ADD COLUMN ... DEFAULT on "LiteLLM_SpendLogs" rewrites existing rows at boot' in rendered
|
||||
|
||||
|
||||
class TestInsert:
|
||||
def test_insert_values_is_bounded_and_passes(self, tmp_path):
|
||||
assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue