feat(db): emit DB service events from native operations

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 crate AGENTS.md records the Prometheus and OTel names a Rust host sink must match

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-09 12:36:02 +09:00
parent 4bb266d02f
commit 90ed6a9bee
13 changed files with 438 additions and 13 deletions

View file

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

View file

@ -54,8 +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 |
| 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 |
## 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,24 @@ 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)?;
crate::execution::run_async(
py,
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

@ -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

@ -7,8 +7,10 @@ from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.connection_budget import RustConnectionLimits
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,
@ -56,7 +58,11 @@ class Database:
if reader_url is None
else await resolve_native_endpoint(reader_url, rust_limits.reader_connection_limit)
)
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

View file

@ -7,7 +7,7 @@ from pydantic import JsonValue
from typing_extensions import disjoint_base
from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
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
@ -392,10 +392,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,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