fix(db): apply the configured connection params to the read replica URL (#37691)

The read replica never received the operator's DB pool settings, so its
Prisma pool fell back to `num_physical_cpus * 2 + 1` and the configured cap
was not enforced. Both startup paths now pass the same params to the reader:
the CLI, and the componentized entrypoints that go through
`DatabaseURLSettings.apply_to_env`.

Only pool and timeout params are inherited, through a single allowlist both
paths share. Anything that decides which tables a query resolves against
stays on the writer, including entries smuggled in through
`database_extra_connection_params`, so a writer `search_path` cannot repoint
reader queries. Params the operator pinned on the replica URL still win.

Co-authored-by: Yassin Kortam <yassin.kortam@gmail.com>
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: 3 people

litellm/proxy/db/db_url_settings.py

@@ -29,12 +29,16 @@
 token auth is gated on the same global toggle as the writer: the chart only
 emits the reader token env vars when the writer also uses token auth.
 Reader-side fields fall back to the writer's user / name / schema / port /
-password when their ``*_READ_REPLICA`` counterpart is unset.
+password when their ``*_READ_REPLICA`` counterpart is unset, and to the
+writer's connection params (pool size, timeouts, pgbouncer mode) for the
+ones the reader URL does not pin itself.
 """
 
 import os
 import urllib.parse
+from collections.abc import Mapping
 from functools import partial
+from types import MappingProxyType
 from typing import Annotated, Final, cast
 
 from pydantic import AliasChoices, BeforeValidator, Field
@@ -62,6 +66,51 @@
 _MISSING_SCHEME: Final = "<missing scheme>"
 
 
+# An allowlist, deliberately not a denylist: only these pool and timeout params
+# follow the writer to the read replica, so nothing that decides which tables a
+# query resolves against (``schema``, or a ``search_path`` inside ``options``)
+# can ever repoint the reader. Without them the reader pool silently falls back
+# to Prisma's default size.
+CONNECTION_PARAM_KEYS: Final[frozenset[str]] = frozenset(
+    {
+        "connection_limit",
+        "pool_timeout",
+        "connect_timeout",
+        "socket_timeout",
+        "pgbouncer",
+    }
+)
+
+
+def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) -> str:
+    """Return ``url`` with the ``params`` it does not already carry appended.
+
+    Params the operator pinned on the URL win, so a hand-tuned replica URL keeps
+    its values. Returns the URL untouched when there is nothing to add, leaving
+    its existing encoding alone.
+    """
+    parsed: Final = urllib.parse.urlsplit(url)
+    existing: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
+    pinned: Final = frozenset(key for key, _ in existing)
+    additions: Final = tuple((key, str(value)) for key, value in params.items() if key not in pinned)
+    if not additions:
+        return url
+    query: Final = urllib.parse.urlencode(existing + additions)
+    return urllib.parse.urlunsplit(parsed._replace(query=query))
+
+
+def reader_shareable_params(params: Mapping[str, str | int | float]) -> Mapping[str, str | int | float]:
+    """Return the subset of ``params`` the read replica is allowed to inherit."""
+    return MappingProxyType({key: value for key, value in params.items() if key in CONNECTION_PARAM_KEYS})
+
+
+def connection_params_from_url(url: str) -> Mapping[str, str | int | float]:
+    """Return the connection params on ``url`` that the read replica shares."""
+    return reader_shareable_params(
+        MappingProxyType({key: value for key, value in urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query)})
+    )
+
+
 def unsupported_db_scheme(database_url: str) -> str | None:
     """Return the connection URL scheme when it is not PostgreSQL, else None.
 
@@ -326,8 +375,14 @@ def apply_to_env(self) -> bool:
         self._raise_for_unsupported_scheme()
         wrote_writer: Final = self.apply_writer_url_to_env()
 
-        reader_url: Final = self.build_reader_url()
+        # The reader inherits the writer's connection params (pool size, timeouts,
+        # pgbouncer mode). Without this the reader pool ignores the configured cap
+        # and falls back to Prisma's `num_physical_cpus * 2 + 1` default.
+        reader_url: Final = self.build_reader_url() or self.database_url_read_replica
         if reader_url is not None:
-            os.environ["DATABASE_URL_READ_REPLICA"] = reader_url
+            os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(
+                reader_url,
+                connection_params_from_url(os.environ.get("DATABASE_URL", "")),
+            )
 
         return wrote_writer

litellm/proxy/proxy_cli.py

@@ -1224,6 +1224,8 @@ def run_server(
 
         if os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None:
             from litellm.proxy.db.db_url_settings import (
+                add_missing_query_params,
+                reader_shareable_params,
                 unsupported_db_scheme,
                 unsupported_db_scheme_message,
             )
@@ -1273,6 +1275,24 @@ def run_server(
                     database_url = os.getenv("DIRECT_URL")
                     modified_url = append_query_params(database_url, connection_url_params)
                     os.environ["DIRECT_URL"] = modified_url
+                # The reader pool is a real pool against the same configured cap, so it
+                # gets the allowlisted pool params. Schema-affecting ones, including any
+                # the operator smuggled in through database_extra_connection_params, stay
+                # on the writer. Anything pinned on the replica URL wins, unlike the
+                # writer where the config is applied on top.
+                read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA")
+                if read_replica_url:
+                    reader_options: Final[str] = _pg_options_with_timeouts(
+                        _url_query_value(read_replica_url, "options"),
+                        db_statement_timeout,
+                        db_lock_timeout,
+                    )
+                    os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params(
+                        _with_query_value(read_replica_url, "options", reader_options)
+                        if reader_options
+                        else read_replica_url,
+                        reader_shareable_params(connection_url_params),
+                    )
                 subprocess.run(["prisma"], capture_output=True)
                 is_prisma_runnable = True
             except FileNotFoundError:

tests/test_litellm/proxy/db/test_db_url_settings.py

@@ -12,6 +12,7 @@
 """
 
 import os
+import urllib.parse
 from unittest.mock import patch
 
 import pytest
@@ -524,6 +525,137 @@ def test_apply_to_env_accepts_pinned_postgres(monkeypatch):
     assert _apply() is False
 
 
+# ---------------------------------------------------------------------------
+# Connection params on the read replica
+# ---------------------------------------------------------------------------
+
+
+def test_reader_inherits_writer_connection_params(monkeypatch):
+    """The reader is a second pool: without the writer's params it sizes itself
+    from Prisma's default and the operator's cap is not enforced."""
+    monkeypatch.setenv(
+        "DATABASE_URL",
+        "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&pool_timeout=20&pgbouncer=true",
+    )
+    monkeypatch.setenv(
+        "DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db"
+    )
+
+    _apply()
+
+    query = urllib.parse.parse_qs(
+        urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query
+    )
+    assert query["connection_limit"] == ["3"]
+    assert query["pool_timeout"] == ["20"]
+    assert query["pgbouncer"] == ["true"]
+
+
+def test_reader_keeps_its_own_pinned_connection_params(monkeypatch):
+    monkeypatch.setenv(
+        "DATABASE_URL",
+        "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&pool_timeout=20",
+    )
+    monkeypatch.setenv(
+        "DATABASE_URL_READ_REPLICA",
+        "postgresql://u:p@reader.example.com:5432/db?connection_limit=50",
+    )
+
+    _apply()
+
+    query = urllib.parse.parse_qs(
+        urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query
+    )
+    assert query["connection_limit"] == ["50"]
+    assert query["pool_timeout"] == ["20"]
+
+
+def test_assembled_reader_url_inherits_writer_connection_params(monkeypatch):
+    """A reader assembled from the discrete DATABASE_*_READ_REPLICA vars must
+    carry the params too, and must not inherit the writer's schema."""
+    monkeypatch.setenv(
+        "DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&schema=writer_schema"
+    )
+    monkeypatch.setenv("DATABASE_USER", "litellm")
+    monkeypatch.setenv("DATABASE_NAME", "litellm_db")
+    monkeypatch.setenv("DATABASE_PASSWORD", "s3cr3t")
+    monkeypatch.setenv("DATABASE_HOST_READ_REPLICA", "reader.example.com")
+
+    _apply()
+
+    reader_url = os.environ["DATABASE_URL_READ_REPLICA"]
+    assert reader_url.startswith("postgresql://litellm:s3cr3t@reader.example.com:5432/litellm_db?")
+    query = urllib.parse.parse_qs(urllib.parse.urlsplit(reader_url).query)
+    assert query["connection_limit"] == ["3"]
+    assert "schema" not in query
+
+
+def test_reader_does_not_inherit_writer_options(monkeypatch):
+    """A writer search_path must not follow the reader, or reader queries resolve
+    against the wrong schema."""
+    monkeypatch.setenv(
+        "DATABASE_URL",
+        "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&options=-c%20search_path%3Dwriter_schema",
+    )
+    monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db")
+
+    _apply()
+
+    query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query)
+    assert query["connection_limit"] == ["3"]
+    assert "options" not in query
+
+
+def test_reader_does_not_inherit_an_unvetted_writer_param(monkeypatch):
+    """Inheritance is an allowlist, so a param nobody vetted for the reader stays
+    on the writer. Flipping this to a denylist would let the next schema-affecting
+    param leak through by default."""
+    monkeypatch.setenv(
+        "DATABASE_URL",
+        "postgresql://u:p@writer.example.com:5432/db?connection_limit=3&application_name=writer&novel_param=x",
+    )
+    monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db")
+
+    _apply()
+
+    query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query)
+    assert query["connection_limit"] == ["3"]
+    assert "application_name" not in query
+    assert "novel_param" not in query
+
+
+def test_reader_keeps_its_own_options_when_writer_params_are_appended(monkeypatch):
+    """Appending the writer's pool params must leave the reader's own search_path
+    intact, since that is what decides which tables its queries resolve against."""
+    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db?connection_limit=3")
+    monkeypatch.setenv(
+        "DATABASE_URL_READ_REPLICA",
+        "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dreader_schema",
+    )
+
+    _apply()
+
+    query = urllib.parse.parse_qs(urllib.parse.urlsplit(os.environ["DATABASE_URL_READ_REPLICA"]).query)
+    assert query["options"] == ["-c search_path=reader_schema"]
+    assert query["connection_limit"] == ["3"]
+
+
+def test_reader_url_left_alone_when_writer_has_no_params(monkeypatch):
+    """No params to inherit must mean the reader URL is not rewritten at all."""
+    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db")
+    monkeypatch.setenv(
+        "DATABASE_URL_READ_REPLICA",
+        "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp",
+    )
+
+    _apply()
+
+    assert (
+        os.environ["DATABASE_URL_READ_REPLICA"]
+        == "postgresql://u:p@reader.example.com:5432/db?options=-c%20search_path%3Dapp"
+    )
+
+
 def test_unsupported_db_scheme_message_names_var_and_scheme():
     msg = unsupported_db_scheme_message("DIRECT_URL", "sqlite")
     assert "DIRECT_URL" in msg