diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py new file mode 100644 index 00000000000..55d325177c6 --- /dev/null +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -0,0 +1,185 @@ +"""One multi-row ``INSERT ... ON CONFLICT DO UPDATE`` per batch of daily spend rows. + +Emitting a statement per aggregated key put every replica's flush on the database as +hundreds of separate statements against the same handful of hot rows, each holding its +row locks for the rest of the enclosing batch transaction. Folding a batch into a single +statement keeps the aggregation identical while collapsing both the statement count and +the window in which those locks are held. +""" + +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import groupby +from types import MappingProxyType +from typing import Final, Literal + +DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"] + +SqlValue = str | int | float | None + +# A queued daily spend transaction, read by column name because the columns are data +# here rather than literals. The concrete TypedDicts in _types.py all satisfy this. +SpendRow = Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class DailySpendTable: + """The physical table behind one entity's daily rollup.""" + + name: str + entity_id_column: str + carries_request_id: bool = False + + +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), + "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), + "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), + "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), + "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), + "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + } +) + +# The unique constraint's columns after the entity id, in constraint order. A NULL can +# never match itself in a unique index, so every one of these is normalized to '': the +# conflict target has to be NULL-free or the row is re-inserted on every single flush. +_KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") + +_COUNTER_COLUMNS: Final = ( + "prompt_tokens", + "completion_tokens", + "api_requests", + "successful_requests", + "failed_requests", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", +) +_SPEND_COLUMNS: Final = ( + "spend", + "compression_savings_spend", + "prompt_caching_savings_spend", + "autorouter_savings_spend", +) + +_CASTS: Final[Mapping[str, str]] = MappingProxyType( + { + **{column: "bigint" for column in _COUNTER_COLUMNS}, + **{column: "double precision" for column in _SPEND_COLUMNS}, + } +) + + +def _quoted(columns: Sequence[str]) -> str: + return ", ".join(f'"{column}"' for column in columns) + + +def _as_text(value: object) -> str: + return "" if value is None else str(value) + + +def _as_int(value: object) -> int: + return int(value) if isinstance(value, (int, float)) else 0 + + +def _as_float(value: object) -> float: + return float(value) if isinstance(value, (int, float)) else 0.0 + + +def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: + """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) + + +def _merge(group: Sequence[SpendRow]) -> SpendRow: + if len(group) == 1: + return group[0] + return { + **group[0], + **{column: sum(_as_int(row.get(column)) for row in group) for column in _COUNTER_COLUMNS}, + **{column: sum(_as_float(row.get(column)) for row in group) for column in _SPEND_COLUMNS}, + } + + +def merge_by_conflict_key( + table: DailySpendTable, + transactions: Sequence[SpendRow], +) -> tuple[tuple[tuple[str, ...], SpendRow], ...]: + """Batch entries keyed by the conflict tuple, in a deterministic order. + + The queue keys transactions by their raw field values, so two entries differing only + in a NULL versus an empty member reach the writer separately while arbitrating to the + same row. Postgres rejects a statement whose values touch one row twice, so they are + summed here into the single row they were always destined to become. Ordering by the + key keeps concurrent writers taking row locks in the same sequence. + """ + ordered: Final = sorted(transactions, key=lambda transaction: conflict_key(table, transaction)) + return tuple((key, _merge(tuple(group))) for key, group in groupby(ordered, key=lambda t: conflict_key(table, t))) + + +def _row_params( + table: DailySpendTable, + key: tuple[str, ...], + transaction: SpendRow, +) -> tuple[SqlValue, ...]: + request_id: Final = transaction.get("request_id") + return ( + str(uuid.uuid4()), + *key, + None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), + *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), + *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), + *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), + ) + + +def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: + return ( + "id", + table.entity_id_column, + *_KEY_COLUMNS, + "model_group", + *_COUNTER_COLUMNS, + *_SPEND_COLUMNS, + *(("request_id",) if table.carries_request_id else ()), + ) + + +def build_bulk_upsert( + table: DailySpendTable, + batch: Sequence[tuple[tuple[str, ...], SpendRow]], +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" + columns: Final = _insert_columns(table) + quoted_table: Final = f'"{table.name}"' + rows: Final = ", ".join( + "(" + + ", ".join( + f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" + for offset, column in enumerate(columns) + ) + + ", (NOW() AT TIME ZONE 'UTC'))" + for row_index in range(len(batch)) + ) + increments: Final = ", ".join( + f'"{column}" = {quoted_table}."{column}" + EXCLUDED."{column}"' + for column in (*_COUNTER_COLUMNS, *_SPEND_COLUMNS) + ) + # request_id names one arbitrary contributing request, so an entry carrying none must + # not blank out the one already recorded. + request_id_update: Final = ( + f', "request_id" = COALESCE(EXCLUDED."request_id", {quoted_table}."request_id")' + if table.carries_request_id + else "" + ) + sql: Final = ( + f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' + f"VALUES {rows}\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f" {increments}{request_id_update},\n" + f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" + ) + return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 385a21976b7..b0130db232a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,9 +12,7 @@ import os import random import time import traceback -from collections.abc import Mapping from datetime import datetime, timedelta, timezone -from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload import litellm @@ -41,6 +39,11 @@ from litellm.proxy._types import ( SpendUpdateQueueItem, ToolDiscoveryQueueItem, ) +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + build_bulk_upsert, + merge_by_conflict_key, +) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) @@ -68,12 +71,6 @@ else: ProxyLogging = Any -# Only tag rows carry a request_id, so the other entity types spread nothing. Built -# once here rather than as an empty literal per transaction, and read-only so it cannot -# be filled in by accident from one of the call sites that spreads it. -_NO_TAG_REQUEST_ID: Final[Mapping[str, Any]] = MappingProxyType({}) - - def _get_llm_router(): """The proxy's router, or None outside a running proxy. @@ -1437,8 +1434,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyUserSpendTransaction], entity_type: Literal["user"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1451,8 +1446,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyTeamSpendTransaction], entity_type: Literal["team"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1465,8 +1458,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyOrganizationSpendTransaction], entity_type: Literal["org"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1479,8 +1470,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyEndUserSpendTransaction], entity_type: Literal["end_user"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1493,8 +1482,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyAgentSpendTransaction], entity_type: Literal["agent"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1507,8 +1494,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyTagSpendTransaction], entity_type: Literal["tag"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... # fmt: on @@ -1526,8 +1511,6 @@ class DBSpendUpdateWriter: | dict[str, DailyAgentSpendTransaction], entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: """ Generic function to update daily spend for any entity type (user, team, org, tag, end_user, agent) @@ -1573,111 +1556,23 @@ class DBSpendUpdateWriter: ) return + table = DAILY_SPEND_TABLES[entity_type] try: - async with prisma_client.db.batch_() as batcher: - for _, transaction in transactions_to_process.items(): - entity_id = transaction.get(entity_id_field) - - # Construct the where clause dynamically - where_clause = { - unique_constraint_name: { - entity_id_field: entity_id, - "date": transaction["date"], - "api_key": transaction["api_key"], - "model": transaction["model"], - "custom_llm_provider": transaction.get("custom_llm_provider") or "", - "mcp_namespaced_tool_name": transaction.get("mcp_namespaced_tool_name") - or "", - "endpoint": transaction.get("endpoint") or "", - } - } - - # Get the table dynamically - table = getattr(batcher, table_name) - - # Additive metrics that older queued rows may omit; one - # enumeration feeds both the create and the increment below - optional_metrics = { - field: value - for field, value in ( - ("cache_read_input_tokens", transaction.get("cache_read_input_tokens")), - ( - "cache_creation_input_tokens", - transaction.get("cache_creation_input_tokens"), - ), - ("compression_saved_tokens", transaction.get("compression_saved_tokens")), - ( - "compression_savings_spend", - transaction.get("compression_savings_spend"), - ), - ( - "prompt_caching_savings_spend", - transaction.get("prompt_caching_savings_spend"), - ), - ("autorouter_savings_spend", transaction.get("autorouter_savings_spend")), - ) - if value is not None - } - - # Only tag rows carry a request_id. Resolved to a spreadable - # value here so both payloads are built in one shot: a dict - # appended to after construction is one nobody can reason about - # by reading its literal. - tag_request_id: Mapping[str, Any] = ( - MappingProxyType({"request_id": transaction["request_id"]}) - if entity_type == "tag" and "request_id" in transaction - else _NO_TAG_REQUEST_ID - ) - - # Common data structure for both create and update - common_data = { - entity_id_field: entity_id, - "date": transaction["date"], - "api_key": transaction["api_key"], - "model": transaction.get("model"), - "model_group": transaction.get("model_group"), - "mcp_namespaced_tool_name": transaction.get("mcp_namespaced_tool_name") or "", - "custom_llm_provider": transaction.get("custom_llm_provider"), - "endpoint": transaction.get("endpoint") or "", - "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction["completion_tokens"], - "spend": transaction["spend"], - "api_requests": transaction["api_requests"], - "successful_requests": transaction["successful_requests"], - "failed_requests": transaction["failed_requests"], - **optional_metrics, - **tag_request_id, - } - - update_data = { - "prompt_tokens": {"increment": transaction["prompt_tokens"]}, - "completion_tokens": {"increment": transaction["completion_tokens"]}, - "spend": {"increment": transaction["spend"]}, - "api_requests": {"increment": transaction["api_requests"]}, - "successful_requests": {"increment": transaction["successful_requests"]}, - "failed_requests": {"increment": transaction["failed_requests"]}, - **{field: {"increment": value} for field, value in optional_metrics.items()}, - # An existing row predating the endpoint column gets it filled in here - "endpoint": transaction.get("endpoint") or "", - **tag_request_id, - } - - table.upsert( - where=where_clause, - data={ - "create": common_data, - "update": update_data, - }, - ) + # One statement per batch rather than per key: the same rows are + # aggregated, but concurrent writers no longer hold a batch's worth + # of row locks across a hundred round trips. + merged_batch = merge_by_conflict_key( + table=table, transactions=tuple(transactions_to_process.values()) + ) + sql, params = build_bulk_upsert(table=table, batch=merged_batch) + await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures # This helps diagnose issues like unique constraint violations spend_log_error( - "Daily %s spend batch upsert failed. " - "Table: %s, Constraint: %s, Batch size: %d, Error: %s", + "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, - table_name, - unique_constraint_name, + table.name, len(transactions_to_process), str(batch_error), exc=batch_error, @@ -1733,8 +1628,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1754,8 +1647,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="team", entity_id_field="team_id", - table_name="litellm_dailyteamspend", - unique_constraint_name="team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1775,8 +1666,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="org", entity_id_field="organization_id", - table_name="litellm_dailyorganizationspend", - unique_constraint_name="organization_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1796,8 +1685,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="end_user", entity_id_field="end_user_id", - table_name="litellm_dailyenduserspend", - unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1817,8 +1704,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="agent", entity_id_field="agent_id", - table_name="litellm_dailyagentspend", - unique_constraint_name="agent_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1838,8 +1723,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="tag", entity_id_field="tag", - table_name="litellm_dailytagspend", - unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) async def _common_add_spend_log_transaction_to_daily_transaction( diff --git a/tests/proxy_unit_tests/test_update_daily_tag_spend.py b/tests/proxy_unit_tests/test_update_daily_tag_spend.py index 80616ade5ef..35e9c6796eb 100644 --- a/tests/proxy_unit_tests/test_update_daily_tag_spend.py +++ b/tests/proxy_unit_tests/test_update_daily_tag_spend.py @@ -91,17 +91,13 @@ async def test_daily_tag_spend_retries_then_succeeds(): prisma_client = MagicMock() proxy_logging_obj = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_batcher.litellm_dailytagspend = mock_table - - # Fail entering batch context 3 times with retryable DB errors, then succeed. - prisma_client.db.batch_.return_value.__aenter__ = AsyncMock( + # Fail the upsert 3 times with retryable DB errors, then succeed. + prisma_client.db.execute_raw = AsyncMock( side_effect=[ httpx.ConnectError("x"), httpx.ConnectError("x"), httpx.ConnectError("x"), - mock_batcher, + 1, ] ) @@ -138,6 +134,10 @@ async def test_daily_tag_spend_retries_then_succeeds(): daily_spend_transactions=daily_spend_transactions, ) - assert prisma_client.db.batch_.return_value.__aenter__.await_count == 4 + assert prisma_client.db.execute_raw.await_count == 4 assert sleep_mock.await_count == 3 - mock_table.upsert.assert_called_once() + # The batch is one statement, so the successful attempt is a single call carrying + # the row rather than one call per key. + final_sql = prisma_client.db.execute_raw.await_args.args[0] + assert final_sql.count("ON CONFLICT") == 1 + assert "prod-tag" in prisma_client.db.execute_raw.await_args.args[1:] diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py new file mode 100644 index 00000000000..c2d0f64461a --- /dev/null +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -0,0 +1,187 @@ +"""Tests for the single-statement daily spend upsert (LIT-5291).""" + +import re + +import pytest + +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + build_bulk_upsert, + conflict_key, + merge_by_conflict_key, +) +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + +TAG_TABLE = DAILY_SPEND_TABLES["tag"] +USER_TABLE = DAILY_SPEND_TABLES["user"] + +# Every nullable member of the unique constraint, so a test that only varied the provider +# cannot pass while a sibling column still leaks a NULL into the conflict target. +NULLABLE_KEY_COLUMNS = ("model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") + + +def tag_txn(**overrides): + return { + "tag": "team-a", + "date": "2026-08-10", + "api_key": "sk-hash", + "model": "gpt-4o-mini", + "model_group": "gpt-4o-mini", + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": "", + "endpoint": "/chat/completions", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.25, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + "request_id": "req-1", + **overrides, + } + + +@pytest.mark.parametrize("column", NULLABLE_KEY_COLUMNS) +def test_conflict_key_normalizes_every_nullable_key_column(column): + """A NULL member can never match itself in a unique index, so the row would be + re-inserted on every flush. Each nullable key column must arrive as ''.""" + key = conflict_key(TAG_TABLE, tag_txn(**{column: None})) + + assert "" in key + assert None not in key + assert key == conflict_key(TAG_TABLE, tag_txn(**{column: ""})) + + +@pytest.mark.parametrize("order", [("null_first"), ("empty_first")]) +def test_null_and_empty_provider_merge_into_one_row(order): + """Two queue entries differing only in NULL versus '' arbitrate to the same row. + Postgres rejects one statement touching a row twice, so they must be folded first. + Asserted under both input orders: a single ordering would prove nothing here.""" + null_entry = tag_txn(custom_llm_provider=None, spend=0.25, api_requests=1) + empty_entry = tag_txn(custom_llm_provider="", spend=0.75, api_requests=3) + transactions = (null_entry, empty_entry) if order == "null_first" else (empty_entry, null_entry) + + merged = merge_by_conflict_key(TAG_TABLE, transactions) + + assert len(merged) == 1 + _, folded = merged[0] + assert folded["spend"] == pytest.approx(1.0) + assert folded["api_requests"] == 4 + + +def test_distinct_keys_are_not_merged_and_are_ordered_deterministically(): + unordered = (tag_txn(tag="z-team"), tag_txn(tag="a-team"), tag_txn(tag="m-team")) + + merged = merge_by_conflict_key(TAG_TABLE, unordered) + + assert [txn["tag"] for _, txn in merged] == ["a-team", "m-team", "z-team"] + assert merged == merge_by_conflict_key(TAG_TABLE, tuple(reversed(unordered))) + + +def test_one_statement_carries_every_row_in_the_batch(): + batch = merge_by_conflict_key(TAG_TABLE, tuple(tag_txn(tag=f"team-{i}") for i in range(100))) + + sql, params = build_bulk_upsert(TAG_TABLE, batch) + + assert sql.count("INSERT INTO") == 1 + assert len(re.findall(r"ON CONFLICT", sql)) == 1 + # 22 bound columns per row plus the inlined updated_at, so the row count is what + # separates one multi-row statement from a hundred single-row ones. + assert len(params) == 100 * 22 + assert "$2200::text" in sql + assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1 + + +def test_conflict_target_is_the_full_unique_constraint(): + sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),))) + + conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + assert conflict_target is not None + assert conflict_target.group(1) == ( + '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' + ) + + +@pytest.mark.parametrize( + "column", + ["prompt_tokens", "completion_tokens", "spend", "api_requests", "successful_requests", "failed_requests"], +) +def test_counters_increment_rather_than_overwrite(column): + """An overwrite would silently discard every earlier flush's spend for that row.""" + sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),))) + + assert f'"{column}" = "LiteLLM_DailyTagSpend"."{column}" + EXCLUDED."{column}"' in sql + + +def test_request_id_is_preserved_when_a_later_batch_carries_none(): + sql, params = build_bulk_upsert( + TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(request_id=None),)) + ) + + assert '"request_id" = COALESCE(EXCLUDED."request_id", "LiteLLM_DailyTagSpend"."request_id")' in sql + assert None in params + + +def test_non_tag_tables_carry_no_request_id_column(): + user_txn = {**tag_txn(), "user_id": "u-1"} + del user_txn["tag"] + + sql, _ = build_bulk_upsert(USER_TABLE, merge_by_conflict_key(USER_TABLE, (user_txn,))) + + assert "request_id" not in sql + assert '"user_id"' in sql + + +class _RecordingDb: + def __init__(self) -> None: + self.statements: list[tuple[str, tuple[object, ...]]] = [] + + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + return len(args) + + +class _RecordingPrismaClient: + def __init__(self) -> None: + self.db = _RecordingDb() + + +@pytest.mark.asyncio +async def test_writer_issues_one_statement_per_batch_not_one_per_key(): + """The whole point of LIT-5291: 250 aggregated keys must not become 250 statements.""" + prisma_client = _RecordingPrismaClient() + transactions = {f"k{i}": tag_txn(tag=f"team-{i}") for i in range(250)} + + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=None, + daily_spend_transactions=transactions, + ) + + # 250 keys at a batch size of 100 is three statements, one per batch. + assert len(prisma_client.db.statements) == 3 + assert [statement.count("ON CONFLICT") for statement, _ in prisma_client.db.statements] == [1, 1, 1] + assert transactions == {} + + +@pytest.mark.asyncio +async def test_writer_survives_a_transaction_whose_key_columns_are_null(): + """A NULL key column used to raise out of prisma and drop the whole batch's spend.""" + prisma_client = _RecordingPrismaClient() + transactions = { + "mcp": tag_txn(model=None, custom_llm_provider=None, mcp_namespaced_tool_name="server/tool"), + "chat": tag_txn(), + } + + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=None, + daily_spend_transactions=transactions, + ) + + assert len(prisma_client.db.statements) == 1 + _, params = prisma_client.db.statements[0] + assert None not in params[:9] + assert transactions == {} diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 191080e3a48..b659ef3321b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2,6 +2,7 @@ import asyncio import copy import json import os +import re import sys sys.path.insert( @@ -9,6 +10,7 @@ sys.path.insert( ) # Adds the parent directory to the system path +from collections.abc import Callable from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch @@ -232,21 +234,49 @@ async def test_update_database_skips_tool_usage_when_spend_logs_disabled(): assert prisma.tool_usage_transactions == [] +Statement = tuple[str, tuple[object, ...]] + + +class _RecordingDb: + """Records the statements the writer sends, in place of a real query engine.""" + + def __init__(self, execute_raw: Callable[[], int] | None = None) -> None: + self.statements: list[Statement] = [] + self._execute_raw = execute_raw + + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + if self._execute_raw is not None: + return self._execute_raw() + return len(args) + + +class _RecordingPrisma: + def __init__(self, execute_raw: Callable[[], int] | None = None) -> None: + self.db = _RecordingDb(execute_raw=execute_raw) + + +def _row_values(statement: Statement, column: str) -> list[object]: + """Every row's value for one column, read out of the flat parameter tuple.""" + sql, params = statement + header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) + assert header is not None, sql + columns = header.group(1).split(", ") + stride = len(columns) - 1 # updated_at is inlined, not bound + offset = columns.index(f'"{column}"') + return [params[row * stride + offset] for row in range(len(params) // stride)] + + @pytest.mark.asyncio async def test_update_daily_spend_with_null_entity_id(): """ - Test that table.upsert is called even when entity_id is null + A null entity_id must still be written, so the 'global view' keeps that spend. - Ensures 'global view' has all daily spend transactions + It is stored as '' rather than NULL: a NULL can never match itself in the unique + index, so such a row would be re-inserted on every flush instead of aggregating. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() - # Create a transaction with null entity_id daily_spend_transactions = { "test_key": { "user_id": None, # null entity_id @@ -263,49 +293,30 @@ async def test_update_daily_spend_with_null_entity_id(): } } - # Call the method await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - # Verify that table.upsert was called - mock_table.upsert.assert_called_once() - - # Verify the where clause contains null entity_id - call_args = mock_table.upsert.call_args[1] - where_clause = call_args["where"][ - "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" - ] - assert where_clause["user_id"] is None - assert where_clause["date"] == "2024-01-01" - assert where_clause["api_key"] == "test-api-key" - assert where_clause["model"] == "gpt-4" - assert where_clause["custom_llm_provider"] == "openai" - assert where_clause["mcp_namespaced_tool_name"] == "" - assert where_clause["endpoint"] == "" - - # Verify the create data contains null entity_id - create_data = call_args["data"]["create"] - assert create_data["user_id"] is None - assert create_data["date"] == "2024-01-01" - assert create_data["api_key"] == "test-api-key" - assert create_data["model"] == "gpt-4" - assert create_data["custom_llm_provider"] == "openai" - assert create_data["mcp_namespaced_tool_name"] == "" - assert create_data["endpoint"] == "" - assert create_data["prompt_tokens"] == 10 - assert create_data["completion_tokens"] == 20 - assert create_data["spend"] == 0.1 - assert create_data["api_requests"] == 1 - assert create_data["successful_requests"] == 1 - assert create_data["failed_requests"] == 0 + assert len(prisma_client.db.statements) == 1 + statement = prisma_client.db.statements[0] + assert _row_values(statement, "user_id") == [""] + assert _row_values(statement, "date") == ["2024-01-01"] + assert _row_values(statement, "api_key") == ["test-api-key"] + assert _row_values(statement, "model") == ["gpt-4"] + assert _row_values(statement, "custom_llm_provider") == ["openai"] + assert _row_values(statement, "mcp_namespaced_tool_name") == [""] + assert _row_values(statement, "endpoint") == [""] + assert _row_values(statement, "prompt_tokens") == [10] + assert _row_values(statement, "completion_tokens") == [20] + assert _row_values(statement, "spend") == [0.1] + assert _row_values(statement, "api_requests") == [1] + assert _row_values(statement, "successful_requests") == [1] + assert _row_values(statement, "failed_requests") == [0] def _daily_txn(user_id: str = "user1") -> dict: @@ -333,24 +344,24 @@ async def test_update_daily_spend_does_not_retry_post_send_ambiguous_errors(): # batch (loudly), never retry it. import httpx - mock_prisma_client = MagicMock() - mock_prisma_client.db.batch_ = MagicMock(side_effect=httpx.ReadTimeout("ambiguous")) + def raise_read_timeout(): + raise httpx.ReadTimeout("ambiguous") + + prisma_client = _RecordingPrisma(execute_raw=raise_read_timeout) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() with pytest.raises(httpx.ReadTimeout): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=3, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=proxy_logging, daily_spend_transactions={"k1": _daily_txn()}, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - mock_prisma_client.db.batch_.assert_called_once() + assert len(prisma_client.db.statements) == 1 @pytest.mark.asyncio @@ -359,12 +370,15 @@ async def test_update_daily_spend_retries_connect_errors(monkeypatch): # the one failure the writer may retry. import httpx - mock_batcher = MagicMock() - good_ctx = MagicMock() - good_ctx.__aenter__ = AsyncMock(return_value=mock_batcher) - good_ctx.__aexit__ = AsyncMock(return_value=None) - mock_prisma_client = MagicMock() - mock_prisma_client.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), good_ctx]) + outcomes = iter([httpx.ConnectError("down"), None]) + + def first_attempt_disconnects(): + outcome = next(outcomes) + if outcome is not None: + raise outcome + return 1 + + prisma_client = _RecordingPrisma(execute_raw=first_attempt_disconnects) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() @@ -374,16 +388,14 @@ async def test_update_daily_spend_retries_connect_errors(monkeypatch): monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", fake_sleep) await DBSpendUpdateWriter._update_daily_spend( n_retry_times=3, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=proxy_logging, daily_spend_transactions={"k1": _daily_txn()}, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - assert mock_prisma_client.db.batch_.call_count == 2 + assert len(prisma_client.db.statements) == 2 @pytest.mark.asyncio @@ -393,19 +405,12 @@ async def test_update_daily_spend_sorting(): Ensures that writes are sorted between transactions to minimize deadlocks """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() - # Create a 50 transactions with out-of-order entity_ids - # In reality we sort using multiple fields, but entity_id is sufficient to test sorting - daily_spend_transactions = {} - upsert_calls = [] - for i in range(50): - daily_spend_transactions[f"test_key_{i}"] = { + # 50 transactions with out-of-order entity_ids. In reality we sort using multiple + # fields, but entity_id is sufficient to test sorting. + daily_spend_transactions = { + f"test_key_{i}": { "user_id": f"user{60-i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", @@ -418,63 +423,22 @@ async def test_update_daily_spend_sorting(): "successful_requests": 1, "failed_requests": 0, } - upsert_calls.append( - call( - where={ - "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": { - "user_id": f"user{i+11}", # user11 ... user60, sorted order - "date": "2024-01-01", - "api_key": "test-api-key", - "model": "gpt-4", - "custom_llm_provider": "openai", - "mcp_namespaced_tool_name": "", - "endpoint": "", - } - }, - data={ - "create": { - "user_id": f"user{i+11}", - "date": "2024-01-01", - "api_key": "test-api-key", - "model": "gpt-4", - "model_group": None, - "mcp_namespaced_tool_name": "", - "custom_llm_provider": "openai", - "endpoint": "", - "prompt_tokens": 10, - "completion_tokens": 20, - "spend": 0.1, - "api_requests": 1, - "successful_requests": 1, - "failed_requests": 0, - }, - "update": { - "prompt_tokens": {"increment": 10}, - "completion_tokens": {"increment": 20}, - "spend": {"increment": 0.1}, - "api_requests": {"increment": 1}, - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 0}, - "endpoint": "", - }, - }, - ) - ) + for i in range(50) + } - # Call the method await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - # Verify that table.upsert was called - mock_table.upsert.assert_has_calls(upsert_calls) + assert len(prisma_client.db.statements) == 1 + written = _row_values(prisma_client.db.statements[0], "user_id") + assert written == sorted(written) + assert written[0] == "user11" and written[-1] == "user60" @pytest.mark.asyncio @@ -485,11 +449,7 @@ async def test_update_daily_spend_drains_all_batches_over_batch_size(): only the first 100 sorted items were upserted then the method returned, silently dropping the remaining entities. """ - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() num_entities = 250 daily_spend_transactions = { @@ -511,17 +471,16 @@ async def test_update_daily_spend_drains_all_batches_over_batch_size(): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - assert mock_table.upsert.call_count == num_entities - assert mock_prisma_client.db.batch_.call_count == 3 + assert len(prisma_client.db.statements) == 3 + all_written = [uid for statement in prisma_client.db.statements for uid in _row_values(statement, "user_id")] + assert sorted(all_written) == sorted(f"user{i:04d}" for i in range(num_entities)) assert daily_spend_transactions == {} @@ -530,14 +489,8 @@ async def test_update_daily_spend_tag_with_request_id(): """ Test that request_id is included in update_data when updating tag transactions. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailytagspend = mock_table + prisma_client = _RecordingPrisma() - # Create a transaction with request_id daily_spend_transactions = { "test_key": { "tag": "prod-tag", @@ -556,26 +509,19 @@ async def test_update_daily_spend_tag_with_request_id(): } } - # Call the method await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="tag", entity_id_field="tag", - table_name="litellm_dailytagspend", - unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", ) - # Verify that table.upsert was called - mock_table.upsert.assert_called_once() - - # Verify request_id is in update_data - call_args = mock_table.upsert.call_args[1] - update_data = call_args["data"]["update"] - assert "request_id" in update_data - assert update_data["request_id"] == "test-request-id-123" + assert len(prisma_client.db.statements) == 1 + sql, _ = prisma_client.db.statements[0] + assert _row_values(prisma_client.db.statements[0], "request_id") == ["test-request-id-123"] + assert '"request_id" = COALESCE(EXCLUDED."request_id", "LiteLLM_DailyTagSpend"."request_id")' in sql @pytest.mark.asyncio @@ -587,12 +533,7 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields(): are None, the sorting doesn't crash with TypeError: '<' not supported between instances of 'NoneType' and 'str'. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() # Create transactions with None values in various sorting fields daily_spend_transactions = { @@ -666,17 +607,20 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields(): # Call the method - this should not raise TypeError await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - # Verify that table.upsert was called (should be called 5 times, once for each transaction) - assert mock_table.upsert.call_count == 5 + # All five distinct rows are written, in one statement, with no NULL anywhere in + # the conflict key. + assert len(prisma_client.db.statements) == 1 + statement = prisma_client.db.statements[0] + assert len(_row_values(statement, "user_id")) == 5 + for column in ("user_id", "date", "api_key", "model", "custom_llm_provider"): + assert None not in _row_values(statement, column) # Tag Spend Tracking Tests @@ -1384,19 +1328,10 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): """ from litellm._logging import verbose_proxy_logger - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_batch_context = MagicMock() - mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher) - mock_batcher.litellm_dailyuserspend = mock_table + def raise_constraint_violation(): + raise Exception("Unique constraint violation") - # Make the batch context manager's exit raise an exception - # This simulates a batch commit failure (e.g., unique constraint violation) - test_exception = Exception("Unique constraint violation") - mock_batch_context.__aexit__ = AsyncMock(side_effect=test_exception) - mock_prisma_client.db.batch_.return_value = mock_batch_context + prisma_client = _RecordingPrisma(execute_raw=raise_constraint_violation) # Create a transaction daily_spend_transactions = { @@ -1427,28 +1362,22 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): with pytest.raises(Exception, match="Unique constraint violation"): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=0, # No retries to make test faster - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=mock_proxy_logging, daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) # Verify that the error was logged with detailed information. # spend_log_error formats the message via ``%`` interpolation, so # render the call args before asserting on substrings. assert mock_error_logger.called - call = mock_error_logger.call_args - formatted = call.args[0] % call.args[1:] + logged = mock_error_logger.call_args + formatted = logged.args[0] % logged.args[1:] assert "Daily user spend batch upsert failed" in formatted - assert "Table: litellm_dailyuserspend" in formatted - assert ( - "Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" - in formatted - ) - assert "Batch size: 1" in formatted + assert "Table: LiteLLM_DailyUserSpend" in formatted + assert "Rows: 1" in formatted assert "Unique constraint violation" in formatted @@ -1458,13 +1387,10 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): Test that when batch upsert fails, the exception is properly re-raised after logging. This ensures that error handling continues to work correctly upstream. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_batch_context = MagicMock() - mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher) - mock_batcher.litellm_dailyuserspend = mock_table + def raise_connection_lost(): + raise ValueError("Database connection lost") + + prisma_client = _RecordingPrisma(execute_raw=raise_connection_lost) # Create a transaction daily_spend_transactions = { @@ -1483,11 +1409,6 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): } } - # Create a custom exception to verify it's re-raised - custom_exception = ValueError("Database connection lost") - mock_batch_context.__aexit__ = AsyncMock(side_effect=custom_exception) - mock_prisma_client.db.batch_.return_value = mock_batch_context - # Create a mock proxy_logging_obj with failure_handler as AsyncMock mock_proxy_logging = MagicMock() mock_proxy_logging.failure_handler = AsyncMock() @@ -1496,13 +1417,11 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): with pytest.raises(ValueError, match="Database connection lost"): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=0, # No retries to make test faster - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=mock_proxy_logging, daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", )