diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c4008a296db..547a10fc8c7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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( diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index d671a4ffc1f..dbf105ce827 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -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