From 4714128da5f59e993a21f53518dc77a6eca5dd9c Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 10 Mar 2026 18:05:13 +0530 Subject: [PATCH 1/3] fix: fail proxy startup if prisma migrate fails Co-Authored-By: Claude Haiku 4.5 --- litellm/proxy/proxy_cli.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 921d86c35c1..5f6ccfd9ba3 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -853,7 +853,14 @@ def run_server( # noqa: PLR0915 ): check_prisma_schema_diff(db_url=None) else: - PrismaManager.setup_database(use_migrate=not use_prisma_db_push) + if not PrismaManager.setup_database(use_migrate=not use_prisma_db_push): + import sys + + print( # noqa + "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " + "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" + ) + sys.exit(1) else: print( # noqa f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa From aae2deb839c8f39f82a58140f2ffa97b42d97a61 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 10 Mar 2026 18:15:11 +0530 Subject: [PATCH 2/3] fix: remove redundant import and add test for startup failure - Remove redundant `import sys` (already imported at module level) - Add test_startup_fails_when_db_setup_fails verifying sys.exit(1) when PrismaManager.setup_database returns False Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/proxy_cli.py | 2 - tests/test_litellm/proxy/test_proxy_cli.py | 57 ++++++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 5f6ccfd9ba3..be2f5ac7c11 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -854,8 +854,6 @@ def run_server( # noqa: PLR0915 check_prisma_schema_diff(db_url=None) else: if not PrismaManager.setup_database(use_migrate=not use_prisma_db_push): - import sys - print( # noqa "\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. " "The proxy cannot start safely. Please check your database connection and migration status.\033[0m" diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index cf6511c18a5..bed77b43c1d 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -664,6 +664,63 @@ class TestHealthAppFactory: ) mock_setup_database.assert_called_with(use_migrate=False) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_startup_fails_when_db_setup_fails( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + ): + """Test that proxy exits with code 1 when PrismaManager.setup_database returns False""" + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + mock_setup_database.return_value = False + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + + with patch.dict( + os.environ, clean_env, clear=True + ), patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + with pytest.raises(SystemExit) as exc_info: + run_server.main( + ["--local", "--skip_server_startup"], standalone_mode=False + ) + assert exc_info.value.code == 1 + # --- Module-level helpers for worker startup hook tests --- From 2f5a553a7d8e8cad9035f3b3c8e22cd6b16d1732 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 10 Mar 2026 18:31:04 +0530 Subject: [PATCH 3/3] test: assert setup_database called with correct args Co-Authored-By: Claude Opus 4.6 --- tests/test_litellm/proxy/test_proxy_cli.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index bed77b43c1d..642d21a42f7 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -720,6 +720,7 @@ class TestHealthAppFactory: ["--local", "--skip_server_startup"], standalone_mode=False ) assert exc_info.value.code == 1 + mock_setup_database.assert_called_once_with(use_migrate=True) # --- Module-level helpers for worker startup hook tests ---