diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..68dcc01c055 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -331,26 +331,14 @@ class ProxyExtrasDBManager: return False @staticmethod - def _resolve_all_migrations( - migrations_dir: str, schema_path: str, mark_all_applied: bool = True - ): - """ - 1. Compare the current database state to schema.prisma and generate a migration for the diff. - 2. Run prisma migrate deploy to apply any pending migrations. - 3. Mark all existing migrations as applied. - """ - database_url = os.getenv("DATABASE_URL") - if not database_url: - logger.error("DATABASE_URL not set") - return - # Prefer DIRECT_URL for schema introspection — pooler URLs (e.g. neon -pooler) - # do not support the extended query protocol required by prisma migrate diff. - diff_url = os.getenv("DIRECT_URL") or database_url + def _write_migration_diff( + diff_url: str, schema_path: str, diff_sql_path: Path + ) -> bool: + """Write the DB-vs-schema drift script, reporting whether it ran to completion. - diff_dir = Path(tempfile.mkdtemp(prefix="litellm_migration_diff_")) - diff_sql_path = diff_dir / "migration.sql" - - # 1. Generate migration SQL for the diff between DB and schema + A killed or failed `prisma migrate diff` still leaves the file behind, so the + caller needs this rather than the file's existence to know the script is whole. + """ try: logger.info("Generating migration diff between DB and schema.prisma...") with open(diff_sql_path, "w") as f: @@ -372,8 +360,40 @@ class ProxyExtrasDBManager: ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") + return False except subprocess.TimeoutExpired: logger.warning("Migration diff generation timed out.") + return False + return True + + @staticmethod + def _resolve_all_migrations( + migrations_dir: str, schema_path: str, mark_all_applied: bool = True + ): + """ + 1. Compare the current database state to schema.prisma and generate a migration for the diff. + 2. Run prisma migrate deploy to apply any pending migrations. + 3. Mark all existing migrations as applied. + """ + database_url = os.getenv("DATABASE_URL") + if not database_url: + logger.error("DATABASE_URL not set") + return + # Prefer DIRECT_URL for schema introspection — pooler URLs (e.g. neon -pooler) + # do not support the extended query protocol required by prisma migrate diff. + diff_url = os.getenv("DIRECT_URL") or database_url + + diff_dir = Path(tempfile.mkdtemp(prefix="litellm_migration_diff_")) + diff_sql_path = diff_dir / "migration.sql" + + # 1. Generate migration SQL for the diff between DB and schema + diff_generated = ProxyExtrasDBManager._write_migration_diff( + diff_url, schema_path, diff_sql_path + ) + + def finish_without_applying() -> None: + if mark_all_applied: + ProxyExtrasDBManager._mark_migrations_applied(migrations_dir) # check if the migration was created if not diff_sql_path.exists(): @@ -414,10 +434,24 @@ class ProxyExtrasDBManager: return logger.info(f"Migration diff created at {diff_sql_path}") - if ProxyExtrasDBManager.spend_logs_is_partitioned(): - filtered_sql = filter_partitioned_spend_logs_diff( - diff_sql_path.read_text() + if not diff_generated: + logger.warning( + "`prisma migrate diff` did not finish, so the drift script may be " + "truncated; not applying it" ) + finish_without_applying() + return + + diff_sql = diff_sql_path.read_text() + if not _without_sql_comments(diff_sql): + logger.info( + "Database already matches schema.prisma; no drift script to apply" + ) + finish_without_applying() + return + + if ProxyExtrasDBManager.spend_logs_is_partitioned(): + filtered_sql = filter_partitioned_spend_logs_diff(diff_sql) diff_sql_path.write_text(filtered_sql) logger.info( "LiteLLM_SpendLogs is partitioned; removed its primary-key " @@ -425,9 +459,7 @@ class ProxyExtrasDBManager: ) if not filtered_sql.strip(): logger.info("Drift script is empty after filtering; nothing to apply") - if not mark_all_applied: - return - ProxyExtrasDBManager._mark_migrations_applied(migrations_dir) + finish_without_applying() return # 2. Run prisma db execute to apply the migration diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 1b95d24c011..b87f2e734b6 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -4,32 +4,84 @@ Migration failures fail the entrypoint by default; set ENFORCE_PRISMA_MIGRATION_ for log-only behavior. A failed 'prisma generate' is always log-only: every shipped image bakes the client at build time, and refreshing it writes into site-packages, which an arbitrary non-root uid or a read-only root filesystem cannot do. + +That same refresh is also skipped outright when the installed client already came from the +schema being generated from, because it can then only rewrite the client with what it +already holds. It is not free: it is the slowest step of the migrations Job (~10s on a +whole CPU, ~40s on a quarter of one), it runs unbounded after the database is already +migrated, and a Job that is still inside it looks identical to a Job still migrating. """ import os import subprocess import sys +from pathlib import Path sys.path.insert(0, os.path.abspath("./")) +from collections.abc import Callable from typing import Final from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_cli import run_server from litellm.secret_managers.main import str_to_bool +SCHEMA_FILE: Final = Path("schema.prisma") -def main() -> int: + +def generated_client_dir() -> Path | None: + """Directory 'prisma generate' writes the client into, or None if prisma is missing.""" + try: + import prisma + except ImportError: + return None + return Path(prisma.__file__).parent if prisma.__file__ else None + + +def client_already_generated_from(schema: Path, client_dir: Path | None) -> bool: + """True when the installed client was generated from this exact schema. + + 'prisma generate' copies its source schema into the generated package, so identical + bytes mean a regeneration would reproduce what is already on disk. + """ + if client_dir is None: + return False + try: + return (client_dir / SCHEMA_FILE.name).read_bytes() == schema.read_bytes() + except OSError: + return False + + +def installed_client_is_current() -> bool: + return client_already_generated_from(SCHEMA_FILE, generated_client_dir()) + + +def generate_prisma_client() -> subprocess.CompletedProcess[str]: + return subprocess.run(("prisma", "generate"), capture_output=True, text=True) + + +def main( + start_server: Callable[..., object] = run_server, + client_is_current: Callable[[], bool] = installed_client_is_current, + generate_client: Callable[[], subprocess.CompletedProcess[str]] = generate_prisma_client, +) -> int: enforce_prisma_migration_check: Final = str_to_bool(os.getenv("ENFORCE_PRISMA_MIGRATION_CHECK")) is not False run_server_args: Final = ( ("--skip_server_startup", "--enforce_prisma_migration_check") if enforce_prisma_migration_check else ("--skip_server_startup",) ) - run_server(run_server_args, standalone_mode=False) + start_server(run_server_args, standalone_mode=False) + + if client_is_current(): + verbose_proxy_logger.info( + "Skipping 'prisma generate': the installed client already comes from %s", + SCHEMA_FILE, + ) + return 0 verbose_proxy_logger.info("Running 'prisma generate'...") - result: Final = subprocess.run(("prisma", "generate"), capture_output=True, text=True) + result: Final = generate_client() verbose_proxy_logger.info("'prisma generate' stdout: %s", result.stdout) if result.returncode != 0: diff --git a/test-quality-budget.json b/test-quality-budget.json index ee33eb581d6..b2b18843bec 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -21,6 +21,6 @@ "limit": 117 }, "TQ008": { - "limit": 11139 + "limit": 11131 } } diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index b3d457707b8..8ea09148356 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -494,6 +494,11 @@ ADD CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id"); DROP TABLE "LiteLLM_SpendLogs_legacy"; """ +_EMPTY_DRIFT_SQL = "-- This is an empty migration.\n\n" + +_TRUNCATED_DRIFT_SQL = """-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COL""" + class TestPartitionedSpendLogsDriftFilter: """A doc-partitioned LiteLLM_SpendLogs (db_scripts/partition_spend_logs.sql) has a @@ -537,15 +542,24 @@ class _FakeCompleted: class TestResolveAllMigrationsLedger: - def _run(self, monkeypatch, tmp_path, partitioned, execute_fails): + def _run( + self, + monkeypatch, + tmp_path, + partitioned, + execute_fails, + diff_sql=None, + diff_times_out=False, + ): import subprocess as subprocess_module import litellm_proxy_extras.utils as utils_module monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") monkeypatch.delenv("DIRECT_URL", raising=False) + is_partitioned = partitioned if callable(partitioned) else lambda: partitioned monkeypatch.setattr( - ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: partitioned) + ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(is_partitioned) ) monkeypatch.setattr( ProxyExtrasDBManager, @@ -557,7 +571,11 @@ class TestResolveAllMigrationsLedger: def fake_run(cmd, **kwargs): calls.append(cmd) if "diff" in cmd: - kwargs["stdout"].write(_PARTITIONED_DRIFT_SQL) + kwargs["stdout"].write( + _PARTITIONED_DRIFT_SQL if diff_sql is None else diff_sql + ) + if diff_times_out: + raise subprocess_module.TimeoutExpired(cmd, 60) return _FakeCompleted() if "execute" in cmd: executed_sql = open(cmd[cmd.index("--file") + 1]).read() @@ -574,6 +592,9 @@ class TestResolveAllMigrationsLedger: def _resolved(self, calls): return [c for c in calls if isinstance(c, list) and "resolve" in c] + def _executed(self, calls): + return [c for c in calls if isinstance(c, list) and "execute" in c] + def _executed_sql(self, calls): return next(c[1] for c in calls if isinstance(c, tuple) and c[0] == "executed_sql") @@ -598,6 +619,44 @@ class TestResolveAllMigrationsLedger: calls = self._run(monkeypatch, tmp_path, partitioned=False, execute_fails=False) assert self._executed_sql(calls) == _PARTITIONED_DRIFT_SQL + def test_an_empty_drift_script_costs_no_prisma_invocation(self, monkeypatch, tmp_path): + calls = self._run( + monkeypatch, + tmp_path, + partitioned=False, + execute_fails=False, + diff_sql=_EMPTY_DRIFT_SQL, + ) + assert self._executed(calls) == [] + assert len(self._resolved(calls)) == 1 + + def test_a_truncated_drift_script_is_never_applied(self, monkeypatch, tmp_path): + calls = self._run( + monkeypatch, + tmp_path, + partitioned=False, + execute_fails=False, + diff_sql=_TRUNCATED_DRIFT_SQL, + diff_times_out=True, + ) + assert self._executed(calls) == [] + assert len(self._resolved(calls)) == 1 + + def test_a_partitioned_database_is_not_queried_for_an_empty_drift_script( + self, monkeypatch, tmp_path + ): + def explode(): + raise AssertionError("spend_logs_is_partitioned should not be reached") + + calls = self._run( + monkeypatch, + tmp_path, + partitioned=explode, + execute_fails=False, + diff_sql=_EMPTY_DRIFT_SQL, + ) + assert self._executed(calls) == [] + class TestPartitionedSpendLogsPushGuard: def _forbid_subprocess(self, monkeypatch): diff --git a/tests/test_litellm/proxy/test_prisma_migration.py b/tests/test_litellm/proxy/test_prisma_migration.py index 729adcfb9e0..47170c816b1 100644 --- a/tests/test_litellm/proxy/test_prisma_migration.py +++ b/tests/test_litellm/proxy/test_prisma_migration.py @@ -1,66 +1,172 @@ import os -from unittest.mock import MagicMock, patch +import subprocess +from pathlib import Path +from unittest.mock import patch import pytest from litellm.proxy import prisma_migration +def _completed(returncode: int = 0) -> subprocess.CompletedProcess[str]: + return subprocess.CompletedProcess(("prisma", "generate"), returncode, stdout="", stderr="") + + +class _RecordingServer: + def __init__(self, exit_code: int | None = None) -> None: + self.calls: list[tuple[tuple[str, ...], bool]] = [] + self.exit_code = exit_code + + def __call__(self, args: tuple[str, ...], standalone_mode: bool) -> None: + self.calls.append((args, standalone_mode)) + if self.exit_code is not None: + raise SystemExit(self.exit_code) + + +class _RecordingGenerator: + def __init__(self, returncode: int = 0) -> None: + self.calls = 0 + self.returncode = returncode + + def __call__(self) -> subprocess.CompletedProcess[str]: + self.calls += 1 + return _completed(self.returncode) + + class TestPrismaMigration: - @patch("litellm.proxy.prisma_migration.subprocess.run") - @patch("litellm.proxy.prisma_migration.run_server") - def test_main_enforces_migration_check_by_default( - self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock - ) -> None: - mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + def test_main_enforces_migration_check_by_default(self) -> None: + server = _RecordingServer() with patch.dict(os.environ, {}, clear=True): - assert prisma_migration.main() == 0 + assert ( + prisma_migration.main( + start_server=server, + client_is_current=lambda: False, + generate_client=_RecordingGenerator(), + ) + == 0 + ) - mock_run_server.assert_called_once_with( - ("--skip_server_startup", "--enforce_prisma_migration_check"), - standalone_mode=False, - ) + assert server.calls == [ + (("--skip_server_startup", "--enforce_prisma_migration_check"), False) + ] - @patch("litellm.proxy.prisma_migration.subprocess.run") - @patch("litellm.proxy.prisma_migration.run_server") - def test_main_disables_migration_check_when_explicitly_false( - self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock - ) -> None: - mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + def test_main_disables_migration_check_when_explicitly_false(self) -> None: + server = _RecordingServer() with patch.dict(os.environ, {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}, clear=True): - assert prisma_migration.main() == 0 + assert ( + prisma_migration.main( + start_server=server, + client_is_current=lambda: False, + generate_client=_RecordingGenerator(), + ) + == 0 + ) - mock_run_server.assert_called_once_with(("--skip_server_startup",), standalone_mode=False) + assert server.calls == [(("--skip_server_startup",), False)] @pytest.mark.parametrize("env", [{}, {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}]) - @patch("litellm.proxy.prisma_migration.subprocess.run") - @patch("litellm.proxy.prisma_migration.run_server") - def test_main_exits_zero_when_only_prisma_generate_fails( - self, - mock_run_server: MagicMock, - mock_subprocess_run: MagicMock, - env: dict[str, str], - ) -> None: - mock_subprocess_run.return_value = MagicMock( - returncode=1, - stdout="", - stderr="PermissionError: [Errno 13] Permission denied: '/app/.venv/lib/python3.13/site-packages/prisma/schema.prisma'", - ) + def test_main_exits_zero_when_only_prisma_generate_fails(self, env: dict[str, str]) -> None: + generator = _RecordingGenerator(returncode=1) with patch.dict(os.environ, env, clear=True): - assert prisma_migration.main() == 0 + assert ( + prisma_migration.main( + start_server=_RecordingServer(), + client_is_current=lambda: False, + generate_client=generator, + ) + == 0 + ) - @patch("litellm.proxy.prisma_migration.subprocess.run") - @patch("litellm.proxy.prisma_migration.run_server") - def test_main_propagates_migration_failure( - self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock - ) -> None: - mock_run_server.side_effect = SystemExit(1) + assert generator.calls == 1 + + def test_main_propagates_migration_failure(self) -> None: + generator = _RecordingGenerator() with patch.dict(os.environ, {}, clear=True): with pytest.raises(SystemExit, match="1"): - prisma_migration.main() + prisma_migration.main( + start_server=_RecordingServer(exit_code=1), + client_is_current=lambda: False, + generate_client=generator, + ) - mock_subprocess_run.assert_not_called() + assert generator.calls == 0 + + def test_main_skips_prisma_generate_when_the_client_is_already_current(self) -> None: + server = _RecordingServer() + generator = _RecordingGenerator() + + with patch.dict(os.environ, {}, clear=True): + assert ( + prisma_migration.main( + start_server=server, + client_is_current=lambda: True, + generate_client=generator, + ) + == 0 + ) + + assert len(server.calls) == 1 + assert generator.calls == 0 + + +class TestClientAlreadyGeneratedFrom: + def _schema(self, tmp_path: Path, body: str) -> Path: + schema = tmp_path / "source" / "schema.prisma" + schema.parent.mkdir() + schema.write_text(body) + return schema + + def _client_dir(self, tmp_path: Path, body: str | None = None) -> Path: + client_dir = tmp_path / "client" + client_dir.mkdir() + if body is not None: + (client_dir / "schema.prisma").write_text(body) + return client_dir + + def test_a_client_generated_from_the_same_schema_is_current(self, tmp_path: Path) -> None: + body = "model Foo {\n id String @id\n}\n" + schema = self._schema(tmp_path, body) + + assert ( + prisma_migration.client_already_generated_from( + schema, self._client_dir(tmp_path, body) + ) + is True + ) + + def test_a_client_generated_from_a_different_schema_is_not_current( + self, tmp_path: Path + ) -> None: + schema = self._schema(tmp_path, "model Foo {\n id String @id\n}\n") + + assert ( + prisma_migration.client_already_generated_from( + schema, self._client_dir(tmp_path, "model Foo {\n id Int @id\n}\n") + ) + is False + ) + + def test_an_ungenerated_client_is_not_current(self, tmp_path: Path) -> None: + schema = self._schema(tmp_path, "model Foo {\n id String @id\n}\n") + + assert ( + prisma_migration.client_already_generated_from(schema, self._client_dir(tmp_path)) + is False + ) + + def test_a_missing_source_schema_is_not_current(self, tmp_path: Path) -> None: + client_dir = self._client_dir(tmp_path, "model Foo {\n id String @id\n}\n") + + assert ( + prisma_migration.client_already_generated_from(tmp_path / "absent.prisma", client_dir) + is False + ) + + def test_an_unimportable_prisma_package_is_not_current(self, tmp_path: Path) -> None: + schema = self._schema(tmp_path, "model Foo {\n id String @id\n}\n") + + assert prisma_migration.client_already_generated_from(schema, None) is False