mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(proxy): self-heal Prisma read paths + harden reconnect state machine
Two related fixes layered on top of the existing reconnect plumbing: 1. Restore reconnect-and-retry on `PrismaClient.get_generic_data` (issue #25143). 1.83.x lost the transport-reconnect-and-retry-once branch that 1.82.6 had on this method, so transient `httpx.ReadError` flaps now surface immediately as `db_exceptions` alerts. `_update_config_from_db` fans out four concurrent `get_generic_data` reads, so a single transport blip used to mark four alerts and a stale config window. Adds `call_with_db_reconnect_retry` to `litellm/proxy/db/exception_handler.py` — a single canonical "try DB read, on transport error reconnect once and retry once" wrapper. Mirrors the inline pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we have one implementation rather than three drifting copies, and gives future read paths a clean opt-in. 2. Fix the `_engine_confirmed_dead` flag-reset bug in `_run_reconnect_cycle`. The flag was cleared before `_do_heavy_reconnect()` ran, so any failure inside the heavy reconnect (timeout, missing DATABASE_URL, recreate failure) left the flag False — and the next attempt could silently demote to the lightweight path even though the engine was genuinely dead. Move the reset into the success branch so the flag stays True across heavy-reconnect failures and the next attempt re-enters the heavy branch. Tests: - `tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py` (new) — 9 tests covering the helper's contract: happy path, retry on transport error, no retry on data-layer errors, propagation when reconnect fails, propagation after second transport error, `hasattr` guard for partial mocks, fresh-coroutine-per-call invariant, explicit timeout override, default timeouts read off the prisma_client. - `tests/test_litellm/proxy/db/test_prisma_self_heal.py` — adds: - `test_get_generic_data_retries_on_transport_error_for_config_table` - `test_get_generic_data_propagates_when_reconnect_fails` - `test_engine_confirmed_dead_persists_across_failed_heavy_reconnect` (regression test for the flag-reset bug). All 16 self-heal tests + 9 helper tests + 535 auth/exception-handler tests pass locally.
This commit is contained in:
Yuneng Jiang
parent
6e6b2ca2d8
commit
1c9c219a74
4 changed files with 501 additions and 12 deletions
|
|
@ -1,5 +1,6 @@
|
|||
from typing import Union
|
||||
from typing import Any, Awaitable, Callable, Optional, Union
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
DB_CONNECTION_ERROR_TYPES,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -123,3 +124,123 @@ class PrismaDBExceptionHandler:
|
|||
):
|
||||
return None
|
||||
raise e
|
||||
|
||||
|
||||
# Default fallback timeouts when neither the caller nor the prisma_client
|
||||
# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`.
|
||||
# Match the auth path's existing defaults so behavior is uniform across read paths.
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1
|
||||
|
||||
|
||||
def _coerce_timeout(value: Any, fallback: float) -> float:
|
||||
"""Return `value` if it is a real int/float, else `fallback`. Guards
|
||||
against tests that mock `prisma_client` and leave the timeout slots as
|
||||
MagicMock instances."""
|
||||
if isinstance(value, (int, float)) and not isinstance(value, bool):
|
||||
return float(value)
|
||||
return fallback
|
||||
|
||||
|
||||
async def call_with_db_reconnect_retry(
|
||||
prisma_client: Any,
|
||||
coro_factory: Callable[[], Awaitable[Any]],
|
||||
*,
|
||||
reason: str,
|
||||
timeout_seconds: Optional[float] = None,
|
||||
lock_timeout_seconds: Optional[float] = None,
|
||||
) -> Any:
|
||||
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
|
||||
|
||||
The canonical "self-heal a transient DB transport blip" wrapper used by
|
||||
`PrismaClient.get_generic_data` and other read paths. Mirrors the inline
|
||||
pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we
|
||||
have a single implementation rather than three drifting copies.
|
||||
|
||||
Behavior:
|
||||
1. Await `coro_factory()`. On success, return its value.
|
||||
2. On exception, if it is NOT a transport error (per
|
||||
`is_database_transport_error`), re-raise — data-layer errors like
|
||||
`UniqueViolationError` mean the DB is reachable, reconnect would be
|
||||
pointless.
|
||||
3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise.
|
||||
This guards against partial stand-ins / older clients in tests.
|
||||
4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns
|
||||
False (cooldown / lock contention / reconnect failure), re-raise.
|
||||
5. Otherwise await `coro_factory()` a second time and return / propagate
|
||||
its result. At-most-one retry by construction — no infinite loop.
|
||||
|
||||
`coro_factory` MUST be a zero-arg callable that returns a fresh awaitable
|
||||
on each call. Passing an already-awaited coroutine would fail on retry
|
||||
with `RuntimeError: cannot reuse already awaited coroutine`.
|
||||
|
||||
`reason` should follow `<subsystem>_<operation>_<table>_failure` so
|
||||
telemetry distinguishes between fan-out callers (e.g.
|
||||
`_update_config_from_db` issues four concurrent reads).
|
||||
|
||||
Args:
|
||||
prisma_client: The `PrismaClient` (or stand-in) that owns
|
||||
`attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults.
|
||||
coro_factory: Zero-arg callable returning the read awaitable.
|
||||
reason: Telemetry tag forwarded to `attempt_db_reconnect`.
|
||||
timeout_seconds: Optional override for the reconnect cycle timeout.
|
||||
Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`,
|
||||
then to 2.0s.
|
||||
lock_timeout_seconds: Optional override for how long the helper will
|
||||
wait to acquire the reconnect lock. Defaults to
|
||||
`prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to
|
||||
0.1s.
|
||||
|
||||
Returns:
|
||||
Whatever `coro_factory()` returns (on first or second attempt).
|
||||
|
||||
Raises:
|
||||
Whatever `coro_factory()` raises if the failure is not a transport
|
||||
error, or if the reconnect attempt does not succeed, or if the retry
|
||||
also fails.
|
||||
"""
|
||||
try:
|
||||
return await coro_factory()
|
||||
except Exception as first_exc:
|
||||
if not PrismaDBExceptionHandler.is_database_transport_error(first_exc):
|
||||
raise
|
||||
if not hasattr(prisma_client, "attempt_db_reconnect"):
|
||||
raise
|
||||
|
||||
resolved_timeout = _coerce_timeout(
|
||||
(
|
||||
timeout_seconds
|
||||
if timeout_seconds is not None
|
||||
else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None)
|
||||
),
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
resolved_lock_timeout = _coerce_timeout(
|
||||
(
|
||||
lock_timeout_seconds
|
||||
if lock_timeout_seconds is not None
|
||||
else getattr(
|
||||
prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None
|
||||
)
|
||||
),
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s",
|
||||
reason,
|
||||
first_exc,
|
||||
)
|
||||
|
||||
did_reconnect = await prisma_client.attempt_db_reconnect(
|
||||
reason=reason,
|
||||
timeout_seconds=resolved_timeout,
|
||||
lock_timeout_seconds=resolved_lock_timeout,
|
||||
)
|
||||
if not did_reconnect:
|
||||
raise
|
||||
|
||||
# At most one retry. If the retry also raises a transport error, we
|
||||
# propagate — repeated reconnect-loops are the watchdog's job, not
|
||||
# this helper's.
|
||||
return await coro_factory()
|
||||
|
|
|
|||
|
|
@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import (
|
|||
should_create_missing_views,
|
||||
)
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
)
|
||||
from litellm.proxy.db.log_db_metrics import log_db_metrics
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
|
|
@ -2779,30 +2782,42 @@ class PrismaClient:
|
|||
table_name: Literal["users", "keys", "config", "spend"],
|
||||
):
|
||||
"""
|
||||
Generic implementation of get data
|
||||
Generic implementation of get data.
|
||||
|
||||
Self-heals across a single transient transport blip via
|
||||
`call_with_db_reconnect_retry`: on `httpx.ReadError` /
|
||||
`ClientNotConnectedError` / similar, attempt one DB reconnect and
|
||||
retry once before surfacing the failure. Restores the 1.82.6 behavior
|
||||
that was lost in 1.83.x — see issue #25143.
|
||||
"""
|
||||
start_time = time.time()
|
||||
try:
|
||||
|
||||
async def _do_query():
|
||||
if table_name == "users":
|
||||
response = await self.db.litellm_usertable.find_first(
|
||||
return await self.db.litellm_usertable.find_first(
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "keys":
|
||||
response = await self.db.litellm_verificationtoken.find_first( # type: ignore
|
||||
return await self.db.litellm_verificationtoken.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "config":
|
||||
response = await self.db.litellm_config.find_first( # type: ignore
|
||||
return await self.db.litellm_config.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "spend":
|
||||
response = await self.db.l.find_first( # type: ignore
|
||||
return await self.db.l.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
import traceback
|
||||
return None
|
||||
|
||||
try:
|
||||
return await call_with_db_reconnect_retry(
|
||||
self,
|
||||
_do_query,
|
||||
reason=f"prisma_get_generic_data_{table_name}_lookup_failure",
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}"
|
||||
verbose_proxy_logger.error(error_msg)
|
||||
error_msg = error_msg + "\nException Type: {}".format(type(e))
|
||||
|
|
@ -4204,7 +4219,6 @@ class PrismaClient:
|
|||
)
|
||||
self._reap_all_zombies()
|
||||
self._cleanup_engine_watcher()
|
||||
self._engine_confirmed_dead = False
|
||||
|
||||
async def _do_heavy_reconnect() -> None:
|
||||
db_url = os.getenv("DATABASE_URL", "")
|
||||
|
|
@ -4217,6 +4231,12 @@ class PrismaClient:
|
|||
await self._start_engine_watcher()
|
||||
|
||||
await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout)
|
||||
# Only clear the "dead engine" flag after the heavy reconnect
|
||||
# actually completed. If `_do_heavy_reconnect()` raises (timeout,
|
||||
# missing DATABASE_URL, recreate failure), the flag stays True so
|
||||
# the next attempt re-enters the heavy branch instead of silently
|
||||
# demoting to the lightweight path.
|
||||
self._engine_confirmed_dead = False
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Performing Prisma DB reconnect (engine alive or unknown)."
|
||||
|
|
|
|||
|
|
@ -0,0 +1,225 @@
|
|||
"""
|
||||
Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read,
|
||||
on transport error reconnect once and retry once" helper.
|
||||
|
||||
Covers the regression in issue #25143 where read paths (e.g.
|
||||
`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in
|
||||
LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient
|
||||
`httpx.ReadError` flaps that used to self-heal in 1.82.6.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
|
||||
|
||||
def _make_client(
|
||||
*,
|
||||
attempt_db_reconnect_return: bool = True,
|
||||
has_attempt_db_reconnect: bool = True,
|
||||
):
|
||||
"""Build a minimal stand-in for PrismaClient that exposes only the surface
|
||||
`call_with_db_reconnect_retry` actually pokes at."""
|
||||
client = MagicMock()
|
||||
if has_attempt_db_reconnect:
|
||||
client.attempt_db_reconnect = AsyncMock(
|
||||
return_value=attempt_db_reconnect_return
|
||||
)
|
||||
else:
|
||||
# `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock
|
||||
# auto-creates attributes, so we wipe it out via `spec`.
|
||||
client = MagicMock(spec=[])
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_returns_value_on_first_success():
|
||||
"""Happy path: factory succeeds first call, no reconnect attempted."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path")
|
||||
|
||||
assert result == {"id": 1}
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_retries_after_transport_error():
|
||||
"""Transport error on first call → reconnect → second call succeeds."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="prisma_get_generic_data_config_lookup_failure"
|
||||
)
|
||||
|
||||
assert result == {"id": 1}
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error():
|
||||
"""Data-layer errors (e.g. UniqueViolationError) are NOT transport errors —
|
||||
propagate immediately, do not reconnect."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
raise UniqueViolationError(
|
||||
data={"user_facing_error": {"meta": {}}},
|
||||
message="Unique constraint failed",
|
||||
)
|
||||
|
||||
with pytest.raises(UniqueViolationError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test")
|
||||
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails():
|
||||
"""Transport error, but reconnect returns False → propagate the original
|
||||
exception. Do not call factory a second time."""
|
||||
client = _make_client(attempt_db_reconnect_return=False)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails")
|
||||
|
||||
assert len(invocations) == 1
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error():
|
||||
"""Transport error, reconnect succeeds, retry also raises transport error →
|
||||
propagate. At most one retry by construction (no infinite loop)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("still failing")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="second_transport_error"
|
||||
)
|
||||
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr():
|
||||
"""Older PrismaClient stand-ins / partial mocks may not expose
|
||||
`attempt_db_reconnect`. The helper must not crash — just propagate the
|
||||
original exception. Mirrors the `hasattr` guard from
|
||||
`auth_checks._fetch_key_object_from_db_with_reconnect`."""
|
||||
client = _make_client(has_attempt_db_reconnect=False)
|
||||
|
||||
async def _factory():
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro():
|
||||
"""Guard against the obvious bug of awaiting the same coroutine twice
|
||||
(`RuntimeError: cannot reuse already awaited coroutine`). The helper must
|
||||
call the factory a fresh time on retry, not cache an awaitable."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
factory_call_count = 0
|
||||
|
||||
async def _factory():
|
||||
nonlocal factory_call_count
|
||||
factory_call_count += 1
|
||||
if factory_call_count == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="fresh_coro_on_retry"
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
assert factory_call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_passes_explicit_timeouts():
|
||||
"""Explicit timeout_seconds / lock_timeout_seconds override the auth
|
||||
defaults read off the prisma_client object."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client,
|
||||
_factory,
|
||||
reason="explicit_timeouts",
|
||||
timeout_seconds=5.5,
|
||||
lock_timeout_seconds=0.25,
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 5.5
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.25
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset():
|
||||
"""When timeouts are not provided, helper reads
|
||||
`_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`
|
||||
off the prisma_client (matching the auth path's existing convention)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
client._db_auth_reconnect_timeout_seconds = 3.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.5
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="defaults")
|
||||
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 3.0
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.5
|
||||
|
|
@ -5,6 +5,7 @@ import sys
|
|||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
|
|
@ -358,3 +359,125 @@ async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
|
|||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
mock_kill.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_generic_data: transport-reconnect-and-retry coverage (issue #25143)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_retries_on_transport_error_for_config_table(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""`get_generic_data(table_name="config")` self-heals on a transient
|
||||
`httpx.ReadError`: reconnect once, retry once, return the row.
|
||||
|
||||
Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry
|
||||
branch that 1.82.6 had on this method. `_update_config_from_db` fans out
|
||||
four concurrent `get_generic_data` calls, so a single transport flap used
|
||||
to surface as four `db_exceptions` alerts and a stale config window.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}}
|
||||
invocations: list[None] = []
|
||||
|
||||
async def _flaky_find_first(**kwargs):
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("simulated transport blip")
|
||||
return expected_row
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
result = await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
assert result == expected_row
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
# The failure_handler telemetry side-effect must NOT fire on the first
|
||||
# transport blip — only if the post-retry call also fails. Drain the
|
||||
# event loop so any spuriously-spawned task would have run by now.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging):
|
||||
"""If reconnect itself does not succeed, propagate the original transport
|
||||
error and let the existing failure_handler / db_exceptions telemetry fire."""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(
|
||||
side_effect=httpx.ReadError("simulated transport blip")
|
||||
)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
# Failure telemetry IS expected here — the read genuinely failed.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_called_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _engine_confirmed_dead flag-reset bug (B2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Regression test for the flag-reset bug.
|
||||
|
||||
Before the fix, `_run_reconnect_cycle` cleared
|
||||
`self._engine_confirmed_dead = False` *before* awaiting
|
||||
`_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout,
|
||||
missing DATABASE_URL, recreate failure), the flag was left cleared and the
|
||||
next attempt could demote to the lightweight path even though the engine
|
||||
was genuinely dead.
|
||||
|
||||
The fix moves the reset into the success branch — the flag must stay True
|
||||
when heavy reconnect raises.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client._engine_confirmed_dead = True
|
||||
client._engine_pid = 0 # so `_is_engine_alive` is not consulted
|
||||
|
||||
# Make the heavy reconnect path raise.
|
||||
client.db.recreate_prisma_client = AsyncMock(
|
||||
side_effect=RuntimeError("simulated heavy reconnect failure")
|
||||
)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
client._cleanup_engine_watcher = MagicMock()
|
||||
client._reap_all_zombies = MagicMock()
|
||||
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
with pytest.raises(Exception):
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
# The flag must STILL be True so the next attempt re-enters the heavy
|
||||
# branch instead of silently demoting to the lightweight path.
|
||||
assert client._engine_confirmed_dead is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue