diff --git a/litellm/proxy/db/db_connection_pool.py b/litellm/proxy/db/db_connection_pool.py new file mode 100644 index 00000000000..ec0efb69040 --- /dev/null +++ b/litellm/proxy/db/db_connection_pool.py @@ -0,0 +1,141 @@ +"""Prisma connection-pool query params, applied consistently across startups. + +The pool knobs (``connection_limit``, ``pool_timeout``, ``connect_timeout``, +``socket_timeout``, ``pgbouncer``, and any custom params) live in the database +URL's query string. They must be appended before Prisma initializes, on every +startup path: the CLI (``proxy_cli.py``), the componentized entrypoints +(``uvicorn gateway.main:app`` / ``backend.main:app``), and the proxy startup +event that all three funnel through. + +Kept deliberately free of ``pydantic_settings`` (unlike ``db_url_settings``) so +``proxy_cli`` can import it at module scope without dragging a proxy-only +dependency into the base ``import litellm`` path. +""" + +import os +import urllib.parse +from collections.abc import Mapping +from typing import Final, cast + +from litellm._logging import verbose_proxy_logger +from litellm.secret_managers.main import str_to_bool + +DEFAULT_DB_CONNECTION_POOL_LIMIT: Final[int] = 10 +DEFAULT_DB_CONNECTION_POOL_TIMEOUT: Final[int] = 60 + +POOL_PARAM_DB_URL_ENV_VARS: Final[tuple[str, ...]] = ( + "DATABASE_URL", + "DIRECT_URL", + "DATABASE_URL_READ_REPLICA", +) + + +def build_db_connection_url_params( + connection_limit: int, + pool_timeout: float | None, + connect_timeout: float | None = None, + socket_timeout: float | None = None, + disable_prepared_statements: bool = False, + extra_params: Mapping[str, object] | None = None, +) -> dict[str, object]: + """Build the Prisma DATABASE_URL query params controlling connection pool behavior. + + ``connect_timeout`` / ``socket_timeout`` map to the Prisma URL params of the + same name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and + are omitted when None so Prisma's defaults apply. ``disable_prepared_statements`` + sets ``pgbouncer=true``, which makes Prisma stop using server-side prepared + statements (pgbouncer transaction-pool compatible; also sidesteps the + "cached plan must not change result type" error during rolling migrations). + ``extra_params`` is an untyped passthrough: keys it provides win over the + named arguments above, so it can override any default set here. + """ + named: dict[str, object] = { + "connection_limit": connection_limit, + **({"pool_timeout": pool_timeout} if pool_timeout is not None else {}), + **({"connect_timeout": connect_timeout} if connect_timeout is not None else {}), + **({"socket_timeout": socket_timeout} if socket_timeout is not None else {}), + **({"pgbouncer": "true"} if disable_prepared_statements else {}), + } + return {**named, **dict(extra_params)} if extra_params else named + + +def append_query_params(url: str | None, params: Mapping[str, object]) -> str: + """Merge ``params`` into ``url``'s query string, params winning on conflict. + + Never logs the URL itself: it embeds the database credential (and, under IAM + auth, a presigned token), so echoing it even at debug level leaks secrets. + """ + if not isinstance(url, str) or url == "": + verbose_proxy_logger.warning("append_query_params received empty or non-string URL, returning empty string") + return "" + parsed_url = urllib.parse.urlparse(url) + merged: dict[str, object] = {**urllib.parse.parse_qs(parsed_url.query), **params} + encoded_query = urllib.parse.urlencode(merged, doseq=True) + return urllib.parse.urlunparse(parsed_url._replace(query=encoded_query)) + + +def _optional_number(value: object) -> float | None: + """Coerce a general_settings value to a number, or None when absent/invalid. + + ``bool`` is rejected explicitly because it is an ``int`` subclass and a stray + ``true`` in YAML must not become a timeout of ``1``. + """ + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return value + return None + + +def pool_params_from_general_settings(general_settings: Mapping[str, object]) -> dict[str, object]: + """Resolve the Prisma pool query params from a config ``general_settings`` block. + + Mirrors the keys the CLI reads so every startup path applies the same pool + ceiling. An empty mapping yields the engine defaults. + """ + raw_limit = _optional_number(general_settings.get("database_connection_pool_limit")) + connection_limit = int(raw_limit) if raw_limit is not None else DEFAULT_DB_CONNECTION_POOL_LIMIT + + pool_timeout = _optional_number(general_settings.get("database_connection_timeout")) + if pool_timeout is None: + pool_timeout = _optional_number(general_settings.get("database_connection_pool_timeout")) + if pool_timeout is None: + pool_timeout = DEFAULT_DB_CONNECTION_POOL_TIMEOUT + + raw_disable = general_settings.get("database_disable_prepared_statements", False) + disable_prepared_statements = ( + str_to_bool(raw_disable) is True if isinstance(raw_disable, str) else bool(raw_disable) + ) + + extra_raw = general_settings.get("database_extra_connection_params") + extra_params = ( + cast("Mapping[str, object]", extra_raw) # cast-ok: isinstance-guarded Mapping, YAML config keys are strings + if isinstance(extra_raw, Mapping) + else None + ) + + return build_db_connection_url_params( + connection_limit=connection_limit, + pool_timeout=pool_timeout, + connect_timeout=_optional_number(general_settings.get("database_connect_timeout")), + socket_timeout=_optional_number(general_settings.get("database_socket_timeout")), + disable_prepared_statements=disable_prepared_statements, + extra_params=extra_params, + ) + + +def apply_pool_params_to_db_urls(general_settings: Mapping[str, object]) -> None: + """Append the resolved pool params to every DB URL env var that is set. + + Idempotent: params overwrite same-named keys, so calling it again (e.g. once + in the CLI pre-migration and once in proxy startup) leaves the URL stable. + Runs for the writer, DIRECT_URL, and the read replica so a componentized + ``uvicorn gateway.main:app`` / ``backend.main:app`` startup gets the same + ceiling the CLI applies, and so the reader URL stops silently keeping + Prisma's default pool size. + """ + params = pool_params_from_general_settings(general_settings) + for env_var in POOL_PARAM_DB_URL_ENV_VARS: + current = os.getenv(env_var) + if current: + os.environ[env_var] = append_query_params(current, params) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index fbccb8a726c..9cdd56bfefe 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -23,20 +23,24 @@ class IAMEndpoint: """Static parts of an RDS IAM-authenticated Postgres connection. The IAM token rotates every ~15 minutes; everything else (host, port, user, - database name, schema) stays fixed. We capture the static fields once so - refresh just regenerates the token and reassembles the URL. + database name, and the full query string) stays fixed. Only the credential + changes on refresh, so ``query`` captures the entire query string verbatim + (``schema``, ``connection_limit``, ``pool_timeout``, ``sslmode``, + ``pgbouncer``, and any custom params) and ``build_url`` reattaches it + unchanged. Dropping it here previously reset every Prisma pool knob back to + the engine default on each refresh. """ host: str port: str user: str name: str - schema: str | None = None + query: str = "" def build_url(self, token: str) -> str: url = f"postgresql://{self.user}:{token}@{self.host}:{self.port}/{self.name}" - if self.schema: - url += f"?schema={self.schema}" + if self.query: + url += f"?{self.query}" return url @@ -44,7 +48,8 @@ def parse_iam_endpoint_from_url(url: str) -> IAMEndpoint: """Parse an IAMEndpoint from a Postgres URL. Used so a reader URL can drive its own IAM refresh without requiring - callers to set parallel DATABASE_HOST_READ_REPLICA / etc. env vars. + callers to set parallel DATABASE_HOST_READ_REPLICA / etc. env vars. The + entire query string is preserved so pool/SSL params survive token rotation. """ parsed = urllib.parse.urlparse(url) if not parsed.hostname or not parsed.username: @@ -53,18 +58,12 @@ def parse_iam_endpoint_from_url(url: str) -> IAMEndpoint: if not name: raise ValueError("Cannot parse IAM endpoint from URL: missing database name") port = str(parsed.port) if parsed.port else "5432" - schema: str | None = None - if parsed.query: - qs = urllib.parse.parse_qs(parsed.query) - schema_vals = qs.get("schema") - if schema_vals: - schema = schema_vals[0] return IAMEndpoint( host=parsed.hostname, port=port, user=parsed.username, name=name, - schema=schema, + query=parsed.query, ) @@ -318,17 +317,34 @@ class PrismaWrapper: db_port = os.getenv("DATABASE_PORT", "5432") db_user = os.getenv("DATABASE_USER") db_name = os.getenv("DATABASE_NAME") - db_schema = os.getenv("DATABASE_SCHEMA") token = generate_iam_auth_token(db_host=db_host, db_port=db_port, db_user=db_user) _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" - if db_schema: - _db_url += f"?schema={db_schema}" + query = self._writer_refresh_query() + if query: + _db_url += f"?{query}" os.environ[self._db_url_env_var] = _db_url return _db_url + def _writer_refresh_query(self) -> str: + """Query string to reattach when the writer refreshes its IAM token. + + Prefers the query already on the live URL so pool/SSL params applied at + startup (``connection_limit``, ``pool_timeout``, ``sslmode``, + ``pgbouncer``, custom params) survive rotation. Falls back to + ``DATABASE_SCHEMA`` for the legacy case where only discrete env vars + were set and no assembled URL exists yet. + """ + current_url = os.getenv(self._db_url_env_var) + if current_url: + existing = urllib.parse.urlparse(current_url).query + if existing: + return existing + db_schema = os.getenv("DATABASE_SCHEMA") + return f"schema={db_schema}" if db_schema else "" + async def recreate_prisma_client( self, new_db_url: str, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 74ec0cc8700..c97a37d2dc8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -5,7 +5,6 @@ import os import random import subprocess import sys -import urllib.parse as urlparse from pathlib import Path from typing import TYPE_CHECKING, Any, Iterable, Optional, Union @@ -15,6 +14,14 @@ from dotenv import load_dotenv import litellm from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY +from litellm.proxy.db.db_connection_pool import ( + DEFAULT_DB_CONNECTION_POOL_LIMIT, + DEFAULT_DB_CONNECTION_POOL_TIMEOUT, + append_query_params, +) +from litellm.proxy.db.db_connection_pool import ( + build_db_connection_url_params as _build_db_connection_url_params, +) from litellm.secret_managers.main import get_secret_bool if TYPE_CHECKING: @@ -35,61 +42,8 @@ telemetry = None class LiteLLMDatabaseConnectionPool(Enum): - database_connection_pool_limit = 10 - database_connection_pool_timeout = 60 - - -def _build_db_connection_url_params( - connection_limit: int, - pool_timeout: Optional[Union[int, float]], - connect_timeout: Optional[Union[int, float]] = None, - socket_timeout: Optional[Union[int, float]] = None, - disable_prepared_statements: bool = False, - extra_params: Optional[dict] = None, -) -> dict: - """Build the Prisma DATABASE_URL query params controlling connection pool behavior. - - `connect_timeout` / `socket_timeout` map to the Prisma URL params of the same - name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are - omitted when None so Prisma's defaults apply. `disable_prepared_statements` - sets `pgbouncer=true`, which makes Prisma stop using server-side prepared - statements (pgbouncer transaction-pool compatible; also sidesteps the - "cached plan must not change result type" error during rolling migrations). - `extra_params` is an untyped passthrough — keys it provides win over the - named arguments above, so it can be used to override any default we set here. - """ - params: dict = { - "connection_limit": connection_limit, - } - if pool_timeout is not None: - params["pool_timeout"] = pool_timeout - if connect_timeout is not None: - params["connect_timeout"] = connect_timeout - if socket_timeout is not None: - params["socket_timeout"] = socket_timeout - if disable_prepared_statements: - params["pgbouncer"] = "true" - if extra_params: - params.update(extra_params) - return params - - -def append_query_params(url: Optional[str], params: dict) -> str: - from litellm._logging import verbose_proxy_logger - - verbose_proxy_logger.debug(f"url: {url}") - verbose_proxy_logger.debug(f"params: {params}") - if not isinstance(url, str) or url == "": - # Preserve previous startup behavior when DATABASE_URL is absent. - # Returning an empty string avoids urlparse type errors in test/dev flows. - verbose_proxy_logger.warning("append_query_params received empty or non-string URL, returning empty string") - return "" - parsed_url = urlparse.urlparse(url) - parsed_query = urlparse.parse_qs(parsed_url.query) - parsed_query.update(params) - encoded_query = urlparse.urlencode(parsed_query, doseq=True) - modified_url = urlparse.urlunparse(parsed_url._replace(query=encoded_query)) - return modified_url # type: ignore + database_connection_pool_limit = DEFAULT_DB_CONNECTION_POOL_LIMIT + database_connection_pool_timeout = DEFAULT_DB_CONNECTION_POOL_TIMEOUT class ProxyInitializationHelpers: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f0ca1f6396f..eb9d689af17 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -913,6 +913,9 @@ async def proxy_startup_event(app: FastAPI): # check if DATABASE_URL in environment - load from there if prisma_client is None: + from litellm.proxy.db.db_connection_pool import apply_pool_params_to_db_urls + + apply_pool_params_to_db_urls(general_settings) _db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore prisma_client = await ProxyStartupEvent._setup_prisma_client( database_url=_db_url, diff --git a/tests/test_litellm/proxy/db/test_db_connection_pool.py b/tests/test_litellm/proxy/db/test_db_connection_pool.py new file mode 100644 index 00000000000..e813316efe3 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_db_connection_pool.py @@ -0,0 +1,147 @@ +import logging +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.db.db_connection_pool import ( + DEFAULT_DB_CONNECTION_POOL_LIMIT, + DEFAULT_DB_CONNECTION_POOL_TIMEOUT, + append_query_params, + apply_pool_params_to_db_urls, + build_db_connection_url_params, + pool_params_from_general_settings, +) + + +def test_build_params_defaults_only_connection_limit_and_pool_timeout(): + assert build_db_connection_url_params(connection_limit=10, pool_timeout=60) == { + "connection_limit": 10, + "pool_timeout": 60, + } + + +def test_build_params_omits_none_timeouts_and_pgbouncer(): + params = build_db_connection_url_params( + connection_limit=10, + pool_timeout=None, + connect_timeout=None, + socket_timeout=None, + disable_prepared_statements=False, + ) + assert params == {"connection_limit": 10} + + +def test_build_params_includes_pgbouncer_and_extra_overrides(): + params = build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + extra_params={"connection_limit": 99, "sslmode": "require"}, + ) + assert params["pgbouncer"] == "true" + assert params["sslmode"] == "require" + assert params["connection_limit"] == 99 + + +def test_append_query_params_merges_and_overwrites(): + merged = append_query_params( + "postgresql://u:p@h:5432/db?schema=public&connection_limit=1", + {"connection_limit": 10, "pool_timeout": 60}, + ) + assert "schema=public" in merged + assert "connection_limit=10" in merged + assert "connection_limit=1&" not in merged and not merged.endswith("connection_limit=1") + assert "pool_timeout=60" in merged + + +def test_append_query_params_missing_url_returns_empty(): + assert append_query_params(None, {"connection_limit": 10}) == "" + assert append_query_params("", {"connection_limit": 10}) == "" + + +def test_append_query_params_never_logs_url(caplog): + """Regression for #33021: the DB URL (which carries credentials) must never + be logged, at any level.""" + secret_url = "postgresql://user:sup3rs3cr3t@h:5432/db?schema=public" + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + append_query_params(secret_url, {"connection_limit": 10}) + assert all("sup3rs3cr3t" not in r.getMessage() for r in caplog.records) + + +def test_pool_params_from_general_settings_defaults_on_empty(): + assert pool_params_from_general_settings({}) == { + "connection_limit": DEFAULT_DB_CONNECTION_POOL_LIMIT, + "pool_timeout": DEFAULT_DB_CONNECTION_POOL_TIMEOUT, + } + + +def test_pool_params_from_general_settings_reads_all_keys(): + params = pool_params_from_general_settings( + { + "database_connection_pool_limit": 25, + "database_connection_pool_timeout": 45, + "database_connect_timeout": 15, + "database_socket_timeout": 20, + "database_disable_prepared_statements": "true", + "database_extra_connection_params": {"application_name": "litellm"}, + } + ) + assert params == { + "connection_limit": 25, + "pool_timeout": 45, + "connect_timeout": 15, + "socket_timeout": 20, + "pgbouncer": "true", + "application_name": "litellm", + } + + +def test_pool_params_connection_timeout_takes_precedence_over_pool_timeout(): + params = pool_params_from_general_settings( + {"database_connection_timeout": 90, "database_connection_pool_timeout": 45} + ) + assert params["pool_timeout"] == 90 + + +def test_apply_pool_params_covers_writer_direct_and_read_replica(monkeypatch): + """Regression for #33021 defect 1: componentized startup must apply pool + params to all three DB URL env vars, including the read replica which was + never receiving them.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer:5432/db?schema=public") + monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@writer:5432/db?schema=public") + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader:5432/db?schema=public") + + apply_pool_params_to_db_urls({"database_connection_pool_limit": 10, "database_connection_pool_timeout": 60}) + + for env_var in ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA"): + value = os.environ[env_var] + assert "connection_limit=10" in value + assert "pool_timeout=60" in value + assert "schema=public" in value + + +def test_apply_pool_params_is_idempotent(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer:5432/db?schema=public") + settings = {"database_connection_pool_limit": 10, "database_connection_pool_timeout": 60} + + apply_pool_params_to_db_urls(settings) + once = os.environ["DATABASE_URL"] + apply_pool_params_to_db_urls(settings) + twice = os.environ["DATABASE_URL"] + + assert once == twice + assert once.count("connection_limit=") == 1 + + +def test_apply_pool_params_skips_unset_env_vars(monkeypatch): + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + + apply_pool_params_to_db_urls({}) + + assert os.getenv("DATABASE_URL") is None + assert os.getenv("DATABASE_URL_READ_REPLICA") is None diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index e5bb8b99507..a891499634f 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -501,7 +501,7 @@ def test_parse_iam_endpoint_from_url_extracts_all_fields(): assert ep.port == "6543" assert ep.user == "litellm_user" assert ep.name == "litellm" - assert ep.schema == "public" + assert ep.query == "schema=public" def test_parse_iam_endpoint_defaults_port_to_5432_and_skips_schema(): @@ -512,7 +512,7 @@ def test_parse_iam_endpoint_defaults_port_to_5432_and_skips_schema(): assert ep.port == "5432" assert ep.user == "u" assert ep.name == "dbname" - assert ep.schema is None + assert ep.query == "" def test_parse_iam_endpoint_rejects_url_without_user_or_dbname(): @@ -530,7 +530,7 @@ def test_iam_endpoint_build_url_inserts_token_verbatim(): # `generate_iam_auth_token` already URL-encodes the presigned token, so # `build_url` must NOT encode again — double-encoding turned `%3D` into # `%253D` and broke RDS auth on the reader path. - ep = IAMEndpoint(host="h", port="5432", user="u", name="db", schema="public") + ep = IAMEndpoint(host="h", port="5432", user="u", name="db", query="schema=public") pre_encoded_token = "token%2Fwith%3Fweird%26chars%3Dyes" url = ep.build_url(pre_encoded_token) assert url == f"postgresql://u:{pre_encoded_token}@h:5432/db?schema=public" @@ -704,6 +704,71 @@ def test_writer_get_rds_iam_token_uses_database_host_env_vars(monkeypatch, unset assert os.environ["DATABASE_URL"] == new_url +def test_writer_iam_refresh_preserves_pool_params_from_live_url(monkeypatch): + """Regression for #33021: a writer IAM refresh must keep the full query + string (pool_timeout, connection_limit, sslmode, custom params) already on + the live DATABASE_URL, replacing only the credential. Previously it rebuilt + the URL with just `schema`, resetting every Prisma pool knob to the engine + default on each ~15-minute rotation.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + monkeypatch.setenv( + "DATABASE_URL", + "postgresql://litellm:stale-token@writer.aurora.local:5432/litellm" + "?schema=public&connection_limit=10&pool_timeout=60&sslmode=require&pgbouncer=true", + ) + monkeypatch.setenv("DATABASE_HOST", "writer.aurora.local") + monkeypatch.setenv("DATABASE_USER", "litellm") + monkeypatch.setenv("DATABASE_NAME", "litellm") + monkeypatch.delenv("DATABASE_SCHEMA", raising=False) + + fake_module = MagicMock() + fake_module.generate_iam_auth_token = lambda db_host=None, db_port=None, db_user=None: "FRESH-TOKEN" + monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_module) + + writer = PrismaWrapper(original_prisma=MagicMock(), iam_token_db_auth=True) + new_url = writer.get_rds_iam_token() + + assert new_url is not None + assert "stale-token" not in new_url + assert new_url.startswith("postgresql://litellm:FRESH-TOKEN@writer.aurora.local:5432/litellm?") + for param in ("schema=public", "connection_limit=10", "pool_timeout=60", "sslmode=require", "pgbouncer=true"): + assert param in new_url + assert os.environ["DATABASE_URL"] == new_url + + +def test_reader_iam_refresh_preserves_pool_params(monkeypatch): + """Regression for #33021: the reader IAM refresh must preserve the full + query string parsed from DATABASE_URL_READ_REPLICA, not collapse it to + `schema` alone.""" + from litellm.proxy.db.prisma_client import PrismaWrapper, parse_iam_endpoint_from_url + + reader_url = ( + "postgresql://lit:stale@reader.aurora.local:6543/litellm" + "?schema=public&connection_limit=5&pool_timeout=30&sslmode=require" + ) + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", reader_url) + + fake_module = MagicMock() + fake_module.generate_iam_auth_token = lambda db_host=None, db_port=None, db_user=None: "READER-TOKEN" + monkeypatch.setitem(sys.modules, "litellm.proxy.auth.rds_iam_token", fake_module) + + reader = PrismaWrapper( + original_prisma=MagicMock(), + iam_token_db_auth=True, + db_url_env_var="DATABASE_URL_READ_REPLICA", + iam_endpoint=parse_iam_endpoint_from_url(reader_url), + recreate_uses_datasource=True, + ) + new_url = reader.get_rds_iam_token() + + assert new_url is not None + assert new_url.startswith("postgresql://lit:READER-TOKEN@reader.aurora.local:6543/litellm?") + for param in ("schema=public", "connection_limit=5", "pool_timeout=30", "sslmode=require"): + assert param in new_url + assert os.environ["DATABASE_URL_READ_REPLICA"] == new_url + + def test_reader_iam_refresh_uses_parsed_endpoint(monkeypatch): """The reader generates fresh tokens against its parsed endpoint and writes the new URL to DATABASE_URL_READ_REPLICA — not DATABASE_URL.""" @@ -730,7 +795,7 @@ def test_reader_iam_refresh_uses_parsed_endpoint(monkeypatch): port="5432", user="lit", name="litellm", - schema=None, + query="", ) reader = PrismaWrapper( original_prisma=MagicMock(),