test_proxy_server_prisma_setup_invalid_url

litellm/proxy/db/prisma_client.py

@@ -7,13 +7,18 @@ import os
 import urllib
 import urllib.parse
 from datetime import datetime, timedelta
-from typing import Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from litellm.secret_managers.main import str_to_bool
 
+if TYPE_CHECKING:
+    from prisma import Prisma
+else:
+    Prisma = Any
+
 
 class PrismaWrapper:
-    def __init__(self, original_prisma: Any, iam_token_db_auth: bool):
+    def __init__(self, original_prisma: Prisma, iam_token_db_auth: bool):
         self._original_prisma = original_prisma
         self.iam_token_db_auth = iam_token_db_auth
 

litellm/proxy/utils.py

@@ -2289,11 +2289,10 @@ class PrismaClient:
             verbose_proxy_logger.debug(
                 "PrismaClient: connect() called Attempting to Connect to DB"
             )
-            if self.db.is_connected() is False:
-                verbose_proxy_logger.debug(
-                    "PrismaClient: DB not connected, Attempting to Connect to DB"
-                )
-                await self.db.connect()
+            verbose_proxy_logger.debug(
+                "PrismaClient: DB not connected, Attempting to Connect to DB"
+            )
+            await self.db.connect()
         except Exception as e:
             import traceback
 

tests/local_testing/test_proxy_server.py

@@ -1921,3 +1921,32 @@ async def test_proxy_server_prisma_setup():
         # Verify our mocked methods were called
         mock_client.connect.assert_called_once()
         mock_client.check_view_exists.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_proxy_server_prisma_setup_invalid_url():
+    """
+    Test that proxy server setup raises an exception with invalid database URL
+
+    This is a PRODUCTION TEST. Think 2-3 times before modifying.
+
+    It's critical that this raises an exception, when unable to connect to the database.
+    """
+    from litellm.proxy.proxy_server import ProxyStartupEvent
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.caching import DualCache
+
+    user_api_key_cache = DualCache()
+    invalid_db_url = "postgresql://invalid:5432/nonexistent"
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+    # Test that it raises an exception with invalid URL
+    with pytest.raises(Exception) as exc_info:
+        await ProxyStartupEvent._setup_prisma_client(
+            database_url=invalid_db_url,
+            proxy_logging_obj=proxy_logging_obj,
+            user_api_key_cache=user_api_key_cache,
+        )
+
+    # Verify the error message indicates a connection failure
+    assert "Could not connect to the query engine" in str(exc_info.value)