(feat) Allow failed DB connection requests to allow virtual keys with allow_failed_db_requests (#6605)

* fix use helper for _handle_failed_db_connection_for_get_key_object

* track ALLOW_FAILED_DB_REQUESTS on prometheus

* fix allow_failed_db_requests check

* fix allow_requests_on_db_unavailable

* fix allow_requests_on_db_unavailable

* docs allow_requests_on_db_unavailable

* identify user_id as litellm_proxy_admin_name when DB is failing

* test_handle_failed_db_connection

* fix test_user_api_key_auth_db_unavailable

* update best practices for prod doc

* update best practices for prod

* fix handle db failure

tests/local_testing/test_auth_checks.py

@@ -12,6 +12,11 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import pytest, litellm
+import httpx
+from litellm.proxy.auth.auth_checks import (
+    _handle_failed_db_connection_for_get_key_object,
+)
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import get_end_user_object
 from litellm.caching.caching import DualCache
 from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_BudgetTable
@@ -60,3 +65,33 @@ async def test_get_end_user_object(customer_spend, customer_budget):
                     customer_spend, customer_budget, str(e)
                 )
             )
+
+
+@pytest.mark.asyncio
+async def test_handle_failed_db_connection():
+    """
+    Test cases:
+    1. When allow_requests_on_db_unavailable=True -> return UserAPIKeyAuth
+    2. When allow_requests_on_db_unavailable=False -> raise original error
+    """
+    from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name
+
+    # Test case 1: allow_requests_on_db_unavailable=True
+    general_settings["allow_requests_on_db_unavailable"] = True
+    mock_error = httpx.ConnectError("Failed to connect to DB")
+
+    result = await _handle_failed_db_connection_for_get_key_object(e=mock_error)
+
+    assert isinstance(result, UserAPIKeyAuth)
+    assert result.key_name == "failed-to-connect-to-db"
+    assert result.token == "failed-to-connect-to-db"
+    assert result.user_id == litellm_proxy_admin_name
+
+    # Test case 2: allow_requests_on_db_unavailable=False
+    general_settings["allow_requests_on_db_unavailable"] = False
+
+    with pytest.raises(httpx.ConnectError) as exc_info:
+        await _handle_failed_db_connection_for_get_key_object(e=mock_error)
+    print("_handle_failed_db_connection_for_get_key_object got exception", exc_info)
+
+    assert str(exc_info.value) == "Failed to connect to DB"

tests/local_testing/test_key_generate_prisma.py

@@ -28,6 +28,7 @@ from datetime import datetime
 from dotenv import load_dotenv
 from fastapi import Request
 from fastapi.routing import APIRoute
+import httpx
 
 load_dotenv()
 import io
@@ -51,6 +52,7 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import (
     user_info,
     user_update,
 )
+from litellm.proxy.auth.auth_checks import get_key_object
 from litellm.proxy.management_endpoints.key_management_endpoints import (
     delete_key_fn,
     generate_key_fn,
@@ -3307,3 +3309,106 @@ async def test_service_accounts(prisma_client):
     print("response from user_api_key_auth", result)
 
     setattr(litellm.proxy.proxy_server, "general_settings", {})
+
+
+@pytest.mark.asyncio
+async def test_user_api_key_auth_db_unavailable():
+    """
+    Test that user_api_key_auth handles DB connection failures appropriately when:
+    1. DB connection fails during token validation
+    2. allow_requests_on_db_unavailable=True
+    """
+    litellm.set_verbose = True
+
+    # Mock dependencies
+    class MockPrismaClient:
+        async def get_data(self, *args, **kwargs):
+            print("MockPrismaClient.get_data() called")
+            raise httpx.ConnectError("Failed to connect to DB")
+
+        async def connect(self):
+            print("MockPrismaClient.connect() called")
+            pass
+
+    class MockDualCache:
+        async def async_get_cache(self, *args, **kwargs):
+            return None
+
+        async def async_set_cache(self, *args, **kwargs):
+            pass
+
+        async def set_cache(self, *args, **kwargs):
+            pass
+
+    # Set up test environment
+    setattr(litellm.proxy.proxy_server, "prisma_client", MockPrismaClient())
+    setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache())
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    setattr(
+        litellm.proxy.proxy_server,
+        "general_settings",
+        {"allow_requests_on_db_unavailable": True},
+    )
+
+    # Create test request
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    # Run test with a sample API key
+    result = await user_api_key_auth(
+        request=request,
+        api_key="Bearer sk-123456789",
+    )
+
+    # Verify results
+    assert isinstance(result, UserAPIKeyAuth)
+    assert result.key_name == "failed-to-connect-to-db"
+    assert result.user_id == litellm.proxy.proxy_server.litellm_proxy_admin_name
+
+
+@pytest.mark.asyncio
+async def test_user_api_key_auth_db_unavailable_not_allowed():
+    """
+    Test that user_api_key_auth raises an exception when:
+    This is default behavior
+
+    1. DB connection fails during token validation
+    2. allow_requests_on_db_unavailable=False (default behavior)
+    """
+
+    # Mock dependencies
+    class MockPrismaClient:
+        async def get_data(self, *args, **kwargs):
+            print("MockPrismaClient.get_data() called")
+            raise httpx.ConnectError("Failed to connect to DB")
+
+        async def connect(self):
+            print("MockPrismaClient.connect() called")
+            pass
+
+    class MockDualCache:
+        async def async_get_cache(self, *args, **kwargs):
+            return None
+
+        async def async_set_cache(self, *args, **kwargs):
+            pass
+
+        async def set_cache(self, *args, **kwargs):
+            pass
+
+    # Set up test environment
+    setattr(litellm.proxy.proxy_server, "prisma_client", MockPrismaClient())
+    setattr(litellm.proxy.proxy_server, "user_api_key_cache", MockDualCache())
+    setattr(litellm.proxy.proxy_server, "general_settings", {})
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+
+    # Create test request
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    # Run test with a sample API key
+    with pytest.raises(litellm.proxy._types.ProxyException):
+        await user_api_key_auth(
+            request=request,
+            api_key="Bearer sk-123456789",
+        )