From e7a3a84dc0ca73abdd214ac81acc7cc2d044caee Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 03:35:31 +0000 Subject: [PATCH 1/3] fix(proxy): bound spend-log cleanup delete transactions to the statement timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 8 +- .../spend_logs_partition_manager.py | 8 +- .../database/test_cleanup_transactions.py | 104 ++++++++++++++++++ .../proxy/test_spend_log_cleanup.py | 36 +++++- 4 files changed, 143 insertions(+), 13 deletions(-) create mode 100644 tests/integration/database/test_cleanup_transactions.py diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index c6f52bf074b..25d22f81c1c 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -27,6 +27,7 @@ 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, ) from litellm.proxy.utils import PrismaClient @@ -294,7 +295,7 @@ class SpendLogCleanup: fault, so the caller stops instead of retrying. """ timeout_ms: Final = self._timeout_ms(deadline) - async with prisma_client.db.tx() as tx: + async with 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}") deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size) @@ -318,9 +319,10 @@ class SpendLogCleanup: LIMIT $2 ) capped """ + timeout_ms: Final = self._timeout_ms(deadline) try: - async with prisma_client.db.tx() as tx: - await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}") + async with bounded_tx(prisma_client, timeout_ms) as tx: + await tx.execute_raw(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) ) diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index b756bfeb6f6..c15784ba12d 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -126,7 +126,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 @@ -159,7 +159,7 @@ class SpendLogsPartitionManager: if budget_ms is None: return False try: - async with _bounded_tx(prisma_client, budget_ms) as tx: + async with 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( """ @@ -194,7 +194,7 @@ class SpendLogsPartitionManager: wait for the lock and statement_timeout bounds the work itself, so a partition this run cannot get is simply left for the next one. """ - async with _bounded_tx(prisma_client, timeout_ms) as tx: + async with 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}") await tx.execute_raw(statement) @@ -231,7 +231,7 @@ class SpendLogsPartitionManager: async def _list_partitions( self, prisma_client: "PrismaClient", timeout_ms: int ) -> list[tuple[str, datetime | None]]: - async with _bounded_tx(prisma_client, timeout_ms) as tx: + async with 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( """ diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py new file mode 100644 index 00000000000..a716cb09b8f --- /dev/null +++ b/tests/integration/database/test_cleanup_transactions.py @@ -0,0 +1,104 @@ +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 + + +@pytest.mark.covers("other.database.cleanup.lock_wait_outlives_transaction_default") +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))) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 72463e17c6b..c0b9b59cb58 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -33,7 +33,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. @@ -47,7 +47,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): @@ -100,6 +102,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 @@ -206,6 +209,27 @@ 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(): + tx_kwargs = [] + + @asynccontextmanager + async def tx(**kwargs): + tx_kwargs.append(kwargs) + transaction = MagicMock() + transaction.execute_raw = AsyncMock(side_effect=[0, 0, 3]) + yield transaction + + prisma_client = MagicMock() + prisma_client.db.tx = tx + cleaner = 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(), _far_deadline()) + + assert tx_kwargs[0]["timeout"] >= timedelta(seconds=12) + + @pytest.mark.asyncio async def test_cleanup_old_spend_logs_batch_deletion(): from unittest.mock import AsyncMock, MagicMock @@ -781,7 +805,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 @@ -981,7 +1005,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): @@ -1190,7 +1214,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): @@ -1242,7 +1266,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): From 4836b598c09b47c5f2051ca2132e0afadd2991fe Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 03:48:37 +0000 Subject: [PATCH 2/3] test(integration): drop legacy covers marker from cleanup lock test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/database/test_cleanup_transactions.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py index a716cb09b8f..f33f819bc23 100644 --- a/tests/integration/database/test_cleanup_transactions.py +++ b/tests/integration/database/test_cleanup_transactions.py @@ -22,7 +22,6 @@ class PartitionConnection: db: Prisma -@pytest.mark.covers("other.database.cleanup.lock_wait_outlives_transaction_default") 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}" From 64be9089fd4fe505c9b06f3690eab2c9f1098b13 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 04:00:44 +0000 Subject: [PATCH 3/3] test(proxy): type the delete-batch transaction timeout capture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/test_spend_log_cleanup.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index c0b9b59cb58..3e51f7a451c 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -210,24 +210,22 @@ async def test_should_delete_spend_logs(): @pytest.mark.asyncio -async def test_delete_batch_transaction_timeout_covers_statement_timeout(): - tx_kwargs = [] +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(**kwargs): - tx_kwargs.append(kwargs) - transaction = MagicMock() - transaction.execute_raw = AsyncMock(side_effect=[0, 0, 3]) + async def tx(timeout: timedelta): yield transaction - prisma_client = MagicMock() - prisma_client.db.tx = tx - cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"}) + 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(), _far_deadline()) - assert tx_kwargs[0]["timeout"] >= timedelta(seconds=12) + assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12) @pytest.mark.asyncio