fix(proxy): isolate poison spend-log rows so one bad record can't drop the whole batch (#31705)

update_spend_logs flushes the queue with a single create_many per batch, so one
row carrying bytes Postgres refuses (a residual NUL byte is the canonical case)
fails the entire insert and drops every good spend log alongside it. PR #29515
strips NUL bytes from the JSON columns, but the scalar string columns (end_user,
model, session_id, ...) still flow through unsanitized, so a poisoned row can
still reach the write and take a batch of up to 1000 good rows down with it.

On a genuine data-layer rejection the batch is now bisected so the good rows
still persist and only the offending row is dropped and logged with its
request_id. The classification lives in PrismaDBExceptionHandler.is_prisma_data_error
(matched by exact type so systemic subclasses like a missing table are not
mistaken for a single poison row), which keeps prisma an in-function import and
litellm.proxy.utils importable without the proxy extra. Transport failures,
including the "can't reach database server" outage that prisma mislabels as a
DataError, are re-raised unchanged so the existing connection-retry path still
runs and a transient outage never turns into silent per-row data loss.

The bisection carries a per-batch isolation budget so an authenticated caller
flooding poisoned rows cannot amplify one failed bulk insert into ~2N failed
inserts and N log lines; once the budget is spent the still-failing remainder
is dropped wholesale under a single log line.

Resolves LIT-4103
This commit is contained in:
Yassin Kortam 2026-06-30 22:21:14 +03:00 • committed by GitHub
parent 6d828e5759
commit 52dc15adfe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 225 additions and 2 deletions

View file

@ -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:
"""

View file

@ -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

View file

@ -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

View file

@ -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: