fix(proxy): bound spend-log cleanup delete transactions to the statement timeout (#43279)

Each spend-log delete batch and the remaining-rows probe ran inside a
prisma interactive transaction with the default 5s timeout while setting a
30s Postgres statement and lock timeout. A batch that waited on a lock for
more than 5s was closed by the query engine with a 422 and rolled back, and
after three such batches the cleanup run aborted with nothing deleted.

Reuse the partition manager's bounded_tx, which sizes the prisma timeout to
the statement bound plus commit slack, for the delete batch and the probe.

Resolves LIT-8686

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py

@@ -28,6 +28,8 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
 from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
     RemainingTimeoutMs,
     SpendLogsPartitionManager,
+    bounded_tx,
+    trusted_sql,
 )
 from litellm.proxy.utils import PrismaClient
 from litellm.types.llms.base import LiteLLMBaseModel
@@ -307,11 +309,15 @@ class SpendLogCleanup:
         fault, so the caller stops instead of retrying.
         """
         timeout_ms: Final = self._timeout_ms(deadline)
-        async with db_span("cleanup_expired_rows", table_name), prisma_client.db.tx() as tx:
-            await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
-            await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
-            deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size)
-        return deleted_result if isinstance(deleted_result, int) else None
+        async with db_span("cleanup_expired_rows", table_name), bounded_tx(prisma_client, timeout_ms) as tx:
+            await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
+            await tx.execute_raw(trusted_sql(f"SET LOCAL lock_timeout = {timeout_ms}"))
+            deleted_result: Final = await tx.execute_raw(trusted_sql(delete_sql), cutoff_date, self.batch_size)
+        return (
+            deleted_result
+            if isinstance(deleted_result, int)  # pyright: ignore[reportUnnecessaryIsInstance]  # driver may return None
+            else None
+        )
 
     async def _count_remaining(
         self, prisma_client: PrismaClient, cutoff_date: Cutoff, table_name: str, time_column: str, deadline: float
@@ -331,11 +337,12 @@ class SpendLogCleanup:
                 LIMIT $2
             ) capped
             """
+        timeout_ms: Final = self._timeout_ms(deadline)
         try:
-            async with db_span("count_expired_rows", table_name), prisma_client.db.tx() as tx:
-                await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}")
+            async with db_span("count_expired_rows", table_name), bounded_tx(prisma_client, timeout_ms) as tx:
+                await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
                 rows: Final = _REMAINING_ROWS.validate_python(
-                    await tx.query_raw(count_sql, cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP)
+                    await tx.query_raw(trusted_sql(count_sql), cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP)
                 )
         except Exception as e:  # noqa: BLE001 - an observability probe must never fail the cleanup run
             verbose_proxy_logger.warning("Could not count remaining %s rows: %s", table_name, e)

litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py

@@ -32,6 +32,7 @@ from litellm.proxy.db.db_span import db_span
 
 if TYPE_CHECKING:
     from prisma.client import TransactionManager
+    from typing_extensions import LiteralString
 
     from litellm.proxy.utils import PrismaClient
 
@@ -127,7 +128,7 @@ def select_partitions_to_drop(partitions: list[tuple[str, datetime | None]], cut
 _TX_COMMIT_SLACK: Final = timedelta(seconds=5)
 
 
-def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager":
+def bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager":
     """
     Open an interactive transaction that outlives the statement bound it
     carries. prisma's default 5s transaction timeout would close it mid
@@ -139,6 +140,12 @@ def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionM
     )
 
 
+def trusted_sql(sql: str) -> "LiteralString":
+    return cast(  # cast-ok: identifiers come from constants or the catalog, values are bound as $n parameters
+        "LiteralString", sql
+    )
+
+
 class SpendLogsPartitionManager:
     def __init__(
         self,
@@ -162,7 +169,7 @@ class SpendLogsPartitionManager:
         try:
             async with (
                 db_span("check_spend_log_partitioning", "LiteLLM_SpendLogs"),
-                _bounded_tx(prisma_client, budget_ms) as tx,
+                bounded_tx(prisma_client, budget_ms) as tx,
             ):
                 await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
                 rows: Final = await tx.query_raw(
@@ -200,7 +207,7 @@ class SpendLogsPartitionManager:
         """
         async with (
             db_span("create_spend_log_partition", "LiteLLM_SpendLogs"),
-            _bounded_tx(prisma_client, timeout_ms) as tx,
+            bounded_tx(prisma_client, timeout_ms) as tx,
         ):
             await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
             await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
@@ -240,7 +247,7 @@ class SpendLogsPartitionManager:
     ) -> list[tuple[str, datetime | None]]:
         async with (
             db_span("list_spend_log_partitions", "LiteLLM_SpendLogs"),
-            _bounded_tx(prisma_client, timeout_ms) as tx,
+            bounded_tx(prisma_client, timeout_ms) as tx,
         ):
             await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
             rows: Final = await tx.query_raw(

tests/integration/database/test_cleanup_transactions.py

@@ -0,0 +1,107 @@
+import asyncio
+import logging
+import os
+import time
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Final
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+import psycopg
+import pytest
+from integration._support.database import read_rows
+from prisma import Prisma
+from psycopg import sql
+
+from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
+
+
+@dataclass(frozen=True)
+class PartitionConnection:
+    db: Prisma
+
+
+async def test_delete_batch_survives_witnessed_lock_past_transaction_default(caplog: pytest.LogCaptureFixture) -> None:
+    caplog.set_level(logging.ERROR)
+    schema: Final = f"integration_{uuid.uuid4().hex}"
+    url: Final = os.environ["DATABASE_URL"]
+    parsed: Final = urlsplit(url)
+    scoped_url: Final = urlunsplit(
+        parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema}))
+    )
+    table: Final = sql.Identifier(schema, "LiteLLM_SpendLogs")
+    now: Final = datetime.now(timezone.utc).replace(tzinfo=None)
+    cutoff_date: Final = now - timedelta(days=30)
+    with psycopg.connect(url, autocommit=True) as setup:
+        setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
+        try:
+            setup.execute(sql.SQL('CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL)').format(table))
+            setup.execute(
+                sql.SQL("INSERT INTO {} VALUES (%s, %s), (%s, %s), (%s, %s), (%s, %s)").format(table),
+                (
+                    "expired-1",
+                    now - timedelta(days=60),
+                    "expired-2",
+                    now - timedelta(days=60),
+                    "expired-3",
+                    now - timedelta(days=60),
+                    "fresh",
+                    now,
+                ),
+            )
+            database: Final = Prisma(datasource={"url": scoped_url})
+            await database.connect()
+            try:
+                cleanup: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
+                with psycopg.connect(url) as blocker:
+                    blocker.execute(sql.SQL("LOCK TABLE {} IN SHARE ROW EXCLUSIVE MODE").format(table))
+                    blocker_pid: Final = blocker.info.backend_pid
+                    operation: Final = asyncio.create_task(
+                        cleanup._delete_old_rows_batched(
+                            PartitionConnection(database),
+                            cutoff_date=cutoff_date,
+                            table_name="LiteLLM_SpendLogs",
+                            key_columns=("request_id",),
+                            time_column="startTime",
+                            deadline=time.monotonic() + 30,
+                        )
+                    )
+                    wait_deadline: Final = time.monotonic() + 3
+                    try:
+                        while True:
+                            witnesses: Final = read_rows(
+                                "SELECT a.pid, extract(epoch FROM "
+                                "clock_timestamp()-a.query_start)::double precision AS age "
+                                "FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) "
+                                "AND a.wait_event_type = 'Lock' AND a.query LIKE '%%DELETE FROM%%'",
+                                (blocker_pid,),
+                            )
+                            if witnesses:
+                                break
+                            assert time.monotonic() < wait_deadline, "Cleanup DELETE never reached the held lock"
+                            await asyncio.sleep(0.02)
+                        assert len(witnesses) == 1
+                        held_at: Final = time.monotonic()
+                        age: Final = float(witnesses[0]["age"])
+                        await asyncio.sleep(max(0, 5.6 - age))
+                        held_seconds: Final = age + time.monotonic() - held_at
+                        assert held_seconds >= 5.5, f"Lock released before the transaction boundary: {held_seconds}"
+                        assert not operation.done(), "Cleanup DELETE completed while its required lock was held"
+                    except BaseException:
+                        operation.cancel()
+                        await asyncio.gather(operation, return_exceptions=True)
+                        raise
+                    finally:
+                        blocker.rollback()
+                    result: Final = await asyncio.wait_for(operation, timeout=5)
+                assert result.rows_deleted == 3
+                assert not any("cleanup batch failed" in record.message for record in caplog.records)
+                assert read_rows(
+                    f"SELECT request_id FROM {table.as_string(setup)} ORDER BY request_id",
+                    (),
+                ) == [{"request_id": "fresh"}]
+            finally:
+                await database.disconnect()
+        finally:
+            setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))

tests/unit/proxy/test_spend_log_cleanup.py

@@ -6,8 +6,8 @@ import asyncio
 import logging
 import math
 import time
-from contextlib import asynccontextmanager
 from collections.abc import Awaitable, Callable
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta, timezone
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock
@@ -24,10 +24,10 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
     SpendLogCleanup,
     TableCleanupResult,
 )
-from tests.unit.proxy.db.fake_prisma_engine import engine_call
 from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
     SpendLogCleanupMetrics,
 )
+from tests.unit.proxy.db.fake_prisma_engine import engine_call
 
 
 def _far_deadline() -> float:
@@ -35,7 +35,7 @@ def _far_deadline() -> float:
     return time.monotonic() + 3600
 
 
-def _wire_tx(db):
+def _wire_tx(db, rejected_timeout=None):
     """
     Model the prisma seam the cleanup job actually uses.
 
@@ -49,7 +49,9 @@ def _wire_tx(db):
     """
 
     @asynccontextmanager
-    async def _tx():
+    async def _tx(**kwargs):
+        if kwargs.get("timeout") == rejected_timeout:
+            raise TypeError("transaction timeout rejected by fake")
         tx = MagicMock()
 
         async def _execute_raw(sql, *args):
@@ -102,6 +104,7 @@ def test_spend_log_cleanup_cron_scheduler_integration():
     a real database connection.
     """
     from unittest.mock import MagicMock
+
     from apscheduler.triggers.cron import CronTrigger
 
     # Mock scheduler
@@ -208,6 +211,50 @@ async def test_should_delete_spend_logs():
     assert cleaner._should_delete_spend_logs() is False
 
 
+@pytest.mark.asyncio
+async def test_delete_batch_transaction_timeout_covers_statement_timeout() -> None:
+    transaction: Final = MagicMock()
+    transaction.execute_raw = AsyncMock(side_effect=[0, 0, 3])
+
+    @asynccontextmanager
+    async def tx(timeout: timedelta):
+        yield transaction
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.tx = MagicMock(side_effect=tx)
+    cleaner: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
+    cleaner.batch_timeout_seconds = 12
+
+    await cleaner._execute_delete_batch(
+        prisma_client, "DELETE FROM table", datetime.now(), "LiteLLM_SpendLogs", _far_deadline()
+    )
+
+    assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12)
+
+
+@pytest.mark.asyncio
+async def test_count_remaining_transaction_timeout_covers_statement_timeout() -> None:
+    transaction: Final = MagicMock()
+    transaction.execute_raw = AsyncMock(return_value=0)
+    transaction.query_raw = AsyncMock(return_value=[{"remaining": 7}])
+
+    @asynccontextmanager
+    async def tx(timeout: timedelta):
+        yield transaction
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.tx = MagicMock(side_effect=tx)
+    cleaner: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
+    cleaner.batch_timeout_seconds = 12
+
+    result: Final = await cleaner._count_remaining(
+        prisma_client, datetime.now(), "LiteLLM_SpendLogs", "startTime", _far_deadline()
+    )
+
+    assert result == 7
+    assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12)
+
+
 @pytest.mark.asyncio
 async def test_cleanup_old_spend_logs_batch_deletion():
     from unittest.mock import AsyncMock, MagicMock
@@ -783,7 +830,7 @@ def _mock_prisma_for_retention(side_effect: list) -> "MagicMock":
     from unittest.mock import AsyncMock, MagicMock
 
     client = MagicMock()
-    _wire_tx(client.db)
+    _wire_tx(client.db, rejected_timeout=timedelta(seconds=10))
     client.db.execute_raw = AsyncMock(side_effect=side_effect)
     return client
 
@@ -1013,7 +1060,7 @@ async def test_each_batch_carries_a_statement_and_lock_timeout():
     mock_db = MagicMock()
 
     @asynccontextmanager
-    async def _tx():
+    async def _tx(**kwargs):
         tx = MagicMock()
 
         async def _execute_raw(sql, *args):
@@ -1222,7 +1269,7 @@ async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
     mock_db = MagicMock()
 
     @asynccontextmanager
-    async def _tx():
+    async def _tx(**kwargs):
         tx = MagicMock()
 
         async def _execute_raw(sql, *args):
@@ -1274,7 +1321,7 @@ async def test_a_statement_timeout_is_clamped_to_the_budget_that_is_left():
     client = MagicMock()
 
     @asynccontextmanager
-    async def _tx():
+    async def _tx(**kwargs):
         tx = MagicMock()
 
         async def _execute_raw(sql, *args):