diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index bbe4b0f5c35..9bf7f6cab96 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -71,12 +71,7 @@ def _as_proxy_exception(e: Exception) -> ProxyException: if isinstance(e, ProxyException): return e if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): - return ProxyException( - message=PrismaDBExceptionHandler.database_unavailable_message(e), - type=ProxyErrorTypes.no_db_connection, - param="None", - code=status.HTTP_503_SERVICE_UNAVAILABLE, - ) + return PrismaDBExceptionHandler.service_unavailable_proxy_exception(e) return ProxyException( message="Authentication Error, " + str(e), type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 99167c1275c..38d6fb9b99d 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,5 +1,6 @@ import re from collections.abc import Awaitable, Callable, Iterator +from http import HTTPStatus from typing import Final, Protocol, TypeVar from pydantic import TypeAdapter, ValidationError @@ -378,6 +379,15 @@ class PrismaDBExceptionHandler: "The proxy deployment needs attention." ) + @staticmethod + def service_unavailable_proxy_exception(e: Exception) -> ProxyException: + return ProxyException( + message=PrismaDBExceptionHandler.database_unavailable_message(e), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=HTTPStatus.SERVICE_UNAVAILABLE.value, + ) + @staticmethod def find_database_service_unavailable_error_in_chain(e: BaseException) -> Exception | None: """The exception in the ``__cause__`` / ``__context__`` chain that diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 60ac7e55eaf..6509977f7ff 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -47,6 +47,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, user_object_permission_id_cache_key, ) +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import ( @@ -2810,6 +2811,9 @@ async def ui_view_users( except HTTPException: raise except Exception as e: + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): + verbose_proxy_logger.warning("Database unavailable during user search: %s", type(e).__name__) + raise PrismaDBExceptionHandler.service_unavailable_proxy_exception(e) from e verbose_proxy_logger.exception("Error searching users: %s", e) raise HTTPException(status_code=500, detail=f"Error searching users: {e}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f8945694051..815fab0dba0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1977,6 +1977,11 @@ async def otel_request_validation_exception_handler(request: Request, exc: Reque async def otel_unhandled_exception_handler(request: Request, exc: Exception): if isinstance(exc, (ProxyException, HTTPException, RequestValidationError)): raise exc + if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): + verbose_proxy_logger.warning("Database unavailable during request: %s", type(exc).__name__) + return await openai_exception_handler( + request=request, exc=PrismaDBExceptionHandler.service_unavailable_proxy_exception(exc) + ) verbose_proxy_logger.exception("Unhandled exception in request: %s", type(exc).__name__) if should_report_bug(exc): verbose_proxy_logger.error( diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 75545a574e3..f684e2040bd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1,8 +1,10 @@ import hashlib import json +import logging from datetime import datetime, timezone from types import SimpleNamespace from typing import Final +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -17,12 +19,15 @@ from litellm.proxy._types import ( LiteLLM_UserTableFiltered, LitellmUserRoles, NewUserRequest, + ProxyErrorTypes, ProxyException, UpdateUserRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.internal_user_endpoints import ( LiteLLM_UserTableWithKeyCount, + _authorize_user_list_request, + _resolve_org_filter_for_user_search, _resolve_user_email_metadata, _update_internal_user_params, get_user_key_counts, @@ -4653,3 +4658,124 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non assert await cache.async_get_cache(key=deleted.user_id, model_type=LiteLLM_UserTable) is None assert await cache.async_get_cache(key=survivor.user_id, model_type=LiteLLM_UserTable) == survivor broadcast.assert_awaited_once_with(cache_key=deleted.user_id) + + +_DB_OUTAGE_503_BODY: Final = { + "error": { + "message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", + "type": "no_db_connection", + "param": "None", + "code": "503", + } +} + + +def _user_read_raising(mocker: MockerFixture, error: Exception) -> tuple[MagicMock, MagicMock]: + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=error) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + mocker.patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True) + return prisma_client, cache + + +def _db_unavailable_fallback_identity(route: str) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID + + return UserAPIKeyAuth( + key_name="failed-to-connect-to-db", + token="failed-to-connect-to-db", + user_id=DB_UNAVAILABLE_FALLBACK_USER_ID, + user_role=LitellmUserRoles.INTERNAL_USER, + request_route=route, + ) + + +@pytest.mark.asyncio +async def test_authorize_user_list_request_propagates_a_db_outage_instead_of_answering_403(mocker): + prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed")) + + with pytest.raises(httpx.ConnectError): + await _authorize_user_list_request( + user_api_key_dict=_db_unavailable_fallback_identity("/user/list"), + organization_ids=None, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_resolve_org_filter_for_user_search_propagates_a_db_outage_instead_of_answering_403(mocker): + prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed")) + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) + + with pytest.raises(httpx.ConnectError): + await _resolve_org_filter_for_user_search( + user_api_key_dict=_db_unavailable_fallback_identity("/user/filter/ui"), + team_id=None, + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_its_own_500(mocker, caplog): + prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed")) + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await ui_view_users( + user_api_key_dict=_db_unavailable_fallback_identity("/user/filter/ui"), + user_id=None, + user_email="lit", + team_id=None, + page=1, + page_size=50, + ) + + assert raised.value.code == "503" + assert raised.value.type == ProxyErrorTypes.no_db_connection + assert isinstance(raised.value.__cause__, httpx.ConnectError) + outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()] + assert outage_logs == ["Database unavailable during user search: ConnectError"] + + +@pytest.mark.parametrize( + ("route", "params"), + [("/user/list", {}), ("/user/filter/ui", {"user_email": "lit"})], + ids=["user_list", "user_filter_ui"], +) +def test_user_routes_answer_503_no_db_connection_when_the_callers_user_read_hits_a_db_outage( + mocker, route: str, params: dict[str, str] +): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + prisma_client, cache = _user_read_raising(mocker, httpx.ConnectError("All connection attempts failed")) + mocker.patch( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_ui_settings_cached", + return_value={"scope_user_search_to_org": True}, + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + app.dependency_overrides[user_api_key_auth] = lambda: _db_unavailable_fallback_identity(route) + try: + response = TestClient(app, raise_server_exceptions=False).get(route, params=params) + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + assert response.status_code == 503, response.text + assert response.json() == _DB_OUTAGE_503_BODY diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index e95359ace8a..8bf9c598b1a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -7,6 +7,7 @@ from collections.abc import Sequence from typing import Final, Optional, cast from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch +import httpx import pytest from fastapi import HTTPException from fastapi.testclient import TestClient @@ -42,6 +43,8 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _STRIP_DELETED_TEAM_FROM_USERS_SQL, GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, + _build_team_list_where_conditions, + _get_org_admin_org_ids, _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, @@ -16569,3 +16572,81 @@ def test_team_member_update_request_rejects_unusable_temp_budget_increase(increa TeamMemberUpdateRequest( team_id="team-1", user_id="user-1", temp_budget_increase=increase, temp_budget_expiry="2030-01-01T00:00:00Z" ) + + +_DB_OUTAGE_503_BODY: Final = { + "error": { + "message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", + "type": "no_db_connection", + "param": "None", + "code": "503", + } +} + + +def _user_read_raising(error: Exception) -> tuple[MagicMock, MagicMock]: + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=error) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + return prisma_client, cache + + +def _db_unavailable_fallback_identity(route: str) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID + + return UserAPIKeyAuth( + key_name="failed-to-connect-to-db", + token="failed-to-connect-to-db", + user_id=DB_UNAVAILABLE_FALLBACK_USER_ID, + user_role=LitellmUserRoles.INTERNAL_USER, + request_route=route, + ) + + +@pytest.mark.asyncio +async def test_get_org_admin_org_ids_propagates_a_db_outage_instead_of_answering_not_an_org_admin(): + prisma_client, cache = _user_read_raising(httpx.ConnectError("All connection attempts failed")) + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(httpx.ConnectError): + await _get_org_admin_org_ids( + user_id="outage-probe-user", + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_build_team_list_where_conditions_propagates_a_db_outage_instead_of_answering_user_not_found(): + prisma_client, cache = _user_read_raising(httpx.ConnectError("All connection attempts failed")) + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(httpx.ConnectError): + await _build_team_list_where_conditions( + prisma_client=prisma_client, + team_id=None, + team_alias=None, + organization_id=None, + user_id="outage-probe-user", + use_deleted_table=False, + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + + +def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hits_a_db_outage(monkeypatch): + prisma_client, cache = _user_read_raising(httpx.ConnectError("All connection attempts failed")) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + app.dependency_overrides[user_api_key_auth] = lambda: _db_unavailable_fallback_identity("/v2/team/list") + try: + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + response = TestClient(app, raise_server_exceptions=False).get("/v2/team/list") + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + assert response.status_code == 503, response.text + assert response.json() == _DB_OUTAGE_503_BODY diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py index b3028ae71dd..089c2d57594 100644 --- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -11,8 +11,10 @@ from __future__ import annotations import json from types import SimpleNamespace +from typing import Final from unittest.mock import MagicMock +import httpx import pytest from fastapi import HTTPException from fastapi.exceptions import RequestValidationError @@ -364,6 +366,39 @@ async def test_otel_unhandled_exception_handler_returns_500_generic_payload(): } +_DB_OUTAGE_503_BODY: Final = { + "error": { + "message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", + "type": "no_db_connection", + "param": "None", + "code": "503", + } +} + + +def _raised_from(outer: Exception, cause: Exception) -> Exception: + try: + raise outer from cause + except Exception as chained: + return chained + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exc", + [ + httpx.ConnectError("All connection attempts failed"), + _raised_from(RuntimeError("user read failed"), httpx.ConnectError("All connection attempts failed")), + ], + ids=["raw_connect_error", "connect_error_as_cause"], +) +async def test_otel_unhandled_exception_handler_answers_a_db_outage_with_503_no_db_connection(exc): + response = await otel_unhandled_exception_handler(request=_make_request(path="/v2/team/list"), exc=exc) + + assert response.status_code == 503 + assert json.loads(response.body) == _DB_OUTAGE_503_BODY + + @pytest.mark.asyncio async def test_otel_unhandled_exception_handler_reraises_proxy_exception_error(): """ProxyException / HTTPException / RequestValidationError are re-raised