From 33a362c4252b5fe9a29f05dfafbe19778f1638b8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 3 Nov 2026 15:09:27 +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() and read(); read() uses the replica, falls back to the writer when no replica is configured or a reader connection cannot be acquired, and skips the replica for 30 seconds after a failure so an outage costs one acquire timeout. The connection budget counts the native reader pool --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/db/AGENTS.md | 6 +- litellm-rust/crates/db/Cargo.toml | 1 + litellm-rust/crates/db/src/endpoint.rs | 25 +-- litellm-rust/crates/db/src/lib.rs | 2 + litellm-rust/crates/db/src/pools.rs | 172 +++++++++++++++++++ litellm-rust/crates/db/src/test_support.rs | 26 ++- litellm-rust/crates/python-bridge/Cargo.toml | 1 - litellm-rust/crates/python-bridge/src/db.rs | 31 +++- litellm/proxy/db/db_url_settings.py | 16 +- litellm/proxy/db/native_database.py | 28 +-- litellm/proxy/proxy_server.py | 4 +- litellm/repositories/database.py | 51 ++++-- litellm/rust_bridge/_native.pyi | 5 +- tests/unit/proxy/db/test_db_url_settings.py | 4 +- tests/unit/proxy/db/test_native_database.py | 29 +++- tests/unit/proxy/test_proxy_cli.py | 2 +- tests/unit/repositories/test_database.py | 31 +++- 18 files changed, 347 insertions(+), 88 deletions(-) create mode 100644 litellm-rust/crates/db/src/pools.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 3a627f6f80f..401ed40a23b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3754,6 +3754,7 @@ name = "litellm-db" version = "0.1.0" dependencies = [ "chrono", + "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 c2965d792f9..4e9ae38452e 100644 --- a/litellm-rust/crates/db/AGENTS.md +++ b/litellm-rust/crates/db/AGENTS.md @@ -36,10 +36,10 @@ Python and SQLx run side by side until every repository operation is native, so | `options` (`statement_timeout`, `lock_timeout`, `idle_in_transaction_session_timeout`) | `proxy_cli.py` | `Endpoint::options`, sent as startup `-c` settings | | `schema` | URL param | `search_path` startup setting when not `public`; a custom schema behind PgBouncer is rejected by PgBouncer | | `pgbouncer=true`, `DATABASE_DISABLE_PREPARED_STATEMENTS` | `db_url_settings.py` | `statement_cache: false` | -| 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 native writer pool | -| Read replica routing | `routing_prisma_wrapper.py` | Planned: optional reader pool with explicit routing | +| 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 | | 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 b59d9747925..a0152f31b1c 100644 --- a/litellm-rust/crates/db/Cargo.toml +++ b/litellm-rust/crates/db/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] chrono.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/endpoint.rs b/litellm-rust/crates/db/src/endpoint.rs index 73534a6f0d3..1f4911f940d 100644 --- a/litellm-rust/crates/db/src/endpoint.rs +++ b/litellm-rust/crates/db/src/endpoint.rs @@ -109,14 +109,14 @@ fn seconds(value: Option) -> Result, Error> { #[cfg(test)] mod tests { - use std::{error::Error as StdError, str::FromStr, time::Duration}; + use std::{error::Error as StdError, time::Duration}; use litellm_secret_value::SecretValue; use rstest::rstest; - use sqlx::postgres::{PgConnectOptions, PgSslMode}; + use sqlx::postgres::PgSslMode; use super::{Endpoint, PoolLimits, SslMode}; - use crate::Error; + use crate::{Error, test_support::database_endpoint}; fn endpoint() -> Endpoint { Endpoint { @@ -133,25 +133,6 @@ mod tests { } } - fn database_endpoint() -> Result> { - let url = PgConnectOptions::from_str(&std::env::var("DATABASE_URL")?)?; - let password = url::Url::parse(&std::env::var("DATABASE_URL")?)? - .password() - .map(SecretValue::new); - Ok(Endpoint { - host: url.get_host().to_owned(), - port: url.get_port(), - username: url.get_username().to_owned(), - password, - database: url.get_database().unwrap_or("postgres").to_owned(), - ssl_mode: SslMode::Prefer, - ssl_root_cert: None, - options: Vec::new(), - statement_cache: true, - pool: PoolLimits::default(), - }) - } - #[test] fn maps_endpoint_to_connect_options() { let options = endpoint().connect_options(); diff --git a/litellm-rust/crates/db/src/lib.rs b/litellm-rust/crates/db/src/lib.rs index b71a6a13548..972b9312cd2 100644 --- a/litellm-rust/crates/db/src/lib.rs +++ b/litellm-rust/crates/db/src/lib.rs @@ -11,6 +11,7 @@ mod end_user; mod endpoint; mod error; mod failure; +mod pools; mod spend; #[cfg(test)] mod test_support; @@ -20,6 +21,7 @@ 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 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..b077e52af0e --- /dev/null +++ b/litellm-rust/crates/db/src/pools.rs @@ -0,0 +1,172 @@ +use std::time::{Duration, Instant}; + +use litellm_circuit_breaker::{Breaker, Config, Failure}; +use sqlx::{PgPool, Postgres, pool::PoolConnection}; + +use crate::{Endpoint, Error}; + +const READER_RECOVERY: Duration = Duration::from_secs(30); + +#[derive(Clone)] +struct Reader { + pool: PgPool, + breaker: Breaker, +} + +#[derive(Clone)] +pub struct Pools { + writer: PgPool, + reader: Option, +} + +impl Pools { + pub fn new(writer: &Endpoint, reader: Option<&Endpoint>) -> Result { + Ok(Self { + writer: writer.pool()?, + reader: reader + .map(|endpoint| { + endpoint.pool().map(|pool| Reader { + pool, + breaker: Breaker::new( + Config { + failure_threshold: 1, + recovery_timeout: READER_RECOVERY, + ..Config::default() + }, + Box::new(|_| {}), + ), + }) + }) + .transpose()?, + }) + } + + pub fn writer(&self) -> &PgPool { + &self.writer + } + + pub fn reader_pool(&self) -> Option<&PgPool> { + self.reader.as_ref().map(|reader| &reader.pool) + } + + pub async fn read(&self) -> Result, Error> { + let Some(reader) = &self.reader else { + return Ok(self.writer.acquire().await?); + }; + let Ok(admission) = reader.breaker.admit(Instant::now()) else { + return Ok(self.writer.acquire().await?); + }; + match reader.pool.acquire().await { + Ok(connection) => { + admission.succeeded(); + Ok(connection) + } + Err(_) => { + admission.failed(Failure::Hard, Instant::now()); + Ok(self.writer.acquire().await?) + } + } + } + + pub async fn close(&self) { + self.writer.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() -> 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.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 reads_use_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?); + 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(()) + } +} diff --git a/litellm-rust/crates/db/src/test_support.rs b/litellm-rust/crates/db/src/test_support.rs index 992be4a4447..99bda232831 100644 --- a/litellm-rust/crates/db/src/test_support.rs +++ b/litellm-rust/crates/db/src/test_support.rs @@ -1,10 +1,32 @@ -use std::error::Error as StdError; +use std::{error::Error as StdError, str::FromStr}; +use litellm_secret_value::SecretValue; use rstest::fixture; -use sqlx::PgPool; +use sqlx::{PgPool, postgres::PgConnectOptions}; + +use crate::{Endpoint, PoolLimits, SslMode}; #[fixture] pub async fn database_pool() -> Result> { let database_url = std::env::var("DATABASE_URL")?; Ok(PgPool::connect(&database_url).await?) } + +pub fn database_endpoint() -> Result> { + let database_url = std::env::var("DATABASE_URL")?; + let options = PgConnectOptions::from_str(&database_url)?; + Ok(Endpoint { + host: options.get_host().to_owned(), + port: options.get_port(), + username: options.get_username().to_owned(), + password: url::Url::parse(&database_url)? + .password() + .map(SecretValue::new), + database: options.get_database().unwrap_or("postgres").to_owned(), + ssl_mode: SslMode::Prefer, + ssl_root_cert: None, + options: Vec::new(), + statement_cache: true, + pool: PoolLimits::default(), + }) +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index f386ad25402..81c7726d6d8 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -62,7 +62,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/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 0b214d9c6de..01033b4e01a 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -259,8 +259,8 @@ def postgres_connection_budget_message(writer_limit: str, reader_limit: str | No 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) and a native writer pool sized from the - writer URL. The limits are read off the final URLs, + ``DATABASE_URL_READ_REPLICA`` is set) and a native pool beside each, sized from + the same URL. 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 @@ -274,13 +274,17 @@ def postgres_connection_budget_message(writer_limit: str, reader_limit: str | No except ValueError: return ( "LiteLLM Proxy: Postgres connection budget per pod = workers x (writer connection_limit + native " - "writer pool + reader connection_limit) " + "writer pool + reader connection_limit + native 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} + native writer pool {writer}" - engines: Final = f"({pools} + reader connection_limit {reader})" if reader_limit is not None else f"({pools})" - per_pod: Final = workers * (2 * writer + reader) + writer_pools: Final = f"writer connection_limit {writer} + native writer pool {writer}" + engines: Final = ( + f"({writer_pools} + reader connection_limit {reader} + native reader pool {reader})" + if reader_limit is not None + else f"({writer_pools})" + ) + per_pod: Final = workers * 2 * (writer + 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 9406dfdd647..5ceebb50630 100644 --- a/litellm/proxy/db/native_database.py +++ b/litellm/proxy/db/native_database.py @@ -13,7 +13,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType -from typing import Final, Protocol +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.proxy.db.pgbouncer import database_url_is_pooled, libpq_option_settings @@ -95,10 +95,6 @@ def native_endpoint(url: str) -> NativeDatabaseEndpoint: ) -class PasswordSink(Protocol): - def set_writer_password(self, password: str) -> None: ... - - @dataclass(frozen=True, slots=True) class TokenRotation: auth: DatabaseTokenAuth @@ -129,11 +125,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 @@ -144,12 +152,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, @@ -163,4 +171,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_server.py b/litellm/proxy/proxy_server.py index 4a4a324f2fe..ac29201ab94 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1520,7 +1520,9 @@ 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"]) + database: Final = await database_runtime.Database.open( + client, os.environ["DATABASE_URL"], os.environ.get("DATABASE_URL_READ_REPLICA") or None + ) resources.push_async_callback(database.aclose) return database diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index 6957b91ff56..4d9f0adabeb 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -1,46 +1,63 @@ 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 -from litellm.proxy.db.native_database import keep_writer_token_fresh, native_endpoint, writer_token_rotation +from litellm.proxy.db.native_database import ( + TokenRotation, + keep_token_fresh, + native_endpoint, + reader_token_rotation, + writer_token_rotation, +) from litellm.rust_bridge._native import NativeDatabase if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +def _token_refresh( + install: Callable[[str], None], rotation: TokenRotation | None, url: str +) -> "asyncio.Task[None] | None": + if rotation is None: + return None + current: Final = urllib.parse.unquote(urllib.parse.urlsplit(url).password or "") + return asyncio.create_task(keep_token_fresh(install, rotation, current)) + + @dataclass(frozen=True, slots=True) class Database: prisma: "PrismaClient" native: NativeDatabase - token_refresh: "asyncio.Task[None] | None" + token_refreshes: "tuple[asyncio.Task[None], ...]" @classmethod - async def open(cls, prisma: "PrismaClient", writer_url: str) -> "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)) + native: Final = await NativeDatabase.connect( + native_endpoint(writer_url), None if reader_url is None else native_endpoint(reader_url) + ) except BaseException: await prisma.aclose() raise - rotation: Final = writer_token_rotation(writer_url) - token_refresh: Final = ( + refreshes: Final = ( + _token_refresh(native.set_writer_password, writer_token_rotation(writer_url), writer_url), None - if rotation is None - else asyncio.create_task( - keep_writer_token_fresh( - native, rotation, urllib.parse.unquote(urllib.parse.urlsplit(writer_url).password or "") - ) - ) + 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, + token_refreshes=tuple(task for task in refreshes if task is not None), ) - return cls(prisma=prisma, native=native, token_refresh=token_refresh) async def aclose(self) -> None: - if self.token_refresh is not None: - self.token_refresh.cancel() - await asyncio.gather(self.token_refresh, return_exceptions=True) + for task in self.token_refreshes: + task.cancel() + await asyncio.gather(*self.token_refreshes, return_exceptions=True) try: await self.prisma.aclose() finally: diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 9925b0a373c..fce6fed1fe7 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -30,8 +30,11 @@ class NativeClickHouseSpendStorage: @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_db_url_settings.py b/tests/unit/proxy/db/test_db_url_settings.py index eb8fe98834b..583edc6c7f3 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): "10", "10", "4", - "4 worker(s) x (writer connection_limit 10 + native writer pool 10 + reader connection_limit 10) = up to 120 connections", + "4 worker(s) x (writer connection_limit 10 + native writer pool 10 + reader connection_limit 10 + native reader pool 10) = up to 160 connections", ), ( "10", "50", "4", - "4 worker(s) x (writer connection_limit 10 + native writer pool 10 + reader connection_limit 50) = up to 280 connections", + "4 worker(s) x (writer connection_limit 10 + native writer pool 10 + reader connection_limit 50 + native reader pool 50) = up to 480 connections", ), ("10", None, "0", "1 worker(s) x (writer connection_limit 10 + native writer pool 10) = up to 20 connections"), ], diff --git a/tests/unit/proxy/db/test_native_database.py b/tests/unit/proxy/db/test_native_database.py index 7688d848729..71d3ee579ec 100644 --- a/tests/unit/proxy/db/test_native_database.py +++ b/tests/unit/proxy/db/test_native_database.py @@ -6,8 +6,9 @@ import pytest from litellm.proxy.db.native_database import ( TokenRotation, - keep_writer_token_fresh, + keep_token_fresh, native_endpoint, + reader_token_rotation, writer_token_rotation, ) from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR @@ -127,6 +128,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") @@ -156,7 +175,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() @@ -187,7 +206,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() @@ -195,7 +214,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() @@ -206,7 +225,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 0d636ea2deb..2fe63353403 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -3392,7 +3392,7 @@ class TestReadReplicaConnectionParams: ) assert ( - "1 worker(s) x (writer connection_limit 3 + native writer pool 3 + reader connection_limit 50) = up to 56 connections" + "1 worker(s) x (writer connection_limit 3 + native writer pool 3 + reader connection_limit 50 + native reader pool 50) = up to 106 connections" in capsys.readouterr().out ) diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py index ab3b251badb..7d33a71dd3f 100644 --- a/tests/unit/repositories/test_database.py +++ b/tests/unit/repositories/test_database.py @@ -7,6 +7,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: @@ -24,7 +25,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) + database: Final = await Database.open(_prisma(aclose), WRITER_URL, None) aclose.assert_not_awaited() await database.aclose() @@ -35,24 +36,25 @@ async def test_closing_the_database_closes_prisma() -> None: @pytest.mark.asyncio @pytest.mark.usefixtures("static_password") async def test_static_password_needs_no_token_refresh() -> None: - database: Final = await Database.open(_prisma(), WRITER_URL) + database: Final = await Database.open(_prisma(), WRITER_URL, None) - assert database.token_refresh is None + assert database.token_refreshes == () 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_writer_and_reader_token_refresh(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) - token_refresh: Final = database.token_refresh - assert token_refresh is not None and not token_refresh.done() + database: Final = await Database.open(_prisma(), WRITER_URL, READER_URL) + refreshes: Final = database.token_refreshes + assert len(refreshes) == 2 + assert not any(task.done() for task in refreshes) await database.aclose() - assert token_refresh.cancelled() + assert all(task.cancelled() for task in refreshes) @pytest.mark.asyncio @@ -60,6 +62,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") + await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", None) + + 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") aclose.assert_awaited_once()