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

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

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))

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)

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;'
+    )