fix(proxy/db): keep prisma predicates from raising TypeError under a mocked prisma module (#39253)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/db/exception_handler.py

@@ -14,6 +14,17 @@ from litellm.secret_managers.main import str_to_bool
 _MAX_EXCEPTION_CHAIN_DEPTH: Final = 20
 
 
+def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]:
+    """Keep only the real exception classes among ``candidates``.
+
+    The predicates below resolve prisma's error classes at call time, so a test
+    that swaps ``sys.modules["prisma"]`` for a ``MagicMock`` hands them mocks,
+    and ``isinstance`` against a mock raises ``TypeError`` instead of answering
+    False. Dropping the non-types lets the call fall through to the other checks.
+    """
+    return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException))
+
+
 class PrismaDBExceptionHandler:
     """
     Class to handle DB Exceptions or Connection Errors
@@ -59,7 +70,7 @@ class PrismaDBExceptionHandler:
 
         if isinstance(e, DB_CONNECTION_ERROR_TYPES):
             return True
-        if isinstance(e, prisma.engine.errors.EngineConnectionError):
+        if isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)):
             return True
         return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection
 
@@ -81,7 +92,7 @@ class PrismaDBExceptionHandler:
         """
         import prisma
 
-        data_layer_errors: Final = (
+        data_layer_errors: Final = _exception_types(
             prisma.errors.DataError,
             prisma.errors.UniqueViolationError,
             prisma.errors.ForeignKeyViolationError,
@@ -94,7 +105,7 @@ class PrismaDBExceptionHandler:
             return False
         if isinstance(e, DB_CONNECTION_ERROR_TYPES):
             return True
-        if isinstance(e, prisma.errors.PrismaError):
+        if isinstance(e, _exception_types(prisma.errors.PrismaError)):
             return True
         if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection:
             return True
@@ -138,13 +149,13 @@ class PrismaDBExceptionHandler:
             return True
         if isinstance(
             e,
-            (
+            _exception_types(
                 prisma.errors.ClientNotConnectedError,
                 prisma.errors.HTTPClientClosedError,
             ),
         ):
             return True
-        if isinstance(e, prisma.errors.PrismaError):
+        if isinstance(e, _exception_types(prisma.errors.PrismaError)):
             error_message: Final = str(e).lower()
             connection_keywords: Final = (
                 "can't reach database server",
@@ -171,7 +182,7 @@ class PrismaDBExceptionHandler:
         """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma."""
         import prisma
 
-        if not isinstance(e, prisma.errors.PrismaError):
+        if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
             return False
         if getattr(e, "code", None) == "P2034":
             return True
@@ -202,7 +213,7 @@ class PrismaDBExceptionHandler:
         """
         import prisma
 
-        if isinstance(e, prisma.errors.PrismaError):
+        if isinstance(e, _exception_types(prisma.errors.PrismaError)):
             return False
         tb = getattr(e, "__traceback__", None)
         while tb is not None:

tests/test_litellm/proxy/db/test_exception_handler.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import sys
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -579,3 +580,40 @@ def test_is_deadlock_error_matches_postgres_deadlock(error):
 def test_is_deadlock_error_excludes_non_deadlocks(error):
     """Non-deadlock prisma errors, connectivity failures, and non-prisma exceptions are not treated as deadlocks."""
     assert PrismaDBExceptionHandler.is_deadlock_error(error) is False
+
+
+MOCKED_PRISMA_PREDICATES: Final = (
+    PrismaDBExceptionHandler.is_database_infrastructure_error,
+    PrismaDBExceptionHandler.is_database_transport_error,
+    PrismaDBExceptionHandler.is_deadlock_error,
+    PrismaDBExceptionHandler.is_prisma_engine_internal_error,
+    PrismaDBExceptionHandler.is_database_service_unavailable_error,
+)
+
+
+@pytest.mark.parametrize("predicate", MOCKED_PRISMA_PREDICATES, ids=lambda p: p.__name__)
+def test_predicates_answer_false_for_a_plain_exception_when_prisma_is_mocked(predicate):
+    """Suites that swap ``sys.modules["prisma"]`` for a ``MagicMock`` hand the
+    predicates mocks in place of prisma's error classes. ``isinstance`` against
+    a mock raises ``TypeError``; the predicate must instead answer for the
+    non-prisma checks it still has."""
+    with patch.dict(sys.modules, {"prisma": MagicMock()}):
+        assert predicate(Exception("db connection dropped")) is False
+
+
+def test_infrastructure_error_still_recognizes_transport_errors_when_prisma_is_mocked():
+    """Skipping the prisma classes must not skip the checks that do not need them."""
+    with patch.dict(sys.modules, {"prisma": MagicMock()}):
+        no_db: Final = ProxyException(message="no db", type=ProxyErrorTypes.no_db_connection, param=None, code=503)
+        assert PrismaDBExceptionHandler.is_database_infrastructure_error(httpx.ConnectError("refused")) is True
+        assert PrismaDBExceptionHandler.is_database_infrastructure_error(no_db) is True
+
+
+def test_connection_error_answers_when_prisma_is_mocked_after_import():
+    """``prisma.engine`` is already loaded in a real process, so a mock parent
+    still resolves ``prisma.engine.errors``; its classes are then mocks too."""
+    import prisma.engine.errors  # noqa: F401
+
+    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