diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py index a032ea7662d..9204cf359b3 100644 --- a/enterprise/enterprise_hooks/blocked_user_list.py +++ b/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} ) ) diff --git a/tests/test_litellm/enterprise/test_blocked_user_list_prisma.py b/tests/test_litellm/enterprise/test_blocked_user_list_prisma.py new file mode 100644 index 00000000000..b1e54d8ddc1 --- /dev/null +++ b/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()