From 2d0c7204252f1b983ab23f18adbaa6396b6af161 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 3 Nov 2026 16:07:00 -0800 Subject: [PATCH] fix(proxy): cancel the Admin UI's raw spend-log reads in Postgres after 30s (#49965) * fix(proxy): cancel the Admin UI's raw spend-log reads in Postgres after 30s Prisma's HTTP client gives up on a statement after 30s, so the Logs, Caching, Cost Optimization and spend report pages already fail at 30s, but Postgres kept running the scan. Every raw LiteLLM_SpendLogs read behind those pages now runs under SET LOCAL statement_timeout = 30000 through one shared helper, which /team/spend/by_user also uses instead of its own copy * fix(proxy): keep the spend report's 503 and time out the session totals read The grouped Logs page's session totals query now runs under the same statement timeout, /global/spend/report re-raises the helper's 503 instead of turning it into a 400, and the hand-written prisma fakes in the endpoint tests expose read_db so every guarded route is exercised * fix(proxy): keep the Logs page when Postgres cancels its session totals read A cancelled session totals read degrades like any other failed enrichment read instead of failing the whole page. The span-name test treats the shared statement-timeout helper as a forwarder, and the spend report test no longer depends on a license in the environment * test(proxy): /global/spend/report keeps the 503 hint when Postgres cancels the read --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../analytics_endpoints/cache_activity.py | 65 +++++--- .../prompt_caching_requests.py | 8 +- .../management_endpoints/team_endpoints.py | 43 ++--- .../spend_management_endpoints.py | 59 +++++-- .../spend_tracking/spend_tracking_utils.py | 5 +- litellm/proxy/utils.py | 34 ++++ .../test_analytics_endpoints.py | 44 ++++- tests/unit/proxy/db/fake_prisma_engine.py | 76 ++++++++- tests/unit/proxy/db/test_prisma_query_span.py | 3 +- .../test_prompt_caching_requests.py | 36 +++++ .../test_team_endpoints.py | 49 +----- .../test_spend_management_endpoints.py | 74 +++++++++ .../test_spend_query_optimization.py | 151 +++++++++++++++--- tests/unit/proxy/test_proxy_utils.py | 52 ++++++ 15 files changed, 554 insertions(+), 146 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 97631b9997a..629cfa34473 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2067,6 +2067,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 +SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS: Final = 30 SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding diff --git a/litellm/proxy/analytics_endpoints/cache_activity.py b/litellm/proxy/analytics_endpoints/cache_activity.py index 8a3f71e656e..17e250a2dbe 100644 --- a/litellm/proxy/analytics_endpoints/cache_activity.py +++ b/litellm/proxy/analytics_endpoints/cache_activity.py @@ -2,31 +2,22 @@ import asyncio import json from collections.abc import Sequence from datetime import datetime -from typing import Final, Protocol +from typing import TYPE_CHECKING, Final from pydantic import TypeAdapter from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.auth.authorization import ReadScope, read_scope_bind_params +from litellm.proxy.utils import query_raw_within_statement_timeout from litellm.types.llms.base import LiteLLMBaseModel +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + UNKNOWN_CALL_TYPE: Final = "Unknown" INFO_ROUTES_JSON: Final = json.dumps(LiteLLMRoutes.info_routes.value) -class _SupportsQueryRaw(Protocol): - """The single database operation the cache-activity queries issue.""" - - async def query_raw(self, query: str, *args: object) -> Sequence[object]: ... - - -class _SupportsRawQueryDb(Protocol): - """A prisma client handle, narrowed to the raw-query surface used here.""" - - @property - def db(self) -> _SupportsQueryRaw: ... - - class CacheActivityGroup(LiteLLMBaseModel): call_type: str api_requests: int @@ -174,7 +165,7 @@ def compute_totals(groups: Sequence[CacheActivityGroup]) -> CacheActivityTotals: async def get_cache_activity( - prisma_client: _SupportsRawQueryDb, + prisma_client: "PrismaClient", read_scope: ReadScope, start_date: datetime, end_date: datetime, @@ -185,14 +176,46 @@ async def get_cache_activity( models_json: Final = json.dumps(list(models)) scope_params: Final = read_scope_bind_params(read_scope) group_rows, error_rows, key_alias_rows, model_rows = await asyncio.gather( - prisma_client.db.query_raw( - GROUPS_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON, *scope_params + query_raw_within_statement_timeout( + prisma_client, + GROUPS_SQL, + start_date, + end_date, + key_aliases_json, + models_json, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", ), - prisma_client.db.query_raw( - ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON, *scope_params + query_raw_within_statement_timeout( + prisma_client, + ERROR_BREAKDOWN_SQL, + start_date, + end_date, + key_aliases_json, + models_json, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", + ), + query_raw_within_statement_timeout( + prisma_client, + KEY_ALIAS_OPTIONS_SQL, + start_date, + end_date, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", + ), + query_raw_within_statement_timeout( + prisma_client, + MODEL_OPTIONS_SQL, + start_date, + end_date, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", ), - prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON, *scope_params), - prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON, *scope_params), ) groups: Final = _groups_adapter.validate_python(group_rows or []) return CacheActivityResponse( diff --git a/litellm/proxy/management_endpoints/prompt_caching_requests.py b/litellm/proxy/management_endpoints/prompt_caching_requests.py index 2f00d263d84..d52c0255b08 100644 --- a/litellm/proxy/management_endpoints/prompt_caching_requests.py +++ b/litellm/proxy/management_endpoints/prompt_caching_requests.py @@ -14,9 +14,7 @@ from litellm.proxy.spend_tracking.savings import ( marks_gateway_injection, prompt_caching_savings_for_request, ) -from litellm.proxy.spend_tracking.spend_tracking_utils import ( - _query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below -) +from litellm.proxy.utils import query_raw_within_statement_timeout from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.management_endpoints.prompt_caching_requests import ( @@ -158,7 +156,7 @@ async def get_prompt_caching_requests( else cursor_start_time ) rows: Final = _REQUEST_ROWS.validate_python( - await _query_raw_rows( + await query_raw_within_statement_timeout( prisma_client, prompt_caching_requests_sql(filter), start.isoformat(), @@ -166,8 +164,8 @@ async def get_prompt_caching_requests( page_size + 1, cursor_time.isoformat() if cursor_time is not None else None, cursor_request_id, + what="Prompt caching requests", ) - or () ) def current_router() -> "Router | None": diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e563944a7b3..54b9b26123e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -140,7 +140,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_reservation_cache_key, ) from litellm.proxy.db.db_span import db_span -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( RoutedModels, @@ -225,7 +224,12 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) -from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy +from litellm.proxy.utils import ( + PrismaClient, + ProxyLogging, + handle_exception_on_proxy, + query_raw_within_statement_timeout, +) from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import delete_many_in, find_many_in from litellm.repositories.daily_activity_buckets import api_key_from_bucket @@ -7474,17 +7478,7 @@ class _TeamUserSpendDbRow(TypedDict): failed_requests: ReadOnly[int] -_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS: Final = 30 - - -async def _team_user_spend_rows( - prisma_client: PrismaClient, sql: str, params: tuple[str, ...] -) -> Sequence[_TeamUserSpendDbRow]: - async with prisma_client.read_db.tx( - timeout=timedelta(seconds=_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS + 5) - ) as tx: - await tx.execute_raw(f"SET LOCAL statement_timeout = {_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS * 1000}") - return await tx.query_raw(sql, *params) +_TEAM_USER_SPEND_ROWS: Final = TypeAdapter(tuple[_TeamUserSpendDbRow, ...]) @router.get( @@ -7536,23 +7530,18 @@ async def get_team_spend_by_user( own_user_only: Final = scope.api_key_filter is not None user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else () - try: - async with db_span("team_user_spend", "LiteLLM_SpendLogs"): - rows: Final = await _team_user_spend_rows( + async with db_span("team_user_spend", "LiteLLM_SpendLogs"): + rows: Final = _TEAM_USER_SPEND_ROWS.validate_python( + await query_raw_within_statement_timeout( prisma_client, _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), - (start_date, end_date, *scoped_team_ids, *user_param), + start_date, + end_date, + *scoped_team_ids, + *user_param, + what="Spend per user", ) - except Exception as e: - if PrismaDBExceptionHandler.is_statement_timeout_error(e): - raise _daily_activity_error( - status_code=503, - message=( - f"Spend per user took longer than {_TEAM_USER_SPEND_STATEMENT_TIMEOUT_SECONDS}s. " - "Pick a shorter date range or fewer teams" - ), - ) from e - raise + ) results: Final = tuple( TeamUserSpendRow( team_id=row["team_id"], diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index fcd719b24b9..d55af56e248 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3,6 +3,7 @@ import collections import json import os from collections.abc import Mapping, Sequence +from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from functools import partial from itertools import groupby @@ -72,7 +73,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team, get_spend_by_team_and_customer, ) -from litellm.proxy.utils import handle_exception_on_proxy +from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository @@ -267,6 +268,30 @@ async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args: return await _query_raw(prisma_client, sql_query, *args) +async def _query_spend_logs(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]: + """``_query_raw`` for raw ``LiteLLM_SpendLogs`` reads, which Postgres cancels at the statement timeout.""" + return cast( # cast-ok: callers declare the row shape their SELECT list produces, as with _query_raw + "Sequence[_RowT]", await query_raw_within_statement_timeout(prisma_client, sql_query, *args) + ) + + +async def _query_spend_logs_or_none( + prisma_client: PrismaClient, sql_query: str, *args: object +) -> Sequence[_RowT] | None: + """``_query_spend_logs`` for the call sites that guard the result against ``None``.""" + return await _query_spend_logs(prisma_client, sql_query, *args) + + +@dataclass(frozen=True, slots=True) +class _SpendLogsReader: + """``read_session_page``'s raw-query handle, each read cancelled by Postgres at the statement timeout.""" + + prisma_client: PrismaClient + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + return await query_raw_within_statement_timeout(self.prisma_client, query, *args) + + class _TeamTable(Protocol): """The subset of the Prisma team table API this module uses.""" @@ -1566,7 +1591,9 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj, api_key) + db_response = await _query_spend_logs_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj, api_key + ) if db_response is None: return [] @@ -1609,7 +1636,7 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await _query_raw_or_none( + db_response = await _query_spend_logs_or_none( prisma_client, sql_query, start_date_obj, end_date_obj, internal_user_id ) if db_response is None: @@ -1676,7 +1703,7 @@ async def get_global_spend_report( group_by_day; """ - db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) + db_response = await _query_spend_logs_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1717,12 +1744,14 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) + db_response = await _query_spend_logs_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] return db_response + except HTTPException: + raise except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -1966,7 +1995,7 @@ async def get_key_spend_report( caller_value=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), scope_name="api_key", ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _scoped_spend_report_sql(scope_column="api_key"), start_date_obj, @@ -2008,7 +2037,7 @@ async def get_user_spend_report( caller_value=user_api_key_dict.user_id, scope_name="internal_user_id", ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _scoped_spend_report_sql(scope_column="user"), start_date_obj, @@ -2050,7 +2079,7 @@ async def get_team_spend_report( caller_value=user_api_key_dict.team_id, scope_name="team_id", ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _scoped_spend_report_sql(scope_column="team_id"), start_date_obj, @@ -2092,7 +2121,7 @@ async def get_organization_spend_report( organization_id=organization_id, prisma_client=prisma_client, ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _ORG_SPEND_REPORT_SQL, start_date_obj, @@ -3017,7 +3046,7 @@ async def ui_view_spend_logs( LIMIT ${p} ) AS bounded_matches """ - count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_raw_or_none( + count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_spend_logs_or_none( prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 ) raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 @@ -3048,7 +3077,7 @@ async def ui_view_spend_logs( ) sql_params.extend([page_size, skip]) - data: Final[Sequence[Mapping[str, object]]] = await _query_raw(prisma_client, sql_query, *sql_params) + data: Final[Sequence[Mapping[str, object]]] = await _query_spend_logs(prisma_client, sql_query, *sql_params) if request_id is not None and not is_v2 and not is_admin_view: await _assert_user_owns_fetched_spend_rows( @@ -3102,7 +3131,7 @@ async def _fetch_session_representatives( ORDER BY {_SESSION_GROUP_KEY_SQL}, {SESSION_REPRESENTATIVE_ORDER_SQL} ) AS session_representatives """ - rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place + rep_rows: Final[Sequence[dict[str, object]]] = await _query_spend_logs( # mutable-ok: rows are enriched in place prisma_client, rep_query, *sql_params, @@ -3130,7 +3159,7 @@ async def _ui_session_grouped_spend_logs( and each is shown by its final agent outcome or newest non-MCP row.""" where_clause: Final = " AND ".join(sql_conditions) if sql_conditions else "TRUE" session_page: Final = await read_session_page( - prisma_client.db, # pyright: ignore[reportArgumentType] # PrismaWrapper forwards raw queries via __getattr__ + _SpendLogsReader(prisma_client), # pyright: ignore[reportArgumentType] # rows are the TypedDict the session SQL selects where_clause=where_clause, params=sql_params, status=status_filter, @@ -4605,7 +4634,7 @@ async def _build_ui_spend_logs_response( if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None } ) - rows: Final[Sequence[_SessionSpendRow]] = await _query_raw( + rows: Final[Sequence[_SessionSpendRow]] = await _query_spend_logs( prisma_client, f""" SELECT s.*, COALESCE(m.session_models, ARRAY[]::text[]) AS session_models @@ -4676,7 +4705,7 @@ async def _build_ui_spend_logs_response( if row.get("session_id") and row.get("api_key") is not None for models in (TypeAdapter(list[str]).validate_python(row.get("session_models") or ()),) } - except PrismaError: + except (PrismaError, HTTPException): verbose_proxy_logger.debug( "Failed to enrich session spend aggregates for spend logs UI", exc_info=True, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a9478a1f889..b1544b1e042 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -75,7 +75,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error -from litellm.proxy.utils import PrismaClient, hash_token +from litellm.proxy.utils import PrismaClient, hash_token, query_raw_within_statement_timeout from litellm.responses.litellm_completion_transformation.compaction_items import ( compaction_summary, summary_only_compaction_content, @@ -1092,7 +1092,8 @@ async def _query_raw_rows( sql_query: str, *args: object, ) -> Sequence[Mapping[str, object]] | None: - return await prisma_client.db.query_raw(sql_query, *args) + """A raw ``LiteLLM_SpendLogs`` read that Postgres cancels at the statement timeout.""" + return await query_raw_within_statement_timeout(prisma_client, sql_query, *args) async def get_spend_by_team( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 837433fd8cc..7bf5e6d28af 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -66,6 +66,7 @@ from litellm.constants import ( REDACTED_BY_LITELLM, REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT, SPEND_LOG_QUEUE_MAX_BYTES, + SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, VLLM_CHOICE_FIELDS, @@ -5309,6 +5310,39 @@ _deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(ma _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60 _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) _PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2) +_SPEND_LOG_READ_STATEMENT_TIMEOUT: Final = timedelta(seconds=SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS) +_STATEMENT_TIMEOUT_TRANSACTION_GRACE: Final = timedelta(seconds=5) + + +class _ReadTransactionClient(Protocol): + @property + def read_db(self) -> PrismaWrapper: ... + + +async def query_raw_within_statement_timeout( + prisma_client: _ReadTransactionClient, + sql: str, + *params: object, + what: str = "Spend logs", + timeout: timedelta = _SPEND_LOG_READ_STATEMENT_TIMEOUT, +) -> Sequence[Mapping[str, object]]: + """One raw read on the read target that Postgres cancels at ``timeout``. + + Prisma's HTTP client gives up on a statement after 30s anyway, so without + this the Admin UI's error is the same but Postgres keeps running the scan. + A cancelled read surfaces as a 503 asking for a narrower date range. + """ + try: + async with prisma_client.read_db.tx(timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE) as transaction: + await transaction.execute_raw(f"SET LOCAL statement_timeout = {int(timeout / timedelta(milliseconds=1))}") + return await transaction.query_raw(sql, *params) + except Exception as e: + if not PrismaDBExceptionHandler.is_statement_timeout_error(e): + raise + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": f"{what} took longer than {int(timeout.total_seconds())}s. Pick a shorter date range"}, + ) from e async def _lookup_deprecated_key( diff --git a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py index 1c706978fed..2d28292a487 100644 --- a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py +++ b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py @@ -26,6 +26,7 @@ from litellm.proxy.analytics_endpoints.cache_activity import ( ADMIN: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") INTERNAL_USER: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1") +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads async def _no_teams(auth: UserAPIKeyAuth) -> tuple[str, ...]: @@ -70,10 +71,15 @@ MODEL_ROWS = [{"model": "gpt-5.1"}] def build_prisma(query_raw: AsyncMock) -> MagicMock: prisma = MagicMock() - prisma.db.query_raw = query_raw + statement_timeout_reads(prisma, query_raw) return prisma +def _spend_log_reads(prisma: MagicMock) -> AsyncMock: + """The raw reads the cache dashboard ran under the statement timeout.""" + return prisma.read_db.tx.return_value.query_raw + + def dispatching_query_raw() -> AsyncMock: async def dispatch(sql: str, *params: object) -> list[dict[str, object]]: if "error_code" in sql: @@ -120,7 +126,9 @@ async def test_filters_are_passed_to_sql_as_json_arrays(mock_prisma: MagicMock): await _activity(ADMIN, key_aliases=["my-key"], models=["gpt-5.1", "claude-opus-4-8"]) filtered_calls = [ - call for call in mock_prisma.db.query_raw.call_args_list if call.args[0] in (GROUPS_SQL, ERROR_BREAKDOWN_SQL) + call + for call in _spend_log_reads(mock_prisma).call_args_list + if call.args[0] in (GROUPS_SQL, ERROR_BREAKDOWN_SQL) ] assert len(filtered_calls) == 2 for call in filtered_calls: @@ -135,7 +143,9 @@ async def test_every_query_excludes_the_same_info_routes(mock_prisma: MagicMock) the SQL does with it is covered against Postgres in tests/integration/spend/test_cache_activity.py.""" await _activity(ADMIN) - exclusions_by_query = {call.args[0]: json.loads(call.args[-4]) for call in mock_prisma.db.query_raw.call_args_list} + exclusions_by_query = { + call.args[0]: json.loads(call.args[-4]) for call in _spend_log_reads(mock_prisma).call_args_list + } assert set(exclusions_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL} for excluded_call_types in exclusions_by_query.values(): assert excluded_call_types == LiteLLMRoutes.info_routes.value @@ -147,7 +157,7 @@ async def test_every_query_excludes_the_same_info_routes(mock_prisma: MagicMock) async def test_admin_queries_grant_all_rows(mock_prisma: MagicMock): await _activity(ADMIN) - scope_by_query = {call.args[0]: call.args[-3:] for call in mock_prisma.db.query_raw.call_args_list} + scope_by_query = {call.args[0]: call.args[-3:] for call in _spend_log_reads(mock_prisma).call_args_list} assert set(scope_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL} for scope_args in scope_by_query.values(): assert scope_args == (True, None, "[]") @@ -157,7 +167,7 @@ async def test_admin_queries_grant_all_rows(mock_prisma: MagicMock): async def test_internal_user_queries_are_scoped_to_own_and_team_rows(mock_prisma: MagicMock): await _activity(INTERNAL_USER, team_lookup=_team_nine) - scope_by_query = {call.args[0]: call.args[-3:] for call in mock_prisma.db.query_raw.call_args_list} + scope_by_query = {call.args[0]: call.args[-3:] for call in _spend_log_reads(mock_prisma).call_args_list} assert set(scope_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL} for scope_args in scope_by_query.values(): assert scope_args == (False, "user-1", json.dumps(["team-9"])) @@ -168,7 +178,7 @@ async def test_key_without_user_id_grants_no_rows(mock_prisma: MagicMock): no_user = UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, user_id=None) await _activity(no_user) - for call in mock_prisma.db.query_raw.call_args_list: + for call in _spend_log_reads(mock_prisma).call_args_list: assert call.args[-3:] == (False, None, "[]") @@ -178,7 +188,27 @@ async def test_rejects_malformed_dates_with_400(mock_prisma: MagicMock): await _activity(ADMIN, start_date="07/01/2026") assert exc_info.value.status_code == 400 - mock_prisma.db.query_raw.assert_not_called() + _spend_log_reads(mock_prisma).assert_not_called() + + +@pytest.mark.asyncio +async def test_every_cache_activity_read_is_cancelled_by_postgres_after_30s(mock_prisma: MagicMock): + await _activity(ADMIN) + + transaction = mock_prisma.read_db.tx.return_value + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) * 4 + + +@pytest.mark.asyncio +async def test_cancelled_cache_activity_read_is_reported_as_503_with_a_date_range_hint(monkeypatch: pytest.MonkeyPatch): + prisma = build_prisma(AsyncMock(side_effect=statement_timeout_error())) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + + with pytest.raises(HTTPException) as exc_info: + await _activity(ADMIN) + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": "Cache activity took longer than 30s. Pick a shorter date range"} def test_totals_ratio_is_zero_without_requests(): diff --git a/tests/unit/proxy/db/fake_prisma_engine.py b/tests/unit/proxy/db/fake_prisma_engine.py index 4221eeced5a..082659a3adc 100644 --- a/tests/unit/proxy/db/fake_prisma_engine.py +++ b/tests/unit/proxy/db/fake_prisma_engine.py @@ -2,8 +2,10 @@ marking the DB I/O witness the way ``_TrackedPrismaEngine`` does, so the producer under test emits its service event.""" -from typing import TypeVar -from unittest.mock import AsyncMock +from typing import Final, TypeVar +from unittest.mock import AsyncMock, MagicMock + +from prisma.errors import DataError from litellm.proxy.db.log_db_metrics import record_db_io @@ -16,3 +18,73 @@ def engine_call(return_value: _T | None = None) -> AsyncMock: return return_value return AsyncMock(side_effect=run) + + +class RecordingReadTransaction: + """The transaction ``query_raw_within_statement_timeout`` opens on ``read_db``, keeping its ``SET LOCAL`` statements.""" + + def __init__(self, query_raw: AsyncMock) -> None: + self.statements: tuple[str, ...] = () + self.query_raw: AsyncMock = query_raw + + async def execute_raw(self, sql: str) -> int: + self.statements = (*self.statements, sql) + return 0 + + async def __aenter__(self) -> "RecordingReadTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> None: + pass + + +def statement_timeout_reads(prisma: MagicMock, query_raw: AsyncMock) -> RecordingReadTransaction: + """Make ``prisma.read_db.tx()`` hand out one recording transaction whose raw reads answer with ``query_raw``.""" + transaction: Final = RecordingReadTransaction(query_raw) + prisma.read_db.tx = MagicMock(return_value=transaction) + return transaction + + +def statement_timeout_error() -> DataError: + """The error prisma raises when Postgres cancels a statement under ``statement_timeout`` (SQLSTATE 57014).""" + message: Final = "canceling statement due to statement timeout" + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": f"Error querying the database: {message}", + "meta": {"code": "57014", "message": message}, + } + } + ) + + +class _DelegatingReadTransaction: + """``read_db.tx()`` of a hand-written prisma fake: raw reads go to the fake's ``db.query_raw``.""" + + def __init__(self, db: object) -> None: + self.db: Final = db + self.statements: tuple[str, ...] = () + + async def execute_raw(self, sql: str) -> int: + self.statements = (*self.statements, sql) + return 0 + + async def query_raw(self, sql: str, *args: object) -> object: + return await getattr(self.db, "query_raw")(sql, *args) + + async def __aenter__(self) -> "_DelegatingReadTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> None: + pass + + +class ReadTargetOf: + """What a prisma fake returns from ``read_db``: ``tx()`` opens a transaction over the fake's ``db``.""" + + def __init__(self, db: object) -> None: + self.db: Final = db + + def tx(self, **kwargs: object) -> _DelegatingReadTransaction: + return _DelegatingReadTransaction(self.db) diff --git a/tests/unit/proxy/db/test_prisma_query_span.py b/tests/unit/proxy/db/test_prisma_query_span.py index 9f2ba28ea7e..55ba03d9f2c 100644 --- a/tests/unit/proxy/db/test_prisma_query_span.py +++ b/tests/unit/proxy/db/test_prisma_query_span.py @@ -19,6 +19,7 @@ from litellm.proxy.db.prisma_query_span import UNKNOWN_PRISMA_QUERY, parse_prism _REPO: Final = Path(__file__).resolve().parents[4] _SOURCE_ROOTS: Final = ("litellm", "enterprise", "litellm-proxy-extras") _RAW_METHODS: Final = frozenset({"query_first", "query_raw", "execute_raw"}) +_CROSS_MODULE_FORWARDERS: Final = frozenset({"query_raw_within_statement_timeout"}) _MODEL_METHODS: Final = frozenset( { "find_unique", @@ -297,7 +298,7 @@ def _parameter_site( """A statement that arrives as a parameter: a ``query_raw`` forwarder adds no round trip of its own, any other helper is named by what its callers in the module hand it.""" function: Final = next(p for p in parents if isinstance(p, (ast.AsyncFunctionDef, ast.FunctionDef))) - if function.name in _RAW_METHODS: + if function.name in _RAW_METHODS or function.name in _CROSS_MODULE_FORWARDERS: return _PrismaCallSite(location, method, "forwarder", f"(callers of {function.name})") callers: Final = tuple( node diff --git a/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py b/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py index 39dcc9630df..ea61177f773 100644 --- a/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py +++ b/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from typing import Final +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -8,6 +9,7 @@ from fastapi import FastAPI from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.prompt_caching_requests import router +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads pytestmark = pytest.mark.usefixtures("local_model_cost_map") @@ -67,3 +69,37 @@ async def test_incomplete_cursor_is_rejected( ) as client: response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params}) assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_the_requests_read_is_cancelled_by_postgres_after_30s(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + prisma: Final = MagicMock() + transaction: Final = statement_timeout_reads(prisma, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END}) + + assert response.status_code == 200, response.text + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) + + +@pytest.mark.asyncio +async def test_a_cancelled_requests_read_is_a_503_with_a_date_range_hint(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + prisma: Final = MagicMock() + statement_timeout_reads(prisma, AsyncMock(side_effect=statement_timeout_error())) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END}) + + assert response.status_code == 503, response.text + assert response.json() == { + "detail": {"error": "Prompt caching requests took longer than 30s. Pick a shorter date range"} + } diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 863be3e2032..697d9810284 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -10,10 +10,10 @@ from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch import httpx import litellm +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads import pytest from fastapi import HTTPException, Request from fastapi.testclient import TestClient -from prisma.errors import DataError from pydantic import ValidationError from litellm._uuid import uuid @@ -17279,28 +17279,6 @@ def _team_spend_by_user_db_row(team_id: str, user_id: str, spend: float, request } -class _RecordingTransaction: - def __init__(self, query_raw: AsyncMock) -> None: - self.statements: tuple[str, ...] = () - self.query_raw: AsyncMock = query_raw - - async def execute_raw(self, sql: str) -> int: - self.statements = (*self.statements, sql) - return 0 - - async def __aenter__(self) -> "_RecordingTransaction": - return self - - async def __aexit__(self, *exc_info: object) -> None: - pass - - -def _team_spend_by_user_query(mock_db_client: MagicMock, query_raw: AsyncMock) -> _RecordingTransaction: - transaction: Final = _RecordingTransaction(query_raw) - mock_db_client.read_db.tx = MagicMock(return_value=transaction) - return transaction - - @pytest.mark.asyncio async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(mock_db_client): from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user @@ -17309,7 +17287,7 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta]) - transaction = _team_spend_by_user_query( + transaction = statement_timeout_reads( mock_db_client, AsyncMock( return_value=[ @@ -17356,7 +17334,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) - transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[])) + transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) ) @@ -17378,7 +17356,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"]) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - transaction = _team_spend_by_user_query( + transaction = statement_timeout_reads( mock_db_client, AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) ) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( @@ -17401,7 +17379,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER) - transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[])) + transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) ) @@ -17430,7 +17408,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error): from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user - transaction = _team_spend_by_user_query(mock_db_client, AsyncMock(return_value=[])) + transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc_info: @@ -17443,19 +17421,6 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids transaction.query_raw.assert_not_called() -def _statement_timeout_error() -> DataError: - message = "canceling statement due to statement timeout" - return DataError( - data={ - "user_facing_error": { - "is_panic": False, - "message": f"Error querying the database: {message}", - "meta": {"code": "57014", "message": message}, - } - } - ) - - @pytest.mark.asyncio async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_client): from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user @@ -17463,7 +17428,7 @@ async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_clie admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) - _team_spend_by_user_query(mock_db_client, AsyncMock(side_effect=_statement_timeout_error())) + statement_timeout_reads(mock_db_client, AsyncMock(side_effect=statement_timeout_error())) with pytest.raises(HTTPException) as exc_info: await get_team_spend_by_user( diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index f8d0ebb0f77..e119acd5ffc 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -20,6 +20,11 @@ import litellm.proxy.proxy_server as ps from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids from litellm.proxy.auth.route_checks import RouteChecks +from tests.unit.proxy.db.fake_prisma_engine import ( + ReadTargetOf, + statement_timeout_error, + statement_timeout_reads, +) def _default_date_range(): @@ -258,6 +263,7 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return [row for row in ordered[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -317,6 +323,7 @@ async def test_is_admin_view_safe_exception(): async def test_can_team_member_view_log_none_team_id(): # team_id=None should immediately return False class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -338,6 +345,7 @@ async def test_can_team_member_view_log_none_team_id(): async def test_can_team_member_view_log_team_not_found(monkeypatch): # Non-existent team should return False class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -377,6 +385,7 @@ async def test_can_team_member_view_log_not_admin(monkeypatch): } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -415,6 +424,7 @@ async def test_can_team_member_view_log_admin(monkeypatch): } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -480,6 +490,7 @@ def _make_owner_lookup_prisma(rows): return _emulate_spend_log_owner_lookup(rows, sql_query, params) class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -985,6 +996,7 @@ async def test_ui_view_spend_logs_sort_by_and_sort_order( return [row for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1046,6 +1058,7 @@ async def test_ui_view_spend_logs_sort_validation_errors(client, monkeypatch, so return 0 class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1122,6 +1135,7 @@ async def test_ui_view_spend_logs_sort_by_request_duration_ms(client, monkeypatc return [row for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1218,6 +1232,7 @@ async def test_ui_view_spend_logs_sort_by_model(client, monkeypatch, sort_order, return [row for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1330,6 +1345,7 @@ async def test_ui_view_spend_logs_sort_by_ttft_ms(client, monkeypatch): return [{k: v for k, v in row.items() if k != "_ttft_ms"} for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -2135,6 +2151,7 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): return [mock_spend_logs[0]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -2187,6 +2204,7 @@ async def test_ui_view_session_spend_logs_rehydrates_metadata_jsonb_text(client, return [dict(raw_row)] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -2235,6 +2253,7 @@ async def test_ui_view_session_spend_logs_includes_permitted_team_logs(client, m ] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -2712,6 +2731,7 @@ async def test_ui_view_request_response_rejects_foreign_row_inserted_after_owner ] class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -2743,6 +2763,7 @@ async def test_ui_view_request_response_custom_logger_denies_foreign_payload_own return [] class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4056,6 +4077,7 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch): ] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4151,6 +4173,7 @@ async def test_view_spend_logs_bounds_row_count(client, monkeypatch): return [{}] * min(kwargs.get("take", 0), self.available_rows) class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4274,6 +4297,7 @@ def _tag_query_prisma_client(assert_query_raw): return assert_query_raw(sql_query, params) class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4493,6 +4517,7 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch): return mock_summarized_response class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4555,6 +4580,7 @@ async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatc return mock_rows class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4620,6 +4646,7 @@ async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch): return [] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4659,6 +4686,7 @@ async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client return mock_rows class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4973,6 +5001,7 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5038,6 +5067,7 @@ async def test_build_ui_spend_logs_response_caps_session_models(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5102,6 +5132,7 @@ async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggre }, ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5155,6 +5186,7 @@ async def test_build_ui_spend_logs_response_empty_api_key_keeps_session_aggregat } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5210,6 +5242,7 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5292,6 +5325,7 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5372,6 +5406,7 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_duration(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5397,6 +5432,35 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_duration(): ) +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_keeps_the_page_when_postgres_cancels_the_session_totals_read(): + """The page already fetched is returned without session totals, like any other failed enrichment read.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + dict_rows = [ + {"request_id": "req-1", "session_id": "sess-slow", "call_type": "completion", "api_key": "k"}, + {"request_id": "req-2", "session_id": "sess-slow", "call_type": "completion", "api_key": "k"}, + ] + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=statement_timeout_error()) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=2, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + assert [row["request_id"] for row in result["data"]] == ["req-1", "req-2"] + assert all(row["session_total_count"] == 1 for row in result["data"]) + + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_session_cache_hit_count(): """ @@ -5430,6 +5494,7 @@ async def test_build_ui_spend_logs_response_session_cache_hit_count(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5476,6 +5541,7 @@ async def test_can_team_member_view_log_with_spend_logs_permission(monkeypatch): } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -5512,6 +5578,7 @@ async def test_can_team_member_view_log_without_spend_logs_permission(monkeypatc } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -6201,6 +6268,7 @@ async def test_ui_view_spend_logs_rehydrates_metadata_jsonb_text(client, monkeyp return [{**raw_row, "total_count": 1}] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -6285,6 +6353,7 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict return [{**raw_row, "total_count": 1}] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -6628,6 +6697,7 @@ _SCOPED_SPEND_REPORT_PATHS = [ def _spend_report_mock_prisma(query_raw_returns=None, team_rows=None, user_row=None): pc = MagicMock() pc.db.query_raw = AsyncMock(return_value=query_raw_returns if query_raw_returns is not None else []) + statement_timeout_reads(pc, pc.db.query_raw) pc.db.litellm_teamtable.find_many = AsyncMock(return_value=team_rows if team_rows is not None else []) pc.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) return pc @@ -7186,6 +7256,7 @@ def _session_grouped_mock_prisma( mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) return mock_prisma @@ -7241,6 +7312,7 @@ def _session_grouped_paginating_prisma(sessions: Sequence[tuple[str, str]]) -> M mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) return mock_prisma @@ -7561,6 +7633,7 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort return [] mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) monkeypatch.setattr( "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", @@ -7611,6 +7684,7 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" diff --git a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py index 7a38c1d8457..aaf4bc5438c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -18,6 +18,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team, get_spend_by_team_and_customer, ) +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads @pytest.mark.asyncio @@ -32,10 +33,8 @@ async def test_spend_query_uses_timestamp_filtering(): """ # Mock prisma client mock_prisma = MagicMock() - mock_db = MagicMock() mock_query_raw = AsyncMock(return_value=[]) - mock_db.query_raw = mock_query_raw - mock_prisma.db = mock_db + statement_timeout_reads(mock_prisma, mock_query_raw) # Use timezone-aware datetime objects start_date = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc) @@ -170,7 +169,7 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + statement_timeout_reads(mock_prisma, AsyncMock(return_value=[])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) @@ -197,22 +196,27 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): user_api_key_dict=auth, ) - assert mock_prisma.db.query_raw.called, "query_raw should have been called" - sql = mock_prisma.db.query_raw.call_args[0][0] + assert _spend_log_reads(mock_prisma).called, "query_raw should have been called" + sql = _spend_log_reads(mock_prisma).call_args[0][0] assert sql.count("AT TIME ZONE 'UTC'") >= 2, ( f"/spend/logs/ui must wrap both `startTime` bounds with `AT TIME ZONE 'UTC'`. SQL was:\n{sql}" ) +def _spend_log_reads(mock_prisma: MagicMock) -> AsyncMock: + """The raw reads ``/spend/logs/ui`` ran under the statement timeout.""" + return mock_prisma.read_db.tx.return_value.query_raw + + def _make_ui_spend_logs_mock(count_total, page_rows): """ - Build a prisma mock whose first `query_raw` (the bounded count) returns - `count_total` and whose second `query_raw` (the page data) returns + Build a prisma mock whose first statement-timeout read (the bounded count) returns + `count_total` and whose second read (the page data) returns `page_rows`. """ mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": count_total}], page_rows]) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": count_total}], page_rows])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) return mock_prisma @@ -262,7 +266,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): mock_prisma.db.litellm_spendlogs.count.assert_not_called() - count_call = mock_prisma.db.query_raw.call_args_list[0] + count_call = _spend_log_reads(mock_prisma).call_args_list[0] count_sql = count_call[0][0] assert "COUNT(*) OVER ()" not in count_sql assert "LIMIT" in count_sql and "FROM (" in count_sql, ( @@ -270,7 +274,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): ) assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1, "the bounded count must probe at most cap+1 rows" - page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0] + page_sql = _spend_log_reads(mock_prisma).call_args_list[1][0][0] assert "COUNT(*) OVER ()" not in page_sql, ( f"the page query must not carry a window count that forces a full-window scan. SQL was:\n{page_sql}" ) @@ -344,7 +348,7 @@ async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): # page. mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": 0}], []]) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": 0}], []])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) @@ -393,7 +397,7 @@ async def test_spend_logs_ui_out_of_range_page_keeps_total(monkeypatch): # out-of-range page (empty). mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": 7}], []]) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": 7}], []])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) @@ -434,9 +438,8 @@ async def test_get_spend_by_team_binds_optional_team_filter(): guard for LIT-4125 (the team query previously carried no team_id predicate). """ mock_prisma = MagicMock() - mock_prisma.db = MagicMock() mock_query_raw = AsyncMock(return_value=[]) - mock_prisma.db.query_raw = mock_query_raw + statement_timeout_reads(mock_prisma, mock_query_raw) start_date = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc) end_date = datetime.datetime(2024, 1, 31, tzinfo=timezone.utc) @@ -486,7 +489,7 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + statement_timeout_reads(mock_prisma, AsyncMock(return_value=[])) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) @@ -501,9 +504,9 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): customer_id=None, ) - assert mock_prisma.db.query_raw.called, "query_raw should have been called" - sql = mock_prisma.db.query_raw.call_args[0][0] - params = mock_prisma.db.query_raw.call_args[0][1:] + assert _spend_log_reads(mock_prisma).called, "query_raw should have been called" + sql = _spend_log_reads(mock_prisma).call_args[0][0] + params = _spend_log_reads(mock_prisma).call_args[0][1:] assert "team_x" in params, "team_id must be forwarded into the DB query params" assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}" @@ -534,7 +537,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=mock_query_raw)) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") @@ -560,7 +563,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): status_filter=None, ) - emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list] + emitted = [call[0] for call in _spend_log_reads(mock_prisma).call_args_list] page_sql = emitted[0][0] assert "GROUP BY session_key, api_key" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}" assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected" @@ -621,10 +624,10 @@ async def test_spend_logs_ui_group_by_session_offset_pages_for_other_sorts(monke ) group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" - count_sql = mock_prisma.db.query_raw.call_args_list[0][0][0] + count_sql = _spend_log_reads(mock_prisma).call_args_list[0][0][0] assert f"GROUP BY {group_key}" in count_sql - page_call = mock_prisma.db.query_raw.call_args_list[1][0] + page_call = _spend_log_reads(mock_prisma).call_args_list[1][0] page_sql = page_call[0] assert f"DISTINCT ON ({group_key})" in page_sql assert "ORDER BY spend DESC" in page_sql @@ -670,8 +673,108 @@ async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(m group_by_session=True, ) - page_call = mock_prisma.db.query_raw.call_args_list[1] + page_call = _spend_log_reads(mock_prisma).call_args_list[1] assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping" assert "req-deep-link" in page_call[0] assert [row["request_id"] for row in response["data"]] == ["req-deep-link"] assert response["total"] == 1 + + +async def _spend_logs_ui_page(mock_prisma: MagicMock, monkeypatch: pytest.MonkeyPatch) -> Mapping[str, object]: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + return await ui_view_spend_logs( + log_team_lookup=await get_log_team_lookup(), + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + search=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_reads_are_cancelled_by_postgres_after_30s(monkeypatch: pytest.MonkeyPatch): + """Prisma's HTTP client gives up at 30s either way; without SET LOCAL statement_timeout the scan kept running.""" + mock_prisma = MagicMock() + transaction = statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": 0}], []])) + + await _spend_logs_ui_page(mock_prisma, monkeypatch) + + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) * 2 + assert mock_prisma.read_db.tx.call_args.kwargs["timeout"] > datetime.timedelta(seconds=30) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_cancelled_read_is_reported_as_503_with_a_date_range_hint(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy._types import ProxyException + + mock_prisma = MagicMock() + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=statement_timeout_error())) + + with pytest.raises(ProxyException) as exc_info: + await _spend_logs_ui_page(mock_prisma, monkeypatch) + + assert str(exc_info.value.code) == "503" + assert "Spend logs took longer than 30s. Pick a shorter date range" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_global_spend_report_returns_the_503_hint_not_its_catch_all_400_when_postgres_cancels( + monkeypatch: pytest.MonkeyPatch, +): + from fastapi import HTTPException + + from litellm.proxy.spend_tracking.spend_management_endpoints import get_global_spend_report + + mock_prisma = MagicMock() + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=statement_timeout_error())) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + with pytest.raises(HTTPException) as exc_info: + await get_global_spend_report( + start_date="2026-02-01", + end_date="2026-02-16", + group_by="api_key", + api_key=None, + internal_user_id=None, + team_id=None, + customer_id=None, + ) + + assert exc_info.value.status_code == 503 + assert "Spend logs took longer than 30s. Pick a shorter date range" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_global_spend_report_reads_are_cancelled_by_postgres_after_30s(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.spend_tracking.spend_management_endpoints import get_global_spend_report + + mock_prisma = MagicMock() + transaction = statement_timeout_reads(mock_prisma, AsyncMock(return_value=[])) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + await get_global_spend_report( + start_date="2026-02-01", + end_date="2026-02-16", + group_by="api_key", + api_key=None, + internal_user_id=None, + team_id=None, + customer_id=None, + ) + + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 1b9cdd003ab..698f4787a42 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -3716,3 +3716,55 @@ async def test_evict_config_param_clears_the_local_layer_and_survives_a_redis_ou await evict_config_param("router_settings", cache=cache) assert await cache.in_memory_cache.async_get_cache(_config_cache_key("router_settings")) is None + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_has_postgres_cancel_the_read_at_30s() -> None: + from datetime import timedelta + + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_reads + + prisma: Final = MagicMock() + transaction: Final = statement_timeout_reads(prisma, AsyncMock(return_value=[{"spend": 1.5}])) + + rows: Final = await query_raw_within_statement_timeout( + prisma, 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE "user" = $1', "u1" + ) + + assert rows == [{"spend": 1.5}] + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) + transaction.query_raw.assert_awaited_once_with('SELECT spend FROM "LiteLLM_SpendLogs" WHERE "user" = $1', "u1") + assert prisma.read_db.tx.call_args.kwargs["timeout"] > timedelta(seconds=30), "Prisma must outlive the Postgres cancel" + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_reports_a_cancelled_read_as_503() -> None: + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads + + prisma: Final = MagicMock() + statement_timeout_reads(prisma, AsyncMock(side_effect=statement_timeout_error())) + + with pytest.raises(HTTPException) as exc_info: + await query_raw_within_statement_timeout(prisma, "SELECT 1", what="Spend per user") + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": "Spend per user took longer than 30s. Pick a shorter date range"} + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_leaves_other_database_errors_alone() -> None: + from prisma.errors import DataError + + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_reads + + prisma: Final = MagicMock() + failure: Final = DataError(data={"user_facing_error": {"is_panic": False, "message": "relation does not exist"}}) + statement_timeout_reads(prisma, AsyncMock(side_effect=failure)) + + with pytest.raises(DataError) as exc_info: + await query_raw_within_statement_timeout(prisma, "SELECT 1") + + assert exc_info.value is failure