diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql index 9a061aaed43..e0ac49d1ecf 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql +++ b/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; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql index 62ad5c42ba7..e0ac49d1ecf 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql +++ b/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; diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 7e3aff75cef..7bd0fcd6a5e 100644 --- a/litellm/proxy/prisma_migration.py +++ b/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 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0477b6c62e9..0ebe8dbc42a 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -921,10 +921,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( @@ -1380,8 +1384,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( diff --git a/tests/test_litellm/proxy/test_prisma_migration.py b/tests/test_litellm/proxy/test_prisma_migration.py index 3fc69b34213..ebcd571d6f7 100644 --- a/tests/test_litellm/proxy/test_prisma_migration.py +++ b/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 diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index c806725d594..c451191605f 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/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 @@ -2064,6 +2065,18 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=False ) + @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") @@ -2076,8 +2089,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) @@ -2094,12 +2112,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", { @@ -2110,26 +2129,20 @@ 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=False - ) + 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=False + ) @patch("subprocess.run") @patch("atexit.register")