mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
264305de23
A stale RoutingPrismaWrapper.writer_unavailable flag made WriterPinnedClient hand back the routed wrapper, where query_raw is classified as a read, so the access-group UPDATE statements behind /key/regenerate, /key/generate with access_group_ids and model rename/delete went to the read replica and failed with SQLSTATE 25006. Route those raw statements through the underlying writer regardless of the flag; a raw SQL write has no replica fallback. WriterPinnedClient keeps yielding to the replica for degraded reads. The model sync's backing-row count stays on the writer too: it runs right after the row delete/update on the writer and a lagging replica could still report the removed row, which would leave the group naming a model nobody serves. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
308 lines
13 KiB
Python
308 lines
13 KiB
Python
"""
|
|
RoutingPrismaWrapper: routes Prisma reads to a read-replica client and writes
|
|
to a writer client. Used when DATABASE_URL_READ_REPLICA is configured;
|
|
otherwise PrismaClient uses the writer-only PrismaWrapper directly.
|
|
"""
|
|
|
|
import os
|
|
from collections.abc import Callable
|
|
from datetime import timedelta
|
|
from typing import TYPE_CHECKING, Any, Final
|
|
|
|
from litellm._logging import verbose_proxy_logger
|
|
from litellm.proxy.db.prisma_client import PrismaWrapper
|
|
|
|
if TYPE_CHECKING:
|
|
from prisma.types import HttpConfig
|
|
|
|
# Per-model action methods that read from the database. These are routed to
|
|
# the read replica when one is configured.
|
|
_MODEL_READ_METHODS: Final = frozenset(
|
|
{
|
|
"find_first",
|
|
"find_first_or_raise",
|
|
"find_many",
|
|
"find_unique",
|
|
"find_unique_or_raise",
|
|
"count",
|
|
"group_by",
|
|
"query_first",
|
|
"query_raw",
|
|
}
|
|
)
|
|
|
|
# Top-level Prisma client methods that read from the database.
|
|
_TOP_LEVEL_READ_METHODS: Final = frozenset({"query_first", "query_raw"})
|
|
|
|
|
|
class _RoutedActions:
|
|
"""Per-model accessor that sends reads to the reader and writes to the writer.
|
|
|
|
`should_use_reader` is consulted on every read dispatch so a mid-call flip
|
|
of the routing wrapper's reader-availability flag (e.g. after the reader
|
|
fails a recreate) is observed without re-fetching the actions accessor.
|
|
"""
|
|
|
|
__slots__ = ("_reader_actions", "_should_use_reader", "_writer_actions")
|
|
|
|
def __init__(
|
|
self,
|
|
writer_actions: object,
|
|
reader_actions: object,
|
|
should_use_reader: Callable[[], bool],
|
|
):
|
|
self._writer_actions = writer_actions
|
|
self._reader_actions = reader_actions
|
|
self._should_use_reader = should_use_reader
|
|
|
|
def __getattr__(self, name: str) -> object:
|
|
if name in _MODEL_READ_METHODS and self._should_use_reader():
|
|
return getattr(self._reader_actions, name)
|
|
return getattr(self._writer_actions, name)
|
|
|
|
|
|
class WriterPinnedClient:
|
|
"""PrismaClient-shaped view whose `.db` resolves to the writer while it is available.
|
|
|
|
Read-after-write paths (e.g. the model reconcile a /model/new triggers to
|
|
verify its own just-committed row) must not read through a lagging read
|
|
replica: the row is not replayed there yet, so the reconcile concludes the
|
|
write is missing and fails the request even though it is durable (#38556).
|
|
|
|
While the writer is degraded (`writer_unavailable`), the pin yields to the
|
|
routed wrapper so reconcile reads keep working from the replica: a proxy
|
|
that starts during a primary outage must still load DB-backed models, and
|
|
no read-after-write hazard exists then because writes are failing anyway.
|
|
"""
|
|
|
|
__slots__ = ("db",)
|
|
|
|
def __init__(self, db: "PrismaWrapper | RoutingPrismaWrapper") -> None:
|
|
self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db
|
|
|
|
|
|
def writer_wrapper(db: "PrismaWrapper | RoutingPrismaWrapper") -> PrismaWrapper:
|
|
"""Unlike `WriterPinnedClient`, ignores `writer_unavailable`: a raw SQL write has no replica fallback."""
|
|
return db.writer if isinstance(db, RoutingPrismaWrapper) else db
|
|
|
|
|
|
class RoutingPrismaWrapper:
|
|
"""
|
|
Routes Prisma operations between a writer and a reader Prisma client.
|
|
|
|
Reads (find_*, count, group_by, query_raw, query_first) go to the reader;
|
|
everything else (writes, transactions, raw execute) goes to the writer.
|
|
Lifecycle methods (connect, disconnect, IAM token refresh) act on both
|
|
clients so callers do not need to know about the split. When
|
|
IAM_TOKEN_DB_AUTH is enabled, both writer and reader refresh their tokens
|
|
independently on their own ~12-minute cadence.
|
|
|
|
Reader degradation: a reader-side failure (failed connect, failed
|
|
recreate) is non-fatal — the wrapper sets `_reader_unavailable=True`, logs
|
|
a warning, and routes subsequent reads to the writer. The next successful
|
|
`connect()` or `recreate_prisma_client()` clears the flag. This keeps the
|
|
proxy serving traffic during transient reader outages instead of failing
|
|
startup or returning errors for read-heavy endpoints.
|
|
|
|
Writer degradation: a writer-side `connect()` failure while the reader
|
|
connects is likewise non-fatal — the wrapper sets
|
|
`_writer_unavailable=True`, logs a warning, and keeps serving reads from
|
|
the reader (key lookups, DB-stored model loads) so a proxy that starts
|
|
during a primary outage still serves inference from the replica. Writes
|
|
fail at call time until the writer recovers; the PrismaClient DB health
|
|
watchdog polls `writer_unavailable` and drives the writer reconnect,
|
|
which clears the flag via `recreate_prisma_client`. Only when BOTH sides
|
|
fail to connect does `connect()` raise (full DB outage — the existing
|
|
`allow_requests_on_db_unavailable` startup handling applies).
|
|
"""
|
|
|
|
def __init__(self, writer: PrismaWrapper, reader: PrismaWrapper):
|
|
self._writer = writer
|
|
self._reader = reader
|
|
# When True, reads fall back to the writer. Flipped on by reader
|
|
# connect/recreate failures and flipped off on the next reader recovery.
|
|
self._reader_unavailable: bool = False
|
|
self._writer_unavailable: bool = False
|
|
|
|
@property
|
|
def writer(self) -> PrismaWrapper:
|
|
return self._writer
|
|
|
|
@property
|
|
def reader(self) -> PrismaWrapper:
|
|
return self._reader
|
|
|
|
@property
|
|
def read_target(self) -> PrismaWrapper:
|
|
"""The wrapper `_TOP_LEVEL_READ_METHODS` dispatch to right now.
|
|
|
|
Callers that need to reason about the engine a read actually ran on
|
|
(e.g. recovering from prepared statements that went stale on it) must
|
|
consult this rather than `writer`, and `__getattr__` routes through it
|
|
so the two cannot drift apart.
|
|
"""
|
|
return self._writer if self._reader_unavailable else self._reader
|
|
|
|
@property
|
|
def reader_unavailable(self) -> bool:
|
|
return self._reader_unavailable
|
|
|
|
@property
|
|
def writer_unavailable(self) -> bool:
|
|
return self._writer_unavailable
|
|
|
|
def mark_writer_recovered(self) -> None:
|
|
"""Clear the degraded-writer flag after an external health probe proved
|
|
the writer reachable. Needed when recovery happens without
|
|
`recreate_prisma_client` (e.g. an IAM token refresh already recreated
|
|
the writer engine), which is otherwise the only runtime path that
|
|
clears the flag — without this, the watchdog would keep firing
|
|
reconnect attempts against an already-healthy writer."""
|
|
self._writer_unavailable = False
|
|
|
|
def _should_use_reader(self) -> bool:
|
|
return not self._reader_unavailable
|
|
|
|
@staticmethod
|
|
async def _try_connect(client: PrismaWrapper, timeout: int | timedelta | None = None) -> Exception | None:
|
|
if client.is_connected() is True:
|
|
return None
|
|
try:
|
|
await client.connect(timeout)
|
|
return None
|
|
except Exception as e:
|
|
return e
|
|
|
|
async def connect(self, timeout: int | timedelta | None = None) -> None:
|
|
writer_error: Final = await self._try_connect(self._writer, timeout)
|
|
if writer_error is None:
|
|
self._writer_unavailable = False
|
|
verbose_proxy_logger.info("[writer] DB connected")
|
|
reader_error: Final = await self._try_connect(self._reader, timeout)
|
|
if reader_error is None:
|
|
self._reader_unavailable = False
|
|
verbose_proxy_logger.info("[reader] DB connected")
|
|
if writer_error is None and reader_error is None:
|
|
return
|
|
if writer_error is not None and reader_error is not None:
|
|
raise writer_error
|
|
if reader_error is not None:
|
|
# Degrade gracefully: the proxy keeps serving traffic with reads
|
|
# routed to the writer until the reader endpoint is reachable.
|
|
# Aborting startup here would tie proxy availability to an
|
|
# opt-in, best-effort reader endpoint.
|
|
self._reader_unavailable = True
|
|
verbose_proxy_logger.warning(
|
|
"Failed to connect to read replica DB: %s. "
|
|
"Falling back to the writer for reads until the reader is reachable.",
|
|
reader_error,
|
|
)
|
|
return
|
|
self._writer_unavailable = True
|
|
verbose_proxy_logger.warning(
|
|
"Failed to connect to primary (writer) DB: %s. "
|
|
"Serving reads from the read replica; writes will fail until the writer recovers.",
|
|
writer_error,
|
|
)
|
|
|
|
async def disconnect(self, timeout: float | timedelta | None = None) -> None:
|
|
first_error: BaseException | None = None
|
|
for client in (self._writer, self._reader):
|
|
try:
|
|
await client.disconnect(timeout)
|
|
except Exception as e:
|
|
if first_error is None:
|
|
first_error = e
|
|
verbose_proxy_logger.warning("Error disconnecting Prisma client: %s", e)
|
|
if first_error is not None:
|
|
raise first_error
|
|
|
|
def is_connected(self) -> bool:
|
|
# Reflects writer health only. The reader is best-effort; its
|
|
# availability is tracked via `_reader_unavailable` and a degraded
|
|
# reader must NOT cause a writer reconnect (would loop indefinitely
|
|
# since recreate_prisma_client only fixes writer-side problems).
|
|
return bool(self._writer.is_connected())
|
|
|
|
async def start_token_refresh_task(self) -> None:
|
|
await self._writer.start_token_refresh_task()
|
|
await self._reader.start_token_refresh_task()
|
|
|
|
async def stop_token_refresh_task(self) -> None:
|
|
await self._writer.stop_token_refresh_task()
|
|
await self._reader.stop_token_refresh_task()
|
|
|
|
async def recreate_prisma_client(
|
|
self,
|
|
new_db_url: str,
|
|
http_client: "HttpConfig | None" = None,
|
|
*,
|
|
expected_generation: int | None = None,
|
|
) -> bool:
|
|
"""Recreate both writer and reader Prisma clients.
|
|
|
|
The writer reconnect path in PrismaClient calls
|
|
`self.db.recreate_prisma_client(...)`. Without this method, a DB-wide
|
|
connectivity event would only re-create the writer; the reader engine
|
|
would stay broken and every routed read would fail. We always recreate
|
|
the writer first (its URL is the one passed in), then best-effort
|
|
recreate the reader. A reader failure flips `_reader_unavailable=True`
|
|
so reads transparently fall through to the writer.
|
|
|
|
`expected_generation` is forwarded to the writer's optimistic-lock
|
|
guard. If the writer recreate is skipped (another path already replaced
|
|
the engine — issue #29176), we skip the reader too rather than churning
|
|
it needlessly, and return ``False``.
|
|
"""
|
|
writer_recreated: Final = await self._writer.recreate_prisma_client(
|
|
new_db_url,
|
|
http_client=http_client,
|
|
expected_generation=expected_generation,
|
|
)
|
|
if not writer_recreated:
|
|
return False
|
|
self._writer_unavailable = False
|
|
try:
|
|
await self._recreate_reader(http_client=http_client)
|
|
self._reader_unavailable = False
|
|
except Exception as e:
|
|
self._reader_unavailable = True
|
|
verbose_proxy_logger.warning(
|
|
"Failed to recreate reader Prisma client: %s. "
|
|
"Reads will fall back to the writer until the reader recovers.",
|
|
e,
|
|
)
|
|
return True
|
|
|
|
async def _recreate_reader(self, http_client: "HttpConfig | None" = None) -> None:
|
|
"""Resolve the reader URL and recreate its Prisma client.
|
|
|
|
Token-authenticated readers regenerate their token (host/port/user came
|
|
from the parsed reader URL at construction time). Password-authenticated
|
|
readers reuse the URL stored in `DATABASE_URL_READ_REPLICA`.
|
|
"""
|
|
if self._reader.iam_token_db_auth:
|
|
new_reader_url: Final = self._reader.get_rds_iam_token()
|
|
if not new_reader_url:
|
|
raise RuntimeError(f"Failed to generate fresh {self._reader.token_label} for read replica")
|
|
await self._reader.recreate_prisma_client(new_reader_url, http_client=http_client)
|
|
return
|
|
reader_url: Final = os.getenv("DATABASE_URL_READ_REPLICA", "")
|
|
if not reader_url:
|
|
raise RuntimeError("DATABASE_URL_READ_REPLICA not set; cannot recreate read replica client")
|
|
await self._reader.recreate_prisma_client(reader_url, http_client=http_client)
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
if name in _TOP_LEVEL_READ_METHODS:
|
|
return getattr(self.read_target, name)
|
|
writer_attr: Final[object] = getattr(self._writer, name)
|
|
# Per-model action accessors are non-callable instances that expose
|
|
# both `find_many` and `create`. Methods like execute_raw / batch_ /
|
|
# tx are callables and stay on the writer untouched.
|
|
if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"):
|
|
try:
|
|
reader_attr: Final[object] = getattr(self._reader, name)
|
|
except AttributeError:
|
|
return writer_attr
|
|
return _RoutedActions(writer_attr, reader_attr, self._should_use_reader)
|
|
return writer_attr
|