fix(proxy): retry lock-timed-out daily spend batches in place so the shutdown flush keeps them

litellm/proxy/db/db_spend_update_writer.py

@@ -2589,9 +2589,11 @@ async def _update_daily_spend(
                             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:

tests/unit/proxy/db/test_db_spend_update_writer.py

@@ -537,6 +537,54 @@ async def fake_sleep(seconds: float) -> None:
     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():
     """