diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index fc761fc1831..25fab8700c2 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -51,6 +51,8 @@ class _PrismaEngine(Protocol): async def rollback_transaction(self, tx_id: str) -> None: ... + async def request(self, method: str, path: str) -> object: ... + class _PrismaClient(Protocol): _Prisma__engine: _PrismaEngine @@ -308,6 +310,10 @@ class PrismaWrapper: await self._original_prisma.connect(timeout) self._active_drain_tracker = self._instrument_prisma_client(self._original_prisma) + async def query_engine_status(self) -> object: + engine: Final = self._read_engine(self._original_prisma) + return await engine.request("GET", "/status") + @staticmethod async def _kill_engine_process(pid: int) -> None: """Force-kill the engine subprocess to prevent DB connection pool leaks. diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c23cde6052e..bcea658427a 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -33,7 +33,6 @@ from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, @@ -1339,21 +1338,11 @@ async def _db_health_readiness_check(): await prisma_client.health_check() db_health_cache = {"status": "connected", "last_updated": datetime.now()} return db_health_cache - except Exception as e: + except Exception: db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} - if PrismaDBExceptionHandler.is_database_transport_error(e): - try: - verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect") - await prisma_client.attempt_db_reconnect(reason="health_readiness_check") - await prisma_client.health_check() - verbose_proxy_logger.info("_db_health_readiness_check: reconnect succeeded") - db_health_cache = { - "status": "connected", - "last_updated": datetime.now(), - } - return db_health_cache - except Exception: - verbose_proxy_logger.error("_db_health_readiness_check: reconnect failed") + verbose_proxy_logger.warning( + "_db_health_readiness_check: health_check failed; returning disconnected without reconnect" + ) return db_health_cache diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8cbf5b685fd..a381da82e01 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5468,7 +5468,7 @@ class PrismaClient: async def start_db_health_watchdog_task(self) -> None: """Start background tasks that monitor DB health: - - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. + - A periodic query-engine status probe that triggers reconnect on local engine failure. - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. """ if self._db_health_watchdog_enabled is not True: @@ -5509,7 +5509,7 @@ class PrismaClient: try: await asyncio.sleep(self._db_health_watchdog_interval_seconds) await asyncio.wait_for( - self.db.query_raw("SELECT 1"), + self.writer_db.query_engine_status(), timeout=self._db_health_watchdog_probe_timeout_seconds, ) if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable: diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 10a48941693..e9f4019fa41 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -269,13 +269,15 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget( @pytest.mark.asyncio -async def test_db_health_watchdog_should_trigger_reconnect_on_db_error( +async def test_db_health_watchdog_should_trigger_reconnect_on_engine_error( mock_proxy_logging, ): client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.query_raw = AsyncMock(side_effect=Exception("db connection dropped")) + client.writer_db.query_engine_status = AsyncMock( + side_effect=Exception("query engine unavailable") + ) client.attempt_db_reconnect = AsyncMock(return_value=True) client._db_health_watchdog_interval_seconds = 1 client._db_watchdog_reconnect_timeout_seconds = 7.0 @@ -306,7 +308,9 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout( client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.query_raw = AsyncMock(side_effect=asyncio.TimeoutError()) + client.writer_db.query_engine_status = AsyncMock( + side_effect=asyncio.TimeoutError() + ) client.attempt_db_reconnect = AsyncMock(return_value=True) client._db_health_watchdog_interval_seconds = 1 client._db_watchdog_reconnect_timeout_seconds = 9.0 @@ -563,7 +567,7 @@ async def test_db_health_watchdog_should_reconnect_degraded_writer( ) writer = MagicMock() reader = MagicMock() - reader.query_raw = AsyncMock(return_value=[{"result": 1}]) + writer.query_engine_status = AsyncMock(return_value={"status": "ok"}) routing = RoutingPrismaWrapper(writer=writer, reader=reader) routing._writer_unavailable = True client.db = routing @@ -596,7 +600,7 @@ async def test_db_health_watchdog_should_not_reconnect_healthy_writer( ) writer = MagicMock() reader = MagicMock() - reader.query_raw = AsyncMock(return_value=[{"result": 1}]) + writer.query_engine_status = AsyncMock(return_value={"status": "ok"}) routing = RoutingPrismaWrapper(writer=writer, reader=reader) client.db = routing client.attempt_db_reconnect = AsyncMock(return_value=True) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 72c6f77a4f7..361727c4e21 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -141,9 +141,7 @@ async def test_db_health_transport_error_never_raises(transport_error): result = await _db_health_readiness_check() assert result["status"] == "disconnected" - mock_prisma.attempt_db_reconnect.assert_called_once_with( - reason="health_readiness_check" - ) + mock_prisma.attempt_db_reconnect.assert_not_called() @pytest.mark.asyncio @@ -155,49 +153,14 @@ async def test_db_health_transport_error_never_raises(transport_error): HTTPClientClosedError(), ], ) -async def test_db_health_transport_error_reconnect_succeeds(transport_error): +async def test_db_health_transport_error_does_not_reconnect(transport_error): """ - When health_check raises a transport error and attempt_db_reconnect - succeeds, the second health_check passes and we return 'connected'. - """ - mock_prisma = MagicMock() - mock_prisma.health_check = AsyncMock(side_effect=[transport_error, None]) - mock_prisma.attempt_db_reconnect = AsyncMock(return_value=True) - - _health_endpoints_module.db_health_cache = { - "status": "connected", - "last_updated": datetime.now() - timedelta(seconds=20), - } - - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): - result = await _db_health_readiness_check() - - assert result["status"] == "connected" - mock_prisma.attempt_db_reconnect.assert_called_once_with( - reason="health_readiness_check" - ) - assert mock_prisma.health_check.call_count == 2 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "transport_error", - [ - httpx.ConnectError("All connection attempts failed"), - ClientNotConnectedError(), - HTTPClientClosedError(), - ], -) -async def test_db_health_transport_error_reconnect_fails(transport_error): - """ - When health_check raises a transport error and attempt_db_reconnect also - fails, return 'disconnected' without raising. + A readiness probe reports a transport error as disconnected without + rebuilding the Prisma query engine or retrying the probe inline. """ mock_prisma = MagicMock() mock_prisma.health_check = AsyncMock(side_effect=transport_error) - mock_prisma.attempt_db_reconnect = AsyncMock( - side_effect=RuntimeError("reconnect failed") - ) + mock_prisma.attempt_db_reconnect = AsyncMock() _health_endpoints_module.db_health_cache = { "status": "connected", @@ -208,6 +171,38 @@ async def test_db_health_transport_error_reconnect_fails(transport_error): result = await _db_health_readiness_check() assert result["status"] == "disconnected" + mock_prisma.attempt_db_reconnect.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transport_error", + [ + httpx.ConnectError("All connection attempts failed"), + ClientNotConnectedError(), + HTTPClientClosedError(), + ], +) +async def test_db_health_transport_error_keeps_disconnected_cache(transport_error): + """ + A failed readiness probe refreshes the disconnected cache timestamp so + callers receive the current state without launching reconnect work. + """ + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=transport_error) + mock_prisma.attempt_db_reconnect = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "connected", + "last_updated": datetime.now() - timedelta(seconds=20), + } + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await _db_health_readiness_check() + + assert result["status"] == "disconnected" + assert _health_endpoints_module.db_health_cache["status"] == "disconnected" + mock_prisma.attempt_db_reconnect.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 719d7cc73f5..0f74d24e653 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -447,58 +447,60 @@ async def test_stop_db_health_watchdog_task_noop_when_no_task( async def test_db_health_watchdog_loop_triggers_reconnect_on_timeout( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: - """The watchdog loop reconnects when ``wait_for`` raises TimeoutError - or a recognized DB connection error. - """ + """A local query-engine status timeout triggers reconnect.""" prisma_client._db_health_watchdog_interval_seconds = 0 prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + prisma_client.writer_db.query_engine_status = AsyncMock( + side_effect=asyncio.TimeoutError() + ) - call_count = {"n": 0} - - async def _timeout_then_cancel(*args: Any, **kwargs: Any) -> None: - call_count["n"] += 1 - if call_count["n"] >= 2: - raise asyncio.CancelledError() - raise asyncio.TimeoutError() - - monkeypatch.setattr("asyncio.wait_for", _timeout_then_cancel) - await prisma_client._db_health_watchdog_loop() + with monkeypatch.context() as patch_context: + patch_context.setattr( + asyncio, + "sleep", + AsyncMock(side_effect=[None, asyncio.CancelledError()]), + ) + await prisma_client._db_health_watchdog_loop() pinned = { "reconnect_called": prisma_client.attempt_db_reconnect.await_count, "reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs[ "reason" ], - "wait_for_calls": call_count["n"], + "status_calls": prisma_client.writer_db.query_engine_status.await_count, "loop_exited_clean": True, } assert pinned == { "reconnect_called": 1, "reconnect_reason": "db_health_watchdog_connection_error", - "wait_for_calls": 2, + "status_calls": 1, "loop_exited_clean": True, } @pytest.mark.asyncio -async def test_db_health_watchdog_loop_swallows_non_db_errors( +async def test_db_health_watchdog_loop_probes_engine_status_without_querying_db( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: - """A non-DB error during the probe should NOT trigger reconnect; the - loop logs and continues until cancellation. - """ + """A healthy local engine does not issue a PostgreSQL query.""" prisma_client._db_health_watchdog_interval_seconds = 0 prisma_client.attempt_db_reconnect = AsyncMock() + prisma_client.writer_db.query_engine_status = AsyncMock( + return_value={"status": "ok"} + ) + prisma_client.db.query_raw = AsyncMock( + side_effect=AssertionError("watchdog must not query PostgreSQL") + ) - call_count = {"n": 0} + with monkeypatch.context() as patch_context: + patch_context.setattr( + asyncio, + "sleep", + AsyncMock(side_effect=[None, asyncio.CancelledError()]), + ) + await prisma_client._db_health_watchdog_loop() - async def _raise_then_cancel(*args: Any, **kwargs: Any) -> None: - call_count["n"] += 1 - if call_count["n"] >= 2: - raise asyncio.CancelledError() - raise ValueError("not a db error") - - monkeypatch.setattr("asyncio.wait_for", _raise_then_cancel) - await prisma_client._db_health_watchdog_loop() + prisma_client.writer_db.query_engine_status.assert_awaited_once() + prisma_client.db.query_raw.assert_not_awaited() assert prisma_client.attempt_db_reconnect.await_count == 0