From ee95ac14a57c856d74dbe4c77131892d426fbcb9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 02:56:56 -0800 Subject: [PATCH] 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. --- .../common_daily_activity.py | 73 +++++---- .../daily_activity_routes.py | 15 +- ruff.toml | 5 +- .../test_common_daily_activity.py | 142 +++++++++++++++++- .../test_daily_activity_routes.py | 92 ++++++++++++ 5 files changed, 276 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 3f5f70e58b4..dedc28e4830 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/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 diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 1a3445aae52..b5af673585c 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/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] diff --git a/ruff.toml b/ruff.toml index 715e5551fad..037b15dbfe9 100644 --- a/ruff.toml +++ b/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"] diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index f49f42a3d83..3e54c0567bc 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/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, diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index d658e59b76c..90690111867 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/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: