diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d7649b524aa..fe5a7baff88 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4748,6 +4748,41 @@ class PrismaClient: else: verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e) + def _prisma_wrappers(self) -> tuple[PrismaWrapper, ...]: + """Underlying writer (and reader, when read-replica routing is on) + wrappers, so callers can inspect their planned-recreate coordination + state regardless of routing.""" + if isinstance(self.db, RoutingPrismaWrapper): + return (self.db.writer, self.db.reader) + if isinstance(self.db, PrismaWrapper): + return (self.db,) + return () + + def _engine_generations(self) -> tuple[int, ...]: + return tuple(w._engine_generation for w in self._prisma_wrappers()) + + def _is_planned_engine_recreate_error(self, e: Exception, generations_before: tuple[int, ...]) -> bool: + """True iff ``e`` is a transient connection error that raced a planned + Prisma engine recreate rather than a real DB outage. + + Every `recreate_prisma_client` (RDS IAM token refresh, guarded + reconnect) kills the query engine on purpose and holds the wrapper's + `_reconnection_lock` while it spawns and connects a replacement, bumping + `_engine_generation` on success. A `SELECT 1` probe that races that + ~0.5-1s window fails with a connection error that should not surface as + a DB exception. We treat the error as planned when a recreate is still + in flight (a lock is held) or when the engine generation moved while the + probe ran. Real outages persist past the recreate window with the lock + free and the generation unchanged, so they are still reported. + """ + if not PrismaDBExceptionHandler.is_database_transport_error(e): + return False + for wrapper in self._prisma_wrappers(): + lock = getattr(wrapper, "_reconnection_lock", None) + if isinstance(lock, asyncio.Lock) and lock.locked(): + return True + return self._engine_generations() != generations_before + @backoff.on_exception( backoff.expo, Exception, @@ -4760,6 +4795,7 @@ class PrismaClient: Health check endpoint for the prisma client """ start_time = time.time() + generations_before = self._engine_generations() try: sql_query = "SELECT 1" @@ -4770,19 +4806,26 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception disconnect(): {str(e)}" + error_msg = f"LiteLLM Prisma Client Exception health_check(): {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() end_time = time.time() _duration = end_time - start_time - asyncio.create_task( - self.proxy_logging_obj.failure_handler( - original_exception=e, - duration=_duration, - call_type="health_check", - traceback_str=error_traceback, + if self._is_planned_engine_recreate_error(e, generations_before): + verbose_proxy_logger.debug( + "health_check: SELECT 1 raced a planned Prisma engine recreate (%s); " + "not reporting as a DB exception.", + str(e), + ) + else: + asyncio.create_task( + self.proxy_logging_obj.failure_handler( + original_exception=e, + duration=_duration, + call_type="health_check", + traceback_str=error_traceback, + ) ) - ) raise e async def _get_spend_logs_row_count(self) -> int: 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 7f723fa3ae0..28af01c7402 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -641,3 +641,82 @@ async def test_direct_reconnect_probe_success_clears_writer_unavailable( writer.query_raw.assert_awaited_once_with("SELECT 1") assert routing.writer_unavailable is False + + +@pytest.fixture +def no_backoff_sleep(): + """Make health_check's @backoff.on_exception retries run instantly.""" + with patch("backoff._async.asyncio.sleep", new=AsyncMock()): + yield + + +@pytest.mark.asyncio +async def test_health_check_not_alerted_while_recreate_lock_held( + mock_proxy_logging, no_backoff_sleep +): + """A SELECT 1 that races a planned engine recreate (wrapper reconnection + lock held) must not be reported to failure_handler as a DB exception.""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + + await client.db._reconnection_lock.acquire() + try: + with pytest.raises(httpx.ConnectError): + await client.health_check() + finally: + client.db._reconnection_lock.release() + + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_health_check_not_alerted_when_engine_generation_moved( + mock_proxy_logging, no_backoff_sleep +): + """A recreate that completes (engine generation bumped) while the probe is + in flight must not be reported, even though the lock is already free.""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + def _bump_then_fail(*args, **kwargs): + client.db._engine_generation += 1 + raise httpx.ConnectError("All connection attempts failed") + + client.db.query_raw = AsyncMock(side_effect=_bump_then_fail) + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_health_check_reports_real_outage_with_correct_label( + mock_proxy_logging, no_backoff_sleep +): + """A persistent connection error with no recreate in flight is a real + outage: it must still be reported, and the traceback must be labeled + health_check() rather than the mislabeled disconnect().""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + await asyncio.sleep(0.05) + mock_proxy_logging.failure_handler.assert_called() + _, kwargs = mock_proxy_logging.failure_handler.call_args + assert kwargs["call_type"] == "health_check" + assert "health_check()" in kwargs["traceback_str"] + assert "disconnect()" not in kwargs["traceback_str"]