From e8c33f35680c222f9a2e0bb02854dc47e2b6ff7f Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 21 Oct 2026 00:00:05 +0000 Subject: [PATCH] fix(proxy): requeue spend logs when a Postgres restart tears down the prisma engine mid-write (#47571) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Backport of #47571 to stable/1.106.x Cherry-picked from ccd0d4ea493d1c5f9c2f6574261002407687663d (main) --- litellm/proxy/utils.py | 6 ++- .../test_proxy_update_spend.py | 50 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b34442378be..f4492f92cbd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -8239,7 +8239,11 @@ MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 def _is_transient_spend_log_write_error(e: Exception) -> bool: - return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) + return ( + PrismaDBExceptionHandler.is_database_transport_error(e) + or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) + or PrismaDBExceptionHandler.is_deadlock_error(e) + ) async def _create_spend_logs_with_poison_isolation( diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index fc3deb0de18..e3b7bb5b49f 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -952,3 +952,53 @@ async def test_update_spend_logs_requeues_batch_when_postgres_is_out_of_connecti assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 sleep.assert_not_awaited() assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r0", "r1", "r2", "r3", "e"] + + +async def _create_many_while_postgres_is_restarting(**_: object) -> int: + from prisma.engine.utils import handle_response_errors + + handle_response_errors( + MagicMock(), + [ + { + "error": "Error in connector: Can't reach database server at `127.0.0.1:5432`", + "user_facing_error": {"error_code": "P1001", "message": "Can't reach database server", "meta": None}, + } + ], + ) + return 0 + + +@pytest.mark.asyncio +async def test_update_spend_logs_retries_and_requeues_batch_when_postgres_restart_tears_down_prisma_engine( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """While Postgres restarts the prisma engine answers the write with a P1001 + payload whose ``meta`` is ``None``, which prisma's own response parser trips + over before it can build a typed error. That outage must be retried like + any other, and once retries exhaust the dequeued batch goes back to the head + of the queue so the rows land when Postgres is back. + """ + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock( + side_effect=_create_many_while_postgres_is_restarting + ) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")] + + with pytest.raises(AttributeError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3