Merge pull request #44206 from BerriAI/litellm_rc_1_104_0_enforce_migration_check_default

fix(proxy): enforce the migration check by default and stop building SpendLogs indexes in migrations (rc/1.104.0)

litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql

@@ -1,2 +1 @@
--- CreateIndex
-CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");
+SELECT 1;

litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql

@@ -1,12 +1 @@
--- CreateIndex (CONCURRENTLY)
---
--- Disclaimer:
--- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
---   single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
--- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
---   interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
--- - Do not edit this file after it has been applied to any database: Prisma checksums
---   migrations; add a new migration instead.
--- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
---   without IF NOT EXISTS if you must support older versions).
-CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
+SELECT 1;

litellm/proxy/prisma_migration.py

@@ -14,15 +14,24 @@ sys.path.insert(0, os.path.abspath("./"))
 
 from typing import Final
 
+import click
 from litellm_proxy_extras.prisma_toolchain import resolve_prisma_argv
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy.proxy_cli import run_server
-from litellm.secret_managers.main import str_to_bool
+
+
+def _enforce_prisma_migration_check(raw: str | None) -> bool:
+    if raw is None or not raw.strip():
+        return True
+    try:
+        return bool(click.BOOL(raw))
+    except click.BadParameter:
+        return True
 
 
 def main() -> int:
-    enforce_prisma_migration_check: Final = str_to_bool(os.getenv("ENFORCE_PRISMA_MIGRATION_CHECK")) is not False
+    enforce_prisma_migration_check: Final = _enforce_prisma_migration_check(os.getenv("ENFORCE_PRISMA_MIGRATION_CHECK"))
     run_server_args: Final = (
         ("--skip_server_startup", "--enforce_prisma_migration_check")
         if enforce_prisma_migration_check

litellm/proxy/proxy_cli.py

@@ -963,10 +963,14 @@ class ProxyInitializationHelpers:
     envvar="LIMIT_CONCURRENCY",
 )
 @click.option(
-    "--enforce_prisma_migration_check",
-    is_flag=True,
-    default=False,
-    help="Exit with error if database migration fails on startup.",
+    "--enforce_prisma_migration_check/--no-enforce_prisma_migration_check",
+    default=True,
+    show_default=True,
+    help=(
+        "Exit when database setup fails on startup instead of serving against a database "
+        "whose schema may be behind the code. Opt out with --no-enforce_prisma_migration_check "
+        "or ENFORCE_PRISMA_MIGRATION_CHECK=false."
+    ),
     envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
 )
 @click.option(
@@ -1434,8 +1438,8 @@ def run_server(
                             sys.exit(1)
                         else:
                             print(
-                                "\033[1;33mLiteLLM Proxy: Database migration failed but continuing startup. "
-                                "Set --enforce_prisma_migration_check or ENFORCE_PRISMA_MIGRATION_CHECK=true to exit on failure.\033[0m"
+                                "\033[1;33mLiteLLM Proxy: Database migration failed but continuing startup because "
+                                "ENFORCE_PRISMA_MIGRATION_CHECK is disabled. The schema may be behind the code.\033[0m"
                             )
             else:
                 print(

tests/test_litellm/proxy/test_prisma_migration.py

@@ -36,6 +36,32 @@ class TestPrismaMigration:
 
         mock_run_server.assert_called_once_with(("--skip_server_startup",), standalone_mode=False)
 
+    @pytest.mark.parametrize(
+        ("value", "enforces"),
+        [
+            ("0", False),
+            ("no", False),
+            ("off", False),
+            ("FALSE", False),
+            ("1", True),
+            ("yes", True),
+            ("", True),
+            ("maybe", True),
+        ],
+    )
+    @patch("litellm.proxy.prisma_migration.subprocess.run")
+    @patch("litellm.proxy.prisma_migration.run_server")
+    def test_main_reads_the_same_opt_out_values_as_the_proxy_flag(
+        self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, value: str, enforces: bool
+    ) -> None:
+        mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="")
+
+        with patch.dict(os.environ, {"ENFORCE_PRISMA_MIGRATION_CHECK": value}, clear=True):
+            assert prisma_migration.main() == 0
+
+        passed: tuple[str, ...] = mock_run_server.call_args.args[0]
+        assert ("--enforce_prisma_migration_check" in passed) is enforces
+
     @pytest.mark.parametrize("env", [{}, {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}])
     @patch("litellm.proxy.prisma_migration.subprocess.run")
     @patch("litellm.proxy.prisma_migration.run_server")
@@ -67,7 +93,9 @@ class TestPrismaMigration:
 
         mock_subprocess_run.assert_not_called()
 
-    @patch("litellm.proxy.prisma_migration.subprocess.run")  # test-quality-ok: the spawned argv is the behavior under test
+    @patch(
+        "litellm.proxy.prisma_migration.subprocess.run"
+    )  # test-quality-ok: the spawned argv is the behavior under test
     @patch("litellm.proxy.prisma_migration.run_server")  # test-quality-ok: run_server boots the whole proxy
     def test_prisma_generate_runs_through_the_module_when_the_cli_is_not_on_path(
         self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path
@@ -81,7 +109,9 @@ class TestPrismaMigration:
 
         assert mock_subprocess_run.call_args.args[0] == (sys.executable, "-m", "prisma", "generate")
 
-    @patch("litellm.proxy.prisma_migration.subprocess.run")  # test-quality-ok: the spawned argv is the behavior under test
+    @patch(
+        "litellm.proxy.prisma_migration.subprocess.run"
+    )  # test-quality-ok: the spawned argv is the behavior under test
     @patch("litellm.proxy.prisma_migration.run_server")  # test-quality-ok: run_server boots the whole proxy
     def test_prisma_generate_runs_the_console_script_when_it_is_on_path(
         self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path

tests/test_litellm/proxy/test_proxy_cli.py

@@ -1,5 +1,6 @@
 import inspect
 import os
+from contextlib import nullcontext
 from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -2266,6 +2267,18 @@ class TestRunServerDbSetup:
             use_migrate=True, use_v2_resolver=True
         )
 
+    @pytest.mark.parametrize(
+        ("arguments", "env", "exits"),
+        [
+            ((), {}, True),
+            (("--enforce_prisma_migration_check",), {}, True),
+            ((), {"ENFORCE_PRISMA_MIGRATION_CHECK": "true"}, True),
+            (("--no-enforce_prisma_migration_check",), {}, False),
+            ((), {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}, False),
+            ((), {"ENFORCE_PRISMA_MIGRATION_CHECK": "0"}, False),
+        ],
+        ids=("default", "flag", "env-true", "no-flag", "env-false", "env-0"),
+    )
     @patch("subprocess.run")
     @patch("atexit.register")
     @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@@ -2278,8 +2291,13 @@ class TestRunServerDbSetup:
         mock_setup_database,
         mock_atexit_register,
         mock_subprocess_run,
+        arguments,
+        env,
+        exits,
+        capsys,
     ):
-        """Test that proxy exits with code 1 when PrismaManager.setup_database returns False and --enforce_prisma_migration_check is set"""
+        """A failed PrismaManager.setup_database exits 1 by default; only an explicit
+        --no-enforce_prisma_migration_check or ENFORCE_PRISMA_MIGRATION_CHECK=false keeps the proxy booting"""
         from litellm.proxy.proxy_cli import run_server
 
         mock_subprocess_run.return_value = MagicMock(returncode=0)
@@ -2296,12 +2314,13 @@ class TestRunServerDbSetup:
         clean_env = {
             k: v
             for k, v in os.environ.items()
-            if k not in ("DATABASE_URL", "DIRECT_URL")
+            if k not in ("DATABASE_URL", "DIRECT_URL", "ENFORCE_PRISMA_MIGRATION_CHECK")
         }
         clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
+        outcome = pytest.raises(SystemExit) if exits else nullcontext()
 
         with (
-            patch.dict(os.environ, clean_env, clear=True),
+            patch.dict(os.environ, {**clean_env, **env}, clear=True),
             patch.dict(
                 "sys.modules",
                 {
@@ -2312,26 +2331,18 @@ class TestRunServerDbSetup:
             patch(
                 "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
             ) as mock_get_args,
+            outcome as exc_info,
         ):
             mock_get_args.return_value = {
                 "app": "litellm.proxy.proxy_server:app",
                 "host": "localhost",
                 "port": 8000,
             }
+            run_server.main(["--local", "--skip_server_startup", *arguments], standalone_mode=False)
 
-            with pytest.raises(SystemExit) as exc_info:
-                run_server.main(
-                    [
-                        "--local",
-                        "--skip_server_startup",
-                        "--enforce_prisma_migration_check",
-                    ],
-                    standalone_mode=False,
-                )
-            assert exc_info.value.code == 1
-            mock_setup_database.assert_called_once_with(
-                use_migrate=True, use_v2_resolver=True
-            )
+        assert (exc_info is not None and exc_info.value.code == 1) is exits
+        assert ("Database migration failed but continuing startup" in capsys.readouterr().out) is not exits
+        mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
 
     @patch("subprocess.run")
     @patch("atexit.register")