feat(db): emit DB service events from native operations
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run

crates/db gains observe(), which records one event per round trip with the
call type, table and operation a Prisma call reports. NativeDatabase forwards
those events to a Python ServiceLoggingEvents that Database.open injects, and
it calls the same ServiceTypes.DB success and failure hooks as db_span on the
proxy's event loop. A native health_check, which pings one pooled connection
without running SQL (call type health_check), is the first observed operation

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

The crate AGENTS.md records the Prometheus and OTel names a Rust host sink must match
This commit is contained in:
Yujong Lee 2026-11-13 11:05:25 +09:00
parent df8a9a09c6
commit 164efe2338
15 changed files with 488 additions and 18 deletions

View file

@ -32,7 +32,7 @@ from litellm.proxy.db.log_db_metrics import ( # noqa: F401 # legacy module exp
_T = TypeVar("_T")
def _service_logging() -> ServiceLogging | None:
def proxy_service_logging() -> ServiceLogging | None:
try:
from litellm.proxy.proxy_server import proxy_logging_obj
except ImportError:
@ -40,7 +40,9 @@ def _service_logging() -> ServiceLogging | None:
return proxy_logging_obj.service_logging_obj
def _event_metadata(table: str | None, operation: str | None) -> dict[str, str]:
def db_event_metadata(
table: str | None, operation: str | None
) -> dict[str, str]: # mutable-ok: ServiceLogging hooks take event_metadata as a dict
pairs: Final = (("table_name", table), ("db_operation", operation))
return {key: value for key, value in pairs if value is not None}
@ -73,9 +75,9 @@ async def db_span(call_type: str, table: str | None, operation: str | None = Non
if db_io_claimed():
yield
return
service_logging: Final = _service_logging()
service_logging: Final = proxy_service_logging()
start_time: Final = datetime.now()
event_metadata: Final = _event_metadata(table, operation)
event_metadata: Final = db_event_metadata(table, operation)
with claim_db_io() as witness:
try:
yield

View file

@ -21,14 +21,16 @@ import struct
import urllib.parse
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final
from typing import Final, Protocol
from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat, pkcs12
from litellm._logging import verbose_proxy_logger
from litellm._service_logger import ServiceTypes
from litellm.proxy.db.db_span import db_event_metadata
from litellm.proxy.db.db_url_settings import NATIVE_APPLICATION_NAME, PRISMA_APPLICATION_NAME
from litellm.proxy.db.pgbouncer import database_url_is_pooled, fold_libpq_option_tokens, libpq_option_setting
from litellm.proxy.db.token_auth import (
@ -287,3 +289,81 @@ async def keep_token_fresh(
while True:
await sleep(delay)
delay = await _refresh(install, rotation, now)
class DatabaseServiceHooks(Protocol):
async def async_service_success_hook(
self,
service: ServiceTypes,
call_type: str,
duration: float,
parent_otel_span: None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
event_metadata: dict[str, str] | None = None, # mutable-ok: matches the ServiceLogging hook signature
) -> None: ...
async def async_service_failure_hook(
self,
service: ServiceTypes,
duration: float,
error: Exception,
call_type: str,
parent_otel_span: None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
event_metadata: dict[str, str] | None = None, # mutable-ok: matches the ServiceLogging hook signature
) -> None: ...
class NativeDatabaseError(Exception):
"""A native (SQLx) query failure, reported to service logging like a Prisma one."""
@dataclass(frozen=True, slots=True)
class ServiceLoggingEvents:
"""Turns native round trips into the ``ServiceTypes.DB`` events ``db_span`` emits.
Rust calls ``record`` from its own threads once per round trip, so the hook is
handed to the proxy's event loop rather than awaited here.
"""
loop: asyncio.AbstractEventLoop
service_logging: Callable[[], DatabaseServiceHooks | None]
def record(
self,
call_type: str,
table: str | None,
operation: str | None,
started_at: float,
duration_seconds: float,
failure: str | None,
) -> None:
service_logging: Final = self.service_logging()
if service_logging is None:
return
start_time: Final = datetime.fromtimestamp(started_at, tz=timezone.utc).astimezone().replace(tzinfo=None)
hook: Final = (
service_logging.async_service_success_hook(
service=ServiceTypes.DB,
call_type=call_type,
parent_otel_span=None,
duration=duration_seconds,
start_time=start_time,
end_time=start_time + timedelta(seconds=duration_seconds),
event_metadata=db_event_metadata(table, operation),
)
if failure is None
else service_logging.async_service_failure_hook(
error=NativeDatabaseError(failure),
service=ServiceTypes.DB,
call_type=call_type,
parent_otel_span=None,
duration=duration_seconds,
start_time=start_time,
end_time=start_time + timedelta(seconds=duration_seconds),
event_metadata=db_event_metadata(table, operation),
)
)
_ = self.loop.call_soon_threadsafe(self.loop.create_task, hook)

View file

@ -6,8 +6,10 @@ from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.db_span import proxy_service_logging
from litellm.proxy.db.native_database import (
NativeEndpointSettings,
ServiceLoggingEvents,
TokenRotation,
keep_token_fresh,
reader_token_rotation,
@ -38,11 +40,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 | None
tasks: "tuple[asyncio.Task[None], ...]" = ()
boot_probe: "asyncio.Task[None] | None" = None
@classmethod
async def open(
@ -55,7 +69,11 @@ class Database:
reader: Final = (
None if reader_url is None else await resolve_native_endpoint(reader_url, native_connections)
)
native: Final = await NativeDatabase.connect(writer.endpoint, None if reader is None else reader.endpoint)
native: Final = await NativeDatabase.connect(
ServiceLoggingEvents(loop=asyncio.get_running_loop(), service_logging=proxy_service_logging),
writer.endpoint,
None if reader is None else reader.endpoint,
)
except BaseException:
await prisma.aclose()
raise
@ -70,12 +88,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, *(() if self.boot_probe is None else (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:

View file

@ -3768,6 +3768,7 @@ dependencies = [
"chrono",
"futures-util",
"litellm-circuit-breaker",
"litellm-error",
"litellm-secret-value",
"macro_rules_attribute",
"rstest",

View file

@ -54,7 +54,21 @@ Python and SQLx run side by side until every repository operation is native, so
| Startup connection budget | `postgres_connection_budget_message` | Counts the Rust writer and reader pools' shares |
| Read replica routing | `routing_prisma_wrapper.py`, `DATABASE_URL_READ_REPLICA` | `Pools`: callers pick `writer()`, `read()` or `read_after_write()`. `read()` prefers the reader and `read_after_write()` the writer (`WriterPinnedClient`); either falls back to the other side when its pool cannot connect and skips that side for 30s, so reads keep working on the reader while the writer is down. Writes have no fallback |
| 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` | 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 |
| Migrations, `check_migration`, views | `prisma_client.py`, `create_views.py` | Prisma keeps them until the SQLx cutover |
## DB event names a Rust host sink must match
Prometheus and OTel dashboards aggregate across the Python proxy and any Rust server, so a Rust sink emits the same names and labels Python does, pinned by inline cases on both sides
- Prometheus
- `litellm_postgres_latency`, a histogram with label `postgres="postgres"`
- `litellm_postgres_total_requests`
- `litellm_postgres_failed_requests`, adding `error_class` (`NativeDatabaseError` for native failures) and `function_name` (the call type)
- OTel span name, from `service_span_name`
- `postgres.{verb} {table}` when the call type maps to a known Prisma helper and its table resolves
- `postgres.{verb}` for targetless verbs, for example `postgres.ping` for `health_check`
- `postgres {call_type}` otherwise

View file

@ -10,6 +10,7 @@ include = ["Cargo.toml", "src/**", "queries/**", ".sqlx/**"]
chrono.workspace = true
futures-util.workspace = true
litellm-circuit-breaker.workspace = true
litellm-error.workspace = true
litellm-secret-value.workspace = true
macro_rules_attribute.workspace = true
serde.workspace = true

View file

@ -0,0 +1,29 @@
use sqlx::{Connection, PgPool};
use crate::Error;
pub async fn health_check(pool: &PgPool) -> Result<(), Error> {
pool.acquire().await?.ping().await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use rstest::rstest;
use sqlx::PgPool;
use super::health_check;
use crate::test_support::database_pool;
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn answers_when_the_database_is_up(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
health_check(&database_pool?).await?;
Ok(())
}
}

View file

@ -12,6 +12,8 @@ mod end_user;
mod endpoint;
mod error;
mod failure;
mod health;
mod observe;
mod pools;
mod spend;
mod statements;
@ -24,6 +26,8 @@ pub use end_user::{EndUser, RateLimits, find_end_user};
pub use endpoint::{ClientIdentity, Endpoint, PoolLimits, SslMode};
pub use error::Error;
pub use failure::DbFailure;
pub use health::health_check;
pub use observe::{Call, Observer, Outcome, QueryEvent, observe};
pub use pools::Pools;
pub use spend::{
DailySpend, DailyTable, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend,

View file

@ -0,0 +1,128 @@
use std::time::{Duration, Instant, SystemTime};
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Call {
pub call_type: &'static str,
pub table: Option<&'static str>,
pub operation: Option<&'static str>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Outcome {
Succeeded,
Failed(String),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct QueryEvent {
pub call: Call,
pub started_at: SystemTime,
pub duration: Duration,
pub outcome: Outcome,
}
pub trait Observer: Send + Sync {
fn record(&self, event: QueryEvent);
}
pub async fn observe<T>(
observer: &dyn Observer,
call: Call,
query: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
let started_at = SystemTime::now();
let clock = Instant::now();
let result = query.await;
let outcome = match &result {
Ok(_) => Some(Outcome::Succeeded),
Err(error @ Error::Query(_)) => Some(Outcome::Failed(
litellm_error::Report::new(error).to_string(),
)),
Err(_) => None,
};
if let Some(outcome) = outcome {
observer.record(QueryEvent {
call,
started_at,
duration: clock.elapsed(),
outcome,
});
}
result
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use super::{Call, Observer, Outcome, QueryEvent, observe};
use crate::Error;
const CALL: Call = Call {
call_type: "add_key_spend",
table: Some("LiteLLM_VerificationToken"),
operation: Some("update"),
};
#[derive(Default)]
struct Recording(Mutex<Vec<QueryEvent>>);
impl Observer for Recording {
fn record(&self, event: QueryEvent) {
self.0.lock().expect("recording lock").push(event);
}
}
impl Recording {
fn outcomes(&self) -> Vec<(Call, Outcome)> {
self.0
.lock()
.expect("recording lock")
.iter()
.map(|event| (event.call, event.outcome.clone()))
.collect()
}
}
#[tokio::test]
async fn records_a_successful_round_trip() {
let recording = Recording::default();
let value = observe(&recording, CALL, async { Ok::<_, Error>(7) }).await;
assert_eq!(value.ok(), Some(7));
assert_eq!(recording.outcomes(), vec![(CALL, Outcome::Succeeded)]);
}
#[tokio::test]
async fn records_a_database_failure() {
let recording = Recording::default();
let result = observe(&recording, CALL, async {
Err::<(), _>(Error::from(sqlx::Error::PoolTimedOut))
})
.await;
assert!(result.is_err());
let outcomes = recording.outcomes();
assert_eq!(outcomes.len(), 1);
assert!(
matches!(&outcomes[0], (call, Outcome::Failed(message)) if *call == CALL && message.contains("pool timed out"))
);
}
#[tokio::test]
async fn skips_failures_that_never_reached_the_database() {
let recording = Recording::default();
let result = observe(&recording, CALL, async {
Err::<(), _>(Error::InvalidSpendBatch { ids: 1, costs: 2 })
})
.await;
assert!(result.is_err());
assert!(recording.outcomes().is_empty());
}
}

View file

@ -1,12 +1,53 @@
use litellm_db::{Endpoint, Pools};
use std::{sync::Arc, time::UNIX_EPOCH};
use litellm_db::{Call, Endpoint, Observer, Outcome, Pools, QueryEvent};
use litellm_secret_value::SecretValue;
use pyo3::{exceptions::PyValueError, prelude::*};
const HEALTH_CHECK: Call = Call {
call_type: "health_check",
table: None,
operation: None,
};
struct PythonEvents(Py<PyAny>);
impl Observer for PythonEvents {
fn record(&self, event: QueryEvent) {
let started_at = event
.started_at
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64();
let failure = match event.outcome {
Outcome::Succeeded => None,
Outcome::Failed(message) => Some(message),
};
Python::attach(|py| {
let events = self.0.bind(py);
if let Err(error) = events.call_method1(
"record",
(
event.call.call_type,
event.call.table,
event.call.operation,
started_at,
event.duration.as_secs_f64(),
failure,
),
) {
error.write_unraisable(py, Some(events));
}
});
}
}
#[pyclass(frozen)]
pub(crate) struct NativeDatabase {
writer_endpoint: Endpoint,
reader_endpoint: Option<Endpoint>,
pools: Pools,
events: Arc<dyn Observer>,
}
fn map_error(error: litellm_db::Error) -> PyErr {
@ -16,9 +57,10 @@ fn map_error(error: litellm_db::Error) -> PyErr {
#[pymethods]
impl NativeDatabase {
#[staticmethod]
#[pyo3(signature = (writer, reader=None))]
#[pyo3(signature = (events, writer, reader=None))]
fn connect(
py: Python<'_>,
events: Py<PyAny>,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] writer: Endpoint,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] reader: Option<Endpoint>,
) -> PyResult<Bound<'_, PyAny>> {
@ -28,6 +70,7 @@ impl NativeDatabase {
writer_endpoint: writer,
reader_endpoint: reader,
pools,
events: Arc::new(PythonEvents(events)),
})
})
}
@ -52,6 +95,26 @@ impl NativeDatabase {
Ok(())
}
fn health_check<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let writer = self.pools.writer().clone();
let events = Arc::clone(&self.events);
let socket_timeout = self.writer_endpoint.socket_timeout().map_err(map_error)?;
// Boxed so each future wrapper between here and the tokio worker holds a pointer instead of a copy of
// the SQLx connect state, which overflowed the worker's 2 MiB stack in unoptimized builds
crate::execution::run_async(
py,
Box::pin(async move {
litellm_db::observe(
events.as_ref(),
HEALTH_CHECK,
litellm_db::within(socket_timeout, litellm_db::health_check(&writer)),
)
.await
}),
map_error,
)
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let pools = self.pools.clone();
crate::execution::run_async(

View file

@ -8,7 +8,7 @@ from typing_extensions import disjoint_base
from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
from litellm.redis.stores.router_session_pins.keys import PinKey
from litellm.rust_bridge.database import NativeDatabaseEndpoint
from litellm.rust_bridge.database import NativeDatabaseEndpoint, NativeDatabaseEvents
from litellm.rust_bridge.execution import Execution
from litellm.rust_bridge.public_call import NativeCall
@ -564,10 +564,11 @@ class NativeSpendCounters:
class NativeDatabase:
@staticmethod
def connect(
writer: NativeDatabaseEndpoint, reader: NativeDatabaseEndpoint | None = None
events: NativeDatabaseEvents, writer: NativeDatabaseEndpoint, reader: NativeDatabaseEndpoint | None = None
) -> Future[NativeDatabase]: ...
def set_writer_password(self, password: str) -> None: ...
def set_reader_password(self, password: str) -> None: ...
def health_check(self) -> Future[None]: ...
def close(self) -> Future[None]: ...
@final

View file

@ -1,4 +1,4 @@
from typing import Literal, TypeAlias
from typing import Literal, Protocol, TypeAlias
from typing_extensions import ReadOnly, TypedDict
@ -31,3 +31,15 @@ class NativeDatabaseEndpoint(TypedDict):
options: ReadOnly[tuple[tuple[str, str], ...]]
statement_cache_capacity: ReadOnly[int]
pool: ReadOnly[NativePoolLimits]
class NativeDatabaseEvents(Protocol):
def record(
self,
call_type: str,
table: str | None,
operation: str | None,
started_at: float,
duration_seconds: float,
failure: str | None,
) -> None: ...

View file

@ -1,7 +1,7 @@
import asyncio
from typing import Final
from litellm.proxy.db.native_database import native_endpoint
from litellm.proxy.db.native_database import ServiceLoggingEvents, native_endpoint
from litellm.rust_bridge._native import NativeDatabase
_NEVER_DIALED_URL: Final = "postgresql://litellm:secret@db.invalid:5432/proxy"
@ -9,7 +9,9 @@ _NEVER_DIALED_URL: Final = "postgresql://litellm:secret@db.invalid:5432/proxy"
async def connect_native_database() -> NativeDatabase:
return await NativeDatabase.connect(
native_endpoint(_NEVER_DIALED_URL, tls_offered=True, max_connections=1).endpoint, None
ServiceLoggingEvents(loop=asyncio.get_running_loop(), service_logging=lambda: None),
native_endpoint(_NEVER_DIALED_URL, tls_offered=True, max_connections=1).endpoint,
None,
)

View file

@ -1,5 +1,6 @@
import asyncio
import datetime as dt
from collections.abc import Mapping
from datetime import datetime, timedelta
from pathlib import Path
from typing import Final
@ -12,7 +13,9 @@ from cryptography.hazmat.primitives.serialization import BestAvailableEncryption
from cryptography.x509.oid import NameOID
from pydantic import TypeAdapter
from litellm._service_logger import ServiceLogging, ServiceTypes
from litellm.proxy.db.native_database import (
ServiceLoggingEvents,
TokenRotation,
keep_token_fresh,
native_endpoint,
@ -23,6 +26,7 @@ from litellm.proxy.db.native_database import (
)
from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR
from litellm.proxy.db.token_auth import IAMEndpoint, RdsIamTokenAuth
from litellm.rust_bridge._native import NativeDatabase
from litellm.rust_bridge.database import NativeDatabaseEndpoint, NativePoolLimits
NOW: Final = datetime(2026, 11, 3, 12, 0, 0)
@ -390,3 +394,86 @@ async def test_keep_token_fresh_retries_a_failed_mint() -> None:
assert sink.passwords == ("recovered",)
assert schedule.delays[:2] == (600.0, 30.0)
class _RecordingServiceLogging(ServiceLogging):
def __init__(self) -> None:
super().__init__(mock_testing=True)
self.events: tuple[tuple[str, str, float, Mapping[str, str] | None, str | None], ...] = ()
self.arrived: Final = asyncio.Event()
async def async_service_success_hook(
self,
service: ServiceTypes,
call_type: str,
duration: float,
parent_otel_span: object = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
event_metadata: dict[str, str] | None = None,
caller: str | None = None,
) -> None:
self.events = (*self.events, (service.value, call_type, duration, event_metadata, None))
self.arrived.set()
async def async_service_failure_hook(
self,
service: ServiceTypes,
duration: float,
error: str | Exception,
call_type: str,
parent_otel_span: object = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
event_metadata: dict[str, str] | None = None,
caller: str | None = None,
) -> None:
self.events = (*self.events, (service.value, call_type, duration, event_metadata, str(error)))
self.arrived.set()
@pytest.mark.asyncio
async def test_service_logging_events_reach_the_loop_from_another_thread() -> None:
logging: Final = _RecordingServiceLogging()
events: Final = ServiceLoggingEvents(loop=asyncio.get_running_loop(), service_logging=lambda: logging)
await asyncio.to_thread(
events.record, "add_key_spend", "LiteLLM_VerificationToken", "update", NOW.timestamp(), 0.25, None
)
await asyncio.wait_for(logging.arrived.wait(), timeout=5)
logging.arrived.clear()
await asyncio.to_thread(events.record, "health_check", None, None, NOW.timestamp(), 1.5, "pool timed out")
await asyncio.wait_for(logging.arrived.wait(), timeout=5)
assert logging.events == (
(
"postgres",
"add_key_spend",
0.25,
{"table_name": "LiteLLM_VerificationToken", "db_operation": "update"},
None,
),
("postgres", "health_check", 1.5, {}, "pool timed out"),
)
@pytest.mark.asyncio
async def test_native_health_check_failure_reports_a_db_event() -> None:
logging: Final = _RecordingServiceLogging()
native: Final = await NativeDatabase.connect(
ServiceLoggingEvents(loop=asyncio.get_running_loop(), service_logging=lambda: logging),
native_endpoint(
"postgresql://litellm:secret@127.0.0.1:1/proxy?pool_timeout=1", tls_offered=True, max_connections=2
).endpoint,
)
with pytest.raises(ValueError, match="database query failed"):
await native.health_check()
await asyncio.wait_for(logging.arrived.wait(), timeout=5)
await native.close()
assert len(logging.events) == 1
service, call_type, duration, metadata, failure = logging.events[0]
assert (service, call_type, metadata) == ("postgres", "health_check", {})
assert duration >= 0.9
assert failure is not None and "database query failed" in failure

View file

@ -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.pgbouncer import PGBOUNCER_POOLED_ENV_VAR
from litellm.repositories.database import Database
@ -94,3 +97,22 @@ async def test_zero_native_connections_open_no_native_pool(monkeypatch: pytest.M
assert (database.native, database.tasks) == (None, ())
await database.aclose()
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, NATIVE_CONNECTIONS
)
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"
]