mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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>
This commit is contained in:
Kerry Lu
GitHub
parent
62f85afacf
commit
2d0c720425
15 changed files with 554 additions and 146 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue