From 15d3ef837fcb0448bcfdf6eb36528122d4e2abec Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 10 Nov 2026 11:14:45 +0900 Subject: [PATCH] feat(db): carry the native error kind across the Python bridge Native failures reached Python as a ValueError, so allow_requests_on_db_unavailable treated a native outage as a permanent fault and a write on a dropped connection was requeued. The bridge now raises NativeDatabaseFailure with the kind and the Rust answers for is_unavailable and is_requeue_safe, and the two Python gates return them --- litellm-rust/crates/db/AGENTS.md | 2 +- litellm-rust/crates/db/src/error.rs | 3 +- litellm-rust/crates/python-bridge/src/db.rs | 20 +++++++- litellm-rust/crates/python-bridge/src/lib.rs | 3 +- litellm/proxy/db/exception_handler.py | 5 ++ litellm/repositories/database.py | 4 +- litellm/rust_bridge/_native.pyi | 6 +++ tests/unit/proxy/db/test_native_database.py | 51 +++++++++++++++++++- 8 files changed, 86 insertions(+), 8 deletions(-) diff --git a/litellm-rust/crates/db/AGENTS.md b/litellm-rust/crates/db/AGENTS.md index 444f7646528..f7d738a9519 100644 --- a/litellm-rust/crates/db/AGENTS.md +++ b/litellm-rust/crates/db/AGENTS.md @@ -53,7 +53,7 @@ Python and SQLx run side by side until every repository operation is native, so | In-container PgBouncer | `pgbouncer.py` | Inherited: the pooled URL carries PgBouncer's credential | | 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` | `Error::kind`, `is_unavailable`, `is_requeue_safe`; `begin` takes the caller's `lock_timeout` | +| Error classification, `lock_timeout` on rollups | `exception_handler.py`, `rollup_lock_timeout.py` | `Error::kind`, `is_unavailable`, `is_requeue_safe`; `begin` takes the caller's `lock_timeout`. The bridge raises `NativeDatabaseFailure` with the kind and both answers, and `is_database_connection_error` and `db_write_failure_is_requeue_safe` return them instead of classifying the error again | | Recreate the client after 25006 (writer became read-only) or a stale cached plan | `exception_handler.py`, `_query_first_with_cached_plan_fallback` | `Pools::recycle_after` drops every pooled connection opened before the error, so the next acquire reconnects and prepares fresh. Callers retry a `StalePlan` once | | `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 | diff --git a/litellm-rust/crates/db/src/error.rs b/litellm-rust/crates/db/src/error.rs index 27c1c06666b..22ddb643d9c 100644 --- a/litellm-rust/crates/db/src/error.rs +++ b/litellm-rust/crates/db/src/error.rs @@ -10,7 +10,8 @@ pub enum Error { InvalidDuration(#[from] std::time::TryFromFloatSecsError), } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr)] +#[strum(serialize_all = "snake_case")] pub enum ErrorKind { Unreachable, Capacity, diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs index 40243b3fe59..1044ba9daee 100644 --- a/litellm-rust/crates/python-bridge/src/db.rs +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -50,8 +50,26 @@ pub(crate) struct NativeDatabase { events: Arc, } +pyo3::create_exception!( + _native, + NativeDatabaseFailure, + pyo3::exceptions::PyException, + "A native database operation failed. The kind attribute is litellm-db's ErrorKind, and unavailable and requeue_safe are the error's is_unavailable and is_requeue_safe." +); + fn map_error(error: litellm_db::Error) -> PyErr { - PyValueError::new_err(litellm_error::Report::new(&error).to_string()) + let kind: &'static str = error.kind().into(); + let failure = NativeDatabaseFailure::new_err(litellm_error::Report::new(&error).to_string()); + let decided = Python::attach(|py| { + let value = failure.value(py); + value.setattr("kind", kind)?; + value.setattr("unavailable", error.is_unavailable())?; + value.setattr("requeue_safe", error.is_requeue_safe()) + }); + match decided { + Ok(()) => failure, + Err(error) => error, + } } #[pymethods] diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 09b787b2c94..96d65648b7f 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -25,7 +25,7 @@ mod _native { use pyo3::{prelude::*, types::PyModule}; #[pymodule_export] - use crate::db::NativeDatabase; + use crate::db::{NativeDatabase, NativeDatabaseFailure}; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -160,6 +160,7 @@ mod tests { "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDatabaseFailure", "trace_encode_error", "bundled_model_cost_map", "TokenCounter", diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index c2d58512fea..473dac84805 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -13,6 +13,7 @@ from litellm.proxy._types import ( ProxyException, ) from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded +from litellm.rust_bridge._native import NativeDatabaseFailure from litellm.secret_managers.main import str_to_bool # Bounds the __cause__/__context__ walk in find_database_service_unavailable_error_in_chain. @@ -152,6 +153,8 @@ class PrismaDBExceptionHandler: """ import prisma.engine.errors + if isinstance(e, NativeDatabaseFailure): + return e.unavailable if isinstance(e, (*DB_CONNECTION_ERROR_TYPES, DBLookupDeadlineExceeded)): return True if isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)): @@ -746,6 +749,8 @@ _DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) def db_write_failure_is_requeue_safe(e: Exception) -> bool: """Whether a failed write provably never applied, so its rows can be sent again without double counting.""" + if isinstance(e, NativeDatabaseFailure): + return e.requeue_safe if isinstance(e, DB_CONNECTION_ERROR_TYPES): return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index b21a8f6a502..93538799aea 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -17,7 +17,7 @@ from litellm.proxy.db.native_database import ( resolve_native_endpoint, writer_token_rotation, ) -from litellm.rust_bridge._native import NativeDatabase +from litellm.rust_bridge._native import NativeDatabase, NativeDatabaseFailure if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -44,7 +44,7 @@ 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: + except NativeDatabaseFailure 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", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index f39dd9a19bd..3903e07b542 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -19,6 +19,11 @@ class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... class NativeRedisStoreError(Exception): ... +class NativeDatabaseFailure(Exception): + kind: str + unavailable: bool + requeue_safe: bool + def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... @@ -590,6 +595,7 @@ __all__ = [ "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDatabaseFailure", "NativeDiagnosticProcessor", "NativePending", "NativeRateLimits", diff --git a/tests/unit/proxy/db/test_native_database.py b/tests/unit/proxy/db/test_native_database.py index dbe24da55c0..575661fab0b 100644 --- a/tests/unit/proxy/db/test_native_database.py +++ b/tests/unit/proxy/db/test_native_database.py @@ -14,6 +14,7 @@ from cryptography.x509.oid import NameOID from pydantic import TypeAdapter from litellm._service_logger import ServiceLogging, ServiceTypes +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe from litellm.proxy.db.native_database import ( ServiceLoggingEvents, TokenRotation, @@ -26,7 +27,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._native import NativeDatabase, NativeDatabaseFailure from litellm.rust_bridge.database import NativeDatabaseEndpoint, NativePoolLimits NOW: Final = datetime(2026, 11, 3, 12, 0, 0) @@ -467,7 +468,7 @@ async def test_native_health_check_failure_reports_a_db_event() -> None: ).endpoint, ) - with pytest.raises(ValueError, match="database query failed"): + with pytest.raises(NativeDatabaseFailure, match="database query failed"): await native.health_check() await asyncio.wait_for(logging.arrived.wait(), timeout=5) await native.close() @@ -477,3 +478,49 @@ async def test_native_health_check_failure_reports_a_db_event() -> None: assert (service, call_type, metadata) == ("postgres", "health_check", {}) assert duration >= 0.9 assert failure is not None and "database query failed" in failure + + +async def _native_failure(url: str) -> NativeDatabaseFailure: + native: Final = await NativeDatabase.connect( + ServiceLoggingEvents(loop=asyncio.get_running_loop(), service_logging=lambda: None), + native_endpoint(url, tls_offered=True, max_connections=2).endpoint, + ) + try: + with pytest.raises(NativeDatabaseFailure) as raised: + await native.health_check() + finally: + await native.close() + return raised.value + + +async def _drop_connection(_: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + writer.close() + + +@pytest.mark.asyncio +async def test_native_outage_passes_the_db_unavailable_gate_and_is_requeue_safe() -> None: + error: Final = await _native_failure("postgresql://litellm:secret@127.0.0.1:1/proxy?pool_timeout=1") + + assert (error.kind, str(error)) == ( + "capacity", + "database query failed: pool timed out while waiting for an open connection", + ) + assert PrismaDBExceptionHandler.is_database_connection_error(error) + assert db_write_failure_is_requeue_safe(error) + + +@pytest.mark.asyncio +async def test_native_dropped_connection_is_unavailable_but_not_requeue_safe() -> None: + server: Final = await asyncio.start_server(_drop_connection, "127.0.0.1", 0) + _, port = _ADDRESS.validate_python(server.sockets[0].getsockname()) + try: + error: Final = await _native_failure( + f"postgresql://litellm:secret@127.0.0.1:{port}/proxy?pool_timeout=5&sslmode=disable" + ) + finally: + server.close() + await server.wait_closed() + + assert error.kind == "unreachable" + assert PrismaDBExceptionHandler.is_database_connection_error(error) + assert not db_write_failure_is_requeue_safe(error)