Merge fbfb0d297f into a6816f0e96

enterprise/enterprise_hooks/blocked_user_list.py

@@ -45,6 +45,22 @@ class _ENTERPRISE_BlockedUserList(CustomLogger):
                     f"An error occurred: {str(e)}, blocked_user_list={blocked_user_list}"
                 )
 
+    def _resolve_prisma_client(self) -> Optional[PrismaClient]:
+        """Use the live proxy Prisma client when init ran before DB setup."""
+        if self.prisma_client is not None:
+            return self.prisma_client
+
+        try:
+            from litellm.proxy.proxy_server import prisma_client as global_prisma_client
+
+            return global_prisma_client
+        except Exception:
+            verbose_proxy_logger.debug(
+                "blocked_user_list: could not resolve global prisma_client",
+                exc_info=True,
+            )
+            return None
+
     def print_verbose(self, print_statement, level: Literal["INFO", "DEBUG"] = "DEBUG"):
         if level == "INFO":
             verbose_proxy_logger.info(print_statement)
@@ -87,10 +103,11 @@ class _ENTERPRISE_BlockedUserList(CustomLogger):
                 end_user_cache_obj: Optional[LiteLLM_EndUserTable] = cache.get_cache(  # type: ignore
                     key=cache_key
                 )
-                if end_user_cache_obj is None and self.prisma_client is not None:
+                prisma_client = self._resolve_prisma_client()
+                if end_user_cache_obj is None and prisma_client is not None:
                     # check db
                     end_user_obj = (
-                        await self.prisma_client.db.litellm_endusertable.find_unique(
+                        await prisma_client.db.litellm_endusertable.find_unique(
                             where={"user_id": user}
                         )
                     )

tests/test_litellm/enterprise/test_blocked_user_list_prisma.py

@@ -0,0 +1,50 @@
+import importlib.util
+from pathlib import Path
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+from fastapi import HTTPException
+
+from litellm.caching.caching import DualCache
+from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth
+
+
+def _load_blocked_user_list_module():
+    module_path = (
+        Path(__file__).resolve().parents[3]
+        / "enterprise"
+        / "enterprise_hooks"
+        / "blocked_user_list.py"
+    )
+    spec = importlib.util.spec_from_file_location("blocked_user_list", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.mark.asyncio
+async def test_blocked_user_check_uses_global_prisma_when_init_had_none(monkeypatch):
+    blocked_user_list = _load_blocked_user_list_module()
+    _ENTERPRISE_BlockedUserList = blocked_user_list._ENTERPRISE_BlockedUserList
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_endusertable.find_unique = AsyncMock(
+        return_value=LiteLLM_EndUserTable(user_id="blocked-user", blocked=True)
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
+
+    hook = _ENTERPRISE_BlockedUserList(prisma_client=None)
+    cache = DualCache()
+
+    with pytest.raises(HTTPException) as exc_info:
+        await hook.async_pre_call_hook(
+            user_api_key_dict=UserAPIKeyAuth(),
+            cache=cache,
+            data={"user": "blocked-user"},
+            call_type="completion",
+        )
+
+    assert "blocked" in str(exc_info.value).lower()
+    mock_prisma.db.litellm_endusertable.find_unique.assert_awaited_once()