Merge 33a362c425 into 786a64ea5f

litellm-rust/Cargo.lock

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

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 |

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

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

@@ -109,14 +109,14 @@ fn seconds(value: Option<f64>) -> Result<Option<Duration>, 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<Endpoint, Box<dyn StdError>> {
-        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();

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,

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<Reader>,
+}
+
+impl Pools {
+    pub fn new(writer: &Endpoint, reader: Option<&Endpoint>) -> Result<Self, Error> {
+        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<PoolConnection<Postgres>, 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<bool, sqlx::Error> {
+        #[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<dyn StdError>> {
+        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<dyn StdError>> {
+        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<dyn StdError>> {
+        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<dyn StdError>> {
+        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(())
+    }
+}

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<PgPool, Box<dyn StdError>> {
     let database_url = std::env::var("DATABASE_URL")?;
     Ok(PgPool::connect(&database_url).await?)
 }
+
+pub fn database_endpoint() -> Result<Endpoint, Box<dyn StdError>> {
+    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(),
+    })
+}

litellm-rust/crates/python-bridge/Cargo.toml

@@ -61,7 +61,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]

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<Endpoint>,
+    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<Endpoint>,
     ) -> PyResult<Bound<'_, PyAny>> {
         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<Bound<'py, PyAny>> {
-        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,

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"

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)

litellm/proxy/proxy_server.py

@@ -1558,7 +1558,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
 

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:

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

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"),
     ],

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()
 

tests/unit/proxy/test_proxy_cli.py

@@ -3500,7 +3500,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
         )
 

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()