perf(spend): write each daily spend batch in one upsert statement (#36448)

The daily spend flush emitted one INSERT ... ON CONFLICT per aggregated key,
so every replica put hundreds of statements on the database each interval,
all contending for the same handful of hot rows and each holding its row
locks for the rest of the enclosing batch transaction. LiteLLM_DailyTagSpend
felt it worst because a request writes one row per tag, and litellm adds two
user-agent tags of its own by default.

A batch now goes out as a single multi-row statement. Rows are folded by the
conflict tuple first, and every nullable member of that tuple is normalized
to '': a NULL can never match itself in a unique index, so such a row was
re-inserted on every flush rather than aggregating, and a NULL model made
prisma reject the whole batch.

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))

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(

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:]

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 == {}

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",
         )