from collections.abc import Awaitable, Callable from typing import Any, Final, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, ProxyErrorTypes, ProxyException, ) from litellm.secret_managers.main import str_to_bool # Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain. # Real exception chains are a few links deep; the cap also makes the walk cycle-safe. _MAX_EXCEPTION_CHAIN_DEPTH: Final = 20 class PrismaDBExceptionHandler: """ Class to handle DB Exceptions or Connection Errors """ @staticmethod def should_allow_request_on_db_unavailable() -> bool: """ Returns True if the request should be allowed to proceed despite the DB connection error """ from litellm.proxy.proxy_server import general_settings _allow_requests_on_db_unavailable: bool | str = general_settings.get("allow_requests_on_db_unavailable", False) if isinstance(_allow_requests_on_db_unavailable, bool): return _allow_requests_on_db_unavailable if str_to_bool(_allow_requests_on_db_unavailable) is True: return True return False @staticmethod def is_database_connection_error(e: Exception) -> bool: """True only for a database that is temporarily unreachable and is expected to come back on its own. This is the gate for ``allow_requests_on_db_unavailable``, which lets the proxy keep serving, and issue fallback identities, without a verified database. Only a transient outage justifies that. A fault that will never resolve by itself, such as a query engine that is missing or version-skewed, a malformed query the client library built, or a transaction used incorrectly, must surface rather than be absorbed into an indefinite degraded mode. Membership is an allowlist, so an unrecognized failure is treated as permanent. A genuine outage reaches the caller as one of ``DB_CONNECTION_ERROR_TYPES``: the engine is a local HTTP server, and an unreachable database surfaces as a transport error against it rather than as a prisma type. Reporting decisions want the opposite breadth; use ``is_database_infrastructure_error`` for those. """ import prisma.engine.errors if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True if isinstance(e, prisma.engine.errors.EngineConnectionError): return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @staticmethod def is_database_infrastructure_error(e: Exception) -> bool: """True when the failure came from the database or its query engine rather than from the caller's request. This answers a reporting question, not a serving one: should the caller be told the service is at fault, or that their request was. It stays deliberately broad, because a permanently faulted engine is still a service problem, and reporting one as a credential failure sends an operator looking in the wrong place. Widening it can only change which error a caller sees; it never grants access. Known data-layer PrismaError subclasses (``UniqueViolationError``, ``RecordNotFoundError``, etc.) are excluded — the DB IS reachable and the request itself is what failed. """ import prisma data_layer_errors: Final = ( prisma.errors.DataError, prisma.errors.UniqueViolationError, prisma.errors.ForeignKeyViolationError, prisma.errors.MissingRequiredValueError, prisma.errors.RawQueryError, prisma.errors.TableNotFoundError, prisma.errors.RecordNotFoundError, ) if isinstance(e, data_layer_errors): return False if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True if isinstance(e, prisma.errors.PrismaError): return True if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True return False @staticmethod def is_prisma_data_error(e: Exception) -> bool: """True iff ``e`` is a base prisma ``DataError``: the database processed the statement and refused the data itself (e.g. ``invalid byte sequence for encoding "UTF8": 0x00``), as opposed to a connectivity failure. Matched by exact type, not ``isinstance``: the specific data-layer subclasses (``UniqueViolationError``, ``TableNotFoundError``, ``MissingRequiredValueError`` ...) all derive from ``DataError`` but carry their own semantics, and a systemic one like a missing table must not be mistaken for a single poison row and bisected away. A raw Postgres execution error with no prisma P-code surfaces as the base ``DataError``. prisma also wraps the P1001 "can't reach database server" outage as a base ``DataError``, so a caller that must not treat an outage as a per-row data rejection has to additionally consult ``is_database_service_unavailable_error`` before acting on a True here. """ import prisma return type(e) is prisma.errors.DataError @staticmethod def is_database_transport_error(e: Exception) -> bool: """ Returns True only for transport/connectivity failures where a reconnect attempt makes sense (e.g. DB is unreachable, connection dropped). Use this for reconnect logic — data-layer errors like UniqueViolationError mean the DB IS reachable, so reconnecting would be pointless. """ import prisma if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True if isinstance( e, ( prisma.errors.ClientNotConnectedError, prisma.errors.HTTPClientClosedError, ), ): return True if isinstance(e, prisma.errors.PrismaError): error_message: Final = str(e).lower() connection_keywords: Final = ( "can't reach database server", "cannot reach database server", "can't connect", "cannot connect", "connection error", "connection closed", "timed out", "timeout", "connection refused", "network is unreachable", "no route to host", "broken pipe", ) if any(keyword in error_message for keyword in connection_keywords): return True if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True return False @staticmethod def is_deadlock_error(e: Exception) -> bool: """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" import prisma if not isinstance(e, prisma.errors.PrismaError): return False if getattr(e, "code", None) == "P2034": return True error_message = str(e).lower() return ( "deadlock detected" in error_message or "40p01" in error_message or "write conflict or a deadlock" in error_message ) @staticmethod def is_prisma_engine_internal_error(e: Exception) -> bool: """True iff ``e`` is a non-``PrismaError`` exception raised from inside prisma-client-py's query-engine layer. During the instant a DB connection is torn down, the query engine can return a malformed error payload (``user_facing_error.meta`` is ``null``). prisma-client-py's ``handle_response_errors`` then crashes with ``AttributeError: 'NoneType' object has no attribute 'get'`` before it can raise the proper P1001 "can't reach database server" error. That AttributeError carries no connection keyword, so it can't be matched by message; identify it by its ``prisma.engine`` origin instead. Recognized ``PrismaError`` subclasses are excluded: connectivity ones are already classified by type/keyword above, and data-layer ones (the DB IS reachable) must stay 401. """ import prisma if isinstance(e, prisma.errors.PrismaError): return False tb = getattr(e, "__traceback__", None) while tb is not None: if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): return True tb = tb.tb_next return False @staticmethod def is_database_service_unavailable_error(e: Exception) -> bool: """True iff the exception means the database could not answer at the infrastructure level (connection refused, socket/interface failure, timeout) rather than a genuine auth failure (key not found) or a data-layer error (the DB IS reachable and rejected the data). Auth must answer 401 only for a key the DB confirms is invalid. When the DB itself is unreachable, the request has to surface as 503 so callers retry instead of treating valid keys as invalid during an outage. Note: prisma-client-py mislabels the P1001 "can't reach database server" connectivity failure as a ``DataError`` (a data-layer type), so a type-only check misses real outages. ``is_database_transport_error`` keyword-matches the connection message and catches that masquerade, while genuine data errors (no connection keyword) correctly stay 401. The Postgres "cached plan must not change result type" error is matched here, not in ``is_database_transport_error``: it is a transient stale-DB- state condition (not an invalid key), but the connection is healthy so it must not trigger a reconnect. A non-``PrismaError`` raised from inside the prisma query engine (e.g. the ``AttributeError`` from ``handle_response_errors`` when the engine returns a malformed error payload mid-tear-down) is also treated as unavailable; see ``is_prisma_engine_internal_error``. """ import asyncio if PrismaDBExceptionHandler.is_database_infrastructure_error(e): return True if PrismaDBExceptionHandler.is_database_transport_error(e): return True if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): return True if "cached plan must not change result type" in str(e).lower(): return True # OSError already covers ConnectionError and (Py3.3+) TimeoutError. # asyncio.TimeoutError is a distinct class before Py3.11. if isinstance(e, (OSError, asyncio.TimeoutError)): return True try: import asyncpg except ImportError: return False return isinstance( e, ( asyncpg.exceptions.PostgresConnectionError, asyncpg.exceptions.InterfaceError, ), ) @staticmethod def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool: """Like ``is_database_service_unavailable_error`` but also walks the ``__cause__`` / ``__context__`` chain. ``is_database_service_unavailable_error`` classifies a single exception by type, which a caller that catches a raw DB failure and re-raises a domain exception of a different type defeats. ``get_user_object`` in ``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps every DB error, a genuine outage included, in a bare ``ValueError`` whose original error survives only as ``__context__``. A type check on the ``ValueError`` misses the outage, so the caller would mistake an infrastructure fault for an auth failure. Walking the chain recovers the real signal, which is the PEP 3134 way to inspect a wrapped cause. The walk is depth-bounded, which also makes it cycle-safe. """ current: BaseException | None = e for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): if not isinstance(current, Exception): return False if PrismaDBExceptionHandler.is_database_service_unavailable_error(current): return True current = current.__cause__ or current.__context__ return False @staticmethod def handle_db_exception(e: Exception): """ Primary handler for `allow_requests_on_db_unavailable` flag. Decides whether to raise a DB Exception or not based on the flag. - If exception is a DB Connection Error, and `allow_requests_on_db_unavailable` is True, - Do not raise an exception, return None - Else, raise the exception """ if ( PrismaDBExceptionHandler.is_database_connection_error(e) and PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() ): return 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: Final = 2.0 _DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS: Final = 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 _ReadResultT: Final = TypeVar("_ReadResultT") async def call_with_db_reconnect_retry( prisma_client: Any, coro_factory: Callable[[], Awaitable[_ReadResultT]], *, reason: str, retry_safe_error_types: tuple[type[Exception], ...] | None = None, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None, ) -> _ReadResultT: """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. Transport errors outside `retry_safe_error_types` are re-raised too. 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`. retry_safe_error_types: Which transport errors may be replayed, or None for every transport error. A non-idempotent write must narrow this to `DB_RETRY_SAFE_ERROR_TYPES`, where the statements provably never reached the database. 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 retry_safe_error_types is not None and not isinstance(first_exc, retry_safe_error_types): raise if not hasattr(prisma_client, "attempt_db_reconnect"): raise resolved_timeout: Final = _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: Final = _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: Final = 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()