From 0b7d6be4a4f1c11bf651e981d7a4d22004002e06 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 9 Nov 2026 12:32:00 +0900 Subject: [PATCH] feat(db): add an optional native reader pool with explicit routing When DATABASE_URL_READ_REPLICA is set, NativeDatabase opens a reader pool resolved from that URL, with its own token rotation. Pools exposes writer(), read() and read_after_write(), matching RoutingPrismaWrapper and WriterPinnedClient. read() prefers the replica and read_after_write() the writer; each falls back to the other side when its pool cannot connect and skips the failed side for 30 seconds, so reads keep working on the replica while the writer is down. The reader pool takes the same native share of the reader's connection_limit, and the startup budget line counts it Co-Authored-By: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/db/AGENTS.md | 6 +- litellm-rust/crates/db/Cargo.toml | 1 + litellm-rust/crates/db/src/lib.rs | 2 + litellm-rust/crates/db/src/pools.rs | 242 ++++++++++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 - litellm-rust/crates/python-bridge/src/db.rs | 31 ++- litellm/proxy/db/connection_budget.py | 2 + litellm/proxy/db/db_url_settings.py | 17 +- litellm/proxy/db/native_database.py | 28 +- litellm/proxy/proxy_cli.py | 24 +- litellm/proxy/proxy_server.py | 7 +- litellm/repositories/database.py | 66 +++-- litellm/rust_bridge/_native.pyi | 5 +- tests/unit/proxy/db/test_connection_budget.py | 7 +- tests/unit/proxy/db/test_db_url_settings.py | 12 +- tests/unit/proxy/db/test_native_database.py | 29 ++- tests/unit/proxy/test_proxy_cli.py | 8 +- tests/unit/repositories/test_database.py | 26 +- 19 files changed, 438 insertions(+), 77 deletions(-) create mode 100644 litellm-rust/crates/db/src/pools.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e48ecf372e6..c833664e345 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3769,6 +3769,7 @@ version = "0.1.0" dependencies = [ "chrono", "futures-util", + "litellm-circuit-breaker", "litellm-secret-value", "macro_rules_attribute", "rstest", diff --git a/litellm-rust/crates/db/AGENTS.md b/litellm-rust/crates/db/AGENTS.md index fd734d56611..1d03b3eb92a 100644 --- a/litellm-rust/crates/db/AGENTS.md +++ b/litellm-rust/crates/db/AGENTS.md @@ -49,10 +49,10 @@ Python and SQLx run side by side until every repository operation is native, so | `schema` | URL param | `search_path` startup setting when not `public`; a custom schema behind PgBouncer is rejected by PgBouncer | | `statement_cache_size`, `pgbouncer=true`, `DATABASE_DISABLE_PREPARED_STATEMENTS` | URL params, `db_url_settings.py` | `statement_cache_capacity`: `statement_cache_size` (Prisma's default 100), 0 under `pgbouncer=true`. Run every query through `Statements::new(executor, endpoint.names_statements())` so a 0 capacity sends unnamed statements like Prisma, until a sqlx release carries transact-rs/sqlx#4337 | | `application_name` | URL param | A value the operator set is used on both pools; Prisma's `litellm-prisma` default becomes `litellm-native` | -| Rotating RDS IAM / Entra token | `token_auth.py`, `PrismaWrapper` | Python mints, `NativeDatabase.set_writer_password` swaps connect options | +| Rotating RDS IAM / Entra token | `token_auth.py`, `PrismaWrapper` | Python mints per side, `NativeDatabase.set_writer_password` / `set_reader_password` swap connect options | | In-container PgBouncer | `pgbouncer.py` | Inherited: the pooled URL carries PgBouncer's credential | -| Startup connection budget | `postgres_connection_budget_message` | Counts the Rust writer pool's share | -| Read replica routing | `routing_prisma_wrapper.py` | Planned: optional reader pool with explicit routing | +| 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 | | 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/Cargo.toml b/litellm-rust/crates/db/Cargo.toml index 527a0927338..c95f982433e 100644 --- a/litellm-rust/crates/db/Cargo.toml +++ b/litellm-rust/crates/db/Cargo.toml @@ -9,6 +9,7 @@ include = ["Cargo.toml", "src/**", "queries/**", ".sqlx/**"] [dependencies] chrono.workspace = true futures-util.workspace = true +litellm-circuit-breaker.workspace = true litellm-secret-value.workspace = true macro_rules_attribute.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index bab2d492c92..cea4e689768 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -12,6 +12,7 @@ mod end_user; mod endpoint; mod error; mod failure; +mod pools; mod spend; mod statements; #[cfg(test)] @@ -23,6 +24,7 @@ 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 pools::Pools; pub use spend::{ DailySpend, DailyTable, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend, add_end_user_spend, add_key_spend, add_team_member_spend, add_team_spend, add_user_spend, diff --git a/litellm-rust/crates/db/src/pools.rs b/litellm-rust/crates/db/src/pools.rs new file mode 100644 index 00000000000..561c1bad914 --- /dev/null +++ b/litellm-rust/crates/db/src/pools.rs @@ -0,0 +1,242 @@ +use std::time::{Duration, Instant}; + +use litellm_circuit_breaker::{Breaker, Config, Failure}; +use sqlx::{PgPool, Postgres, pool::PoolConnection}; + +use crate::{Endpoint, Error}; + +const RECOVERY: Duration = Duration::from_secs(30); + +type Connection = PoolConnection; + +#[derive(Clone)] +struct Side { + pool: PgPool, + breaker: Breaker, +} + +impl Side { + fn new(endpoint: &Endpoint) -> Result { + Ok(Self { + pool: endpoint.pool()?, + breaker: Breaker::new( + Config { + failure_threshold: 1, + recovery_timeout: RECOVERY, + ..Config::default() + }, + Box::new(|_| {}), + ), + }) + } + + async fn acquire_unless_skipped(&self) -> Option { + let admission = self.breaker.admit(Instant::now()).ok()?; + match self.pool.acquire().await { + Ok(connection) => { + admission.succeeded(); + Some(connection) + } + Err(_) => { + admission.failed(Failure::Hard, Instant::now()); + None + } + } + } + + async fn acquire(&self) -> Result { + let result = self.pool.acquire().await; + if result.is_err() { + self.breaker.failed(Failure::Hard, Instant::now()); + } + Ok(result?) + } +} + +async fn acquire_preferring( + preferred: &Side, + fallback: Option<&Side>, +) -> Result { + let Some(fallback) = fallback else { + return preferred.acquire().await; + }; + match preferred.acquire_unless_skipped().await { + Some(connection) => Ok(connection), + None => fallback.acquire().await, + } +} + +#[derive(Clone)] +pub struct Pools { + writer: Side, + reader: Option, +} + +impl Pools { + pub fn new(writer: &Endpoint, reader: Option<&Endpoint>) -> Result { + Ok(Self { + writer: Side::new(writer)?, + reader: reader.map(Side::new).transpose()?, + }) + } + + pub fn writer(&self) -> &PgPool { + &self.writer.pool + } + + pub fn reader_pool(&self) -> Option<&PgPool> { + self.reader.as_ref().map(|reader| &reader.pool) + } + + pub async fn read(&self) -> Result { + match &self.reader { + Some(reader) => acquire_preferring(reader, Some(&self.writer)).await, + None => self.writer.acquire().await, + } + } + + pub async fn read_after_write(&self) -> Result { + acquire_preferring(&self.writer, self.reader.as_ref()).await + } + + pub async fn close(&self) { + self.writer.pool.close().await; + if let Some(reader) = &self.reader { + reader.pool.close().await; + } + } +} + +#[cfg(test)] +mod tests { + use std::{error::Error as StdError, time::Instant}; + + use sqlx::PgConnection; + + use super::Pools; + use crate::{Endpoint, PoolLimits, test_support::database_endpoint}; + + fn read_only(endpoint: &Endpoint) -> Endpoint { + Endpoint { + options: vec![("default_transaction_read_only".to_owned(), "on".to_owned())], + ..endpoint.clone() + } + } + + fn unreachable(endpoint: &Endpoint) -> Endpoint { + Endpoint { + host: "127.0.0.1".to_owned(), + port: 1, + pool: PoolLimits { + acquire_timeout_seconds: Some(2.0), + ..endpoint.pool + }, + ..endpoint.clone() + } + } + + async fn is_read_only(connection: &mut PgConnection) -> Result { + #[expect( + clippy::disallowed_methods, + reason = "Session settings are not part of the compile-checked schema" + )] + let setting: String = sqlx::query_scalar("SHOW default_transaction_read_only") + .fetch_one(connection) + .await?; + Ok(setting == "on") + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn reads_go_to_the_reader_and_read_after_write_to_the_writer() + -> Result<(), Box> { + let writer = database_endpoint()?; + let pools = Pools::new(&writer, Some(&read_only(&writer)))?; + + assert!(is_read_only(&mut *pools.read().await?).await?); + assert!(!is_read_only(&mut *pools.read_after_write().await?).await?); + assert!(!is_read_only(&mut *pools.writer().acquire().await?).await?); + Ok(()) + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn writes_sent_to_the_reader_fail_read_only() -> Result<(), Box> { + let writer = database_endpoint()?; + let pools = Pools::new(&writer, Some(&read_only(&writer)))?; + let mut connection = pools.read().await?; + + #[expect(clippy::disallowed_methods, reason = "A scratch table needs no schema")] + let error = sqlx::query("CREATE TABLE native_reader_write_probe (id int)") + .execute(&mut *connection) + .await + .expect_err("a read-only session must reject writes"); + + assert_eq!( + error.as_database_error().and_then(|error| error.code()), + Some("25006".into()) + ); + Ok(()) + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn every_read_uses_the_writer_without_a_reader() -> Result<(), Box> { + let pools = Pools::new(&database_endpoint()?, None)?; + + assert!(pools.reader_pool().is_none()); + assert!(!is_read_only(&mut *pools.read().await?).await?); + assert!(!is_read_only(&mut *pools.read_after_write().await?).await?); + Ok(()) + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn unreachable_reader_falls_back_once_then_skips_it() -> Result<(), Box> { + let writer = database_endpoint()?; + let pools = Pools::new(&writer, Some(&unreachable(&writer)))?; + + assert!(!is_read_only(&mut *pools.read().await?).await?); + let started = Instant::now(); + assert!(!is_read_only(&mut *pools.read().await?).await?); + + assert!(started.elapsed().as_secs_f64() < 1.0); + Ok(()) + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn reads_stay_on_the_reader_while_the_writer_is_down() -> Result<(), Box> { + let reader = read_only(&database_endpoint()?); + let pools = Pools::new(&unreachable(&reader), Some(&reader))?; + + assert!(is_read_only(&mut *pools.read().await?).await?); + assert!(pools.writer().acquire().await.is_err()); + Ok(()) + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn read_after_write_yields_to_the_reader_while_the_writer_is_down() + -> Result<(), Box> { + let reader = read_only(&database_endpoint()?); + let pools = Pools::new(&unreachable(&reader), Some(&reader))?; + + assert!(is_read_only(&mut *pools.read_after_write().await?).await?); + let started = Instant::now(); + assert!(is_read_only(&mut *pools.read_after_write().await?).await?); + + assert!(started.elapsed().as_secs_f64() < 1.0); + Ok(()) + } + + #[tokio::test] + #[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"] + async fn read_after_write_fails_when_the_only_database_is_down() -> Result<(), Box> + { + let pools = Pools::new(&unreachable(&database_endpoint()?), None)?; + + assert!(pools.read_after_write().await.is_err()); + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 8406a1b89a1..28884b5ce40 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -70,7 +70,6 @@ serde_json.workspace = true strum.workspace = true veil.workspace = true thiserror.workspace = true -sqlx.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] diff --git a/litellm-rust/crates/python-bridge/src/db.rs b/litellm-rust/crates/python-bridge/src/db.rs index ca1ccc8fb2b..d17f5fd8099 100644 --- a/litellm-rust/crates/python-bridge/src/db.rs +++ b/litellm-rust/crates/python-bridge/src/db.rs @@ -1,12 +1,12 @@ -use litellm_db::Endpoint; +use litellm_db::{Endpoint, Pools}; use litellm_secret_value::SecretValue; use pyo3::{exceptions::PyValueError, prelude::*}; -use sqlx::PgPool; #[pyclass(frozen)] pub(crate) struct NativeDatabase { writer_endpoint: Endpoint, - writer: PgPool, + reader_endpoint: Option, + pools: Pools, } fn map_error(error: litellm_db::Error) -> PyErr { @@ -16,33 +16,48 @@ fn map_error(error: litellm_db::Error) -> PyErr { #[pymethods] impl NativeDatabase { #[staticmethod] + #[pyo3(signature = (writer, reader=None))] fn connect( py: Python<'_>, #[pyo3(from_py_with = litellm_host_python::from_py_argument)] writer: Endpoint, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] reader: Option, ) -> PyResult> { crate::execution::run_async_value(py, async move { - let pool = writer.pool().map_err(map_error)?; + let pools = Pools::new(&writer, reader.as_ref()).map_err(map_error)?; Ok(NativeDatabase { writer_endpoint: writer, - writer: pool, + reader_endpoint: reader, + pools, }) }) } fn set_writer_password(&self, password: String) { - self.writer.set_connect_options( + self.pools.writer().set_connect_options( self.writer_endpoint .with_password(SecretValue::new(password)) .connect_options(), ); } + fn set_reader_password(&self, password: String) -> PyResult<()> { + let (Some(endpoint), Some(pool)) = (&self.reader_endpoint, self.pools.reader_pool()) else { + return Err(PyValueError::new_err("no read replica is configured")); + }; + pool.set_connect_options( + endpoint + .with_password(SecretValue::new(password)) + .connect_options(), + ); + Ok(()) + } + fn close<'py>(&self, py: Python<'py>) -> PyResult> { - let writer = self.writer.clone(); + let pools = self.pools.clone(); crate::execution::run_async( py, async move { - writer.close().await; + pools.close().await; Ok(()) }, |error: PyErr| error, diff --git a/litellm/proxy/db/connection_budget.py b/litellm/proxy/db/connection_budget.py index 541208c5dcb..b55a22d4383 100644 --- a/litellm/proxy/db/connection_budget.py +++ b/litellm/proxy/db/connection_budget.py @@ -26,6 +26,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict RUST_RESERVED_CONNECTIONS: Final = 2 RUST_CONNECTION_LIMIT_ENV_PREFIX: Final = "LITELLM_DATABASE_RUST_" RUST_WRITER_CONNECTION_LIMIT_ENV_VAR: Final = f"{RUST_CONNECTION_LIMIT_ENV_PREFIX}WRITER_CONNECTION_LIMIT" +RUST_READER_CONNECTION_LIMIT_ENV_VAR: Final = f"{RUST_CONNECTION_LIMIT_ENV_PREFIX}READER_CONNECTION_LIMIT" @dataclass(frozen=True, slots=True) @@ -51,3 +52,4 @@ class RustConnectionLimits(BaseSettings): model_config = SettingsConfigDict(env_prefix=RUST_CONNECTION_LIMIT_ENV_PREFIX, extra="ignore", frozen=True) writer_connection_limit: PositiveInt = RUST_RESERVED_CONNECTIONS + reader_connection_limit: PositiveInt = RUST_RESERVED_CONNECTIONS diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index fce08a0518f..d93e9e4b1b2 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -259,14 +259,14 @@ NATIVE_APPLICATION_NAME: Final = "litellm-native" def postgres_connection_budget_message( - writer_limit: str, rust_writer_limit: int, reader_limit: str | None, num_workers: str + writer_limit: str, rust_writer_limit: int, reader_limit: str | None, rust_reader_limit: int, num_workers: str ) -> str: """The startup line that states this pod's worst-case Postgres connection demand. Prisma's ``connection_limit`` is per query engine, and every uvicorn worker owns one engine per configured database (writer, plus the reader when - ``DATABASE_URL_READ_REPLICA`` is set) plus the Rust writer pool's share - (``connection_budget``). The limits are read off the final URLs, + ``DATABASE_URL_READ_REPLICA`` is set) plus a Rust pool beside each, holding that + pool's share (``connection_budget``). The limits are read off the final URLs, so a reader that pins its own ``connection_limit`` or an operator override in ``database_extra_connection_params`` is counted at its real value. The server-side cap is shared by every pod, so the number an operator has to keep @@ -280,13 +280,18 @@ def postgres_connection_budget_message( except ValueError: return ( "LiteLLM Proxy: Postgres connection budget per pod = workers x (writer connection_limit + rust " - "writer pool + reader connection_limit) " + "writer pool + reader connection_limit + rust reader pool) " f"(workers={num_workers!r}, writer={writer_limit!r}, reader={reader_limit!r}); " "keep pods x that figure under max_connections minus superuser_reserved_connections" ) pools: Final = f"writer connection_limit {writer} + rust writer pool {rust_writer_limit}" - engines: Final = f"({pools} + reader connection_limit {reader})" if reader_limit is not None else f"({pools})" - per_pod: Final = workers * (writer + rust_writer_limit + reader) + engines: Final = ( + f"({pools} + reader connection_limit {reader} + rust reader pool {rust_reader_limit})" + if reader_limit is not None + else f"({pools})" + ) + rust_reader: Final = rust_reader_limit if reader_limit is not None else 0 + per_pod: Final = workers * (writer + rust_writer_limit + reader + rust_reader) return ( f"LiteLLM Proxy: Postgres connection budget per pod = {workers} worker(s) x {engines} = up to {per_pod} " "connections; keep pods x that figure under max_connections minus superuser_reserved_connections" diff --git a/litellm/proxy/db/native_database.py b/litellm/proxy/db/native_database.py index eb5f5ff40b8..430755652fb 100644 --- a/litellm/proxy/db/native_database.py +++ b/litellm/proxy/db/native_database.py @@ -24,7 +24,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from types import MappingProxyType -from typing import Final, Protocol +from typing import Final from cryptography.hazmat.primitives.serialization import Encoding, NoEncryption, PrivateFormat, pkcs12 @@ -210,10 +210,6 @@ async def resolve_native_endpoint(url: str, max_connections: int) -> NativeEndpo return native_endpoint(url, tls_offered=tls_offered, max_connections=max_connections) -class PasswordSink(Protocol): - def set_writer_password(self, password: str) -> None: ... - - @dataclass(frozen=True, slots=True) class TokenRotation: auth: DatabaseTokenAuth @@ -244,11 +240,23 @@ def writer_token_rotation(url: str) -> TokenRotation | None: return TokenRotation(auth=auth, endpoint=parse_iam_endpoint_from_url(url)) +def reader_token_rotation(url: str) -> TokenRotation | None: + """The token rotation the read replica needs, or None for a static password. + + The reader never goes through the in-container PgBouncer, so only the token toggles + decide, with the reader's own region, the same as its Prisma client. + """ + auth: Final = resolve_database_token_auth(read_replica=True) + if auth is None: + return None + return TokenRotation(auth=auth, endpoint=parse_iam_endpoint_from_url(url)) + + def _utcnow() -> datetime: return datetime.now(timezone.utc).replace(tzinfo=None) -async def _refresh(sink: PasswordSink, rotation: TokenRotation, now: Callable[[], datetime]) -> float: +async def _refresh(install: Callable[[str], None], rotation: TokenRotation, now: Callable[[], datetime]) -> float: try: token: Final = await asyncio.to_thread(rotation.mint) except Exception as error: # noqa: BLE001 # boto3 and azure-identity raise many types; any failure means retry @@ -259,12 +267,12 @@ async def _refresh(sink: PasswordSink, rotation: TokenRotation, now: Callable[[] TOKEN_RETRY_SECONDS, ) return TOKEN_RETRY_SECONDS - sink.set_writer_password(token) + install(token) return rotation.seconds_until_refresh(token, now()) -async def keep_writer_token_fresh( - sink: PasswordSink, +async def keep_token_fresh( + install: Callable[[str], None], rotation: TokenRotation, current_token: str, now: Callable[[], datetime] = _utcnow, @@ -278,4 +286,4 @@ async def keep_writer_token_fresh( delay = rotation.seconds_until_refresh(current_token, now()) # rebind-ok: each refresh schedules the next while True: await sleep(delay) - delay = await _refresh(sink, rotation, now) + delay = await _refresh(install, rotation, now) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 8b35e1ae860..1a1e443748d 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1526,6 +1526,7 @@ def run_server( if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None: from litellm.proxy.db.connection_budget import ( + RUST_READER_CONNECTION_LIMIT_ENV_VAR, RUST_WRITER_CONNECTION_LIMIT_ENV_VAR, ConnectionBudget, RustConnectionLimits, @@ -1616,7 +1617,17 @@ def run_server( # on the writer. Anything pinned on the replica URL wins, unlike the # writer where the config is applied on top. read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA") - if read_replica_url: + reader_budget: Final = ( + ConnectionBudget.split( + TypeAdapter(PositiveInt).validate_python( + _url_query_value(read_replica_url, "connection_limit") or writer_budget.total + ) + ) + if read_replica_url + else None + ) + if read_replica_url and reader_budget is not None: + os.environ[RUST_READER_CONNECTION_LIMIT_ENV_VAR] = str(reader_budget.rust) reader_options: Final[str] = _pg_options_with_timeouts( _url_query_value(read_replica_url, "options"), db_statement_timeout, @@ -1626,9 +1637,13 @@ def run_server( os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( add_missing_query_params( add_missing_query_params( - _with_query_value(read_replica_url, "options", reader_options) - if reader_options - else read_replica_url, + _with_query_value( + _with_query_value(read_replica_url, "options", reader_options) + if reader_options + else read_replica_url, + "connection_limit", + str(reader_budget.python), + ), reader_shareable_params(prisma_url_params), ), prisma_default_params, @@ -1645,6 +1660,7 @@ def run_server( ) if read_replica_url else None, + rust_reader_limit=reader_budget.rust if reader_budget is not None else 0, num_workers=f"{1 if run_hypercorn else num_workers}", ) ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 138b688a316..133523efa10 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1584,7 +1584,12 @@ async def _connect_to_count_stored_values(resources: AsyncExitStack) -> Supports async def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None: if client is None: return None - database: Final = await database_runtime.Database.open(client, os.environ["DATABASE_URL"], RustConnectionLimits()) + database: Final = await database_runtime.Database.open( + client, + os.environ["DATABASE_URL"], + os.environ.get("DATABASE_URL_READ_REPLICA") or None, + RustConnectionLimits(), + ) resources.push_async_callback(database.aclose) return database diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index 971a2b806a8..8b09ec8983d 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -1,6 +1,6 @@ import asyncio import urllib.parse -from collections.abc import AsyncGenerator, Coroutine +from collections.abc import AsyncGenerator, Callable, Coroutine from contextlib import asynccontextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Final @@ -8,7 +8,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.native_database import ( - keep_writer_token_fresh, + NativeEndpointSettings, + TokenRotation, + keep_token_fresh, + reader_token_rotation, resolve_native_endpoint, writer_token_rotation, ) @@ -18,6 +21,24 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +def _token_refresh( + install: Callable[[str], None], rotation: TokenRotation | None, url: str +) -> "tuple[asyncio.Task[None], ...]": + if rotation is None: + return () + current: Final = urllib.parse.unquote(urllib.parse.urlsplit(url).password or "") + return (asyncio.create_task(keep_token_fresh(install, rotation, current)),) + + +def _log_ignored_options(env_var: str, settings: NativeEndpointSettings) -> None: + if settings.ignored_options: + verbose_proxy_logger.error( + "Native database pool cannot send the %s options %s; only -c name=value settings carry over", + env_var, + settings.ignored_options, + ) + + @dataclass(frozen=True, slots=True) class Database: prisma: "PrismaClient" @@ -25,31 +46,32 @@ class Database: tasks: "tuple[asyncio.Task[None], ...]" @classmethod - async def open(cls, prisma: "PrismaClient", writer_url: str, rust_limits: RustConnectionLimits) -> "Database": + async def open( + cls, prisma: "PrismaClient", writer_url: str, reader_url: str | None, rust_limits: RustConnectionLimits + ) -> "Database": try: - settings: Final = await resolve_native_endpoint(writer_url, rust_limits.writer_connection_limit) - native: Final = await NativeDatabase.connect(settings.endpoint) + writer: Final = await resolve_native_endpoint(writer_url, rust_limits.writer_connection_limit) + reader: Final = ( + None + 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) except BaseException: await prisma.aclose() raise - if settings.ignored_options: - verbose_proxy_logger.error( - "Native database pool cannot send the DATABASE_URL options %s; only -c name=value settings carry over", - settings.ignored_options, - ) - rotation: Final = writer_token_rotation(writer_url) - token_refresh: Final = ( - () - if rotation is None - else ( - asyncio.create_task( - keep_writer_token_fresh( - native, rotation, urllib.parse.unquote(urllib.parse.urlsplit(writer_url).password or "") - ) - ), - ) + _log_ignored_options("DATABASE_URL", writer) + if reader is not None: + _log_ignored_options("DATABASE_URL_READ_REPLICA", reader) + writer_refresh: Final = _token_refresh( + native.set_writer_password, writer_token_rotation(writer_url), writer_url ) - return cls(prisma=prisma, native=native, tasks=token_refresh) + reader_refresh: Final = ( + () + 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)) async def aclose(self) -> None: for task in self.tasks: diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 4a5abac7afa..5faf1dea248 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -391,8 +391,11 @@ class NativeSpendCounters: @final class NativeDatabase: @staticmethod - def connect(writer: NativeDatabaseEndpoint) -> Future[NativeDatabase]: ... + def connect( + 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 close(self) -> Future[None]: ... @final diff --git a/tests/unit/proxy/db/test_connection_budget.py b/tests/unit/proxy/db/test_connection_budget.py index 3415f280480..e741e01358b 100644 --- a/tests/unit/proxy/db/test_connection_budget.py +++ b/tests/unit/proxy/db/test_connection_budget.py @@ -25,7 +25,10 @@ def test_split_takes_the_rust_reserve_out_of_the_total(total: int, rust_reserved assert (budget.python, budget.rust) == (python, rust) -def test_workers_read_the_rust_share_proxy_cli_exported(monkeypatch: pytest.MonkeyPatch) -> None: +def test_workers_read_the_rust_shares_proxy_cli_exported(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_DATABASE_RUST_WRITER_CONNECTION_LIMIT", "3") + monkeypatch.setenv("LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT", "1") - assert RustConnectionLimits().writer_connection_limit == 3 + limits: Final = RustConnectionLimits() + + assert (limits.writer_connection_limit, limits.reader_connection_limit) == (3, 1) diff --git a/tests/unit/proxy/db/test_db_url_settings.py b/tests/unit/proxy/db/test_db_url_settings.py index 019a11b3287..719c6080a1a 100644 --- a/tests/unit/proxy/db/test_db_url_settings.py +++ b/tests/unit/proxy/db/test_db_url_settings.py @@ -1285,13 +1285,13 @@ def test_reader_keeps_its_own_pinned_idle_lifetime(monkeypatch): "8", "8", "4", - "4 worker(s) x (writer connection_limit 8 + rust writer pool 2 + reader connection_limit 8) = up to 72 connections", + "4 worker(s) x (writer connection_limit 8 + rust writer pool 2 + reader connection_limit 8 + rust reader pool 2) = up to 80 connections", ), ( "8", "50", "4", - "4 worker(s) x (writer connection_limit 8 + rust writer pool 2 + reader connection_limit 50) = up to 240 connections", + "4 worker(s) x (writer connection_limit 8 + rust writer pool 2 + reader connection_limit 50 + rust reader pool 2) = up to 248 connections", ), ("8", None, "0", "1 worker(s) x (writer connection_limit 8 + rust writer pool 2) = up to 10 connections"), ], @@ -1303,7 +1303,11 @@ def test_connection_budget_message_sums_the_engines_each_worker_owns( from litellm.proxy.db.db_url_settings import postgres_connection_budget_message message: Final = postgres_connection_budget_message( - writer_limit=writer_limit, rust_writer_limit=2, reader_limit=reader_limit, num_workers=num_workers + writer_limit=writer_limit, + rust_writer_limit=2, + reader_limit=reader_limit, + rust_reader_limit=2, + num_workers=num_workers, ) assert expected in message assert "max_connections minus superuser_reserved_connections" in message @@ -1313,6 +1317,6 @@ def test_connection_budget_message_survives_an_unparseable_limit() -> None: from litellm.proxy.db.db_url_settings import postgres_connection_budget_message message: Final = postgres_connection_budget_message( - writer_limit="ten", rust_writer_limit=2, reader_limit=None, num_workers="2" + writer_limit="ten", rust_writer_limit=2, reader_limit=None, rust_reader_limit=0, num_workers="2" ) assert "writer='ten'" in message diff --git a/tests/unit/proxy/db/test_native_database.py b/tests/unit/proxy/db/test_native_database.py index dad278ea13c..edd4cf7fa2d 100644 --- a/tests/unit/proxy/db/test_native_database.py +++ b/tests/unit/proxy/db/test_native_database.py @@ -14,9 +14,10 @@ from pydantic import TypeAdapter from litellm.proxy.db.native_database import ( TokenRotation, - keep_writer_token_fresh, + keep_token_fresh, native_endpoint, needs_tls_probe, + reader_token_rotation, server_offers_tls, writer_token_rotation, ) @@ -284,6 +285,24 @@ def test_writer_token_rotation_follows_token_auth(monkeypatch: pytest.MonkeyPatc ) +def test_reader_token_rotation_uses_the_reader_region_even_when_the_writer_is_pooled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + url: Final = "postgresql://reader:tok@replica.internal:5433/proxy" + monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False) + monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "true") + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "false") + assert reader_token_rotation(url) is None + + monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") + monkeypatch.setenv("AWS_RDS_REGION", "us-west-2") + monkeypatch.setenv("AWS_RDS_READ_REPLICA_REGION", "eu-west-1") + assert reader_token_rotation(url) == TokenRotation( + auth=RdsIamTokenAuth(region="eu-west-1"), + endpoint=IAMEndpoint(host="replica.internal", port="5433", user="reader", name="proxy"), + ) + + def test_pooled_writer_needs_no_token_rotation(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "true") @@ -313,7 +332,7 @@ class _Sink: self._wanted: Final = wanted self.done: Final = asyncio.Event() - def set_writer_password(self, password: str) -> None: + def install(self, password: str) -> None: self.passwords = (*self.passwords, password) if len(self.passwords) >= self._wanted: self.done.set() @@ -344,7 +363,7 @@ def _minting(*outcomes: str | Exception) -> TokenRotation: async def _run_until(sink: _Sink, rotation: TokenRotation, schedule: _Schedule) -> None: task: Final = asyncio.create_task( - keep_writer_token_fresh(sink, rotation, "opaque", now=lambda: NOW, sleep=schedule.sleep) + keep_token_fresh(sink.install, rotation, "opaque", now=lambda: NOW, sleep=schedule.sleep) ) await asyncio.wait_for(sink.done.wait(), timeout=5) task.cancel() @@ -352,7 +371,7 @@ async def _run_until(sink: _Sink, rotation: TokenRotation, schedule: _Schedule) @pytest.mark.asyncio -async def test_keep_writer_token_fresh_installs_each_minted_token() -> None: +async def test_keep_token_fresh_installs_each_minted_token() -> None: sink: Final = _Sink(wanted=2) schedule: Final = _Schedule() @@ -363,7 +382,7 @@ async def test_keep_writer_token_fresh_installs_each_minted_token() -> None: @pytest.mark.asyncio -async def test_keep_writer_token_fresh_retries_a_failed_mint() -> None: +async def test_keep_token_fresh_retries_a_failed_mint() -> None: sink: Final = _Sink(wanted=1) schedule: Final = _Schedule() diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index d7bcdd94189..0b5099fc800 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -3654,6 +3654,7 @@ _CAPTURED_DB_ENV_VARS = ( "DIRECT_URL", "DATABASE_URL_READ_REPLICA", "LITELLM_DATABASE_RUST_WRITER_CONNECTION_LIMIT", + "LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT", ) @@ -3832,7 +3833,8 @@ class TestReadReplicaConnectionParams: def test_operator_pinned_replica_params_win(self, tmp_path): """The documented workaround (params pinned on the replica URL) must keep - working, so an operator who tuned the reader separately is not overridden. + working, so an operator who tuned the reader separately is not overridden. A pinned + connection_limit is the reader's whole budget, shared with the Rust reader pool. """ import yaml @@ -3855,7 +3857,7 @@ class TestReadReplicaConnectionParams: ) query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query) - assert query["connection_limit"] == ["50"] + assert int(query["connection_limit"][0]) + int(captured["LITELLM_DATABASE_RUST_READER_CONNECTION_LIMIT"]) == 50 assert query["pool_timeout"] == ["20"] def test_connection_budget_line_counts_the_limits_the_final_urls_carry( @@ -3876,7 +3878,7 @@ class TestReadReplicaConnectionParams: ) assert ( - "1 worker(s) x (writer connection_limit 1 + rust writer pool 2 + reader connection_limit 50) = up to 53 connections" + "1 worker(s) x (writer connection_limit 1 + rust writer pool 2 + reader connection_limit 48 + rust reader pool 2) = up to 53 connections" in capsys.readouterr().out ) diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py index 25a4973dcde..702742c3eac 100644 --- a/tests/unit/repositories/test_database.py +++ b/tests/unit/repositories/test_database.py @@ -8,6 +8,7 @@ from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR from litellm.repositories.database import Database WRITER_URL: Final = "postgresql://litellm:secret@db.invalid:5432/proxy?connection_limit=3" +READER_URL: Final = "postgresql://reader:secret@replica.invalid:5432/proxy?connection_limit=3" def _prisma(aclose: AsyncMock | None = None) -> MagicMock: @@ -25,7 +26,7 @@ def static_password(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.usefixtures("static_password") async def test_closing_the_database_closes_prisma() -> None: aclose: Final = AsyncMock() - database: Final = await Database.open(_prisma(aclose), WRITER_URL, RustConnectionLimits()) + database: Final = await Database.open(_prisma(aclose), WRITER_URL, None, RustConnectionLimits()) aclose.assert_not_awaited() await database.aclose() @@ -40,24 +41,24 @@ def _task_names(database: Database) -> tuple[str, ...]: @pytest.mark.asyncio @pytest.mark.usefixtures("static_password") async def test_static_password_runs_no_background_task() -> None: - database: Final = await Database.open(_prisma(), WRITER_URL, RustConnectionLimits()) + database: Final = await Database.open(_prisma(), WRITER_URL, READER_URL, RustConnectionLimits()) assert _task_names(database) == () await database.aclose() @pytest.mark.asyncio -async def test_closing_the_database_stops_the_token_refresh(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_closing_the_database_stops_both_token_refreshes(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False) monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR, raising=False) - database: Final = await Database.open(_prisma(), WRITER_URL, RustConnectionLimits()) - assert _task_names(database) == ("keep_writer_token_fresh",) + database: Final = await Database.open(_prisma(), WRITER_URL, READER_URL, RustConnectionLimits()) + assert _task_names(database) == ("keep_token_fresh", "keep_token_fresh") await database.aclose() assert all(task.done() for task in database.tasks) - assert database.tasks[0].cancelled() + assert database.tasks[0].cancelled() and database.tasks[1].cancelled() @pytest.mark.asyncio @@ -65,6 +66,17 @@ async def test_rejected_native_settings_still_close_prisma() -> None: aclose: Final = AsyncMock() with pytest.raises(ValueError, match="must name a host and a user"): - await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", RustConnectionLimits()) + await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", None, RustConnectionLimits()) + + aclose.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("static_password") +async def test_rejected_reader_settings_still_close_prisma() -> None: + aclose: Final = AsyncMock() + + with pytest.raises(ValueError, match="must name a host and a user"): + await Database.open(_prisma(aclose), WRITER_URL, "postgresql://replica.invalid/proxy", RustConnectionLimits()) aclose.assert_awaited_once()