mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(db): report an unreachable native pool at boot
The native pool is lazy, so a pool that could not connect stayed silent until the first ported query failed on a request path. Database.open now pings it once in the background and logs a warning on failure, without failing or delaying boot
This commit is contained in:
parent
90ed6a9bee
commit
fc7a3029e2
3 changed files with 44 additions and 3 deletions
|
|
@ -56,6 +56,7 @@ Python and SQLx run side by side until every repository operation is native, so
|
|||
| Error classification, `lock_timeout` on rollups | `exception_handler.py`, `rollup_lock_timeout.py` | Planned |
|
||||
| `ServiceTypes.DB` events per round trip | `db_span.py`, `log_db_metrics.py`, `prometheus_services.py`, `otel/model/spans.py` | Python host only. Wrap each native operation in `observe` with the call type and table its Prisma version uses, and the bridge forwards events to `ServiceLoggingEvents`. Nesting under an outer Python `db_span` is not detected. A Rust host has no sink yet and must emit the names below |
|
||||
| Auth lookup gate, fail open on DB outage, reader choice per query | `db_lookup_gate.py`, `allow_requests_on_db_unavailable` | With each ported auth read |
|
||||
| Connect at boot, so a database that cannot be reached is reported before traffic | `PrismaClient.connect` | The pool stays lazy and never fails boot. `Database.open` pings it once in the background and logs a warning when it cannot connect, so a mistranslated endpoint is not first found by a ported query |
|
||||
| Health watchdog, engine reaper, drain tracker, engine recreation | `utils.py`, `query_engine_reaper.py`, `prisma_client.py` | Not needed: SQLx pools reconnect per connection and `close` waits for checked-out connections |
|
||||
| Columns Prisma fills itself: `@updatedAt`, `uuid()` / `cuid()` ids, `now()` timestamps | `schema.prisma` | Every native write sets them in its SQL. Postgres has no default for `@updatedAt` or the ids, and its `now()` default follows the session time zone while Prisma and the Python raw inserts write UTC, so timestamps are `NOW() AT TIME ZONE 'UTC'` |
|
||||
| Migrations, `check_migration`, views | `prisma_client.py`, `create_views.py` | Prisma keeps them until the SQLx cutover |
|
||||
|
|
|
|||
|
|
@ -41,11 +41,23 @@ def _log_ignored_options(env_var: str, settings: NativeEndpointSettings) -> None
|
|||
)
|
||||
|
||||
|
||||
async def _warn_if_unreachable(native: NativeDatabase) -> None:
|
||||
try:
|
||||
await native.health_check()
|
||||
except ValueError as error:
|
||||
verbose_proxy_logger.warning(
|
||||
"Native database pool cannot reach the database, so operations ported to it will fail until it"
|
||||
" connects. Prisma is checked separately. Cause: %s",
|
||||
error,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Database:
|
||||
prisma: "PrismaClient"
|
||||
native: NativeDatabase
|
||||
tasks: "tuple[asyncio.Task[None], ...]"
|
||||
boot_probe: "asyncio.Task[None]"
|
||||
|
||||
@classmethod
|
||||
async def open(
|
||||
|
|
@ -77,12 +89,18 @@ class Database:
|
|||
if reader_url is None
|
||||
else _token_refresh(native.set_reader_password, reader_token_rotation(reader_url), reader_url)
|
||||
)
|
||||
return cls(prisma=prisma, native=native, tasks=(*writer_refresh, *reader_refresh))
|
||||
return cls(
|
||||
prisma=prisma,
|
||||
native=native,
|
||||
tasks=(*writer_refresh, *reader_refresh),
|
||||
boot_probe=asyncio.create_task(_warn_if_unreachable(native)),
|
||||
)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
for task in self.tasks:
|
||||
background: Final = (*self.tasks, self.boot_probe)
|
||||
for task in background:
|
||||
task.cancel()
|
||||
await asyncio.gather(*self.tasks, return_exceptions=True)
|
||||
await asyncio.gather(*background, return_exceptions=True)
|
||||
try:
|
||||
await self.prisma.aclose()
|
||||
finally:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
import asyncio
|
||||
import logging
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.db.connection_budget import RustConnectionLimits
|
||||
from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR
|
||||
from litellm.repositories.database import Database
|
||||
|
|
@ -80,3 +83,22 @@ async def test_rejected_reader_settings_still_close_prisma() -> None:
|
|||
await Database.open(_prisma(aclose), WRITER_URL, "postgresql://replica.invalid/proxy", RustConnectionLimits())
|
||||
|
||||
aclose.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("static_password")
|
||||
async def test_an_unreachable_native_pool_is_reported_without_failing_boot(caplog: pytest.LogCaptureFixture) -> None:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name):
|
||||
database: Final = await Database.open(
|
||||
_prisma(), "postgresql://litellm:secret@127.0.0.1:1/proxy?pool_timeout=1", None, RustConnectionLimits()
|
||||
)
|
||||
await asyncio.wait_for(database.boot_probe, timeout=10)
|
||||
await database.aclose()
|
||||
|
||||
assert [
|
||||
record.getMessage() for record in caplog.records if "Native database pool cannot reach" in record.getMessage()
|
||||
] == [
|
||||
"Native database pool cannot reach the database, so operations ported to it will fail until it connects."
|
||||
" Prisma is checked separately. Cause: database query failed: pool timed out while waiting for an open"
|
||||
" connection"
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue