From 53d8f0d72ef0ad8f3fa3234a8f93962b6db6f893 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:51:08 +0000 Subject: [PATCH] feat(proxy): roll back each migration of the dry run on its own and baseline-check databases without a ledger --- .../litellm_proxy_extras/migration_dry_run.py | 139 ++++++++++++++---- .../litellm_proxy_extras/utils.py | 36 ++++- .../migrations/test_request_log_indexes.py | 65 +++++++- .../test_migration_dry_run.py | 31 +++- 4 files changed, 231 insertions(+), 40 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py index 09c1e11f537..f008be5144d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py @@ -1,6 +1,7 @@ -"""`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.""" +"""`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 @@ -8,6 +9,7 @@ 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 @@ -19,12 +21,16 @@ if TYPE_CHECKING: _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) @@ -34,14 +40,27 @@ class Rejected: seconds: float +@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.""" + + diff: str | None + + @property + def matches(self) -> bool: + return self.diff 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 + return self.rejected is None and (self.baseline is None or self.baseline.matches) def transactional(migration_sql: str) -> str: @@ -49,61 +68,117 @@ def transactional(migration_sql: str) -> str: return _CONCURRENTLY.sub(r"\1", migration_sql) -def _run( +def _attempt( connection: psycopg.Connection[tuple[object, ...]], - done: DryRun, - migration: tuple[str, str], + schema: str, + chain: tuple[Migration, ...], clock: Callable[[], float], -) -> DryRun: +) -> tuple[float, psycopg.Error | None]: import psycopg + from psycopg import sql - if done.rejected is not None: - return done - name, migration_sql = migration started: Final = clock() try: - connection.execute(transactional(migration_sql)) + 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 DryRun(done.applied, Rejected(name, str(exc).strip(), clock() - started)) - return DryRun(done.applied + (Applied(name, clock() - started),), None) + 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, + schema_diff: Callable[[], str | None] | None = None, 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 - + """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())) 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 + return reduce( + lambda done, position: _run(connection, schema, pending, position, done, clock), + range(len(pending)), + DryRun((), None), + ) -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 _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.""" - timings: Final = "".join(f"\n {applied.name}: {applied.seconds:.1f} s" for applied in result.applied) + 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 ( + "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}" + ) + 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 in a dry run and rolled back, nothing changed:{timings}" + 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}" diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index f7a188bb3bd..36907b0c3fc 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -1425,6 +1425,33 @@ class ProxyExtrasDBManager: ProxyExtrasDBManager.apply_replica_identity_full_if_requested() 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.""" + from litellm_proxy_extras.migration_lock import migration_environment + + 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: + return str(exc.stdout or exc.stderr or exc).strip() + 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 @@ -1442,10 +1469,15 @@ class ProxyExtrasDBManager: 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" + prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) try: with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection: - result: Final = dry_run(connection, schema, migrations_dir) + result: Final = dry_run( + connection, + schema, + prisma_dir / "migrations", + schema_diff=lambda: ProxyExtrasDBManager._schema_diff_against_build(prisma_dir), + ) except psycopg.OperationalError as exc: return (False, f"Could not reach the database: {_redact_credentials(str(exc))}") return (result.ok, dry_run_report(result)) diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 4df0e73c03b..65b20db06fc 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -1173,7 +1173,11 @@ def test_the_dry_run_applies_every_pending_migration_in_order_and_rolls_all_of_i 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) + 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 @@ -1218,8 +1222,61 @@ def test_the_cli_summary_turns_from_postgres_error_to_timings_once_the_rows_are_ 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 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" + assert ( + _index_validity(scratch_database, "region_idx") == {} and _index_validity(scratch_database, "alias_idx") == {} + ) + + +@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, +) -> 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) + differing: Final = migration_dry_run.dry_run( + conn, + "public", + PACKAGE / "migrations", + schema_diff=lambda: "-- AlterTable\nALTER TABLE x ADD COLUMN y TEXT;", + ) + + 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) diff --git a/tests/unit/litellm_proxy_extras/test_migration_dry_run.py b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py index 6352532982e..38bd6a2d6ac 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_dry_run.py +++ b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py @@ -1,7 +1,7 @@ from typing import Final import pytest -from litellm_proxy_extras.migration_dry_run import Applied, DryRun, Rejected, report, transactional +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")' @@ -27,7 +27,9 @@ def test_only_concurrent_index_statements_change_so_the_migration_can_run_in_a_t 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 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 @@ -48,3 +50,28 @@ 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_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;' + )