diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f469587ab8e..2cee5128c66 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00ccad33b6d..13cba9a25ce 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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( diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index c685d778c0e..3f009137a1c 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -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, ) diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 10a48941693..e6914213939 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -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 diff --git a/tests/test_litellm/proxy/test_prisma_engine_watchdog.py b/tests/test_litellm/proxy/test_prisma_engine_watchdog.py index d73f74c5cd2..ae750b0770a 100644 --- a/tests/test_litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/test_litellm/proxy/test_prisma_engine_watchdog.py @@ -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() diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index dc30798df55..117c5aa3081 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -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 diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 41b0eb3cf95..cec6dd99ce8 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -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