From e49e6bc660f7490a485d6cc1f6a66514b8ca7e7b Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 19 Sep 2026 19:54:11 +0000 Subject: [PATCH] fix(proxy): stop re-sending un-resendable spend batches from the Redis buffer --- litellm/proxy/db/db_spend_update_writer.py | 150 ++++++++-- litellm/proxy/db/exception_handler.py | 13 +- .../proxy/db/test_db_spend_update_writer.py | 267 ++++++++++++++++++ .../proxy/db/test_exception_handler.py | 19 +- 4 files changed, 424 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 165486a4669..7a996a86d59 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,7 @@ import os import random import time import traceback -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload @@ -166,13 +166,67 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): _DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) -def _daily_spend_commit_failure_is_requeue_safe(e: Exception) -> bool: +def _spend_commit_failure_is_requeue_safe(e: Exception) -> bool: if isinstance(e, DB_CONNECTION_ERROR_TYPES): return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES +_SpendTableName = Literal[ + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", + "org_member_list_transactions", + "project_list_transactions", + "tag_list_transactions", + "model_access_group_list_transactions", + "agent_list_transactions", +] +_SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = ( + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", + "org_member_list_transactions", + "project_list_transactions", + "tag_list_transactions", + "model_access_group_list_transactions", + "agent_list_transactions", +) + + +def _spend_tables_left_to_send( + transactions: DBSpendUpdateTransactions, + committed: Sequence[_SpendTableName], + failure: Exception, +) -> DBSpendUpdateTransactions | None: + in_flight: Final[_SpendTableName | None] = ( + _SPEND_TABLE_COMMIT_ORDER[len(committed)] if len(committed) < len(_SPEND_TABLE_COMMIT_ORDER) else None + ) + dropped: Final[frozenset[_SpendTableName]] = ( + frozenset() if in_flight is None or _spend_commit_failure_is_requeue_safe(failure) else frozenset({in_flight}) + ) + if dropped and in_flight is not None: + spend_log_error( + "Spend tracking - dropped %d %s increments: the failed statement may have applied or the " + "database refused the data, so re-sending it is not safe. Error: %s", + len(cast(dict[str, dict[str, float] | None], transactions).get(in_flight) or ()), + in_flight, + str(failure), + exc=failure, + ) + remaining: Final = { + name: (None if name in committed or name in dropped else txns) for name, txns in transactions.items() + } + return cast(DBSpendUpdateTransactions, remaining) if any(remaining.values()) else None + + def _timed_request_duration_ms( payload: dict | SpendLogsPayload, request_status: Literal["success", "failure"], @@ -1284,6 +1338,7 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("acquired lock for spend updates") uncommitted: dict[str, Any] = {} # mutable-ok: tracks popped categories still needing commit + committed_spend_tables: Final[list[_SpendTableName]] = [] # mutable-ok: filled as each table lands try: ( @@ -1323,12 +1378,19 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("agent_list_transactions") or ()), len(db_spend_update_transactions.get("model_access_group_list_transactions") or ()), ) - await self._commit_spend_updates_to_db( - prisma_client=prisma_client, - n_retry_times=n_retry_times, - proxy_logging_obj=proxy_logging_obj, - db_spend_update_transactions=db_spend_update_transactions, - ) + try: + await self._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=db_spend_update_transactions, + on_table_committed=committed_spend_tables.append, + ) + except Exception as e: + uncommitted["db_spend_update_transactions"] = _spend_tables_left_to_send( + db_spend_update_transactions, committed_spend_tables, e + ) + raise uncommitted.pop("db_spend_update_transactions", None) if daily_spend_update_transactions is not None: @@ -1376,10 +1438,22 @@ class DBSpendUpdateWriter: ) uncommitted.pop("daily_agent_spend_update_transactions", None) if window_spend_update_transactions is not None: - await DBSpendUpdateWriter._commit_window_spend_updates( - prisma_client=prisma_client, - window_spend_transactions=window_spend_update_transactions, - ) + try: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + except Exception as e: + if not _spend_commit_failure_is_requeue_safe(e): + uncommitted.pop("window_spend_update_transactions", None) + spend_log_error( + "Spend tracking - dropped %d budget window increments: the failed statement may have " + "applied or the database refused the data, so re-sending it is not safe. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) + raise uncommitted.pop("window_spend_update_transactions", None) except Exception as e: spend_log_error( @@ -1523,14 +1597,23 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run - spend_log_error( - "Spend tracking - failed to commit budget window spend updates. " - "Re-queued %d window increments for retry on next tick. Error: %s", - len(window_spend_update_transactions), - str(e), - exc=e, - ) - await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions) + if _spend_commit_failure_is_requeue_safe(e): + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. " + "Re-queued %d window increments for retry on next tick. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) + await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions) + else: + spend_log_error( + "Spend tracking - dropped %d budget window increments: the failed statement may have " + "applied or the database refused the data, so re-sending it is not safe. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1688,6 +1771,7 @@ class DBSpendUpdateWriter: n_retry_times: int, proxy_logging_obj: ProxyLogging, db_spend_update_transactions: DBSpendUpdateTransactions, + on_table_committed: Callable[[_SpendTableName], None] | None = None, ): """ Commits all the spend `UPDATE` transactions to the Database @@ -1721,6 +1805,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("user_list_transactions") ### UPDATE END-USER TABLE ### end_user_list_transactions: Final = db_spend_update_transactions["end_user_list_transactions"] @@ -1732,6 +1818,8 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, end_user_list_transactions=end_user_list_transactions, ) + if on_table_committed is not None: + on_table_committed("end_user_list_transactions") ### UPDATE KEY TABLE ### key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) @@ -1761,6 +1849,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("key_list_transactions") ### UPDATE TEAM TABLE ### team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] @@ -1789,6 +1879,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("team_list_transactions") ### UPDATE TEAM Membership TABLE with spend ### team_member_list_transactions: Final = db_spend_update_transactions["team_member_list_transactions"] @@ -1817,6 +1909,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("team_member_list_transactions") # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database @@ -1829,6 +1923,8 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug( "Invalidated team membership cache for user_id=%s, team_id=%s", user_id, team_id ) + elif on_table_committed is not None: + on_table_committed("team_member_list_transactions") ### UPDATE ORG TABLE ### org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] @@ -1854,6 +1950,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("org_list_transactions") org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) @@ -1877,6 +1975,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("org_member_list_transactions") ### UPDATE PROJECT TABLE ### project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") @@ -1889,6 +1989,8 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("project_list_transactions") await DBSpendUpdateWriter._invalidate_project_caches( project_ids=tuple(project_list_transactions or ()), proxy_logging_obj=proxy_logging_obj, @@ -1905,6 +2007,8 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("tag_list_transactions") ### UPDATE MODEL ACCESS GROUP TABLE ### model_access_group_list_transactions: Final = db_spend_update_transactions.get( @@ -1919,6 +2023,8 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("model_access_group_list_transactions") ### UPDATE AGENT TABLE ### agent_list_transactions: Final = db_spend_update_transactions["agent_list_transactions"] @@ -1931,6 +2037,8 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) + if on_table_committed is not None: + on_table_committed("agent_list_transactions") @staticmethod async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None: @@ -2140,7 +2248,7 @@ class DBSpendUpdateWriter: sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: - if _daily_spend_commit_failure_is_requeue_safe(batch_error): + if _spend_commit_failure_is_requeue_safe(batch_error): spend_log_error( "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 460bf5db3b1..9146f234570 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,3 +1,4 @@ +import re from collections.abc import Awaitable, Callable, Iterator from typing import Any, Final, TypeVar @@ -20,6 +21,7 @@ _TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = ( ) _DATABASE_ERROR_META: Final = TypeAdapter(dict[str, object]) +_BATCH_POSTGRES_ERROR_CODE: Final = re.compile(r'PostgresError \{ code: "([0-9A-Z]{5})"') def _exception_chain(e: BaseException) -> Iterator[BaseException]: @@ -40,6 +42,13 @@ def _database_service_unavailable_errors(e: BaseException) -> tuple[Exception, . ) +def _batch_postgres_sqlstate(e: Exception) -> str | None: + """The SQLSTATE a batched statement failed with: prisma reports those without a + ``meta`` payload and only prints the connector error into the message.""" + match: Final = _BATCH_POSTGRES_ERROR_CODE.search(str(e)) + return match.group(1) if match is not None else None + + def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]: """Keep only the real exception classes among ``candidates``. @@ -235,9 +244,9 @@ class PrismaDBExceptionHandler: try: meta: Final = _DATABASE_ERROR_META.validate_python(getattr(e, "meta", None)) except ValidationError: - return None + return _batch_postgres_sqlstate(e) code: Final = meta.get("code") - return code if isinstance(code, str) else None + return code if isinstance(code, str) else _batch_postgres_sqlstate(e) @staticmethod def is_read_only_transaction_error(e: Exception) -> bool: 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 d08ff77f364..d65acb1f9b7 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 @@ -14,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest +from prisma.errors import DataError as PrismaDataError from prisma.errors import RawQueryError from redis.exceptions import DataError @@ -24,6 +25,8 @@ from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, _TEAM_MEMBER_SPEND_SQL, DBSpendUpdateWriter, + _SpendTableName, + _spend_tables_left_to_send, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( @@ -3007,6 +3010,29 @@ def _postgres_rejection(sqlstate: str) -> RawQueryError: ) +def _batched_postgres_rejection(sqlstate: str) -> PrismaDataError: + return PrismaDataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error occurred during query execution:\nConnectorError(ConnectorError { " + f'user_facing_error: None, kind: QueryError(PostgresError {{ code: "{sqlstate}", ' + 'message: "db error", severity: "ERROR" }) })', + "batch_request_idx": 0, + } + } + ) + + +_REQUEUE_SAFETY_CASES: Final = [ + pytest.param(httpx.ReadTimeout("no reply"), False, id="reply lost after the statement was sent"), + pytest.param(httpx.ConnectError("refused"), True, id="statement never reached the database"), + pytest.param(_postgres_rejection("23514"), False, id="postgres refused a constraint violation"), + pytest.param(_batched_postgres_rejection("23514"), False, id="postgres refused a batched constraint violation"), + pytest.param(_postgres_rejection("42P01"), True, id="table missing"), +] + + @pytest.mark.parametrize( ("failure", "lands_on_the_next_tick"), [ @@ -3175,6 +3201,169 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): db_writer.pod_lock_manager.release_lock.assert_awaited_once() +@pytest.mark.parametrize(("failure", "safe_to_resend"), _REQUEUE_SAFETY_CASES) +@pytest.mark.asyncio +async def test_failed_per_entity_increment_from_redis_restores_only_what_may_still_be_sent( + failure: Exception, safe_to_resend: bool +): + db_writer = DBSpendUpdateWriter() + db_spend_transactions = _empty_spend_transactions( + user_list_transactions={"user-1": 1.5}, + key_list_transactions={"key-1": 1.5}, + team_list_transactions={"team-1": 1.5}, + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(db_spend_transactions, None, None, None, None, None, None) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + + mock_batcher = MagicMock() + for table_name in ( + "litellm_usertable", + "litellm_verificationtoken", + "litellm_teamtable", + "litellm_teammembership", + "litellm_organizationtable", + "litellm_organizationmembership", + "litellm_projecttable", + "litellm_tagtable", + "litellm_modelaccessgroupbudgettable", + "litellm_agentstable", + ): + setattr(mock_batcher, table_name, MagicMock()) + mock_batcher.litellm_verificationtoken.update_many.side_effect = failure + + class _BatchContext: + async def __aenter__(self): + return mock_batcher + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + class _Transaction: + def batch_(self): + return _BatchContext() + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=_Transaction()) + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + with patch( # test-quality-ok: retry sleeps are disabled to exercise ConnectError without waiting + "litellm.proxy.db.db_spend_update_writer.asyncio.sleep", + new_callable=AsyncMock, + ): + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=proxy_logging_obj, + ) + + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() + restored = mock_redis_update_buffer.restore_transactions_to_redis.call_args.kwargs[ + "db_spend_update_transactions" + ] + assert restored["user_list_transactions"] is None + assert restored["team_list_transactions"] == {"team-1": 1.5} + assert restored["key_list_transactions"] == ({"key-1": 1.5} if safe_to_resend else None) + mock_batcher.litellm_usertable.update_many.assert_called_once() + + +@pytest.mark.parametrize(("failure", "safe_to_resend"), _REQUEUE_SAFETY_CASES) +@pytest.mark.asyncio +async def test_failed_window_spend_commit_from_redis_is_restored_only_when_safe_to_resend( + failure: Exception, safe_to_resend: bool +): + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=failure) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + assert _window_spend_upserts(db) == [] + if safe_to_resend: + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + window_spend_update_transactions=window_transactions + ) + else: + mock_redis_update_buffer.restore_transactions_to_redis.assert_not_awaited() + db_writer.pod_lock_manager.release_lock.assert_awaited_once() + + +@pytest.mark.parametrize(("failure", "safe_to_resend"), _REQUEUE_SAFETY_CASES) +@pytest.mark.asyncio +async def test_failed_window_spend_commit_is_requeued_only_when_the_rows_are_provably_uncommitted( + failure: Exception, safe_to_resend: bool +): + class _WindowSpendFailureDB(_WindowSpendFakeDB): + def __init__(self, failure: Exception | None) -> None: + super().__init__() + self.failure = failure + + async def query_raw(self, query, *args): + if self.failure is not None: + raise self.failure + return await super().query_raw(query, *args) + + db_writer = DBSpendUpdateWriter() + transaction = build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + ) + await db_writer.window_spend_update_queue.add_update(transaction) + db = _WindowSpendFailureDB(failure) + db_writer._flush_tool_discovery_queue = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + db.failure = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + assert len(_window_spend_upserts(db)) == (1 if safe_to_resend else 0) + assert db_writer.window_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at(): """Spend flushes must leave settings_updated_at alone, or it decays into @@ -3239,6 +3428,84 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at assert call_kwargs["data"]["spend"] == {"increment": response_cost} +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_reports_each_completed_table(): + db_writer = DBSpendUpdateWriter() + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(MagicMock())) + proxy_logging_obj = MagicMock() + proxy_logging_obj.call_details = {} + on_table_committed = MagicMock() + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=_empty_spend_transactions( + org_member_list_transactions={}, + project_list_transactions={}, + model_access_group_list_transactions={}, + ), + on_table_committed=on_table_committed, + ) + + assert on_table_committed.call_args_list == [ + call("user_list_transactions"), + call("end_user_list_transactions"), + call("key_list_transactions"), + call("team_list_transactions"), + call("team_member_list_transactions"), + call("org_list_transactions"), + call("org_member_list_transactions"), + call("project_list_transactions"), + call("tag_list_transactions"), + call("model_access_group_list_transactions"), + call("agent_list_transactions"), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "transactions, table", + [ + pytest.param( + {"team_member_list_transactions": {"team_id::t1::user_id::u1": 0.5}}, + "team_member_list_transactions", + id="team membership spend landed before its cache invalidation failed", + ), + pytest.param( + {"project_list_transactions": {"p1": 0.5}}, + "project_list_transactions", + id="project spend landed before its cache invalidation failed", + ), + ], +) +async def test_commit_spend_updates_to_db_reports_table_committed_before_cache_invalidation( + transactions: dict[str, dict[str, float]], table: _SpendTableName +): + db_writer = DBSpendUpdateWriter() + mock_prisma_client = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(MagicMock())) + user_api_key_cache = MagicMock() + user_api_key_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down")) + proxy_logging_obj = MagicMock() + proxy_logging_obj.call_details = {"user_api_key_cache": user_api_key_cache} + committed = [] + + with pytest.raises(ConnectionError): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=_empty_spend_transactions(**transactions), + on_table_committed=committed.append, + ) + + user_api_key_cache.async_delete_cache.assert_awaited_once() + assert committed[-1] == table + assert _spend_tables_left_to_send(_empty_spend_transactions(**transactions), committed, ConnectionError()) is None + + @pytest.mark.asyncio async def test_daily_transaction_internal_call_keeps_spend_but_not_request_counts(): """Internal sub-calls (auto-router classifier, shadow eval's shadow and judge) bill diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index f7cc5e3ed83..613ca847115 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -677,13 +677,28 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"message": "m"}}}), None), (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"code": 42, "message": "m"}}}), None), (prisma_errors.DataError(data={"user_facing_error": {"meta": None}}), None), + ( + prisma_errors.DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error occurred during query execution:\nConnectorError(ConnectorError { " + 'user_facing_error: None, kind: QueryError(PostgresError { code: "23514", ' + 'message: "new row violates check constraint", severity: "ERROR" }) })', + "batch_request_idx": 0, + } + } + ), + "23514", + ), (PrismaError("db error"), None), (httpx.ReadTimeout("no reply"), None), ], ) def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error: Exception, sqlstate: str | None): - """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a - codeless or malformed payload, an engine-level error, and a transport error yield None.""" + """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE, whether in ``meta`` + or, for a batched statement, only in the message; a codeless or malformed payload, an engine-level + error, and a transport error yield None.""" assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate