fix(proxy): requeue spend logs when a Postgres restart tears down the prisma engine mid-write (#47571)

Backport of #47571 to stable/1.103.x

Cherry-picked from ccd0d4ea49 (main)
This commit is contained in:
devin-ai-integration[bot] 2026-10-21 00:23:12 +00:00
parent 6120ddf786
commit d7e344440c
2 changed files with 55 additions and 1 deletions

View file

@ -7672,7 +7672,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(

View file

@ -883,3 +883,53 @@ def test_disable_spend_updates_error_when_general_settings_unavailable(
monkeypatch.delattr(proxy_server_mod, "general_settings", raising=False)
with pytest.raises(ImportError):
ProxyUpdateSpend.disable_spend_updates()
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