Merge db10e18d07 into 8cb628ce6a

litellm-rust/Cargo.lock

@@ -3712,6 +3712,7 @@ version = "0.1.0"
 dependencies = [
  "chrono",
  "litellm-circuit-breaker",
+ "litellm-error",
  "litellm-secret-value",
  "macro_rules_attribute",
  "rstest",

litellm-rust/crates/db/.sqlx/query-626aedc76c92e336c0ef536f58d102cb2fb823560b341be1f01ca31de3341e7b.json

@@ -0,0 +1,21 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT 1::int4 AS \"ok!\"\n",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "ok!",
+        "type_info": "Int4",
+        "origin": "Expression"
+      }
+    ],
+    "parameters": {
+      "Left": []
+    },
+    "nullable": [
+      null
+    ]
+  },
+  "hash": "626aedc76c92e336c0ef536f58d102cb2fb823560b341be1f01ca31de3341e7b"
+}

litellm-rust/crates/db/AGENTS.md

@@ -41,7 +41,20 @@ Python and SQLx run side by side until every repository operation is native, so
 | Startup connection budget | `postgres_connection_budget_message` | Counts the native writer and reader pools |
 | Read replica routing | `routing_prisma_wrapper.py`, `DATABASE_URL_READ_REPLICA` | `Pools`: callers pick `writer()` or `read()`; a failed reader acquire falls back to the writer and skips the reader for 30s |
 | 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 |
 | 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

litellm-rust/crates/db/Cargo.toml

@@ -8,6 +8,7 @@ repository.workspace = true
 [dependencies]
 chrono.workspace = true
 litellm-circuit-breaker.workspace = true
+litellm-error.workspace = true
 litellm-secret-value.workspace = true
 macro_rules_attribute.workspace = true
 serde.workspace = true

litellm-rust/crates/db/queries/health/health_check.sql

@@ -0,0 +1 @@
+SELECT 1::int4 AS "ok!"

litellm-rust/crates/db/src/health.rs

@@ -0,0 +1,31 @@
+use sqlx::PgExecutor;
+
+use crate::Error;
+
+pub async fn health_check(executor: impl PgExecutor<'_>) -> Result<(), Error> {
+    sqlx::query_file_scalar!("queries/health/health_check.sql")
+        .fetch_one(executor)
+        .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(())
+    }
+}

litellm-rust/crates/db/src/lib.rs

@@ -11,6 +11,8 @@ mod end_user;
 mod endpoint;
 mod error;
 mod failure;
+mod health;
+mod observe;
 mod pools;
 mod spend;
 #[cfg(test)]
@@ -21,6 +23,8 @@ pub use end_user::{EndUser, RateLimits, find_end_user};
 pub use endpoint::{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,

litellm-rust/crates/db/src/observe.rs

@@ -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());
+    }
+}

litellm-rust/crates/python-bridge/src/db.rs

@@ -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,23 @@ 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);
+        crate::execution::run_async(
+            py,
+            async move {
+                litellm_db::observe(
+                    events.as_ref(),
+                    HEALTH_CHECK,
+                    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(

litellm/proxy/db/db_span.py

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

litellm/proxy/db/native_database.py

@@ -11,11 +11,13 @@ import shlex
 import urllib.parse
 from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from types import MappingProxyType
-from typing import Final
+from typing import Final, Protocol
 
 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.pgbouncer import database_url_is_pooled, libpq_option_settings
 from litellm.proxy.db.token_auth import (
     DEFAULT_POSTGRES_PORT,
@@ -172,3 +174,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)

litellm/repositories/database.py

@@ -5,7 +5,9 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Final
 
+from litellm.proxy.db.db_span import proxy_service_logging
 from litellm.proxy.db.native_database import (
+    ServiceLoggingEvents,
     TokenRotation,
     keep_token_fresh,
     native_endpoint,
@@ -37,7 +39,9 @@ class Database:
     async def open(cls, prisma: "PrismaClient", writer_url: str, reader_url: str | None) -> "Database":
         try:
             native: Final = await NativeDatabase.connect(
-                native_endpoint(writer_url), None if reader_url is None else native_endpoint(reader_url)
+                ServiceLoggingEvents(loop=asyncio.get_running_loop(), service_logging=proxy_service_logging),
+                native_endpoint(writer_url),
+                None if reader_url is None else native_endpoint(reader_url),
             )
         except BaseException:
             await prisma.aclose()

litellm/rust_bridge/_native.pyi

@@ -5,7 +5,7 @@ from typing import Generic, Literal, Never, NotRequired, TypedDict, TypeVar, fin
 import httpx
 from pydantic import JsonValue
 
-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
 
@@ -31,10 +31,11 @@ class NativeClickHouseSpendStorage:
 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

litellm/rust_bridge/database.py

@@ -1,4 +1,4 @@
-from typing import Literal, TypeAlias
+from typing import Literal, Protocol, TypeAlias
 
 from typing_extensions import ReadOnly, TypedDict
 
@@ -22,3 +22,15 @@ class NativeDatabaseEndpoint(TypedDict):
     options: ReadOnly[tuple[tuple[str, str], ...]]
     statement_cache: ReadOnly[bool]
     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: ...

tests/unit/proxy/db/test_native_database.py

@@ -1,10 +1,13 @@
 import asyncio
+from collections.abc import Mapping
 from datetime import datetime, timedelta
 from typing import Final
 
 import pytest
 
+from litellm._service_logger import ServiceLogging, ServiceTypes
 from litellm.proxy.db.native_database import (
+    ServiceLoggingEvents,
     TokenRotation,
     keep_token_fresh,
     native_endpoint,
@@ -13,6 +16,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)
@@ -233,3 +237,84 @@ 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"),
+    )
+
+    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