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

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

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-10-28 23:41:34 +00:00
parent 96fe78f96e
commit 634f755164
8 changed files with 319 additions and 461 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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."

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

@ -0,0 +1,111 @@
"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real,
inside one transaction that is always rolled back, so Postgres itself says whether this
database's rows and schema take them. It holds the same locks the real upgrade will."""
from __future__ import annotations
import re
import time
from collections.abc import Callable
from dataclasses import dataclass
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
)
@dataclass(frozen=True, slots=True)
class Applied:
name: str
seconds: float
@dataclass(frozen=True, slots=True)
class Rejected:
name: str
error: str
seconds: float
@dataclass(frozen=True, slots=True)
class DryRun:
applied: tuple[Applied, ...]
rejected: Rejected | None
@property
def ok(self) -> bool:
return self.rejected is None
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 _run(
connection: psycopg.Connection[tuple[object, ...]],
done: DryRun,
migration: tuple[str, str],
clock: Callable[[], float],
) -> DryRun:
import psycopg
if done.rejected is not None:
return done
name, migration_sql = migration
started: Final = clock()
try:
connection.execute(transactional(migration_sql))
except psycopg.Error as exc:
return DryRun(done.applied, Rejected(name, str(exc).strip(), clock() - started))
return DryRun(done.applied + (Applied(name, clock() - started),), None)
def dry_run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
migrations_dir: Path,
clock: Callable[[], float] = time.monotonic,
) -> DryRun:
"""Apply every pending migration in order on `connection`, then roll all of it back,
whatever happened. The first migration Postgres rejects ends the run with its error."""
from psycopg import sql
pending: Final = pending_migrations(connection, schema, migrations_dir)
with connection.transaction(force_rollback=True):
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
result: Final = _fold(connection, pending, clock)
return result
def _fold(
connection: psycopg.Connection[tuple[object, ...]],
pending: tuple[tuple[str, str], ...],
clock: Callable[[], float],
) -> DryRun:
from functools import reduce
return reduce(lambda done, migration: _run(connection, done, migration, clock), pending, DryRun((), None))
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."""
timings: Final = "".join(f"\n {applied.name}: {applied.seconds:.1f} s" 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 in a dry run and rolled back, 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 "")
)

67
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,29 @@ 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 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)
migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations"
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, migrations_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 +1700,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

@ -245,18 +245,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
@ -1028,8 +1029,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",

174
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,102 @@ 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"
assert "pending migration(s) applied in a dry run and rolled back" 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 in a dry run and rolled back, nothing changed:"), (
summary_after
)
assert f"\n {REQUIRE_TEAM_ALIAS}: " in summary_after
assert _team_alias_is_nullable(scratch_database)

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

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

@ -0,0 +1,50 @@
from typing import Final
import pytest
from litellm_proxy_extras.migration_dry_run import Applied, 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 in a dry run and rolled back, nothing changed:")
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"')