fix(usage): keep spend when metadata lookups hit database failures

Treat transport errors and the query engine crash on a closed connection as database
failures in the key, user and guardrail lookups of the aggregated, keys, users,
model_top_keys and cache_leakage_keys routes, and keep prisma out of module-level imports
so a proxy install without prisma still imports the usage endpoints.

litellm/proxy/management_endpoints/common_daily_activity.py

@@ -5,16 +5,16 @@ from dataclasses import dataclass, replace
 from datetime import date, datetime, timedelta
 from itertools import groupby
 from types import MappingProxyType
-from typing import Final, Literal, NoReturn, Protocol
+from typing import Final, Literal, NoReturn, Protocol, TypeVar
 
 from fastapi import HTTPException, status
-from prisma.errors import PrismaError
 from typing_extensions import ReadOnly, TypedDict, assert_never
 
 from litellm import constants
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import PTU_SENTINEL_API_KEY, SUBSCRIPTION_BACKED_PROVIDERS
 from litellm.proxy._types import CommonProxyErrors
+from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 from litellm.proxy.spend_tracking.key_metadata_recovery import (
     attach_user_details,
     recover_cli_session_key_metadata,
@@ -57,6 +57,8 @@ from litellm.types.repositories.daily_activity import (
     SpendLogsWindow,
 )
 
+_T: Final = TypeVar("_T")
+
 
 @dataclass(frozen=True, slots=True)
 class ScopeDenied:
@@ -920,24 +922,40 @@ def _aggregate_grouping_sets_records_sync(
     return {"results": results, "totals": total_metrics}
 
 
+async def unless_database_fails(lookup: Awaitable[_T], fallback: _T, left_out: str) -> _T:
+    try:
+        return await lookup
+    except Exception as exc:
+        if not (
+            PrismaDBExceptionHandler.is_prisma_error(exc)
+            or PrismaDBExceptionHandler.is_database_service_unavailable_error(exc)
+        ):
+            raise
+        verbose_proxy_logger.warning("Returning daily activity without %s: %s", left_out, exc)
+        return fallback
+
+
+async def _key_metadata_unless_database_fails(
+    repository: DailyActivityRepository, rows: Sequence[GroupingSetsRow] | Sequence[EntityRollupRow]
+) -> Mapping[str, KeyMetadataRow]:
+    api_keys: Final = frozenset(row.api_key for row in rows if row.api_key and row.api_key != PTU_SENTINEL_API_KEY)
+    return await unless_database_fails(
+        repository.key_metadata(api_keys, spend_logs_window(frozenset(row.date for row in rows))),
+        MappingProxyType({}),
+        "its key metadata",
+    )
+
+
 async def _aggregate_grouping_sets_records(
     *,
     repository: DailyActivityRepository,
     records: Sequence[GroupingSetsRow],
 ) -> _AggregatedSpendData:
     """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread."""
-    api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY}
-
-    api_key_metadata: Final[Mapping[str, KeyMetadataRow]] = (
-        await repository.key_metadata(frozenset(api_keys), spend_logs_window(frozenset(r.date for r in records)))
-        if api_keys
-        else MappingProxyType({})
-    )
-
     return await asyncio.to_thread(
         _aggregate_grouping_sets_records_sync,
         records=records,
-        api_key_metadata=api_key_metadata,
+        api_key_metadata=await _key_metadata_unless_database_fails(repository, records),
     )
 
 
@@ -1105,24 +1123,16 @@ def _with_guardrail_breakdown(
 async def _guardrail_spend_or_empty(
     repository: DailyActivityRepository, scope: DailyActivityScope
 ) -> tuple[GuardrailSpendRow, ...]:
-    from prisma.errors import PrismaError
-
-    try:
-        return await repository.guardrail_spend(scope)
-    except PrismaError as exc:
-        verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc)
-        return ()
+    return await unless_database_fails(repository.guardrail_spend(scope), (), "its guardrail breakdown")
 
 
 async def _user_entity_metadata(
     repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow]
 ) -> Mapping[str, Mapping[str, object]]:
     user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id)
-    try:
-        users: Final = await repository.user_metadata(user_ids)
-    except PrismaError as exc:
-        verbose_proxy_logger.warning("Failed daily activity entity user metadata recovery: %s", exc)
-        return MappingProxyType({})
+    users: Final = await unless_database_fails(
+        repository.user_metadata(user_ids), MappingProxyType({}), "its entity user metadata"
+    )
     return MappingProxyType(
         {
             user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias})
@@ -1131,21 +1141,6 @@ async def _user_entity_metadata(
     )
 
 
-async def _entity_key_metadata(
-    repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow]
-) -> Mapping[str, KeyMetadataRow]:
-    api_keys: Final = frozenset(
-        row.api_key for row in entity_rows if row.api_key and row.api_key != PTU_SENTINEL_API_KEY
-    )
-    if not api_keys:
-        return MappingProxyType({})
-    try:
-        return await repository.key_metadata(api_keys, spend_logs_window(frozenset(row.date for row in entity_rows)))
-    except PrismaError as exc:
-        verbose_proxy_logger.warning("Failed daily activity entity key metadata recovery: %s", exc)
-        return MappingProxyType({})
-
-
 async def get_daily_activity_aggregated(
     repository: DailyActivityRepository,
     scope: DailyActivityScope,
@@ -1180,7 +1175,7 @@ async def get_daily_activity_aggregated(
         )
         if aggregated_rows.entity_rows:
             entity_records: Final = aggregated_rows.entity_rows
-            entity_key_metadata: Final = await _entity_key_metadata(repository, entity_records)
+            entity_key_metadata: Final = await _key_metadata_unless_database_fails(repository, entity_records)
             user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = (
                 await _user_entity_metadata(repository, entity_records)
                 if scope.table is DailyActivityTable.USER

litellm/proxy/management_endpoints/daily_activity_routes.py

@@ -4,6 +4,7 @@ import json
 from collections.abc import AsyncIterator, Mapping, Sequence
 from dataclasses import asdict, fields, replace
 from datetime import date, datetime
+from types import MappingProxyType
 from typing import Annotated, Final, Literal
 
 from fastapi import APIRouter, Depends, HTTPException, Query
@@ -23,6 +24,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
     raise_public,
     spend_logs_window,
     sum_spend_metrics,
+    unless_database_fails,
 )
 from litellm.proxy.management_endpoints.daily_activity_scopes import (
     AGENT_RESOLVER,
@@ -149,9 +151,10 @@ async def _key_activity_rows(
     spend_window: Final[tuple[datetime, datetime] | None] = spend_logs_window(
         frozenset((resolved_scope.scope.start_date, resolved_scope.scope.end_date))
     )
-    metadata: Final[Mapping[str, KeyMetadataRow]] = await repository.key_metadata(
-        frozenset(row.api_key for row in rows if row.api_key),
-        spend_window,
+    metadata: Final[Mapping[str, KeyMetadataRow]] = await unless_database_fails(
+        repository.key_metadata(frozenset(row.api_key for row in rows if row.api_key), spend_window),
+        MappingProxyType({}),
+        "its key metadata",
     )
     return [_key_activity_row(row, metadata) for row in rows]
 
@@ -176,8 +179,10 @@ async def _user_activity_rows(
     repository: DailyActivityRepository,
     rows: Sequence[UserSpendRow],
 ) -> Sequence[UserActivityRow]:
-    metadata: Final[Mapping[str, UserMetadataRow]] = await repository.user_metadata(
-        frozenset(user_id for row in rows if (user_id := row.user_id) is not None)
+    metadata: Final[Mapping[str, UserMetadataRow]] = await unless_database_fails(
+        repository.user_metadata(frozenset(user_id for row in rows if (user_id := row.user_id) is not None)),
+        MappingProxyType({}),
+        "its user metadata",
     )
     return [_user_activity_row(row, metadata) for row in rows]
 

ruff.toml

@@ -9,7 +9,7 @@ lint.extend-select = [
     "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", "RUF022", "RUF023", "RUF051", "S113",
     "SIM114", "SIM118", "SIM201", "SIM211", "SIM222", "TC005", "UP006", "UP007", "UP008",
     "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045",
-    "C404", "C419",
+    "C404", "C419", "TID253",
 ]
 # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip
 # `# noqa` directives that protect rules enforced elsewhere. List those codes as external
@@ -35,6 +35,9 @@ format.exclude = ["**/enterprise/**"]
 lint.exclude = ["litellm/__init__.py", "litellm/proxy/example_config_yaml/*", "tests/*"]
 
 
+[lint.flake8-tidy-imports]
+banned-module-level-imports = ["prisma"]
+
 [lint.per-file-ignores]
 "litellm/main.py" = ["F401"]
 "litellm/types/caching.py" = ["F401"]

tests/unit/proxy/management_endpoints/test_common_daily_activity.py

@@ -1,13 +1,15 @@
 from collections.abc import Mapping, Sequence
-from dataclasses import asdict
+from dataclasses import asdict, fields
 from datetime import date, datetime
 from types import SimpleNamespace
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock
 
+import httpx
 import pytest
 from fastapi import HTTPException
-from prisma.errors import PrismaError
+from prisma.engine.utils import handle_response_errors
+from prisma.errors import DataError, PrismaError
 
 import litellm.proxy.management_endpoints.common_daily_activity as common_daily_activity_module
 from litellm.constants import USAGE_TOP_API_KEYS_DEFAULT
@@ -38,16 +40,65 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import (
     SpendAnalyticsPaginatedResponse,
     SpendMetrics,
 )
-from litellm.types.repositories.daily_activity import GroupingSetsRow, KeyMetadataRow
+from litellm.types.repositories.daily_activity import (
+    AggregatedRows,
+    DailyActivityScope,
+    EntityRollupRow,
+    GroupingSetsRow,
+    KeyMetadataRow,
+)
+
+
+def _query_engine_crash_on_closed_connection() -> AttributeError:
+    with pytest.raises(AttributeError) as crash:
+        handle_response_errors(
+            MagicMock(),
+            [
+                {
+                    "error": "Server has closed the connection.",
+                    "user_facing_error": {
+                        "is_panic": False,
+                        "message": "Server has closed the connection.",
+                        "meta": None,
+                        "error_code": "P1017",
+                    },
+                }
+            ],
+        )
+    return crash.value
+
+
+DATABASE_FAILURES: Final = pytest.mark.parametrize(
+    "failure",
+    [
+        DataError(
+            {"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server at `db:5432`"}}
+        ),
+        DataError({"user_facing_error": {"error_code": "P2010", "message": "Raw query failed"}}),
+        PrismaError("metadata unavailable"),
+        httpx.ReadTimeout("the query engine stopped answering"),
+        httpx.ConnectError("the query engine refused the connection"),
+        _query_engine_crash_on_closed_connection(),
+    ],
+    ids=[
+        "unreachable-database",
+        "rejected-query",
+        "prisma-error",
+        "engine-read-timeout",
+        "engine-connect-error",
+        "engine-crash-on-closed-connection",
+    ],
+)
 
 
 @pytest.mark.asyncio
+@DATABASE_FAILURES
 @pytest.mark.parametrize(
     ("key_failure", "user_failure", "empty"),
     [(False, False, False), (True, False, False), (False, True, False), (True, True, False), (False, False, True)],
 )
 async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata(
-    key_failure: bool, user_failure: bool, empty: bool
+    failure: Exception, key_failure: bool, user_failure: bool, empty: bool
 ) -> None:
     prisma: Final = MagicMock()
     key: Final = SimpleNamespace(token="key", key_alias="Alias", team_id=None, user_id=None, metadata=None)
@@ -62,11 +113,11 @@ async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata(
     )
     prisma.db.query_raw = AsyncMock(side_effect=[main_rows, entity_rows])
     prisma.db.litellm_verificationtoken.find_many = AsyncMock(
-        side_effect=[[key], PrismaError("key metadata unavailable") if key_failure else ([] if empty else [key])]
+        side_effect=[[key], failure if key_failure else ([] if empty else [key])]
     )
     prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
     prisma.db.litellm_usertable.find_many = AsyncMock(
-        side_effect=PrismaError("user metadata unavailable") if user_failure else None,
+        side_effect=failure if user_failure else None,
         return_value=[] if empty else [user],
     )
 
@@ -95,6 +146,85 @@ async def test_entity_metadata_outages_preserve_daily_spend_and_other_metadata(
     )
 
 
+def _usage_scope() -> DailyActivityScope:
+    return daily_activity_scope(
+        "litellm_dailyuserspend", "user_id", None, None, None, "2024-01-01", "2024-01-01", None, None
+    )
+
+
+def _entity_row() -> EntityRollupRow:
+    grouping: Final = asdict(_grouping_row(63, api_key="key", spend=2.0))
+    return EntityRollupRow(
+        **{
+            **{field.name: grouping[field.name] for field in fields(EntityRollupRow) if field.name in grouping},
+            "entity_id": "user",
+            "api_key_rolled": 0,
+        }
+    )
+
+
+def _usage_repository() -> MagicMock:
+    repository: Final = MagicMock()
+    repository.aggregated = AsyncMock(
+        return_value=AggregatedRows(
+            grouping_rows=tuple(
+                _grouping_row(level, api_key=api_key, spend=2.0)
+                for level, api_key in ((127, None), (63, None), (31, "key"))
+            ),
+            entity_rows=(_entity_row(),),
+            distinct_api_keys=1,
+        )
+    )
+    repository.key_metadata = AsyncMock(return_value={})
+    repository.user_metadata = AsyncMock(return_value={})
+    repository.guardrail_spend = AsyncMock(return_value=())
+    return repository
+
+
+@pytest.mark.asyncio
+@DATABASE_FAILURES
+async def test_usage_keeps_spend_when_the_first_key_lookup_hits_a_database_failure(failure: Exception) -> None:
+    repository: Final = _usage_repository()
+    repository.key_metadata = AsyncMock(side_effect=failure)
+
+    result: Final = await _get_daily_activity_aggregated(repository, _usage_scope())
+
+    day: Final = result.results[0]
+    assert result.metadata.total_spend == day.metrics.spend == 2.0
+    assert day.breakdown.api_keys["key"].metrics.spend == 2.0
+    assert day.breakdown.api_keys["key"].metadata == KeyMetadata(key_alias=None, key_exists=False)
+
+
+@pytest.mark.asyncio
+@DATABASE_FAILURES
+async def test_usage_keeps_spend_when_the_guardrail_lookup_hits_a_database_failure(failure: Exception) -> None:
+    repository: Final = _usage_repository()
+    repository.guardrail_spend = AsyncMock(side_effect=failure)
+
+    result: Final = await _get_daily_activity_aggregated(repository, _usage_scope(), include_guardrail_breakdown=True)
+
+    assert result.metadata.total_spend == result.results[0].metrics.spend == 2.0
+    assert result.results[0].breakdown.guardrails == {}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("lookup", ["key_metadata", "user_metadata", "guardrail_spend"])
+@pytest.mark.parametrize("bug", [TypeError, AttributeError])
+async def test_usage_still_reports_a_lookup_bug_that_is_not_a_database_failure(
+    lookup: str, bug: type[Exception]
+) -> None:
+    repository: Final = _usage_repository()
+    setattr(repository, lookup, AsyncMock(side_effect=bug("not a database failure")))
+
+    with pytest.raises(HTTPException) as error:
+        await _get_daily_activity_aggregated(
+            repository, _usage_scope(), include_entity_breakdown=True, include_guardrail_breakdown=True
+        )
+
+    assert error.value.status_code == 500
+    assert error.value.detail == {"error": "Failed to fetch analytics: not a database failure"}
+
+
 async def _run_aggregated_daily_activity(
     *,
     prisma_client: PrismaClient,

tests/unit/proxy/management_endpoints/test_daily_activity_routes.py

@@ -7,9 +7,11 @@ from types import MappingProxyType, SimpleNamespace
 from typing import Final
 from unittest.mock import AsyncMock
 
+import httpx
 import pytest
 from fastapi import FastAPI, Request
 from fastapi.testclient import TestClient
+from prisma.errors import DataError
 
 from litellm import constants
 from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
@@ -1015,6 +1017,96 @@ def test_user_page_route_ranks_users_by_spend_and_attaches_metadata(
     assert repository.user_metadata.call_args.args[0] == frozenset(("user-a", "user-b"))
 
 
+_LOOKUP_DATABASE_FAILURES: Final = pytest.mark.parametrize(
+    "failure",
+    [
+        DataError({"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server"}}),
+        httpx.ConnectError("the query engine refused the connection"),
+    ],
+    ids=["unreachable-database", "engine-connect-error"],
+)
+_KEY_ROWS_ROUTES: Final = pytest.mark.parametrize(
+    ("suffix", "extra_params"),
+    [
+        ("keys", {}),
+        ("search", {"search": "key"}),
+        ("model_top_keys", {"model_group": "rare-group"}),
+        ("cache_leakage_keys", {}),
+    ],
+    ids=["keys", "search", "model_top_keys", "cache_leakage_keys"],
+)
+
+
+@_LOOKUP_DATABASE_FAILURES
+@_KEY_ROWS_ROUTES
+def test_key_routes_keep_spend_when_the_key_lookup_hits_a_database_failure(
+    daily_activity_client: tuple[TestClient, _FakeRepository],
+    suffix: str,
+    extra_params: Mapping[str, str],
+    failure: Exception,
+) -> None:
+    client, repository = daily_activity_client
+    repository.key_metadata.side_effect = failure
+
+    response: Final = client.get(
+        f"/user/daily/activity/aggregated/{suffix}",
+        params={**_entity_params("user_id", "user-a"), **extra_params},
+    )
+
+    assert response.status_code == 200, response.text
+    rows: Final = response.json()["api_keys"]
+    assert rows
+    assert all(row["metrics"]["spend"] > 0 for row in rows)
+    assert {row["metadata"]["key_alias"] for row in rows} == {None}
+
+
+@_KEY_ROWS_ROUTES
+def test_key_routes_still_report_a_lookup_bug_that_is_not_a_database_failure(
+    daily_activity_client: tuple[TestClient, _FakeRepository],
+    suffix: str,
+    extra_params: Mapping[str, str],
+) -> None:
+    client, repository = daily_activity_client
+    repository.key_metadata.side_effect = AttributeError("not a database failure")
+
+    response: Final = client.get(
+        f"/user/daily/activity/aggregated/{suffix}",
+        params={**_entity_params("user_id", "user-a"), **extra_params},
+    )
+
+    assert response.status_code == 500, response.text
+    assert "not a database failure" in response.text
+
+
+@_LOOKUP_DATABASE_FAILURES
+def test_user_page_route_keeps_spend_when_the_user_lookup_hits_a_database_failure(
+    daily_activity_client: tuple[TestClient, _FakeRepository],
+    failure: Exception,
+) -> None:
+    client, repository = daily_activity_client
+    repository.user_metadata.side_effect = failure
+
+    response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS)
+
+    assert response.status_code == 200, response.text
+    assert [(user["user_id"], user["spend"], user["user_email"]) for user in response.json()["users"]] == [
+        ("user-b", 100.0, None),
+        ("user-a", 14.5, None),
+    ]
+
+
+def test_user_page_route_still_reports_a_lookup_bug_that_is_not_a_database_failure(
+    daily_activity_client: tuple[TestClient, _FakeRepository],
+) -> None:
+    client, repository = daily_activity_client
+    repository.user_metadata.side_effect = AttributeError("not a database failure")
+
+    response: Final = client.get("/user/daily/activity/aggregated/users", params=_DATE_PARAMS)
+
+    assert response.status_code == 500, response.text
+    assert "not a database failure" in response.text
+
+
 def test_user_page_route_forwards_offset_and_limit(
     daily_activity_client: tuple[TestClient, _FakeRepository],
 ) -> None: