From 2b5db2b8bf131cf44ed6a025b78c4315d41b4502 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 5 Nov 2024 18:31:50 -0800
Subject: [PATCH] test_proxy_server_prisma_setup_invalid_url

---
 litellm/proxy/db/prisma_client.py        |  9 ++++++--
 litellm/proxy/utils.py                   |  9 ++++----
 tests/local_testing/test_proxy_server.py | 29 ++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py
index 76e425bf2cb..7b7b3002676 100644
--- a/litellm/proxy/db/prisma_client.py
+++ b/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
 
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 44243cab0b9..53923106c30 100644
--- a/litellm/proxy/utils.py
+++ b/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
 
diff --git a/tests/local_testing/test_proxy_server.py b/tests/local_testing/test_proxy_server.py
index 808b10db3ad..40f02be4f5d 100644
--- a/tests/local_testing/test_proxy_server.py
+++ b/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)