mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): apply DB pool params on componentized startup and preserve IAM URL params
Componentized entrypoints (gateway.main:app / backend.main:app) only assembled the DB URLs and never applied the general_settings connection-pool params, and the read replica never received them even via the CLI. Move the pool-param logic into a shared litellm/proxy/db/db_connection_pool module and apply it in proxy_startup_event, before Prisma is built, so every startup path shares one ceiling across DATABASE_URL, DIRECT_URL, and DATABASE_URL_READ_REPLICA. IAM token refresh rebuilt the writer/reader URL with only the schema query param, discarding connection_limit, pool_timeout, connect_timeout, socket_timeout, sslmode, pgbouncer, and custom params on every rotation. IAMEndpoint now keeps the full query string and reattaches it verbatim; the legacy writer path prefers the live URL query and falls back to DATABASE_SCHEMA. Also stop logging the full DB URL in append_query_params, which leaked the credential (and IAM presigned token) into debug logs.
This commit is contained in:
Devin AI
parent
3e9e52042a
commit
a9e06f8337
6 changed files with 402 additions and 76 deletions
141
litellm/proxy/db/db_connection_pool.py
Normal file
141
litellm/proxy/db/db_connection_pool.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
147
tests/test_litellm/proxy/db/test_db_connection_pool.py
Normal file
147
tests/test_litellm/proxy/db/test_db_connection_pool.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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(),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue