diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7fcecbc940..9f9a88d38a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -51,6 +51,7 @@ from litellm._service_logger import ServiceLogging from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( _cache_key_object, + _handle_failed_db_connection_for_get_key_object, allowed_routes_check, can_key_call_model, common_checks, @@ -802,7 +803,9 @@ async def user_api_key_auth( # 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.") + return await _handle_failed_db_connection_for_get_key_object( + e=Exception("No connected db.") + ) ## check for cache hit (In-Memory Cache) _user_role = None diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 31daa358af..9cd52e591c 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -14,7 +14,7 @@ import pytest from starlette.datastructures import URL import litellm -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth, UserAPIKeyAuth class Request: @@ -387,3 +387,34 @@ def test_is_api_route_allowed(route, user_role, expected_result): pass else: raise e + + +@pytest.mark.asyncio +async def test_auth_not_connected_to_db(): + """ + ensure requests don't fail when `prisma_client` = None + """ + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import hash_token, user_api_key_cache + + user_key = "sk-12345678" + + setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr(litellm.proxy.proxy_server, "prisma_client", None) + setattr( + litellm.proxy.proxy_server, + "general_settings", + {"allow_requests_on_db_unavailable": True}, + ) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + valid_token = await user_api_key_auth(request=request, api_key="Bearer " + user_key) + print("got valid token", valid_token) + assert valid_token.key_name == "failed-to-connect-to-db" + assert valid_token.token == "failed-to-connect-to-db"