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

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

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-10-29 01:38:56 +00:00
parent 883cbf259d
commit 6f9ce77aa9
5 changed files with 69 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
View file

@ -42,14 +42,15 @@ class Rejected:
@dataclass(frozen=True, slots=True)
class Baseline:
"""A database with tables but no `_prisma_migrations`: the boot baselines it when the live
schema matches this build, so the dry run compares schemas instead of replaying history."""
"""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)."""
diff: str | None
refusal: str | None
@property
def matches(self) -> bool:
return self.diff is None
return self.refusal is None
@dataclass(frozen=True, slots=True)
@ -135,15 +136,15 @@ def dry_run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
migrations_dir: Path,
schema_diff: Callable[[], str | None] | None = None,
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.
`schema_diff` (None when the live schema matches this build's schema.prisma) stands in for the
replay on a database that has tables but no migration history, which the boot baselines."""
if schema_diff is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema):
return DryRun((), None, Baseline(schema_diff()))
`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),
@ -166,10 +167,7 @@ def report(result: DryRun) -> str:
"No migration history and the schema matches this build, "
"the proxy baselines it at boot, nothing to dry-run."
)
return (
"No migration history and the schema differs from this build, so the proxy cannot baseline it "
f"at boot. prisma migrate diff:\n{result.baseline.diff}"
)
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)

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(

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

@ -1426,11 +1426,18 @@ class ProxyExtrasDBManager:
return True
@staticmethod
def _schema_diff_against_build(prisma_dir: Path) -> str | None:
"""`prisma migrate diff` between the live database and this build's schema.prisma, the
same check `baseline_current_schema` runs at boot: None when they match, the diff otherwise."""
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(
(
@ -1447,7 +1454,8 @@ class ProxyExtrasDBManager:
env=migration_environment(_get_prisma_env()),
)
except subprocess.CalledProcessError as exc:
return str(exc.stdout or exc.stderr or exc).strip()
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
@ -1476,7 +1484,7 @@ class ProxyExtrasDBManager:
connection,
schema,
prisma_dir / "migrations",
schema_diff=lambda: ProxyExtrasDBManager._schema_diff_against_build(prisma_dir),
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir),
)
except psycopg.OperationalError as exc:
return (False, f"Could not reach the database: {_redact_credentials(str(exc))}")

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

@ -1256,27 +1256,44 @@ def test_a_migration_that_depends_on_the_one_before_it_is_replayed_with_it_while
("20991231000002_index_team_alias", ()),
)
assert region is None, "every transaction must have been rolled back"
assert (
_index_validity(scratch_database, "region_idx") == {} and _index_validity(scratch_database, "alias_idx") == {}
)
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,
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", schema_diff=lambda: None)
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",
schema_diff=lambda: "-- AlterTable\nALTER TABLE x ADD COLUMN y TEXT;",
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)

13
tests/unit/litellm_proxy_extras/test_migration_dry_run.py
View file

@ -69,9 +69,12 @@ def test_a_database_without_history_that_matches_this_build_has_nothing_to_dry_r
)
def test_a_database_without_history_that_differs_from_this_build_reports_the_diff() -> None:
result: Final = DryRun((), None, Baseline('-- AlterTable\nALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "x" TEXT;'))
assert result.ok is False
assert report(result).endswith(
'prisma migrate diff:\n-- AlterTable\nALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "x" TEXT;'
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"'
)