diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 158a6c1a3c4..9fb8cc20981 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -2589,9 +2589,11 @@ class DBSpendUpdateWriter: PrismaDBExceptionHandler, ) - is_retryable = isinstance( - e, DB_RETRY_SAFE_ERROR_TYPES - ) or PrismaDBExceptionHandler.is_deadlock_error(e) + is_retryable = ( + isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + or PrismaDBExceptionHandler.is_deadlock_error(e) + or PrismaDBExceptionHandler.is_lock_timeout_error(e) + ) if not is_retryable: raise if i >= n_retry_times: diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index dc9b8fd395b..4a508fcd272 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -537,6 +537,54 @@ async def test_update_daily_spend_retries_connect_errors(monkeypatch): assert len(prisma_client.db.statements) == 2 +@pytest.mark.asyncio +async def test_update_daily_spend_retries_lock_timeout_errors(monkeypatch: pytest.MonkeyPatch) -> None: + def _lock_timeout_error() -> PrismaDataError: + return PrismaDataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error querying the database: canceling statement due to lock timeout", + "meta": {"code": "55P03", "message": "canceling statement due to lock timeout"}, + } + } + ) + + outcomes: Final = iter([_lock_timeout_error(), None]) + + def first_attempt_locks_out() -> int: + outcome: Final = next(outcomes) + if outcome is not None: + raise outcome + return 1 + + prisma_client: Final = _RecordingPrisma(execute_raw=first_attempt_locks_out) + proxy_logging: Final = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + async def fake_sleep(seconds: float) -> None: + return None + + monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", fake_sleep) + daily_spend_transactions: Final = {"k1": _daily_txn()} + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=3, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging, + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + ) + + assert len(prisma_client.db.statements) == 2, ( + "a 55P03 lock_timeout cancels the upsert before it applies, so the writer must " + "resend it in place instead of only requeueing it for a next tick a shutdown " + "flush never gets" + ) + assert daily_spend_transactions == {}, "the retried batch must drain the transactions dict" + proxy_logging.failure_handler.assert_not_called() + + @pytest.mark.asyncio async def test_update_daily_spend_sorting(): """