diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index cfa90a48ee..ab9d341aa5 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,5 +1,6 @@ -from typing import Union +from typing import Any, Awaitable, Callable, Optional, Union +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, ProxyErrorTypes, @@ -123,3 +124,138 @@ class PrismaDBExceptionHandler: ): return None raise e + + +# Default fallback timeouts when neither the caller nor the prisma_client +# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`. +# Match the auth path's existing defaults so behavior is uniform across read paths. +_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0 +_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1 + + +def _coerce_timeout(value: Any, fallback: float) -> float: + """Return `value` if it is a real int/float, else `fallback`. Guards + against tests that mock `prisma_client` and leave the timeout slots as + MagicMock instances.""" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return fallback + + +async def call_with_db_reconnect_retry( + prisma_client: Any, + coro_factory: Callable[[], Awaitable[Any]], + *, + reason: str, + timeout_seconds: Optional[float] = None, + lock_timeout_seconds: Optional[float] = None, +) -> Any: + """Run a Prisma read coroutine with one transport-reconnect-and-retry. + + The canonical "self-heal a transient DB transport blip" wrapper used by + `PrismaClient.get_generic_data` and other read paths. Mirrors the inline + pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we + have a single implementation rather than three drifting copies. + + Behavior: + 1. Await `coro_factory()`. On success, return its value. + 2. On exception, if it is NOT a transport error (per + `is_database_transport_error`), re-raise — data-layer errors like + `UniqueViolationError` mean the DB is reachable, reconnect would be + pointless. + 3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise. + This guards against partial stand-ins / older clients in tests. + 4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns + False (cooldown / lock contention / reconnect failure), re-raise. + 5. Otherwise await `coro_factory()` a second time and return / propagate + its result. At-most-one retry by construction — no infinite loop. + + `coro_factory` MUST be a zero-arg callable that returns a fresh awaitable + on each call. Passing an already-awaited coroutine would fail on retry + with `RuntimeError: cannot reuse already awaited coroutine`. + + `reason` should follow `___failure` so + telemetry distinguishes between fan-out callers (e.g. + `_update_config_from_db` issues four concurrent reads). + + Args: + prisma_client: The `PrismaClient` (or stand-in) that owns + `attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults. + coro_factory: Zero-arg callable returning the read awaitable. + reason: Telemetry tag forwarded to `attempt_db_reconnect`. + timeout_seconds: Optional override for the reconnect cycle timeout. + Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`, + then to 2.0s. + lock_timeout_seconds: Optional override for how long the helper will + wait to acquire the reconnect lock. Defaults to + `prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to + 0.1s. + + Returns: + Whatever `coro_factory()` returns (on first or second attempt). + + Raises: + Whatever `coro_factory()` raises if the failure is not a transport + error, or if the reconnect attempt does not succeed, or if the retry + also fails. + """ + try: + return await coro_factory() + except Exception as first_exc: + if not PrismaDBExceptionHandler.is_database_transport_error(first_exc): + raise + if not hasattr(prisma_client, "attempt_db_reconnect"): + raise + + resolved_timeout = _coerce_timeout( + ( + timeout_seconds + if timeout_seconds is not None + else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None) + ), + _DEFAULT_RECONNECT_TIMEOUT_SECONDS, + ) + resolved_lock_timeout = _coerce_timeout( + ( + lock_timeout_seconds + if lock_timeout_seconds is not None + else getattr( + prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None + ) + ), + _DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS, + ) + + verbose_proxy_logger.warning( + "DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s", + reason, + first_exc, + ) + + # Preserve the original transport error in telemetry. If + # `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer + # error, unexpected internal failure), surfacing that exception + # instead of `first_exc` would mask the actual DB transport problem + # in `failure_handler` / `db_exceptions` alerts. Chain the reconnect + # error as the cause for debuggability without losing the original. + try: + did_reconnect = await prisma_client.attempt_db_reconnect( + reason=reason, + timeout_seconds=resolved_timeout, + lock_timeout_seconds=resolved_lock_timeout, + ) + except Exception as reconnect_exc: + verbose_proxy_logger.warning( + "DB reconnect attempt raised; preserving original transport error. " + "reason=%s reconnect_error=%s", + reason, + reconnect_exc, + ) + raise first_exc from reconnect_exc + if not did_reconnect: + raise + + # At most one retry. If the retry also raises a transport error, we + # propagate — repeated reconnect-loops are the watchdog's job, not + # this helper's. + return await coro_factory() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3a1184c434..93cc8bfd22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import ( should_create_missing_views, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + call_with_db_reconnect_retry, +) from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( @@ -2779,30 +2782,42 @@ class PrismaClient: table_name: Literal["users", "keys", "config", "spend"], ): """ - Generic implementation of get data + Generic implementation of get data. + + Self-heals across a single transient transport blip via + `call_with_db_reconnect_retry`: on `httpx.ReadError` / + `ClientNotConnectedError` / similar, attempt one DB reconnect and + retry once before surfacing the failure. Restores the 1.82.6 behavior + that was lost in 1.83.x — see issue #25143. """ start_time = time.time() - try: + + async def _do_query(): if table_name == "users": - response = await self.db.litellm_usertable.find_first( + return await self.db.litellm_usertable.find_first( where={key: value} # type: ignore ) elif table_name == "keys": - response = await self.db.litellm_verificationtoken.find_first( # type: ignore + return await self.db.litellm_verificationtoken.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "config": - response = await self.db.litellm_config.find_first( # type: ignore + return await self.db.litellm_config.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "spend": - response = await self.db.l.find_first( # type: ignore + return await self.db.l.find_first( # type: ignore where={key: value} # type: ignore ) - return response - except Exception as e: - import traceback + return None + try: + return await call_with_db_reconnect_retry( + self, + _do_query, + reason=f"prisma_get_generic_data_{table_name}_lookup_failure", + ) + except Exception as e: error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}" verbose_proxy_logger.error(error_msg) error_msg = error_msg + "\nException Type: {}".format(type(e)) @@ -4204,7 +4219,6 @@ class PrismaClient: ) self._reap_all_zombies() self._cleanup_engine_watcher() - self._engine_confirmed_dead = False async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") @@ -4217,6 +4231,12 @@ class PrismaClient: await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + # Only clear the "dead engine" flag after the heavy reconnect + # actually completed. If `_do_heavy_reconnect()` raises (timeout, + # missing DATABASE_URL, recreate failure), the flag stays True so + # the next attempt re-enters the heavy branch instead of silently + # demoting to the lightweight path. + self._engine_confirmed_dead = False else: verbose_proxy_logger.debug( "Performing Prisma DB reconnect (engine alive or unknown)." diff --git a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py new file mode 100644 index 0000000000..ae0e1f845b --- /dev/null +++ b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py @@ -0,0 +1,255 @@ +""" +Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read, +on transport error reconnect once and retry once" helper. + +Covers the regression in issue #25143 where read paths (e.g. +`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in +LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient +`httpx.ReadError` flaps that used to self-heal in 1.82.6. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from prisma.errors import UniqueViolationError + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry + + +def _make_client( + *, + attempt_db_reconnect_return: bool = True, + has_attempt_db_reconnect: bool = True, +): + """Build a minimal stand-in for PrismaClient that exposes only the surface + `call_with_db_reconnect_retry` actually pokes at.""" + client = MagicMock() + if has_attempt_db_reconnect: + client.attempt_db_reconnect = AsyncMock( + return_value=attempt_db_reconnect_return + ) + else: + # `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock + # auto-creates attributes, so we wipe it out via `spec`. + client = MagicMock(spec=[]) + client._db_auth_reconnect_timeout_seconds = 2.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.1 + return client + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_returns_value_on_first_success(): + """Happy path: factory succeeds first call, no reconnect attempted.""" + client = _make_client() + + async def _factory(): + return {"id": 1} + + result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path") + + assert result == {"id": 1} + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_retries_after_transport_error(): + """Transport error on first call → reconnect → second call succeeds.""" + client = _make_client(attempt_db_reconnect_return=True) + + invocations = [] + + async def _factory(): + invocations.append(None) + if len(invocations) == 1: + raise httpx.ReadError("transport blip") + return {"id": 1} + + result = await call_with_db_reconnect_retry( + client, _factory, reason="prisma_get_generic_data_config_lookup_failure" + ) + + assert result == {"id": 1} + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure" + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error(): + """Data-layer errors (e.g. UniqueViolationError) are NOT transport errors — + propagate immediately, do not reconnect.""" + client = _make_client() + + async def _factory(): + raise UniqueViolationError( + data={"user_facing_error": {"meta": {}}}, + message="Unique constraint failed", + ) + + with pytest.raises(UniqueViolationError): + await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test") + + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails(): + """Transport error, but reconnect returns False → propagate the original + exception. Do not call factory a second time.""" + client = _make_client(attempt_db_reconnect_return=False) + + invocations = [] + + async def _factory(): + invocations.append(None) + raise httpx.ReadError("transport blip") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails") + + assert len(invocations) == 1 + client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error(): + """Transport error, reconnect succeeds, retry also raises transport error → + propagate. At most one retry by construction (no infinite loop).""" + client = _make_client(attempt_db_reconnect_return=True) + + invocations = [] + + async def _factory(): + invocations.append(None) + raise httpx.ReadError("still failing") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry( + client, _factory, reason="second_transport_error" + ) + + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr(): + """Older PrismaClient stand-ins / partial mocks may not expose + `attempt_db_reconnect`. The helper must not crash — just propagate the + original exception. Mirrors the `hasattr` guard from + `auth_checks._fetch_key_object_from_db_with_reconnect`.""" + client = _make_client(has_attempt_db_reconnect=False) + + async def _factory(): + raise httpx.ReadError("transport blip") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr") + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro(): + """Guard against the obvious bug of awaiting the same coroutine twice + (`RuntimeError: cannot reuse already awaited coroutine`). The helper must + call the factory a fresh time on retry, not cache an awaitable.""" + client = _make_client(attempt_db_reconnect_return=True) + + factory_call_count = 0 + + async def _factory(): + nonlocal factory_call_count + factory_call_count += 1 + if factory_call_count == 1: + raise httpx.ReadError("transport blip") + return "ok" + + result = await call_with_db_reconnect_retry( + client, _factory, reason="fresh_coro_on_retry" + ) + + assert result == "ok" + assert factory_call_count == 2 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_passes_explicit_timeouts(): + """Explicit timeout_seconds / lock_timeout_seconds override the auth + defaults read off the prisma_client object.""" + client = _make_client(attempt_db_reconnect_return=True) + + async def _factory(): + if not hasattr(_factory, "_called"): + _factory._called = True # type: ignore[attr-defined] + raise httpx.ReadError("transport blip") + return "ok" + + result = await call_with_db_reconnect_retry( + client, + _factory, + reason="explicit_timeouts", + timeout_seconds=5.5, + lock_timeout_seconds=0.25, + ) + + assert result == "ok" + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["timeout_seconds"] == 5.5 + assert call_kwargs["lock_timeout_seconds"] == 0.25 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset(): + """When timeouts are not provided, helper reads + `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds` + off the prisma_client (matching the auth path's existing convention).""" + client = _make_client(attempt_db_reconnect_return=True) + client._db_auth_reconnect_timeout_seconds = 3.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.5 + + async def _factory(): + if not hasattr(_factory, "_called"): + _factory._called = True # type: ignore[attr-defined] + raise httpx.ReadError("transport blip") + return "ok" + + await call_with_db_reconnect_retry(client, _factory, reason="defaults") + + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["timeout_seconds"] == 3.0 + assert call_kwargs["lock_timeout_seconds"] == 0.5 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises(): + """If `attempt_db_reconnect` itself raises (lock cancellation, timer + error, unexpected internal failure), the helper must surface the + *original* transport error to telemetry — not the reconnect exception. + Otherwise `failure_handler` / `db_exceptions` alerts log the wrong + error string and the actual DB transport problem becomes invisible. + + The reconnect error is chained as the `__cause__` for debuggability.""" + client = MagicMock() + reconnect_exc = RuntimeError("simulated reconnect lock cancellation") + client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc) + client._db_auth_reconnect_timeout_seconds = 2.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + original_exc = httpx.ReadError("transport blip") + + async def _factory(): + raise original_exc + + with pytest.raises(httpx.ReadError) as exc_info: + await call_with_db_reconnect_retry( + client, _factory, reason="reconnect_itself_raises" + ) + + assert exc_info.value is original_exc + assert exc_info.value.__cause__ is reconnect_exc + client.attempt_db_reconnect.assert_awaited_once() 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 fb215e5477..af65cad09e 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -5,6 +5,7 @@ import sys import time from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest sys.path.insert( @@ -358,3 +359,125 @@ async def test_lightweight_reconnect_skips_kill_on_successful_disconnect( await client._run_reconnect_cycle(timeout_seconds=5.0) mock_kill.assert_not_called() + + +# --------------------------------------------------------------------------- +# get_generic_data: transport-reconnect-and-retry coverage (issue #25143) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_generic_data_retries_on_transport_error_for_config_table( + mock_proxy_logging, +): + """`get_generic_data(table_name="config")` self-heals on a transient + `httpx.ReadError`: reconnect once, retry once, return the row. + + Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry + branch that 1.82.6 had on this method. `_update_config_from_db` fans out + four concurrent `get_generic_data` calls, so a single transport flap used + to surface as four `db_exceptions` alerts and a stale config window. + """ + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}} + invocations: list[None] = [] + + async def _flaky_find_first(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise httpx.ReadError("simulated transport blip") + return expected_row + + client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first) + client.attempt_db_reconnect = AsyncMock(return_value=True) + + result = await client.get_generic_data( + key="param_name", + value="general_settings", + table_name="config", + ) + + assert result == expected_row + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure" + + # The failure_handler telemetry side-effect must NOT fire on the first + # transport blip — only if the post-retry call also fails. Drain the + # event loop so any spuriously-spawned task would have run by now. + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging): + """If reconnect itself does not succeed, propagate the original transport + error and let the existing failure_handler / db_exceptions telemetry fire.""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + client.db.litellm_config.find_first = AsyncMock( + side_effect=httpx.ReadError("simulated transport blip") + ) + client.attempt_db_reconnect = AsyncMock(return_value=False) + + with pytest.raises(httpx.ReadError): + await client.get_generic_data( + key="param_name", + value="general_settings", + table_name="config", + ) + + client.attempt_db_reconnect.assert_awaited_once() + # Failure telemetry IS expected here — the read genuinely failed. + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_called_once() + + +# --------------------------------------------------------------------------- +# _engine_confirmed_dead flag-reset bug (B2) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect( + mock_proxy_logging, +): + """Regression test for the flag-reset bug. + + Before the fix, `_run_reconnect_cycle` cleared + `self._engine_confirmed_dead = False` *before* awaiting + `_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout, + missing DATABASE_URL, recreate failure), the flag was left cleared and the + next attempt could demote to the lightweight path even though the engine + was genuinely dead. + + The fix moves the reset into the success branch — the flag must stay True + when heavy reconnect raises. + """ + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client._engine_confirmed_dead = True + client._engine_pid = 0 # so `_is_engine_alive` is not consulted + + # Make the heavy reconnect path raise. + client.db.recreate_prisma_client = AsyncMock( + side_effect=RuntimeError("simulated heavy reconnect failure") + ) + client._start_engine_watcher = AsyncMock() + client._cleanup_engine_watcher = MagicMock() + client._reap_all_zombies = MagicMock() + + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + with pytest.raises(Exception): + await client._run_reconnect_cycle(timeout_seconds=5.0) + + # The flag must STILL be True so the next attempt re-enters the heavy + # branch instead of silently demoting to the lightweight path. + assert client._engine_confirmed_dead is True