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 f008be5144d..c24742ce145 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py @@ -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) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py index 5a55b35b255..1c16eb5e6b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -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( diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 36907b0c3fc..8b19576d86e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -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))}") diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 65b20db06fc..d06f0b13cff 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -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) 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 38bd6a2d6ac..75a7fea6c89 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_dry_run.py +++ b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py @@ -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"' )