diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index ef1c4a66203..3a9e505275a 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,4 +1,6 @@ from collections.abc import Awaitable, Callable, Iterator +from importlib import import_module +from types import ModuleType from typing import Any, Final, TypeVar from litellm._logging import verbose_proxy_logger @@ -47,6 +49,28 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]: return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException)) +def _optional_prisma() -> ModuleType | None: + """The ``prisma`` package, or None when it is not installed. + + A DB-less proxy (config + master key, no ``DATABASE_URL``) is installed + without the optional ``prisma`` package, yet the predicates below still run + on every error path, auth failures that never touch a database included. An + unguarded import there raises ``ModuleNotFoundError`` from inside the error + handler itself, which replaces the intended response (e.g. a 401 for a + missing API key) with a generic 500. Without prisma installed no exception + can be a prisma one, so callers fall back to the non-prisma checks. + + ``prisma.engine.errors`` is imported alongside it because that submodule is + not loaded by importing the top-level package. + """ + try: + prisma: Final = import_module("prisma") + import_module("prisma.engine.errors") + except ImportError: + return None + return prisma + + class PrismaDBExceptionHandler: """ Class to handle DB Exceptions or Connection Errors @@ -88,11 +112,11 @@ class PrismaDBExceptionHandler: Reporting decisions want the opposite breadth; use ``is_database_infrastructure_error`` for those. """ - import prisma.engine.errors + prisma: Final = _optional_prisma() if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True - if isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)): + if prisma is not None and isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)): return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @@ -112,23 +136,24 @@ class PrismaDBExceptionHandler: ``RecordNotFoundError``, etc.) are excluded — the DB IS reachable and the request itself is what failed. """ - import prisma + prisma: Final = _optional_prisma() - data_layer_errors: Final = _exception_types( - 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 prisma is not None: + data_layer_errors: Final = _exception_types( + 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, _exception_types(prisma.errors.PrismaError)): + return True if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True - if isinstance(e, _exception_types(prisma.errors.PrismaError)): - return True if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True return False @@ -152,9 +177,9 @@ class PrismaDBExceptionHandler: per-row data rejection has to additionally consult ``is_database_service_unavailable_error`` before acting on a True here. """ - import prisma + prisma: Final = _optional_prisma() - return type(e) is prisma.errors.DataError + return prisma is not None and type(e) is prisma.errors.DataError @staticmethod def is_database_transport_error(e: Exception) -> bool: @@ -165,36 +190,37 @@ class PrismaDBExceptionHandler: Use this for reconnect logic — data-layer errors like UniqueViolationError mean the DB IS reachable, so reconnecting would be pointless. """ - import prisma + prisma: Final = _optional_prisma() if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True - if isinstance( - e, - _exception_types( - prisma.errors.ClientNotConnectedError, - prisma.errors.HTTPClientClosedError, - ), - ): - return True - if isinstance(e, _exception_types(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): + if prisma is not None: + if isinstance( + e, + _exception_types( + prisma.errors.ClientNotConnectedError, + prisma.errors.HTTPClientClosedError, + ), + ): return True + if isinstance(e, _exception_types(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 @@ -202,9 +228,9 @@ class PrismaDBExceptionHandler: @staticmethod def is_deadlock_error(e: Exception) -> bool: """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" - import prisma + prisma: Final = _optional_prisma() - if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + if prisma is None or not isinstance(e, _exception_types(prisma.errors.PrismaError)): return False if getattr(e, "code", None) == "P2034": return True @@ -233,9 +259,9 @@ class PrismaDBExceptionHandler: are already classified by type/keyword above, and data-layer ones (the DB IS reachable) must stay 401. """ - import prisma + prisma: Final = _optional_prisma() - if isinstance(e, _exception_types(prisma.errors.PrismaError)): + if prisma is None or isinstance(e, _exception_types(prisma.errors.PrismaError)): return False tb = getattr(e, "__traceback__", None) while tb is not None: diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 21e0b83791f..25bf05791cb 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -1,5 +1,6 @@ import asyncio import json +import sys from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -871,3 +872,32 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( assert records[0].levelname == expect_level expected_logger_name = "LiteLLM Proxy.stdout" if expect_level == "WARNING" else "LiteLLM Proxy" assert records[0].name == expected_logger_name + + +@pytest.mark.asyncio +async def test_missing_api_key_returns_401_when_prisma_is_not_installed(): + """ + On a DB-less proxy (config + master key, no DATABASE_URL) the optional + prisma package is absent. Classifying the "no api key passed in" failure + must not crash on that missing import, or the intended 401 is replaced by a + generic 500 from the catch-all handler. + """ + prisma_absent = dict.fromkeys( + ("prisma", "prisma.errors", "prisma.engine", "prisma.engine.errors"), None + ) + handler = UserAPIKeyAuthExceptionHandler() + + with patch.dict(sys.modules, prisma_absent): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + Exception("No api key passed in."), + MagicMock(), + {}, + "/v1/chat/completions", + None, + None, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == str(status.HTTP_401_UNAUTHORIZED) + assert "No api key passed in." in exc_info.value.message diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index c685d778c0e..3ec06b3b1be 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -700,3 +700,71 @@ def test_connection_error_answers_when_prisma_is_mocked_after_import(): with patch.dict(sys.modules, {"prisma": MagicMock()}): assert PrismaDBExceptionHandler.is_database_connection_error(Exception("x")) is False assert PrismaDBExceptionHandler.is_database_connection_error(httpx.ConnectError("refused")) is True + + +_PRISMA_ABSENT: Final = dict.fromkeys( + ("prisma", "prisma.errors", "prisma.engine", "prisma.engine.errors"), None +) + + +@pytest.mark.parametrize( + "predicate", + [ + PrismaDBExceptionHandler.is_database_connection_error, + PrismaDBExceptionHandler.is_database_infrastructure_error, + PrismaDBExceptionHandler.is_prisma_data_error, + PrismaDBExceptionHandler.is_database_transport_error, + PrismaDBExceptionHandler.is_deadlock_error, + PrismaDBExceptionHandler.is_prisma_engine_internal_error, + PrismaDBExceptionHandler.is_database_service_unavailable_error, + PrismaDBExceptionHandler.is_permanent_database_fault, + PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain, + ], + ids=lambda predicate: predicate.__name__, +) +def test_predicates_answer_false_when_prisma_is_not_installed(predicate): + """ + A DB-less proxy install ships without the optional prisma package, so no + exception it sees can be a prisma one. Every predicate must answer False + instead of raising ModuleNotFoundError out of the error path it is being + used to handle. + """ + with patch.dict(sys.modules, _PRISMA_ABSENT): + assert predicate(Exception("No api key passed in.")) is False + + +def test_chain_helpers_answer_when_prisma_is_not_installed(): + """ + The chain walk and the 503 wording go through the same predicates, so they + have to survive a missing prisma too. + """ + auth_failure = Exception("No api key passed in.") + + with patch.dict(sys.modules, _PRISMA_ABSENT): + assert PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(auth_failure) is None + assert PrismaDBExceptionHandler.database_unavailable_message(httpx.ConnectError("refused")) + + +def test_non_prisma_db_signals_still_classified_when_prisma_is_not_installed(): + """ + Guarding the prisma import must not blind the predicates to the signals + that do not need prisma: httpx transport failures and the proxy's own + no_db_connection exception. + """ + connect_error = httpx.ConnectError("All connection attempts failed") + no_db_connection = ProxyException( + message="No connected db.", + type=ProxyErrorTypes.no_db_connection, + param="None", + code=500, + ) + + with patch.dict(sys.modules, _PRISMA_ABSENT): + assert PrismaDBExceptionHandler.is_database_connection_error(connect_error) is True + assert PrismaDBExceptionHandler.is_database_infrastructure_error(connect_error) is True + assert PrismaDBExceptionHandler.is_database_transport_error(connect_error) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(connect_error) is True + + assert PrismaDBExceptionHandler.is_database_connection_error(no_db_connection) is True + assert PrismaDBExceptionHandler.is_database_infrastructure_error(no_db_connection) is True + assert PrismaDBExceptionHandler.is_database_transport_error(no_db_connection) is True