diff --git a/litellm/proxy/db/connection_budget.py b/litellm/proxy/db/connection_budget.py index 71b839b91ef..26fd7d91f33 100644 --- a/litellm/proxy/db/connection_budget.py +++ b/litellm/proxy/db/connection_budget.py @@ -45,8 +45,12 @@ class ConnectionBudget: return None return cls.split(int(digits), rust_reserved) + @property + def prisma_limit(self) -> int | None: + """Prisma's share once the Rust pool takes some of the total, or None to leave Prisma's limit untouched.""" + return None if self.rust == 0 else self.python + def prisma_params(self, params: Mapping[str, str | int | float]) -> Mapping[str, str | int | float]: """``params`` with Prisma's share as its ``connection_limit``, untouched while the Rust pool takes none.""" - if self.rust == 0: - return params - return MappingProxyType({**params, "connection_limit": self.python}) + limit: Final = self.prisma_limit + return params if limit is None else MappingProxyType({**params, "connection_limit": limit}) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 231904d4d78..217c2dca8c4 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -268,14 +268,14 @@ def prisma_default_params( 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``) when it is above 0. The limits are read off the final URLs, + ``DATABASE_URL_READ_REPLICA`` is set) plus a Rust pool beside each whose + share (``connection_budget``) is above 0. 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 @@ -289,17 +289,20 @@ def postgres_connection_budget_message( except ValueError: return ( "LiteLLM Proxy: Postgres connection budget per pod = workers x (writer connection_limit + " - f"{'rust writer pool + ' if rust_writer_limit else ''}reader connection_limit) " + f"{'rust writer pool + ' if rust_writer_limit else ''}reader connection_limit" + f"{' + rust reader pool' if rust_reader_limit else ''}) " 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" ) + rust_reader: Final = rust_reader_limit if reader_limit is not None else 0 terms: Final = ( f"writer connection_limit {writer}", *((f"rust writer pool {rust_writer_limit}",) if rust_writer_limit else ()), *((f"reader connection_limit {reader}",) if reader_limit is not None else ()), + *((f"rust reader pool {rust_reader}",) if rust_reader else ()), ) engines: Final = terms[0] if len(terms) == 1 else f"({' + '.join(terms)})" - per_pod: Final = workers * (writer + rust_writer_limit + reader) + 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 cef24e441c1..a56f6d9f00d 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1718,6 +1718,14 @@ 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") + reader_budget: Final = ( + ConnectionBudget.of( + _url_query_value(read_replica_url, "connection_limit") or connection_url_params["connection_limit"] + ) + if read_replica_url + else None + ) + reader_prisma_limit: Final = None if reader_budget is None else reader_budget.prisma_limit if read_replica_url: reader_options: Final[str] = _pg_options_with_timeouts( _url_query_value(read_replica_url, "options"), @@ -1725,12 +1733,17 @@ def run_server( db_lock_timeout, db_idle_in_transaction_timeout, ) + reader_url: Final = ( + _with_query_value(read_replica_url, "options", reader_options) + if reader_options + else read_replica_url + ) 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, + reader_url + if reader_prisma_limit is None + else _with_query_value(reader_url, "connection_limit", str(reader_prisma_limit)), reader_shareable_params(prisma_url_params), ), prisma_defaults, @@ -1747,6 +1760,7 @@ def run_server( ) if read_replica_url else None, + rust_reader_limit=0 if reader_budget is None else reader_budget.rust, 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 65853f61549..cdcca47284a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1603,7 +1603,10 @@ async def _own_database(resources: AsyncExitStack, client: PrismaClient | None) if client is None: return None database: Final = await database_runtime.Database.open( - client, os.environ["DATABASE_URL"], RUST_RESERVED_CONNECTIONS + client, + os.environ["DATABASE_URL"], + os.environ.get("DATABASE_URL_READ_REPLICA") or None, + RUST_RESERVED_CONNECTIONS, ) resources.push_async_callback(database.aclose) return database diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index b577ba2bd67..0f4dd33a2f8 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -1,13 +1,16 @@ 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._logging import verbose_proxy_logger 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, ) @@ -17,6 +20,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" @@ -24,33 +45,32 @@ class Database: tasks: "tuple[asyncio.Task[None], ...]" = () @classmethod - async def open(cls, prisma: "PrismaClient", writer_url: str, native_connections: int) -> "Database": + async def open( + cls, prisma: "PrismaClient", writer_url: str, reader_url: str | None, native_connections: int + ) -> "Database": if native_connections == 0: return cls(prisma=prisma, native=None) try: - settings: Final = await resolve_native_endpoint(writer_url, native_connections) - native: Final = await NativeDatabase.connect(settings.endpoint) + writer: Final = await resolve_native_endpoint(writer_url, native_connections) + reader: Final = ( + None if reader_url is None else await resolve_native_endpoint(reader_url, native_connections) + ) + 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/Cargo.lock b/litellm/rust/Cargo.lock index ddd0eaf0870..a4672290e1e 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -3767,6 +3767,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 1fedf9c9a2b..2ffecedd789 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 504706def5f..f8bff4d2569 100644 --- a/litellm/rust/crates/python-bridge/Cargo.toml +++ b/litellm/rust/crates/python-bridge/Cargo.toml @@ -72,7 +72,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/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 964862db17a..5e3583f2279 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -563,8 +563,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/_support/native_database.py b/tests/_support/native_database.py index bade92d79f8..8e7ea8feb18 100644 --- a/tests/_support/native_database.py +++ b/tests/_support/native_database.py @@ -8,7 +8,9 @@ _NEVER_DIALED_URL: Final = "postgresql://litellm:secret@db.invalid:5432/proxy" async def connect_native_database() -> NativeDatabase: - return await NativeDatabase.connect(native_endpoint(_NEVER_DIALED_URL, tls_offered=True, max_connections=1).endpoint) + return await NativeDatabase.connect( + native_endpoint(_NEVER_DIALED_URL, tls_offered=True, max_connections=1).endpoint, None + ) def native_database() -> NativeDatabase: diff --git a/tests/unit/proxy/db/test_db_url_settings.py b/tests/unit/proxy/db/test_db_url_settings.py index a42894efa89..abd4ed14b2e 100644 --- a/tests/unit/proxy/db/test_db_url_settings.py +++ b/tests/unit/proxy/db/test_db_url_settings.py @@ -1296,14 +1296,14 @@ def test_reader_keeps_its_own_pinned_idle_lifetime(monkeypatch): 2, "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", 2, "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", 2, None, "0", "1 worker(s) x (writer connection_limit 8 + rust writer pool 2) = up to 10 connections"), ], @@ -1325,6 +1325,7 @@ def test_connection_budget_message_sums_the_engines_each_worker_owns( writer_limit=writer_limit, rust_writer_limit=rust_writer_limit, reader_limit=reader_limit, + rust_reader_limit=rust_writer_limit, num_workers=num_workers, ) assert expected in message @@ -1335,7 +1336,7 @@ 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 4e2ddfe859d..f964782939d 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -3650,6 +3650,23 @@ class TestReadReplicaConnectionParams: assert query["connection_limit"] == ["50"] assert query["pool_timeout"] == ["20"] + @pytest.mark.parametrize("pinned_limit", ["0", "abc"]) + def test_replica_limit_that_is_not_a_positive_integer_is_left_to_prisma( + self, tmp_path: Path, pinned_limit: str + ) -> None: + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {}})) + + captured: Final = _CAPTURED_URLS.validate_python( + _run_server_and_capture_urls( + str(config_path), + read_replica_url=f"postgresql://t:t@reader:5432/t?connection_limit={pinned_limit}", + ) + ) + + query: Final = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query) + assert query["connection_limit"] == [pinned_limit] + def test_connection_budget_line_counts_the_limits_the_final_urls_carry( self, tmp_path: Path, diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py index 798bdc39325..2497c9237e9 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" NATIVE_CONNECTIONS: Final = 2 @@ -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, NATIVE_CONNECTIONS) + database: Final = await Database.open(_prisma(aclose), WRITER_URL, None, NATIVE_CONNECTIONS) 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, NATIVE_CONNECTIONS) + database: Final = await Database.open(_prisma(), WRITER_URL, READER_URL, NATIVE_CONNECTIONS) 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, NATIVE_CONNECTIONS) - assert _task_names(database) == ("keep_writer_token_fresh",) + database: Final = await Database.open(_prisma(), WRITER_URL, READER_URL, NATIVE_CONNECTIONS) + 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,7 +66,18 @@ 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", NATIVE_CONNECTIONS) + await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", None, NATIVE_CONNECTIONS) + + 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", NATIVE_CONNECTIONS) aclose.assert_awaited_once() @@ -75,7 +87,9 @@ async def test_zero_native_connections_open_no_native_pool(monkeypatch: pytest.M monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true") aclose: Final = AsyncMock() - database: Final = await Database.open(_prisma(aclose), "postgresql://db.invalid/proxy", 0) + database: Final = await Database.open( + _prisma(aclose), "postgresql://db.invalid/proxy", "postgresql://replica.invalid/proxy", 0 + ) assert (database.native, database.tasks) == (None, ()) await database.aclose()