diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 48066945131..3a93896a206 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -66,6 +66,29 @@ class PrismaDBExceptionHandler: return True return False + @staticmethod + def is_prisma_data_error(e: Exception) -> bool: + """True iff ``e`` is a base prisma ``DataError``: the database processed + the statement and refused the data itself (e.g. ``invalid byte sequence + for encoding "UTF8": 0x00``), as opposed to a connectivity failure. + + Matched by exact type, not ``isinstance``: the specific data-layer + subclasses (``UniqueViolationError``, ``TableNotFoundError``, + ``MissingRequiredValueError`` ...) all derive from ``DataError`` but + carry their own semantics, and a systemic one like a missing table must + not be mistaken for a single poison row and bisected away. A raw + Postgres execution error with no prisma P-code surfaces as the base + ``DataError``. + + prisma also wraps the P1001 "can't reach database server" outage as a + base ``DataError``, so a caller that must not treat an outage as a + per-row data rejection has to additionally consult + ``is_database_service_unavailable_error`` before acting on a True here. + """ + import prisma + + return type(e) is prisma.errors.DataError + @staticmethod def is_database_transport_error(e: Exception) -> bool: """ diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 154a17bc4db..4433a35f5d0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -23,7 +23,9 @@ from typing import ( Dict, List, Literal, + Mapping, Optional, + Sequence, Tuple, Union, cast, @@ -5194,8 +5196,10 @@ class ProxyUpdateSpend: for j in range(0, len(logs_to_process), BATCH_SIZE): batch = logs_to_process[j : j + BATCH_SIZE] batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch] - await SpendLogsRepository(prisma_client).table.create_many( - data=batch_with_dates, skip_duplicates=True + await _create_spend_logs_with_poison_isolation( + SpendLogsRepository(prisma_client), + batch_with_dates, + MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH, ) verbose_proxy_logger.debug(f"Flushed {len(batch)} logs to the DB.") # Explicitly clear batch memory @@ -5462,6 +5466,65 @@ async def _monitor_spend_logs_queue( await asyncio.sleep(current_interval) +MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH = 256 + + +async def _create_spend_logs_with_poison_isolation( + repo: SpendLogsRepository, + rows: Sequence[Mapping[str, object]], + attempts_left: int, +) -> int: + """Write spend-log rows, isolating any row Postgres rejects on its data. + + ``create_many`` writes the whole batch in a single statement, so one row + carrying bytes Postgres refuses (a residual NUL byte is the canonical case) + fails the entire insert and drops every good row alongside it. On a genuine + data-layer rejection the batch is bisected so the good rows still persist + and only the offending row is dropped and logged. Transport failures, + including the "can't reach database server" outage that prisma mislabels as + a ``DataError``, are re-raised unchanged so the caller's connection-retry + path still runs. + + ``attempts_left`` is a hard ceiling on the number of ``create_many`` calls + the isolation may issue for this batch, so an authenticated caller flooding + poisoned rows cannot amplify one failed bulk insert into unbounded failed + inserts and log lines. It is checked before any insert (so an exhausted + budget never even attempts a write), decremented once per ``create_many`` + call, and threaded through the recursion so the whole bisection shares one + allowance; total inserts are therefore bounded by the initial value + regardless of how many rows are poisoned. When it runs out the still-failing + remainder is dropped wholesale (the pre-existing drop-the-batch behavior) + under one log line. Returns the budget left after this subtree. + """ + if attempts_left <= 0: + spend_log_error( + "Spend tracking - dropping %d spend log rows without per-row isolation; " + "isolation attempt budget exhausted for this flush", + len(rows), + ) + return 0 + try: + await repo.table.create_many(data=rows, skip_duplicates=True) + return attempts_left - 1 + except Exception as e: + if not PrismaDBExceptionHandler.is_prisma_data_error(e): + raise + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise + if len(rows) == 1: + request_id = rows[0].get("request_id") + spend_log_error( + "Spend tracking - dropping spend log row Postgres rejected. request_id=%s error=%s", + request_id, + str(e), + exc=e, + ) + return attempts_left - 1 + mid = len(rows) // 2 + remaining = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], attempts_left - 1) + return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining) + + def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_logging_obj: ProxyLogging): """ Raise an exception for failed update spend logs diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 6021c221426..0634a01326c 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -148,6 +148,37 @@ def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_datae ) +def test_is_prisma_data_error_only_true_for_dataerror(): + """The spend-log poison-row isolation gates on this: only a prisma + ``DataError`` (the DB refused the data, e.g. a NUL byte) may be bisected + into a per-row drop. A connectivity failure or any non-prisma exception + must not be treated as a data rejection, so the whole batch surfaces.""" + import httpx + + data_error = DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}) + assert PrismaDBExceptionHandler.is_prisma_data_error(data_error) is True + + for non_data in ( + httpx.ConnectError("conn refused"), + PrismaError("can't reach database server"), + UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}), + RuntimeError("boom"), + ): + assert PrismaDBExceptionHandler.is_prisma_data_error(non_data) is False + + +def test_is_prisma_data_error_true_for_connection_masquerade_dataerror(): + """The P1001 outage prisma mislabels as a ``DataError`` is still a + ``DataError`` by type, so this returns True; the spend-log helper relies on + ``is_database_service_unavailable_error`` (not this check) to keep that + outage on the retry path instead of dropping rows.""" + p1001_as_dataerror = DataError( + data={"user_facing_error": {"message": "Can't reach database server at `127.0.0.1`:`5499`"}} + ) + assert PrismaDBExceptionHandler.is_prisma_data_error(p1001_as_dataerror) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(p1001_as_dataerror) is True + + def test_is_database_service_unavailable_error_cached_plan_escapes_as_503(): """Composes with the cached-plan retry: when that recovery fails and the Postgres "cached plan must not change result type" error escapes (raised by 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 6a4fd516c9b..d5d4de7f2cf 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 @@ -231,6 +231,112 @@ async def test_update_spend_logs_failure_raises_after_retries( ) +def _data_error(message: str) -> Any: + from prisma.errors import DataError + + return DataError({"user_facing_error": {"message": message}}) + + +@pytest.mark.asyncio +async def test_update_spend_logs_isolates_poison_row_and_persists_good_rows( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """One row Postgres rejects (22P05) must not drop the whole batch. + + The good rows still persist and only the offending row is dropped, with no + exception bubbling up. On the unfixed single-shot ``create_many`` the first + write raises and the entire batch is lost. + """ + poison_id = "r1" + written: List[str] = [] + + async def _create_many(*, data: Any, skip_duplicates: bool) -> None: + ids = [row["request_id"] for row in data] + if poison_id in ids: + raise _data_error( + "Inconsistent column data: 22P05 invalid byte sequence for encoding UTF8: 0x00" + ) + written.extend(ids) + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + logs = [make_spend_log_row(request_id=f"r{i}") for i in range(4)] + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=logs, + ) + assert sorted(written) == ["r0", "r2", "r3"] + + +@pytest.mark.asyncio +async def test_update_spend_logs_reraises_connection_masquerade_dataerror( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A P1001 "can't reach database server" outage that prisma mislabels as a + ``DataError`` is transient, not a poison row: it must propagate so the batch + is surfaced/retried rather than bisected into silent per-row drops. + """ + err = _data_error("Can't reach database server at db-host:5432") + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(type(err)): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + 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"), + ], + ) + + +@pytest.mark.asyncio +async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A flood of poisoned rows must not amplify one failed bulk insert into + unbounded failed inserts. The per-batch attempt budget hard-caps the number + of ``create_many`` calls regardless of how many rows are poisoned, so the DB + work stays bounded and well below the input row count, and the helper still + completes without raising. + """ + import litellm.proxy.utils as utils_mod + + attempt_cap = utils_mod.MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH + # single create_many batch (< BATCH_SIZE) whose row count exceeds the attempt + # cap, so the bound bites and attempts stay below the input row count + n_rows = attempt_cap * 3 + + async def _always_poison(*, data: Any, skip_duplicates: bool) -> None: + raise _data_error("invalid byte sequence for encoding UTF8: 0x00") + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_always_poison) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + logs = [make_spend_log_row(request_id=f"r{i}") for i in range(n_rows)] + + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=logs, + ) + + attempts = mock_prisma_client.db.litellm_spendlogs.create_many.await_count + assert attempts <= attempt_cap + assert attempts < n_rows + + def test_disable_spend_updates_reflects_general_settings( monkeypatch: pytest.MonkeyPatch, ) -> None: