scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 249

feat(proxy): replace the pre-migration row check with a rolled-back dry run of the pending migrations (#48961)

Browse source 
* feat(proxy): replace the pre-migration row check with a rolled-back dry run of the pending migrations

* feat(proxy): roll back each migration of the dry run on its own and baseline-check databases without a ledger

* test(proxy): cover the dry run CLI and PrismaManager wrapper without a database

* fix(proxy): refuse the no-ledger dry run for a custom migration folder like boot does, and probe the rolled-back indexes directly

* ci(codecov): stop carrying forward the roi-postgres and redis-compat flags

Their carried-forward sessions were measured against older proxy_cli.py revisions (633 and 689 line maps against 710 today) and mark docstring lines of the merged file as missed, which failed codecov/patch on this PR at 70% while every unit shard reports the lines hit

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-29 09:39:26 -07:00 • committed by GitHub
parent c45b264024
commit 4bc157e281
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 591 additions and 471 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
codecov.yaml
View file

@ -15,9 +15,10 @@ ignore:
flag_management:
default_rules:
carryforward: true
# Dead flags no CI job uploads anymore: their carried-forward sessions were
# measured against old revisions, and the stale line maps mark comment lines
# of since-edited files as missed, sinking patch coverage on unrelated PRs.
# Dead flags no CI job uploads anymore, and flags whose job is path-gated on
# PRs: their carried-forward sessions were measured against old revisions, and
# the stale line maps mark comment lines of since-edited files as missed,
# sinking patch coverage on unrelated PRs.
individual_flags:
- name: proxy-mgmt-behavior
carryforward: false
@ -99,6 +100,10 @@ flag_management:
carryforward: false
- name: proxy-db-endpoints-and-responses
carryforward: false
- name: roi-postgres
carryforward: false
- name: redis-compat
carryforward: false
component_management:
individual_components:

252
litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py
View file

@ -14,11 +14,9 @@ either, so the migration keeps today's behavior for it.
"""
import os
import re
from collections.abc import Callable, Iterator
from collections.abc import Callable
from dataclasses import dataclass
from functools import reduce
from itertools import accumulate
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
@ -43,7 +41,7 @@ MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIME
DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0
REFUSAL: Final = "Existing rows would make a pending schema change fail:"
Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"]
Kind = Literal["too_long"]
@dataclass(frozen=True, slots=True)
@ -54,9 +52,6 @@ class Check:
statement: str
source: str
index: str | None = None
referenced_table: str | None = None
referenced_columns: tuple[str, ...] = ()
expression: str | None = None
@dataclass(frozen=True, slots=True)
@ -100,33 +95,6 @@ class ColumnType:
ColumnTypeLookup = Callable[[str, str], ColumnType | None]
_IDENT: Final = r'"?(\w+)"?'
_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?'
_CREATE_INDEX: Final = re.compile(
rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?"
rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)",
re.IGNORECASE | re.DOTALL,
)
_ALTER_TABLE: Final = re.compile(
rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL
)
_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE)
_ADD_COLUMN: Final = re.compile(
rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$",
re.IGNORECASE | re.DOTALL,
)
_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?"
_ADD_FOREIGN_KEY: Final = re.compile(
rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$",
re.IGNORECASE | re.DOTALL,
)
_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL)
_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL)
_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE)
_COMMENT: Final = re.compile(r"--[^\n]*")
_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE)
_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE)
def migration_data_check_timeout_seconds() -> float:
"""Seconds one count query may run before the preflight gives up on verifying it."""
@ -145,117 +113,6 @@ def migration_data_check_timeout_seconds() -> float:
return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
def statements(migration_sql: str) -> tuple[str, ...]:
"""The statements of one migration file, comments stripped, empty ones dropped."""
parts: Final = _COMMENT.sub("", migration_sql).split(";")
return tuple(part.strip() for part in parts if part.strip())
def _split_top_level(text: str) -> tuple[str, ...]:
"""Split on the commas outside parentheses."""
depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text)
cuts: Final = tuple(
position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0
)
starts: Final = (0, *(cut + 1 for cut in cuts))
stops: Final = (*cuts, len(text))
pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True))
return tuple(piece for piece in pieces if piece)
def _plain_columns(column_list: str) -> tuple[str, ...] | None:
"""The column names of an index or constraint column list, None when one entry is an
expression the derivation cannot turn into a count."""
matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list))
if any(match is None for match in matches):
return None
return tuple(match.group(1) for match in matches if match is not None)
def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]:
unique, name, table, method, column_list = match.groups()
columns: Final = _plain_columns(column_list)
if columns is None:
return
label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"'
if method is None or method.lower() == "btree":
yield Check("too_long", table, columns, label, source, index=name)
if unique:
yield Check("duplicates", table, columns, label, source, index=name)
def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]:
if (set_not_null := _SET_NOT_NULL.match(action)) is not None:
yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source)
return
if (add_column := _ADD_COLUMN.match(action)) is not None:
column, body = add_column.groups()
if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None:
yield Check(
"any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source
)
return
if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None:
constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups()
own: Final = _plain_columns(columns)
theirs: Final = _plain_columns(referenced_columns)
if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper():
return
yield Check(
"orphans",
table,
own,
_constraint_label(table, constraint, "FOREIGN KEY"),
source,
referenced_table=referenced_table,
referenced_columns=theirs,
)
return
if (unique := _ADD_UNIQUE.match(action)) is not None:
constraint, columns = unique.groups()
unique_columns: Final = _plain_columns(columns)
if unique_columns is not None:
yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source)
return
if (check := _ADD_CHECK.match(action)) is not None:
constraint, expression, not_valid = check.groups()
if not_valid is None:
yield Check(
"check_violations",
table,
(),
_constraint_label(table, constraint, "CHECK"),
source,
expression=expression,
)
def _constraint_label(table: str, constraint: str | None, kind: str) -> str:
named: Final = f'CONSTRAINT "{constraint}" ' if constraint else ""
return f'ALTER TABLE "{table}" ADD {named}{kind}'
def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]:
"""The count queries one SQL statement needs before it can be applied, empty for a
statement that cannot fail on existing rows or that the derivation cannot read."""
if (index := _CREATE_INDEX.match(statement_text)) is not None:
return tuple(_index_checks(index, statement_text, source))
if (alter := _ALTER_TABLE.match(statement_text)) is not None:
table, actions = alter.groups()
return tuple(
check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source)
) # comprehension-ok: flattening the per-action checks of one statement
return ()
def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]:
return tuple(
check
for statement_text in statements(migration_sql)
for check in checks_for_statement(statement_text, f"migration {name}")
) # comprehension-ok: flattening the per-statement checks of one migration
def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]:
"""The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`."""
return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),)
@ -424,42 +281,7 @@ def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...],
def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
from psycopg import sql
table: Final = sql.Identifier(schema, check.table)
columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns)
match check.kind:
case "too_long":
return _entry_size_query(check, schema, types, limit_bytes)
case "duplicates":
not_null: Final = sql.SQL(" AND ").join(
sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d"
).format(t=table, nn=not_null, cols=columns)
case "nulls":
return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format(
t=table, col=sql.Identifier(check.columns[0])
)
case "any_row":
return sql.SQL("SELECT count(*) FROM {t}").format(t=table)
case "orphans":
pairs: Final = sql.SQL(" AND ").join(
sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own))
for own, theirs in zip(check.columns, check.referenced_columns, strict=True)
)
present: Final = sql.SQL(" AND ").join(
sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})"
).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs)
case "check_violations":
return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format(
t=table,
expr=sql.SQL(check.expression or "true"), # pyright: ignore[reportArgumentType] # copied from our own migration file
)
return _entry_size_query(check, schema, types, limit_bytes)
def _applies(
@ -474,15 +296,7 @@ def _applies(
return False
if check.index is not None and _index_is_valid(connection, schema, check.index):
return False
if check.kind == "any_row":
return types is None
if types is None:
return False
if check.kind == "too_long":
return any(column.can_overflow_btree for column in types)
if check.kind == "orphans":
return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table)
return True
return types is not None and any(column.can_overflow_btree for column in types)
def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
@ -515,21 +329,7 @@ def _run_check(
rows: Final = _int(values[0], 0) if values else 0
if rows == 0:
return None
match check.kind:
case "too_long":
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
case "duplicates":
return BlockedChange(check, rows, "duplicated value groups")
case "nulls":
return BlockedChange(check, rows, "NULL")
case "any_row":
return BlockedChange(
check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows"
)
case "orphans":
return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"')
case "check_violations":
return BlockedChange(check, rows, f"failing CHECK ({check.expression})")
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
def find_blockers(
@ -571,19 +371,7 @@ def find_blockers(
def _remedy(check: Check) -> str:
joined: Final = ", ".join(check.columns)
match check.kind:
case "too_long":
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
case "duplicates":
return f'Remove or merge the duplicate "{check.table}" ({joined}) values.'
case "nulls":
return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.'
case "any_row":
return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards."
case "orphans":
return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.'
case "check_violations":
return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.'
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
def _describe(finding: Finding) -> str:
@ -607,30 +395,4 @@ def _describe(finding: Finding) -> str:
def report(findings: tuple[Finding, ...]) -> str:
"""The operator-facing refusal: one block per finding, then what to do next."""
lines: Final = "\n".join(_describe(finding) for finding in findings)
return (
f"{REFUSAL}\n{lines}\n"
"No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again."
)
def pending_change_checks(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
migrations_dir: Path | None,
indexes: tuple[tuple[str, str, tuple[str, ...]], ...],
) -> tuple[Check, ...]:
"""Every check the pending migrations and the request-log indexes need. `migrations_dir`
is None when the schema is pushed rather than migrated, so only the indexes are checked."""
migration_checks: Final = (
tuple(
check
for name, migration_sql in pending_migrations(connection, schema, migrations_dir)
for check in checks_for_migration(name, migration_sql)
) # comprehension-ok: flattening the per-migration checks
if migrations_dir is not None
else ()
)
index_checks: Final = tuple(
check for table, name, columns in indexes for check in checks_for_index(table, name, columns)
) # comprehension-ok: flattening the per-index checks
return migration_checks + index_checks
return f"{REFUSAL}\n{lines}\nNo index was built. Fix the rows and the next index build picks it up."

184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py Normal file
View file

@ -0,0 +1,184 @@
"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real, each
inside its own transaction that is always rolled back, so Postgres itself says whether this
database's rows and schema take them. Each migration holds the locks the real upgrade will,
and releases them before the next one runs."""
from __future__ import annotations
import re
import time
from collections.abc import Callable
from dataclasses import dataclass
from functools import reduce
from pathlib import Path
from typing import TYPE_CHECKING, Final
from litellm_proxy_extras.migration_data_checks import pending_migrations
if TYPE_CHECKING:
import psycopg
_CONCURRENTLY: Final = re.compile(
r"\b((?:CREATE\s+(?:UNIQUE\s+)?INDEX|DROP\s+INDEX|REINDEX\s+\w+)\s+)CONCURRENTLY\s+", re.IGNORECASE
)
UNDEFINED_SQLSTATES: Final = frozenset({"42703", "42P01", "42704", "42883"})
Migration = tuple[str, str]
@dataclass(frozen=True, slots=True)
class Applied:
name: str
seconds: float
replayed: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class Rejected:
name: str
error: str
seconds: float
@dataclass(frozen=True, slots=True)
class Baseline:
"""A database with tables but no `_prisma_migrations`: the boot baselines it instead of replaying
history, so the dry run asks the same question the boot does. `refusal` is None when the boot
will baseline, otherwise its reason (a custom migration folder, or the schema diff)."""
refusal: str | None
@property
def matches(self) -> bool:
return self.refusal is None
@dataclass(frozen=True, slots=True)
class DryRun:
applied: tuple[Applied, ...]
rejected: Rejected | None
baseline: Baseline | None = None
@property
def ok(self) -> bool:
return self.rejected is None and (self.baseline is None or self.baseline.matches)
def transactional(migration_sql: str) -> str:
"""The migration as it can run inside a transaction: `CONCURRENTLY` index builds become plain ones."""
return _CONCURRENTLY.sub(r"\1", migration_sql)
def _attempt(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
chain: tuple[Migration, ...],
clock: Callable[[], float],
) -> tuple[float, psycopg.Error | None]:
import psycopg
from psycopg import sql
started: Final = clock()
try:
with connection.transaction(force_rollback=True):
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
for _, migration_sql in chain:
connection.execute(transactional(migration_sql))
except psycopg.Error as exc:
return clock() - started, exc
return clock() - started, None
def _run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
pending: tuple[Migration, ...],
position: int,
done: DryRun,
clock: Callable[[], float],
) -> DryRun:
"""One migration in its own rolled-back transaction. A migration that fails because something
an earlier pending migration creates does not exist yet is rerun with those replayed in front."""
if done.rejected is not None:
return done
name: Final = pending[position][0]
seconds, error = _attempt(connection, schema, (pending[position],), clock)
if error is None:
return DryRun(done.applied + (Applied(name, seconds),), None, done.baseline)
if error.sqlstate not in UNDEFINED_SQLSTATES or position == 0:
return DryRun(done.applied, Rejected(name, str(error).strip(), seconds), done.baseline)
earlier: Final = pending[:position]
chained_seconds, chained_error = _attempt(connection, schema, earlier + (pending[position],), clock)
if chained_error is None:
replayed: Final = tuple(earlier_name for earlier_name, _ in earlier)
return DryRun(done.applied + (Applied(name, chained_seconds, replayed),), None, done.baseline)
return DryRun(done.applied, Rejected(name, str(chained_error).strip(), chained_seconds), done.baseline)
def _has_ledger(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
from psycopg import sql
row: Final = connection.execute(
"SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(connection),)
).fetchone()
return row is not None and row[0] is not None
def _has_litellm_tables(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
row: Final = connection.execute(
"SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name LIKE 'LiteLLM\\_%%' LIMIT 1",
(schema,),
).fetchone()
return row is not None
def dry_run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
migrations_dir: Path,
baseline_check: Callable[[], str | None] | None = None,
clock: Callable[[], float] = time.monotonic,
) -> DryRun:
"""Apply every pending migration in order, each in its own transaction that is rolled back
before the next one starts. The first migration Postgres rejects ends the run with its error.
`baseline_check` (None when the boot would baseline, otherwise why it refuses) stands in for
the replay on a database that has tables but no migration history, which the boot baselines."""
if baseline_check is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema):
return DryRun((), None, Baseline(baseline_check()))
pending: Final = pending_migrations(connection, schema, migrations_dir)
return reduce(
lambda done, position: _run(connection, schema, pending, position, done, clock),
range(len(pending)),
DryRun((), None),
)
def _timing(applied: Applied) -> str:
replay: Final = f" (depends on {', '.join(applied.replayed)}, replayed in front of it)" if applied.replayed else ""
return f"\n {applied.name}: {applied.seconds:.1f} s{replay}"
def report(result: DryRun) -> str:
"""What the operator reads: the per-migration time the real upgrade will hold its locks,
or the migration Postgres rejected with Postgres's own words. Nothing was applied either way."""
if result.baseline is not None:
if result.baseline.matches:
return (
"No migration history and the schema matches this build, "
"the proxy baselines it at boot, nothing to dry-run."
)
return f"No migration history and the proxy cannot baseline it at boot: {result.baseline.refusal}"
timings: Final = "".join(_timing(applied) for applied in result.applied)
if result.rejected is None:
count: Final = len(result.applied)
if count == 0:
return "No pending migrations, this database is on this version's schema."
return (
f"{count} pending migration(s) applied one at a time in a dry run, each rolled back before the next "
f"so its locks are released, nothing changed:{timings}"
)
return (
f"Migration {result.rejected.name} fails on this database after {result.rejected.seconds:.1f} s, "
f"nothing was applied.\n{result.rejected.error}"
+ (f"\nApplied before it in the dry run:{timings}" if timings else "")
)

21
litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py
View file

@ -86,6 +86,18 @@ def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
)
def has_packaged_history(migrations_dir: Path) -> bool:
"""Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships,
the only history the boot baselines an existing database onto."""
packaged_dir: Final = Path(__file__).parent
migrations: Final = migration_files(migrations_dir)
return (
bool(migrations)
and migrations == migration_files(packaged_dir)
and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes()
)
def baseline_current_schema(
coordinator: MigrationCoordinator,
schema: str,
@ -95,14 +107,9 @@ def baseline_current_schema(
) -> None:
from psycopg import sql
packaged_dir: Final = Path(__file__).parent
migrations: Final = migration_files(migrations_dir)
if (
not migrations
or migrations != migration_files(packaged_dir)
or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
):
if not has_packaged_history(migrations_dir):
raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
migrations: Final = migration_files(migrations_dir)
coordinator.acquire_prisma_lock()
existing: Final = coordinator.connection.execute(

107
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -11,18 +11,13 @@ from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
from typing import TYPE_CHECKING, Final, Protocol
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.migration_data_checks import (
Finding,
find_blockers,
migration_data_check_timeout_seconds,
pending_change_checks,
report,
)
from litellm_proxy_extras.migration_dry_run import dry_run
from litellm_proxy_extras.migration_dry_run import report as dry_run_report
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@ -36,7 +31,6 @@ from litellm_proxy_extras.replica_identity import (
apply_replica_identity_full,
)
from litellm_proxy_extras.request_log_indexes import (
REQUEST_LOG_INDEXES,
ensure_request_log_indexes,
filter_request_log_index_diff,
)
@ -52,7 +46,7 @@ class LensCheckConnect(Protocol):
) -> "psycopg.Connection[tuple[object, ...]]": ...
def str_to_bool(value: Optional[str]) -> bool:
def str_to_bool(value: str | None) -> bool:
if value is None:
return False
return value.lower() in ("true", "1", "t", "y", "yes")
@ -169,7 +163,7 @@ def _without_sql_comments(statement: str) -> str:
).strip()
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
def _without_spend_logs_pk_clauses(statement: str) -> str | None:
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
if not prefix_match:
return statement
@ -217,7 +211,7 @@ def _max_migration_timestamp(names) -> int:
return max(_migration_timestamp(n) for n in names)
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
def _redacted_command(command: object) -> str | tuple[str, ...] | list[str]:
if isinstance(command, tuple):
return tuple(_redact_credentials(str(argument)) for argument in command)
if isinstance(command, list):
@ -404,7 +398,7 @@ class ProxyExtrasDBManager:
return tuple(row) if row is not None else ()
@staticmethod
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
def _failed_migration_logs(migration_name: str, started_at: str) -> str | None:
row: Final = ProxyExtrasDBManager._read_migration_ledger(
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
"AND finished_at IS NULL AND rolled_back_at IS NULL",
@ -721,10 +715,10 @@ class ProxyExtrasDBManager:
return row is not None
@staticmethod
def _prisma_schema_param(url: str) -> Optional[str]:
def _prisma_schema_param(url: str) -> str | None:
"""The `schema` query param Prisma uses to pick its target schema,
or None when the URL does not set one."""
from urllib.parse import urlparse, parse_qsl
from urllib.parse import parse_qsl, urlparse
return next(
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
@ -1424,7 +1418,6 @@ class ProxyExtrasDBManager:
bool: True if setup was successful, False otherwise
"""
ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env())
ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate)
migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver)
if not migrated:
return False
@ -1433,46 +1426,69 @@ class ProxyExtrasDBManager:
return True
@staticmethod
def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]:
"""What in the existing rows would make a pending migration or request-log index
fail, read with count queries only: no write, no table lock. Empty when nothing
blocks, when no database is configured, or when it cannot be reached, so the
migration step that follows reports the connection problem the way it does today."""
def _baseline_refusal(prisma_dir: Path) -> str | None:
"""Why `baseline_current_schema` would refuse this database at boot, None when it would baseline:
the migrations folder must be the packaged one, then `prisma migrate diff` must find the live
database and this build's schema.prisma identical."""
from litellm_proxy_extras.migration_lock import migration_environment
from litellm_proxy_extras.migration_recovery import has_packaged_history
if not has_packaged_history(prisma_dir):
return (
f"{prisma_dir} is a custom migration history (LITELLM_MIGRATION_DIR), it differs from the "
"migrations this build ships, and the proxy only baselines onto the packaged history."
)
try:
prisma_toolchain.run_prisma(
(
_get_prisma_command(),
"migrate",
"diff",
"--from-schema-datasource",
str(prisma_dir / "schema.prisma"),
"--to-schema-datamodel",
str(prisma_dir / "schema.prisma"),
"--exit-code",
),
timeout=prisma_toolchain.prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
except subprocess.CalledProcessError as exc:
diff: Final = str(exc.stdout or exc.stderr or exc).strip()
return f"the schema differs from this build. prisma migrate diff:\n{diff}"
except subprocess.TimeoutExpired as exc:
return f"prisma migrate diff did not finish within {exc.timeout} s"
return None
@staticmethod
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
"""`--check_migrations_against_existing_data`: the pending migrations applied for real on
DIRECT_URL or DATABASE_URL inside one transaction that is rolled back, so Postgres reports
what the upgrade would hit. (True, the per-migration timings) when they all apply,
(False, Postgres's own error) when one does not."""
import psycopg
if not use_migrate:
return (True, "The schema is pushed rather than migrated, there is nothing to dry-run.")
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return ()
return (False, "DATABASE_URL is not set.")
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
os.environ.get("DIRECT_URL") or database_url
)
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None
indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES)
prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir())
try:
with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection:
checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes)
return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds())
result: Final = dry_run(
connection,
schema,
prisma_dir / "migrations",
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir),
)
except psycopg.OperationalError as exc:
logger.warning(
"Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc))
)
return ()
@staticmethod
def migration_data_check_report(use_migrate: bool = True) -> Optional[str]:
"""The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks."""
findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate)
return report(findings) if findings else None
@staticmethod
def raise_if_rows_block_pending_changes(use_migrate: bool) -> None:
"""Refuse to run migrations that the existing rows would make fail, before any is
applied. The RuntimeError carries the report and reaches the operator through the
proxy's "Database migration cannot proceed" exit."""
refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)
if refusal is not None:
raise RuntimeError(refusal)
return (False, f"Could not reach the database: {_redact_credentials(str(exc))}")
return (result.ok, dry_run_report(result))
@staticmethod
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
@ -1724,5 +1740,4 @@ class ProxyExtrasDBManager:
time.sleep(random.randrange(5, 15))
finally:
os.chdir(original_dir)
pass
return False

12
litellm/proxy/db/prisma_client.py
View file

@ -1127,16 +1127,16 @@ class PrismaManager:
return ProxyExtrasDBManager.build_request_log_indexes()
@staticmethod
def migration_data_check_report(use_migrate: bool = True) -> str | None:
"""The read-only pre-migration row check on its own, for `--check_migrations_against_existing_data`: the
refusal the migration step would print, or None when the pending schema changes
can be applied to the rows the database has."""
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
"""`--check_migrations_against_existing_data` on its own: the pending migrations applied inside a
transaction that is rolled back, as (whether they all applied, what the operator reads)."""
from litellm_proxy_extras.utils import (
ProxyExtrasDBManager, # pyright: ignore[reportUnknownVariableType] # litellm-proxy-extras is an optional install
)
from pydantic import TypeAdapter
refusal: Final[object] = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # litellm-proxy-extras is an optional install
return refusal if isinstance(refusal, str) else None
summary: Final[object] = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=use_migrate) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # litellm-proxy-extras is an optional install
return TypeAdapter(tuple[bool, str]).validate_python(summary)
@staticmethod
def start_request_log_index_build() -> None:

18
litellm/proxy/proxy_cli.py
View file

@ -276,18 +276,19 @@ class ProxyInitializationHelpers:
@staticmethod
def run_migration_data_check(use_migrate: bool) -> int:
"""`--check_migrations_against_existing_data`: the pre-migration row check against DATABASE_URL, nothing applied.
0 when the pending schema changes fit the rows, 2 with the refusal when they do not."""
"""`--check_migrations_against_existing_data`: the pending migrations applied to DATABASE_URL inside one
transaction that is rolled back, so Postgres reports what the upgrade would hit. 0 when they all apply,
2 with Postgres's error when one does not."""
if os.getenv("DATABASE_URL") is None:
raise click.UsageError("--check_migrations_against_existing_data requires DATABASE_URL")
from litellm.proxy.db.prisma_client import PrismaManager
refusal: Final[str | None] = PrismaManager.migration_data_check_report(use_migrate=use_migrate) # pyright: ignore[reportUnknownVariableType] # litellm-proxy-extras is an optional install
if refusal is None:
print("LiteLLM Proxy: the pending migrations and indexes can be applied to the rows this database has.")
ok, summary = PrismaManager.migration_dry_run_summary(use_migrate=use_migrate)
if ok:
print(f"LiteLLM Proxy: {summary}")
return 0
print(
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {refusal}\033[0m", file=sys.stderr, flush=True
f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {summary}\033[0m", file=sys.stderr, flush=True
)
return 2
@ -1059,8 +1060,9 @@ class ProxyInitializationHelpers:
"--check_migrations_against_existing_data",
is_flag=True,
default=False,
help="Check, read-only, that the rows in DATABASE_URL can take this version's pending migrations and indexes, "
"print what would block them, then exit without applying anything. Exit code 2 when something blocks.",
help="Dry-run this version's pending migrations on DATABASE_URL: the real SQL inside one transaction that is "
"rolled back, so Postgres reports the row or schema that would make the upgrade fail, then exit. Holds the same "
"locks the upgrade will. Exit code 2 when a migration fails.",
)
@click.option(
"--keepalive_timeout",

248
tests/integration/database/migrations/test_request_log_indexes.py
View file

@ -16,7 +16,7 @@ from typing import Final
import psycopg
import pytest
from litellm_proxy_extras import migration_data_checks, request_log_indexes
from litellm_proxy_extras import migration_data_checks, migration_dry_run, request_log_indexes
from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lock
from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration
from litellm_proxy_extras.request_log_indexes import (
@ -967,33 +967,25 @@ _HASH_LEGACY_KEYS: Final = (
@requires_db
@RESOLVERS
@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",))
def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_keys_are_too_long_for_the_index(
def test_the_migration_job_applies_every_migration_and_only_the_index_build_waits_for_too_long_keys(
scratch_database: str, use_v2_resolver: bool
) -> None:
"""Rows never stop an upgrade. The migrations go through, and the one request-log index
whose keys do not fit is left for the next build with the rows named in the log."""
with psycopg.connect(scratch_database, autocommit=True) as conn:
_insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token())
_insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token())
_insert_spend_log_with_key(conn, "hashed-1", "a" * 64)
_insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000)
ledger_before: Final = _ledger(scratch_database)
with pytest.raises(RuntimeError) as refused:
ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver)
refusal: Final = str(refused.value)
assert (
'"LiteLLM_SpendLogs".api_key, startTime: 2 row(s) with an index entry over 2704 bytes, largest 3232 bytes '
"(api_key up to 3207 bytes)"
) in refusal, refusal
assert f'blocks: CREATE INDEX "{API_KEY_INDEX}"' in refusal, refusal
assert _ledger(scratch_database) == ledger_before, "a refused job must not apply a single migration"
assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is False
assert len(_ledger(scratch_database)) > len(ledger_before), "the rows must not hold back a single migration"
assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx")
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute(_HASH_LEGACY_KEYS)
assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
assert len(_ledger(scratch_database)) > len(ledger_before)
@requires_db
@ -1017,49 +1009,6 @@ def test_a_serving_proxy_index_build_that_meets_too_long_keys_reports_the_rows_i
assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx")
_PROBE_STATEMENTS: Final = (
'CREATE UNIQUE INDEX "probe_b_key" ON "preflight_probe" ("b")',
'ALTER TABLE "preflight_probe" ALTER COLUMN "b" SET NOT NULL',
'ALTER TABLE "preflight_probe" ADD COLUMN "c" INTEGER NOT NULL',
'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_r_fkey" FOREIGN KEY ("r") REFERENCES "preflight_parent"("id")',
'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_a_check" CHECK ("a" <> \'bad\')',
'CREATE INDEX "probe_a_idx" ON "preflight_probe" ("a")',
)
@requires_db
def test_each_check_kind_counts_exactly_the_rows_postgres_would_reject(scratch_database: str) -> None:
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('CREATE TABLE "preflight_parent" ("id" INTEGER PRIMARY KEY)')
conn.execute('CREATE TABLE "preflight_probe" ("a" TEXT, "b" INTEGER, "r" INTEGER)')
conn.execute('INSERT INTO "preflight_parent" VALUES (1)')
conn.execute(
'INSERT INTO "preflight_probe" VALUES (%s, 1, 1), (%s, 1, 1), (%s, NULL, 2), (%s, 2, NULL), (%s, 3, 1)',
("ok", "ok", "bad", "bad", secrets.token_hex(1400)),
)
checks: Final = tuple(
check
for statement in _PROBE_STATEMENTS
for check in migration_data_checks.checks_for_statement(statement, "probe")
)
findings: Final = migration_data_checks.find_blockers(conn, "public", checks)
for statement in _PROBE_STATEMENTS:
with conn.transaction(force_rollback=True), pytest.raises(psycopg.Error):
conn.execute(statement) # pyright: ignore[reportArgumentType] # test SQL
counted: Final = {
(f.check.kind, f.check.columns): f.rows for f in findings if isinstance(f, migration_data_checks.BlockedChange)
}
assert counted == {
("duplicates", ("b",)): 1,
("nulls", ("b",)): 1,
("any_row", ("c",)): 5,
("orphans", ("r",)): 1,
("check_violations", ()): 2,
("too_long", ("a",)): 1,
}
@requires_db
@pytest.mark.parametrize(
("index_columns", "s_value", "t_value", "longest_fitting_key", "described_as"),
@ -1098,7 +1047,7 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does(
),
)
findings: Final = migration_data_checks.find_blockers(
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
conn, "public", migration_data_checks.checks_for_index("entry_probe", "entry_probe_idx", index_columns)
)
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"):
conn.execute(create_index)
@ -1122,7 +1071,7 @@ def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random
conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)')
conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500)))
findings: Final = migration_data_checks.find_blockers(
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
conn, "public", migration_data_checks.checks_for_index("compress_probe", "compress_probe_idx", ("k",))
)
with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"):
conn.execute(create_index)
@ -1154,7 +1103,6 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent
check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it."""
noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40))
keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500))
create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")')
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)')
conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),))
@ -1162,7 +1110,7 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent
'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704'
).fetchone()
findings: Final = migration_data_checks.find_blockers(
conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe")
conn, "public", migration_data_checks.checks_for_index("stored_probe", "stored_probe_idx", ("k",))
)
rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key))
@ -1176,14 +1124,176 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent
@requires_db
def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None:
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('CREATE TABLE "preflight_slow" ("a" INTEGER)')
conn.execute('INSERT INTO "preflight_slow" SELECT generate_series(1, 50)')
(check,) = migration_data_checks.checks_for_statement(
'ALTER TABLE "preflight_slow" ADD CONSTRAINT "slow" CHECK (pg_sleep(0.1) IS NULL)', "probe"
conn.execute('CREATE TABLE "preflight_slow" ("k" TEXT)')
conn.execute(
"INSERT INTO \"preflight_slow\" SELECT md5(i::text) || repeat('a', 1500) FROM generate_series(1, 200000) AS i"
)
findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.5)
(check,) = migration_data_checks.checks_for_index("preflight_slow", "preflight_slow_idx", ("k",))
findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.01)
timeout_afterwards: Final = conn.execute("SHOW statement_timeout").fetchone()
assert findings == (migration_data_checks.CouldNotVerify(check, 0.5),)
assert timeout_afterwards == ("0",), "the preflight must hand the session back with the timeout it found"
assert "could not be checked within 0.5 s" in migration_data_checks.report(findings)
assert findings == (migration_data_checks.CouldNotVerify(check, 0.01),)
assert timeout_afterwards == ("0",), "the check must hand the session back with the timeout it found"
assert "could not be checked within 0.01 s" in migration_data_checks.report(findings)
CURRENT_RELEASE: Final = "99999999999999"
REQUIRE_TEAM_ALIAS: Final = "20991231000000_require_team_alias"
def _extras_with_scratch_migration(tmp_path: Path, name: str, migration_sql: str) -> Path:
"""A copy of the shipped package dir with one more migration folder, the shape `LITELLM_MIGRATION_DIR` points at."""
extras: Final = tmp_path / "extras"
shutil.copytree(PACKAGE / "migrations", extras / "migrations")
shutil.copy(PACKAGE / "schema.prisma", extras / "schema.prisma")
(extras / "migrations" / name).mkdir()
(extras / "migrations" / name / "migration.sql").write_text(migration_sql)
return extras
def _team_alias_is_nullable(database_url: str) -> bool:
with psycopg.connect(database_url) as conn:
row: Final = conn.execute(
"SELECT is_nullable FROM information_schema.columns "
"WHERE table_name = 'LiteLLM_TeamTable' AND column_name = 'team_alias'"
).fetchone()
return row == ("YES",)
@requires_db
@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",))
def test_the_dry_run_applies_every_pending_migration_in_order_and_rolls_all_of_it_back(scratch_database: str) -> None:
ledger_before: Final = _ledger(scratch_database)
with psycopg.connect(scratch_database, autocommit=True) as conn:
pending: Final = migration_data_checks.pending_migrations(conn, "public", PACKAGE / "migrations")
result: Final = migration_dry_run.dry_run(conn, "public", PACKAGE / "migrations")
assert result.rejected is None, result
assert tuple(applied.name for applied in result.applied) == tuple(name for name, _ in pending)
assert len(result.applied) > 1
assert _ledger(scratch_database) == ledger_before, "a dry run must leave the ledger alone"
assert _index_validity(scratch_database, "api_key_startTime_idx") == {}, "a dry run must leave no index behind"
names: Final = tuple(name for name, _ in pending)
assert all(set(applied.replayed) <= set(names[: names.index(applied.name)]) for applied in result.applied), (
"a migration is only ever replayed behind a later one that depends on it"
)
assert "applied one at a time in a dry run" in migration_dry_run.report(result)
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_the_dry_run_reports_postgres_own_error_for_the_row_a_migration_cannot_take(
scratch_database: str, tmp_path: Path
) -> None:
extras: Final = _extras_with_scratch_migration(
tmp_path, REQUIRE_TEAM_ALIAS, 'ALTER TABLE "LiteLLM_TeamTable" ALTER COLUMN "team_alias" SET NOT NULL;\n'
)
ledger_before: Final = _ledger(scratch_database)
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('INSERT INTO "LiteLLM_TeamTable" ("team_id") VALUES (%s)', ("no-alias",))
result: Final = migration_dry_run.dry_run(conn, "public", extras / "migrations")
assert result.rejected is not None, result
assert result.rejected.name == REQUIRE_TEAM_ALIAS
assert 'column "team_alias" of relation "LiteLLM_TeamTable" contains null values' in result.rejected.error
assert result.applied == ()
assert _team_alias_is_nullable(scratch_database), "a rejected dry run must leave the schema as it was"
assert _ledger(scratch_database) == ledger_before
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_the_cli_summary_turns_from_postgres_error_to_timings_once_the_rows_are_fixed(
scratch_database: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
extras: Final = _extras_with_scratch_migration(
tmp_path, REQUIRE_TEAM_ALIAS, 'ALTER TABLE "LiteLLM_TeamTable" ALTER COLUMN "team_alias" SET NOT NULL;\n'
)
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(extras))
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('INSERT INTO "LiteLLM_TeamTable" ("team_id") VALUES (%s)', ("no-alias",))
ok, summary = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=True)
assert ok is False
assert summary.startswith(f"Migration {REQUIRE_TEAM_ALIAS} fails on this database after "), summary
assert "contains null values" in summary
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute('UPDATE "LiteLLM_TeamTable" SET "team_alias" = \'filled\' WHERE "team_alias" IS NULL')
ok_after, summary_after = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=True)
assert ok_after is True
assert summary_after.startswith("1 pending migration(s) applied one at a time in a dry run"), summary_after
assert f"\n {REQUIRE_TEAM_ALIAS}: " in summary_after
assert _team_alias_is_nullable(scratch_database)
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_a_migration_that_depends_on_the_one_before_it_is_replayed_with_it_while_independent_ones_run_alone(
scratch_database: str, tmp_path: Path
) -> None:
extras: Final = _extras_with_scratch_migration(
tmp_path, "20991231000000_add_team_region", 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "region" TEXT;\n'
)
(extras / "migrations" / "20991231000001_index_team_region").mkdir()
(extras / "migrations" / "20991231000001_index_team_region" / "migration.sql").write_text(
'CREATE INDEX CONCURRENTLY "LiteLLM_TeamTable_region_idx" ON "LiteLLM_TeamTable"("region");\n'
)
(extras / "migrations" / "20991231000002_index_team_alias").mkdir()
(extras / "migrations" / "20991231000002_index_team_alias" / "migration.sql").write_text(
'CREATE INDEX "LiteLLM_TeamTable_alias_idx" ON "LiteLLM_TeamTable"("team_alias");\n'
)
with psycopg.connect(scratch_database, autocommit=True) as conn:
result: Final = migration_dry_run.dry_run(conn, "public", extras / "migrations")
region: Final = conn.execute(
"SELECT 1 FROM information_schema.columns WHERE table_name = 'LiteLLM_TeamTable' AND column_name = 'region'"
).fetchone()
assert result.rejected is None, result
assert tuple((applied.name, applied.replayed) for applied in result.applied) == (
("20991231000000_add_team_region", ()),
("20991231000001_index_team_region", ("20991231000000_add_team_region",)),
("20991231000002_index_team_alias", ()),
)
assert region is None, "every transaction must have been rolled back"
with psycopg.connect(scratch_database) as conn:
probe_indexes: Final = conn.execute(
"SELECT indexname FROM pg_indexes "
"WHERE indexname IN ('LiteLLM_TeamTable_region_idx', 'LiteLLM_TeamTable_alias_idx')"
).fetchall()
assert probe_indexes == [], probe_indexes
@requires_db
@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",))
def test_a_database_with_tables_but_no_ledger_is_judged_by_the_schema_diff_instead_of_replaying_history(
scratch_database: str, tmp_path: Path
) -> None:
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute("DROP TABLE _prisma_migrations")
matching: Final = migration_dry_run.dry_run(conn, "public", PACKAGE / "migrations", baseline_check=lambda: None)
differing: Final = migration_dry_run.dry_run(
conn,
"public",
PACKAGE / "migrations",
baseline_check=lambda: (
"the schema differs from this build. prisma migrate diff:\nALTER TABLE x ADD COLUMN y TEXT;"
),
)
custom: Final = migration_dry_run.dry_run(
conn,
"public",
PACKAGE / "migrations",
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(
_extras_with_scratch_migration(
tmp_path, "20991231000000_custom", 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "c" TEXT;\n'
)
),
)
assert matching.ok is True and matching.applied == ()
assert "nothing to dry-run" in migration_dry_run.report(matching)
assert differing.ok is False
assert "ALTER TABLE x ADD COLUMN y TEXT;" in migration_dry_run.report(differing)
assert custom.ok is False
assert "custom migration history" in migration_dry_run.report(custom)

96
tests/unit/litellm_proxy_extras/test_migration_data_checks.py
View file

@ -12,98 +12,15 @@ from litellm_proxy_extras.migration_data_checks import (
_count_query,
btree_max_entry_bytes,
checks_for_index,
checks_for_migration,
checks_for_statement,
migration_data_check_timeout_seconds,
report,
statements,
)
MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0: Final = 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime")'
def _shape(checks: tuple[Check, ...]) -> tuple[tuple[str, str, tuple[str, ...]], ...]:
return tuple((check.kind, check.table, check.columns) for check in checks)
@pytest.mark.parametrize(
("statement", "expected"),
(
(
API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,
(("too_long", "LiteLLM_SpendLogs", ("api_key", "startTime")),),
),
(
'CREATE UNIQUE INDEX CONCURRENTLY "T_a_b_key" ON "public"."T" USING btree ("a", "b" DESC)',
(("too_long", "T", ("a", "b")), ("duplicates", "T", ("a", "b"))),
),
('CREATE INDEX "T_a_idx" ON "T" USING gin ("a")', ()),
('CREATE INDEX "T_a_idx" ON "T" (lower("a"))', ()),
('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', (("nulls", "T", ("a",)),)),
('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL', (("any_row", "T", ("a",)),)),
('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL DEFAULT \'x\'', ()),
('ALTER TABLE "T" ADD COLUMN "a" TEXT', ()),
(
'ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") ON DELETE CASCADE',
(("orphans", "T", ("r",)),),
),
('ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") NOT VALID', ()),
('ALTER TABLE "T" ADD CONSTRAINT "T_a_key" UNIQUE ("a", "b")', (("duplicates", "T", ("a", "b")),)),
(
'ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0 AND "b" IS NOT NULL)',
(("check_violations", "T", ()),),
),
('ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0) NOT VALID', ()),
(
'ALTER TABLE "T" ADD COLUMN "c" TEXT, ALTER COLUMN "b" SET NOT NULL, ADD CONSTRAINT "u" UNIQUE ("c")',
(("nulls", "T", ("b",)), ("duplicates", "T", ("c",))),
),
('CREATE TABLE "T" ("id" TEXT NOT NULL, CONSTRAINT "T_pkey" PRIMARY KEY ("id"))', ()),
('DROP INDEX IF EXISTS "T_a_idx"', ()),
("SELECT 1", ()),
),
)
def test_a_statement_gets_one_check_per_way_existing_rows_can_make_it_fail(
statement: str, expected: tuple[tuple[str, str, tuple[str, ...]], ...]
) -> None:
assert _shape(checks_for_statement(statement, "migration x")) == expected
def test_a_foreign_key_check_keeps_the_referenced_table_and_columns() -> None:
(check,) = checks_for_statement(
'ALTER TABLE "T" ADD CONSTRAINT "fk" FOREIGN KEY ("r", "s") REFERENCES "R"("id", "sub")', "migration x"
)
assert (check.referenced_table, check.referenced_columns) == ("R", ("id", "sub"))
def test_a_check_constraint_keeps_its_expression_for_the_count_query() -> None:
(check,) = checks_for_statement('ALTER TABLE "T" ADD CONSTRAINT "c" CHECK ("a" > 0)', "migration x")
assert check.expression == '"a" > 0'
def test_comments_and_blank_statements_are_dropped_before_matching() -> None:
migration_sql: Final = "-- CreateIndex\n" + API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0 + ";\n\n-- done; really\n"
assert statements(migration_sql) == (API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,)
assert _shape(checks_for_migration("x", migration_sql)) == _shape(
checks_for_statement(API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0, "x")
)
def test_every_shipped_migration_is_readable_and_the_later_indexes_on_text_columns_get_a_check() -> None:
derived: Final = {
folder.name: checks_for_migration(folder.name, (folder / "migration.sql").read_text())
for folder in MIGRATIONS.iterdir()
if folder.is_dir()
}
assert len(derived) > 200
everywhere: Final = tuple(check for checks in derived.values() for check in checks) # comprehension-ok: flattening
assert any(
check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "session_id" in check.columns
for check in everywhere
)
def test_a_request_log_index_gets_one_entry_size_check_over_all_its_columns() -> None:
assert _shape(
checks_for_index("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime"))
@ -142,14 +59,19 @@ def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_wa
assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text
assert 'fix: Shorten or hash the long "LiteLLM_SpendLogs" (api_key, startTime) values' in text
assert text.startswith("Existing rows would make a pending schema change fail:")
assert "No schema change was made" in text
assert "No index was built" in text
def test_a_scan_that_timed_out_is_reported_as_unverified_with_the_knob_to_turn() -> None:
(check,) = checks_for_statement('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', "migration 20270101000000_x")
(check,) = checks_for_index(
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")
)
text: Final = report((CouldNotVerify(check, 600.0),))
assert '"T".a: could not be checked within 600 s, refusing rather than guessing' in text
assert 'blocks: ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL (migration 20270101000000_x)' in text
assert (
'"LiteLLM_SpendLogs".api_key, startTime: could not be checked within 600 s, refusing rather than guessing'
in text
)
assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text
assert MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR in text

80
tests/unit/litellm_proxy_extras/test_migration_dry_run.py Normal file
View file

@ -0,0 +1,80 @@
from typing import Final
import pytest
from litellm_proxy_extras.migration_dry_run import Applied, Baseline, DryRun, Rejected, report, transactional
INDEX: Final = 'CREATE INDEX IF NOT EXISTS "a_idx" ON "T"("c")'
@pytest.mark.parametrize(
("migration_sql", "expected"),
(
('CREATE INDEX CONCURRENTLY IF NOT EXISTS "a_idx" ON "T"("c")', INDEX),
('create unique index concurrently "u" on "T"("d")', 'create unique index "u" on "T"("d")'),
('DROP INDEX CONCURRENTLY "x"', 'DROP INDEX "x"'),
('REINDEX INDEX CONCURRENTLY "y"', 'REINDEX INDEX "y"'),
(INDEX, INDEX),
('ALTER TABLE "T" ALTER COLUMN "c" SET NOT NULL', 'ALTER TABLE "T" ALTER COLUMN "c" SET NOT NULL'),
("-- built CONCURRENTLY by the index job\n" + INDEX, "-- built CONCURRENTLY by the index job\n" + INDEX),
),
ids=("concurrently", "unique_lowercase", "drop", "reindex", "plain_index", "alter", "word_in_comment"),
)
def test_only_concurrent_index_statements_change_so_the_migration_can_run_in_a_transaction(
migration_sql: str, expected: str
) -> None:
assert transactional(migration_sql) == expected
def test_a_clean_dry_run_lists_each_migration_with_the_time_the_upgrade_will_hold_its_locks() -> None:
text: Final = report(DryRun((Applied("20270101000000_a", 0.04), Applied("20270102000000_b", 12.345)), None))
assert text.startswith(
"2 pending migration(s) applied one at a time in a dry run, each rolled back before the next"
)
assert "\n 20270101000000_a: 0.0 s\n 20270102000000_b: 12.3 s" in text
def test_no_pending_migration_is_said_plainly() -> None:
assert report(DryRun((), None)) == "No pending migrations, this database is on this version's schema."
def test_a_rejected_migration_is_reported_with_postgres_own_words_and_what_ran_before_it() -> None:
error: Final = 'column "team_alias" of relation "LiteLLM_TeamTable" contains null values'
text: Final = report(DryRun((Applied("20270101000000_a", 0.5),), Rejected("20270102000000_b", error, 3.25)))
assert text.startswith(
"Migration 20270102000000_b fails on this database after 3.2 s, nothing was applied.\n" + error
)
assert text.endswith("Applied before it in the dry run:\n 20270101000000_a: 0.5 s")
def test_a_rejection_on_the_first_migration_lists_nothing_before_it() -> None:
text: Final = report(DryRun((), Rejected("20270101000000_a", 'syntax error at or near "NOPE"', 0.0)))
assert "Applied before it" not in text
assert text.endswith('syntax error at or near "NOPE"')
def test_a_migration_that_needed_the_one_before_it_says_which_were_replayed_in_front() -> None:
text: Final = report(
DryRun((Applied("20270101000000_a", 0.1), Applied("20270102000000_b", 0.4, ("20270101000000_a",))), None)
)
assert text.endswith(
"\n 20270101000000_a: 0.1 s\n 20270102000000_b: 0.4 s (depends on 20270101000000_a, replayed in front of it)"
)
def test_a_database_without_history_that_matches_this_build_has_nothing_to_dry_run() -> None:
result: Final = DryRun((), None, Baseline(None))
assert result.ok is True
assert report(result) == (
"No migration history and the schema matches this build, the proxy baselines it at boot, nothing to dry-run."
)
def test_a_database_without_history_the_boot_refuses_to_baseline_reports_the_refusal() -> None:
result: Final = DryRun(
(), None, Baseline('the schema differs from this build. prisma migrate diff:\n-- AlterTable\nALTER TABLE "x"')
)
assert result.ok is False
assert report(result) == (
"No migration history and the proxy cannot baseline it at boot: "
'the schema differs from this build. prisma migrate diff:\n-- AlterTable\nALTER TABLE "x"'
)

16
tests/unit/proxy/db/test_prisma_client.py
View file

@ -586,3 +586,19 @@ async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain
await second
await asyncio.wait_for(drained, timeout=1)
assert engine.queries_finished == 2
def test_migration_dry_run_summary_without_migrations_has_nothing_to_dry_run(unset_database_url):
"""`prisma db push` applies no migration files, so the dry run reports that instead of connecting."""
from litellm.proxy.db.prisma_client import PrismaManager
assert PrismaManager.migration_dry_run_summary(use_migrate=False) == (
True,
"The schema is pushed rather than migrated, there is nothing to dry-run.",
)
def test_migration_dry_run_summary_needs_database_url(unset_database_url):
from litellm.proxy.db.prisma_client import PrismaManager
assert PrismaManager.migration_dry_run_summary(use_migrate=True) == (False, "DATABASE_URL is not set.")

17
tests/unit/proxy/test_proxy_cli.py
View file

@ -4393,3 +4393,20 @@ def test_granian_flag_notices_are_json_records_in_json_mode(
"LiteLLM: --ciphers is not applied when using --run_granian.",
]
assert [record["level"] for record in records] == ["INFO", "WARNING", "WARNING"]
def test_migration_data_check_exits_zero_when_there_is_nothing_to_dry_run(monkeypatch, capfd):
monkeypatch.setenv("DATABASE_URL", "postgresql://user:pass@db.example/litellm")
assert ProxyInitializationHelpers.run_migration_data_check(use_migrate=False) == 0
out, err = capfd.readouterr()
assert out == "LiteLLM Proxy: The schema is pushed rather than migrated, there is nothing to dry-run.\n"
assert err == ""
def test_migration_data_check_requires_database_url(monkeypatch):
monkeypatch.delenv("DATABASE_URL", raising=False)
with pytest.raises(click.UsageError, match="requires DATABASE_URL"):
ProxyInitializationHelpers.run_migration_data_check(use_migrate=True)