From 2a9843e649a4336927646c237632b97acc451e59 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 5 Aug 2026 12:27:49 -0700 Subject: [PATCH] fix(proxy): keep the connected DB client when a startup health check fails (#35837) `_setup_prisma_client` ran `connect()`, then a `SELECT 1` health check, then armed the DB health watchdog. Any failure fell into one handler that, with `allow_requests_on_db_unavailable` set, swallowed the error and returned None, which the caller assigns to the module-level `prisma_client`. A single transient timeout on that health check therefore discarded a client that had already connected, for the life of the process, and skipped the watchdog that exists to reconnect it. The watchdog now starts before the health check, and a swallowed post-connect failure returns the connected client instead of None. A client whose `connect()` failed is still discarded, and startup still hard-fails when `allow_requests_on_db_unavailable` is not set. The same check also misreported its own failure. `health_check()` labelled its error `disconnect()`, a copy-paste from the real `disconnect()` below it, so grepping the logs for the health check turned up nothing and read as "the check never ran". Both it and the sibling `connect()` failure reported through `print_verbose`, which reaches `verbose_proxy_logger.debug` and otherwise prints only under the deprecated `litellm.set_verbose`, leaving a startup-blocking database fault invisible at the verbosity operators actually run. Both now log at warning under their own names. The proxy logger's handler carries the secret redaction filter, so a connection string in the exception text is redacted exactly as it was on the old print path. --- litellm/proxy/proxy_server.py | 72 ++++++----- litellm/proxy/utils.py | 6 +- tests/test_litellm/proxy/test_proxy_server.py | 120 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 76 +++++++++++ 4 files changed, 239 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61c6ce22a91..8d6b930591b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8664,48 +8664,56 @@ class ProxyStartupEvent: - Sets up prisma client - Adds necessary views to proxy """ + connected_client: PrismaClient | None = None try: - prisma_client: PrismaClient | None = None - if database_url is not None: - try: - prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) - except Exception as e: - raise e + if database_url is None: + return None - try: - await prisma_client.connect() - except Exception as e: - if "P3018" in str(e) or "P3009" in str(e): - verbose_proxy_logger.debug("CRITICAL: DATABASE MIGRATION FAILED") - verbose_proxy_logger.debug("Your database is in a 'dirty' state.") - verbose_proxy_logger.debug("FIX: Run 'prisma migrate resolve --applied '") - raise e + prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) - ## Start RDS IAM token refresh background task if enabled ## - # This proactively refreshes IAM tokens before they expire, - # preventing the 15-minute connection failure bug (#16220) - if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): - await prisma_client.db.start_token_refresh_task() + try: + await prisma_client.connect() + except Exception as e: + if "P3018" in str(e) or "P3009" in str(e): + verbose_proxy_logger.debug("CRITICAL: DATABASE MIGRATION FAILED") + verbose_proxy_logger.debug("Your database is in a 'dirty' state.") + verbose_proxy_logger.debug("FIX: Run 'prisma migrate resolve --applied '") + raise e - ## Add necessary views to proxy ## - asyncio.create_task( - prisma_client.check_view_exists() - ) # check if all necessary views exist. Don't block execution + connected_client = prisma_client - asyncio.create_task( - prisma_client._set_spend_logs_row_count_in_proxy_state() - ) # set the spend logs row count in proxy state. Don't block execution + ## Start RDS IAM token refresh background task if enabled ## + # This proactively refreshes IAM tokens before they expire, + # preventing the 15-minute connection failure bug (#16220) + if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): + await prisma_client.db.start_token_refresh_task() - # run a health check to ensure the DB is ready - if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True: - await prisma_client.health_check() + ## Add necessary views to proxy ## + asyncio.create_task( + prisma_client.check_view_exists() + ) # check if all necessary views exist. Don't block execution + + asyncio.create_task( + prisma_client._set_spend_logs_row_count_in_proxy_state() + ) # set the spend logs row count in proxy state. Don't block execution + + if hasattr(prisma_client, "start_db_health_watchdog_task"): + await prisma_client.start_db_health_watchdog_task() + + # run a health check to ensure the DB is ready + if get_secret_bool("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", False) is not True: + await prisma_client.health_check() - if hasattr(prisma_client, "start_db_health_watchdog_task"): - await prisma_client.start_db_health_watchdog_task() return prisma_client except Exception as e: PrismaDBExceptionHandler.handle_db_exception(e) - return None + if connected_client is not None: + verbose_proxy_logger.warning( + "Retaining the connected Prisma client after a post-connect startup step failed: %s. " + "The DB health watchdog keeps probing and reconnects once the database recovers.", + e, + ) + return connected_client @classmethod def _init_dd_tracer(cls): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8d638dedff8..7717b4da1af 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4269,7 +4269,7 @@ class PrismaClient: import traceback error_msg: Final = f"LiteLLM Prisma Client Exception connect(): {e}" - print_verbose(error_msg) + verbose_proxy_logger.warning(error_msg) error_traceback: Final = error_msg + "\n" + traceback.format_exc() end_time: Final = time.time() _duration: Final = end_time - start_time @@ -4987,8 +4987,8 @@ class PrismaClient: except Exception as e: import traceback - error_msg: Final = f"LiteLLM Prisma Client Exception disconnect(): {e}" - print_verbose(error_msg) + error_msg: Final = f"LiteLLM Prisma Client Exception health_check(): {e}" + verbose_proxy_logger.warning(error_msg) error_traceback: Final = error_msg + "\n" + traceback.format_exc() end_time: Final = time.time() _duration: Final = end_time - start_time diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ede93dc0c58..4a491ec0cff 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11018,3 +11018,123 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={}) assert MOCK_TESTING_CONFIG_KEY not in caplog.text + + +def _mock_startup_prisma_client(health_check_error=None, connect_error=None): + client = MagicMock() + client.connect = AsyncMock(side_effect=connect_error) + client.db.start_token_refresh_task = AsyncMock() + client.check_view_exists = AsyncMock() + client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + client.start_db_health_watchdog_task = AsyncMock() + client.health_check = AsyncMock(side_effect=health_check_error) + return client + + +async def _run_setup_prisma_client(mock_client): + from litellm.proxy.proxy_server import ProxyStartupEvent + + with patch.object(proxy_server_module, "PrismaClient", return_value=mock_client): + result = await ProxyStartupEvent._setup_prisma_client( + database_url="postgresql://litellm:litellm@localhost:5432/litellm", + proxy_logging_obj=MagicMock(), + user_api_key_cache=DualCache(), + ) + await asyncio.sleep(0.05) + return result + + +@pytest.mark.asyncio +async def test_setup_prisma_client_retains_connected_client_when_startup_health_check_fails( + monkeypatch, +): + """A transient failure of the startup ``SELECT 1`` must not discard a client + whose ``connect()`` already succeeded. + + Discarding it assigns ``None`` to the module-level ``prisma_client`` for the + life of the process, so a database that came back a second later is never + used again until the proxy is restarted.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + + mock_client = _mock_startup_prisma_client( + health_check_error=httpx.ReadTimeout("startup health check timed out") + ) + result = await _run_setup_prisma_client(mock_client) + + assert mock_client.connect.await_count == 1 + assert mock_client.health_check.await_count == 1 + assert result is mock_client + + +@pytest.mark.asyncio +async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_check( + monkeypatch, +): + """The health watchdog is the only thing that reconnects a dropped DB, so it + has to be armed before the startup health check can fail. + + Armed after, the single failure it exists to recover from is exactly the one + that skips it, and recovery never happens.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + + mock_client = _mock_startup_prisma_client( + health_check_error=httpx.ReadTimeout("startup health check timed out") + ) + call_order = MagicMock() + call_order.attach_mock(mock_client.start_db_health_watchdog_task, "watchdog") + call_order.attach_mock(mock_client.health_check, "health_check") + + await _run_setup_prisma_client(mock_client) + + assert mock_client.start_db_health_watchdog_task.await_count == 1 + assert [call[0] for call in call_order.mock_calls] == ["watchdog", "health_check"] + + +@pytest.mark.asyncio +async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(monkeypatch): + """Without ``allow_requests_on_db_unavailable`` a failed startup health check + must still hard-fail startup. Retaining the client is a fallback for + operators who opted into serving traffic without a database, never a way to + boot a proxy whose DB never answered.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": False}, + ) + + mock_client = _mock_startup_prisma_client( + health_check_error=httpx.ReadTimeout("startup health check timed out") + ) + with pytest.raises(httpx.ReadTimeout): + await _run_setup_prisma_client(mock_client) + + +@pytest.mark.asyncio +async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkeypatch): + """Retaining only ever applies to a client that connected. If ``connect()`` + failed there is no usable client and no watchdog to recover it, so the caller + must still get ``None``.""" + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + monkeypatch.setattr( + proxy_server_module, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + + mock_client = _mock_startup_prisma_client(connect_error=httpx.ConnectError("connection refused")) + result = await _run_setup_prisma_client(mock_client) + + assert result is None + assert mock_client.start_db_health_watchdog_task.await_count == 0 + assert mock_client.health_check.await_count == 0 diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 3421751d962..abd6220144b 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1085,3 +1085,79 @@ async def test_post_mcp_call_hook_propagates_guardrail_block(restore_callbacks): request_data={"mcp_tool_name": "echo"}, user_api_key_dict=None, ) + + +@pytest.mark.asyncio +async def test_prisma_health_check_failure_names_itself_at_operator_visible_level(caplog): + """A failing DB health check has to name the check that failed, at a level + operators actually run at. + + Reporting it as ``disconnect()`` sends anyone grepping the logs to the wrong + function and reads as "the check never ran", and reporting it only at debug + level hides a database fault behind a flag nobody enables in production.""" + import logging + from unittest.mock import AsyncMock + + from litellm.proxy.utils import PrismaClient + + client = MagicMock() + client.db.query_raw = AsyncMock(side_effect=Exception("connection refused")) + client.proxy_logging_obj.failure_handler = AsyncMock() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(Exception, match="connection refused"): + await PrismaClient.health_check(client) + + assert "health_check()" in caplog.text + assert "disconnect()" not in caplog.text + assert "connection refused" in caplog.text + + +@pytest.mark.asyncio +async def test_prisma_connect_failure_is_reported_at_operator_visible_level(caplog): + """The sibling connect failure is labelled correctly but was equally + invisible. A database the proxy could not connect to at startup must not be + a debug-only record.""" + import logging + from unittest.mock import AsyncMock + + from litellm.proxy.utils import PrismaClient + + client = MagicMock() + client.db.is_connected = MagicMock(return_value=False) + client.db.connect = AsyncMock(side_effect=Exception("could not reach database")) + client.proxy_logging_obj.failure_handler = AsyncMock() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(Exception, match="could not reach database"): + await PrismaClient.connect(client) + + assert "connect()" in caplog.text + assert "could not reach database" in caplog.text + + +@pytest.mark.asyncio +async def test_prisma_health_check_failure_redacts_database_credentials(caplog): + """Raising the level must not widen what reaches the logs. The exception + text can carry a full connection string, so the credential has to be gone + from the emitted record.""" + import logging + from unittest.mock import AsyncMock + + from litellm.proxy.utils import PrismaClient + + client = MagicMock() + client.db.query_raw = AsyncMock( + side_effect=Exception("could not connect to postgresql://admin:hunter2@db.internal:5432/litellm") + ) + client.proxy_logging_obj.failure_handler = AsyncMock() + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(Exception): + await PrismaClient.health_check(client) + + emitted = [record.getMessage() for record in caplog.records if record.name == "LiteLLM Proxy"] + + assert emitted + assert all("hunter2" not in message for message in emitted) + assert any("postgresql://REDACTED@db.internal" in message for message in emitted)