scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix(proxy): recreate the Prisma client when the writer session turns read-only (#40610)

Browse source 
The writer health probe only ran SELECT 1, which a read-only Postgres
session answers fine, so a pooled connection left pointing at a demoted
primary kept failing every write with SQLSTATE 25006 until the pod was
restarted. Probe transaction_read_only instead, treat a 25006 on the
request path as a signal to recreate the client, and back off
exponentially while the database as a whole stays read-only so a replica
or an in-progress failover does not get its engine killed every cycle.

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-10 13:53:42 -07:00 • committed by GitHub
parent 9cd1c4c29c
commit 0e35c8fee9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 295 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/db/exception_handler.py
View file

@ -221,6 +221,18 @@ class PrismaDBExceptionHandler:
or "write conflict or a deadlock" in error_message
)
@staticmethod
def is_read_only_transaction_error(e: Exception) -> bool:
"""True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the
pooled session answers reads but rejects writes, so the connection is
poisoned until the client is recreated."""
import prisma
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
return False
error_message: Final = str(e).lower()
return '"25006"' in error_message or "read-only transaction" in error_message
@staticmethod
def is_prisma_engine_internal_error(e: Exception) -> bool:
"""True iff ``e`` is a non-``PrismaError`` exception raised from inside

86
litellm/proxy/utils.py
View file

@ -68,6 +68,7 @@ except ImportError:
raise ImportError("backoff is not installed. Please install it via 'pip install backoff'")
from fastapi import HTTPException, status
from pydantic import TypeAdapter
import litellm
import litellm.litellm_core_utils
@ -3902,6 +3903,11 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam
)
_WRITER_WRITABILITY_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only"
_WRITER_WRITABILITY_PROBE_ROWS: Final = TypeAdapter(list[dict[str, object]])
_READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS: Final = 600
class _ForcedRecreateDeclined(Exception):
"""A forced recreate was declined by the engine-generation guard.
@ -4079,6 +4085,8 @@ class PrismaClient:
self._db_health_watchdog_task: asyncio.Task | None = None
self._db_last_reconnect_attempt_ts: float = 0.0
self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15")))
self._db_read_only_recreate_ts: float = 0.0
self._db_read_only_recreate_streak: int = 0
self._db_health_watchdog_interval_seconds: int = max(
5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30"))
)
@ -5796,15 +5804,20 @@ class PrismaClient:
writer: Final = self.writer_db
if force_recreate is False:
try:
await writer.query_raw("SELECT 1")
verbose_proxy_logger.info(
"Writer healthy on probe; skipping recreate (engine "
"likely already replaced by a token refresh)."
)
if isinstance(self.db, RoutingPrismaWrapper):
self.db.mark_writer_recovered()
await self._start_engine_watcher()
return
if await self._writer_is_read_only(writer):
verbose_proxy_logger.warning(
"Writer answers the probe but its session is read-only "
"(writes fail with SQLSTATE 25006); recreating Prisma client."
)
else:
verbose_proxy_logger.info(
"Writer healthy on probe; skipping recreate (engine "
"likely already replaced by a token refresh)."
)
if isinstance(self.db, RoutingPrismaWrapper):
self.db.mark_writer_recovered()
await self._start_engine_watcher()
return
except Exception as probe_err:
verbose_proxy_logger.warning(
"Writer probe failed (%s); recreating Prisma client.",
@ -6124,6 +6137,18 @@ class PrismaClient:
reason="db_health_watchdog_writer_unavailable",
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
)
continue
if await asyncio.wait_for(
self._writer_is_read_only(self.writer_db),
timeout=self._db_health_watchdog_probe_timeout_seconds,
):
await self.recreate_read_only_writer(
reason="db_health_watchdog_writer_read_only",
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
)
continue
self._db_read_only_recreate_streak = 0
self._db_read_only_recreate_ts = 0.0
except asyncio.CancelledError:
break
except Exception as e:
@ -6135,6 +6160,39 @@ class PrismaClient:
else:
verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e)
async def recreate_read_only_writer(self, reason: str, timeout_seconds: float | None = None) -> bool:
"""Force-recreate the client behind a writer session that rejects writes
(SQLSTATE 25006). Each recreate doubles the wait before the next one
until the watchdog sees a writable session again, so a database that is
read-only as a whole (replica, failover in progress) does not get its
engine killed on every watchdog cycle or failed write."""
backoff_seconds: Final = min(
self._db_reconnect_cooldown_seconds * 2 ** min(self._db_read_only_recreate_streak, 10),
_READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS,
)
if time.time() - self._db_read_only_recreate_ts < backoff_seconds:
verbose_proxy_logger.debug(
"Writer session still read-only after %s recreate(s); backing off %ss. reason=%s",
self._db_read_only_recreate_streak,
backoff_seconds,
reason,
)
return False
verbose_proxy_logger.warning(
"Writer session is read-only (writes fail with SQLSTATE 25006); recreating Prisma client. reason=%s",
reason,
)
self._db_read_only_recreate_ts = time.time()
self._db_read_only_recreate_streak += 1
return await self.attempt_db_reconnect(reason=reason, timeout_seconds=timeout_seconds, force_recreate=True)
async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool:
"""True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006)."""
rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python(
await writer.query_raw(_WRITER_WRITABILITY_PROBE_SQL)
)
return any(row.get("transaction_read_only") == "on" for row in rows)
def _probe_target_wrapper(self) -> PrismaWrapper:
"""The Prisma wrapper a `SELECT 1` health probe actually reaches.
@ -7547,6 +7605,12 @@ def _get_openapi_url() -> str | None:
return "/openapi.json"
def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | None") -> None:
if prisma_client is None:
return
asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction"))
def handle_exception_on_proxy(e: Exception) -> ProxyException:
"""
Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible
@ -7554,6 +7618,10 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
from fastapi import status
verbose_proxy_logger.exception("Exception: %s", e)
if PrismaDBExceptionHandler.is_read_only_transaction_error(e):
from litellm.proxy.proxy_server import prisma_client
_recreate_writer_on_read_only_transaction(prisma_client)
if isinstance(e, HTTPException):
return ProxyException(

37
tests/test_litellm/proxy/db/test_exception_handler.py
View file

@ -665,10 +665,47 @@ def test_is_deadlock_error_excludes_non_deadlocks(error):
assert PrismaDBExceptionHandler.is_deadlock_error(error) is False
READ_ONLY_CONNECTOR_ERROR: Final = (
"Error occurred during query execution:\nConnectorError(ConnectorError { user_facing_error: None, "
'kind: QueryError(PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction", '
'severity: "ERROR", detail: None, column: None, hint: None }), transient: false })'
)
@pytest.mark.parametrize(
"error",
[
DataError(data={"user_facing_error": {"message": READ_ONLY_CONNECTOR_ERROR}}),
RawQueryError(data={"user_facing_error": {"message": "cannot execute INSERT in a read-only transaction"}}),
PrismaError(
'PostgresError { code: "25006", message: "kann DELETE in einer Read-Only-Transaktion nicht ausführen" }'
),
],
)
def test_is_read_only_transaction_error_matches_sqlstate_25006(error):
assert PrismaDBExceptionHandler.is_read_only_transaction_error(error) is True
@pytest.mark.parametrize(
"error",
[
UniqueViolationError(data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "t"}}}),
PrismaError("can't reach database server"),
RawQueryError(data={"user_facing_error": {"message": "deadlock detected", "meta": {"table": "t"}}}),
httpx.ConnectError("connection refused"),
RuntimeError("cannot execute UPDATE in a read-only transaction"),
ValueError('"25006"'),
],
)
def test_is_read_only_transaction_error_excludes_other_failures(error):
assert PrismaDBExceptionHandler.is_read_only_transaction_error(error) is False
MOCKED_PRISMA_PREDICATES: Final = (
PrismaDBExceptionHandler.is_database_infrastructure_error,
PrismaDBExceptionHandler.is_database_transport_error,
PrismaDBExceptionHandler.is_deadlock_error,
PrismaDBExceptionHandler.is_read_only_transaction_error,
PrismaDBExceptionHandler.is_prisma_engine_internal_error,
PrismaDBExceptionHandler.is_database_service_unavailable_error,
)

4
tests/test_litellm/proxy/db/test_prisma_self_heal.py
View file

@ -626,7 +626,7 @@ async def test_direct_reconnect_probe_success_clears_writer_unavailable(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
writer = MagicMock()
writer.query_raw = AsyncMock(return_value=[{"result": 1}])
writer.query_raw = AsyncMock(return_value=[{"transaction_read_only": "off"}])
reader = MagicMock()
routing = RoutingPrismaWrapper(writer=writer, reader=reader)
routing._writer_unavailable = True
@ -636,5 +636,5 @@ async def test_direct_reconnect_probe_success_clears_writer_unavailable(
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await client._run_reconnect_cycle(timeout_seconds=5.0)
writer.query_raw.assert_awaited_once_with("SELECT 1")
writer.query_raw.assert_awaited_once_with("SELECT current_setting('transaction_read_only') AS transaction_read_only")
assert routing.writer_unavailable is False

9
tests/test_litellm/proxy/test_prisma_engine_watchdog.py
View file

@ -18,12 +18,15 @@ import asyncio
import os
import threading
import time
from typing import Final
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import pytest
from litellm.proxy.utils import PrismaClient, ProxyLogging
WRITER_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only"
@pytest.fixture(autouse=True)
def mock_prisma_binary():
@ -260,7 +263,7 @@ async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
"""Direct reconnect (engine alive) probes the writer first and skips the
recreate when the probe is healthy.
The engine-alive path now runs a SELECT 1 probe before recreating. A
The engine-alive path now runs a writability probe before recreating. A
healthy probe means the connection is fine — e.g. an IAM token refresh
already replaced the engine (issue #29176) — so recreating would kill a
working engine. Recreate happens only when the probe fails (covered in
@ -278,7 +281,7 @@ async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.query_raw.assert_awaited_once_with(WRITER_PROBE_SQL)
engine_client.db.disconnect.assert_not_awaited()
engine_client._start_engine_watcher.assert_awaited_once()
@ -296,7 +299,7 @@ async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.query_raw.assert_awaited_once_with(WRITER_PROBE_SQL)
engine_client.db.disconnect.assert_not_awaited()
engine_client._start_engine_watcher.assert_awaited_once()

45
tests/test_litellm/proxy/utils/helpers/test_error_helpers.py
View file

@ -1,7 +1,10 @@
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
from prisma.errors import DataError
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.utils import get_error_message_str, handle_exception_on_proxy
@ -171,3 +174,45 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500():
"code": "500",
"type": ProxyErrorTypes.internal_server_error.value,
}
@pytest.mark.asyncio
async def test_handle_exception_on_proxy_read_only_transaction_forces_writer_recreate(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prisma_client = MagicMock()
prisma_client.recreate_read_only_writer = AsyncMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
exc = DataError(
data={
"user_facing_error": {
"message": 'PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction" }'
}
}
)
result = handle_exception_on_proxy(exc)
await asyncio.sleep(0)
snapshot = {
"code": result.code,
"recreate_kwargs": prisma_client.recreate_read_only_writer.await_args.kwargs,
}
assert snapshot == {
"code": "500",
"recreate_kwargs": {"reason": "postgres_read_only_transaction"},
}
@pytest.mark.asyncio
async def test_handle_exception_on_proxy_leaves_writer_alone_for_other_db_errors(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prisma_client = MagicMock()
prisma_client.recreate_read_only_writer = AsyncMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
handle_exception_on_proxy(DataError(data={"user_facing_error": {"message": "deadlock detected"}}))
await asyncio.sleep(0)
assert prisma_client.recreate_read_only_writer.await_count == 0

117
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py
View file

@ -29,7 +29,7 @@ from __future__ import annotations
import asyncio
import time
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, call
import pytest
@ -111,6 +111,34 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails(
}
@pytest.mark.asyncio
async def test_run_reconnect_cycle_direct_path_recreates_when_writer_is_read_only(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("DATABASE_URL", "postgres://x:y@h:5432/db")
prisma_client._engine_confirmed_dead = False
prisma_client._engine_pid = 0
prisma_client._start_engine_watcher = AsyncMock()
prisma_client._cleanup_engine_watcher = MagicMock()
writer: Final = MagicMock()
writer.query_raw = AsyncMock(side_effect=[[{"transaction_read_only": "on"}], [{"?column?": 1}]])
writer.recreate_prisma_client = AsyncMock()
prisma_client.db = writer
await prisma_client._run_reconnect_cycle(timeout_seconds=5)
pinned = {
"recreate_called": writer.recreate_prisma_client.await_count,
"start_watcher_called": prisma_client._start_engine_watcher.await_count,
"cleanup_called": prisma_client._cleanup_engine_watcher.call_count,
}
assert pinned == {
"recreate_called": 1,
"start_watcher_called": 1,
"cleanup_called": 1,
}
@pytest.mark.asyncio
async def test_run_reconnect_cycle_force_recreate_skips_probe_and_recreates(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
@ -488,6 +516,93 @@ async def test_db_health_watchdog_loop_swallows_non_db_errors(
assert prisma_client.attempt_db_reconnect.await_count == 0
def _routing_db_with_writer_sessions(*transaction_read_only: str) -> tuple[RoutingPrismaWrapper, MagicMock]:
"""One watchdog cycle per value, then the loop is cancelled."""
writer: Final = MagicMock()
writer.query_raw = AsyncMock(side_effect=[[{"transaction_read_only": value}] for value in transaction_read_only])
reader: Final = MagicMock()
reader.query_raw = AsyncMock(
side_effect=[[{"?column?": 1}] for _ in transaction_read_only] + [asyncio.CancelledError()]
)
return RoutingPrismaWrapper(writer=writer, reader=reader), writer
@pytest.mark.asyncio
async def test_db_health_watchdog_loop_forces_recreate_when_writer_is_read_only(
prisma_client: PrismaClient,
) -> None:
prisma_client._db_health_watchdog_interval_seconds = 0
prisma_client.attempt_db_reconnect = AsyncMock(side_effect=asyncio.CancelledError())
prisma_client.db, _ = _routing_db_with_writer_sessions("on")
await prisma_client._db_health_watchdog_loop()
assert prisma_client.attempt_db_reconnect.await_args is not None
assert prisma_client.attempt_db_reconnect.await_args.kwargs == {
"reason": "db_health_watchdog_writer_read_only",
"timeout_seconds": prisma_client._db_watchdog_reconnect_timeout_seconds,
"force_recreate": True,
}
@pytest.mark.asyncio
async def test_db_health_watchdog_loop_leaves_writable_writer_alone(
prisma_client: PrismaClient,
) -> None:
prisma_client._db_health_watchdog_interval_seconds = 0
prisma_client.attempt_db_reconnect = AsyncMock()
prisma_client.db, writer = _routing_db_with_writer_sessions("off")
await prisma_client._db_health_watchdog_loop()
assert (prisma_client.attempt_db_reconnect.await_count, writer.query_raw.await_count) == (0, 1)
@pytest.mark.asyncio
async def test_db_health_watchdog_loop_backs_off_while_database_stays_read_only(
prisma_client: PrismaClient,
) -> None:
prisma_client._db_health_watchdog_interval_seconds = 0
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
prisma_client.db, writer = _routing_db_with_writer_sessions("on", "on", "on")
await prisma_client._db_health_watchdog_loop()
assert (prisma_client.attempt_db_reconnect.await_count, writer.query_raw.await_count) == (1, 3)
@pytest.mark.asyncio
async def test_db_health_watchdog_loop_recreates_again_once_writer_was_writable_in_between(
prisma_client: PrismaClient,
) -> None:
prisma_client._db_health_watchdog_interval_seconds = 0
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
prisma_client.db, _ = _routing_db_with_writer_sessions("on", "off", "on")
await prisma_client._db_health_watchdog_loop()
assert prisma_client.attempt_db_reconnect.await_count == 2
@pytest.mark.asyncio
async def test_recreate_read_only_writer_retries_after_backoff_elapses(
prisma_client: PrismaClient,
) -> None:
prisma_client._db_reconnect_cooldown_seconds = 15
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
first: Final = await prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")
within_backoff: Final = await prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")
prisma_client._db_read_only_recreate_ts -= 30
after_backoff: Final = await prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")
prisma_client._db_read_only_recreate_ts -= 30
still_within_doubled_backoff: Final = await prisma_client.recreate_read_only_writer(
reason="postgres_read_only_transaction"
)
assert (first, within_backoff, after_backoff, still_within_doubled_backoff) == (True, False, True, False)
assert prisma_client.attempt_db_reconnect.await_args_list == [
call(reason="postgres_read_only_transaction", timeout_seconds=None, force_recreate=True),
call(reason="postgres_read_only_transaction", timeout_seconds=None, force_recreate=True),
]
@pytest.mark.asyncio
async def test_iam_refresh_racing_reconnect_recreates_engine_only_once(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch