From f8caebc54b28dda058ec42ea638d344a5662a784 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 26 Mar 2025 14:55:40 -0700
Subject: [PATCH] is_database_connection_error

---
 litellm/proxy/_types.py                      |  4 ++++
 litellm/proxy/auth/auth_exception_handler.py | 10 +++++++---
 litellm/proxy/auth/user_api_key_auth.py      |  9 ++++++---
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index c0fb0750eb..104d1e7338 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2070,6 +2070,10 @@ class ProxyErrorTypes(str, enum.Enum):
     """
     Object was over budget
     """
+    no_db_connection = "no_db_connection"
+    """
+    No database connection
+    """
 
     token_not_found_in_db = "token_not_found_in_db"
     """
diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py
index 5b314ef931..c1a546b569 100644
--- a/litellm/proxy/auth/auth_exception_handler.py
+++ b/litellm/proxy/auth/auth_exception_handler.py
@@ -144,6 +144,10 @@ class UserAPIKeyAuthExceptionHandler:
         """
         import prisma
 
-        return isinstance(e, DB_CONNECTION_ERROR_TYPES) or isinstance(
-            e, prisma.errors.PrismaError
-        )
+        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
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 06fa560866..85ee76ff31 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -675,7 +675,12 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
         if (
             prisma_client is None
         ):  # if both master key + user key submitted, and user key != master key, and no db connected, raise an error
-            raise Exception("No connected db.")
+            raise ProxyException(
+                message="No connected db.",
+                type=ProxyErrorTypes.no_db_connection,
+                code=400,
+                param=None,
+            )
 
         ## check for cache hit (In-Memory Cache)
         _user_role = None
@@ -1001,8 +1006,6 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 route=route,
                 start_time=start_time,
             )
-        else:
-            raise Exception()
     except Exception as e:
         return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
             e=e,