feat(proxy): let the in-container pgbouncer follow rotating RDS IAM and Azure Entra tokens (#40623)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-10 16:18:00 -07:00 committed by GitHub
parent ef1a37795c
commit 3b93da414e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 703 additions and 131 deletions

View file

@ -4,8 +4,9 @@
is fine for a plain Postgres URL but not for the pooler: PgBouncer must be
started exactly once per pod, before the workers fork, and the workers must be
handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
``DatabaseURLSettings.apply_to_env`` under password auth, so setting it here is
enough for every worker to pick the pooled URL up unchanged.
``DatabaseURLSettings.apply_to_env`` under password auth, and one marked pooled
wins under token auth too, so exporting it here is enough for every worker to
pick the pooled URL up unchanged.
Run with:
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
@ -19,7 +20,12 @@ from typing import Final
from uvicorn.main import main as uvicorn_main
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerSettings,
export_pooled_database_url,
start_in_container_pgbouncer,
)
GATEWAY_APP: Final = "gateway.main:app"
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
@ -41,15 +47,15 @@ def pool_database_url(
"""Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Token auth is
rejected by the pooler itself, since it holds one password for its lifetime.
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
auth the pooler mints and renews the upstream token itself.
"""
if not pgbouncer.enabled:
return None
upstream_url: Final = environ.get("DATABASE_URL")
if upstream_url is None:
return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled")
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth_enabled=settings.token_auth() is not None)
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())
def _serve(argv: Sequence[str]) -> None:
@ -63,7 +69,7 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -
if isinstance(pooled_url, PgBouncerError):
sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}")
if pooled_url is not None:
os.environ["DATABASE_URL"] = pooled_url
export_pooled_database_url(pooled_url)
serve(uvicorn_argv(argv, os.environ))

View file

@ -361,12 +361,9 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end -}}
{{/*
In-container PgBouncer env for the gateway container. Fails at render time under IAM or Entra auth: the pooler holds one static password for the life of the pod.
In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself.
*/}}
{{- define "litellm.connectionPoolEnv" -}}
{{- if or .Values.database.writer.useIAMAuth .Values.database.writer.useAzureEntraAuth }}
{{- fail "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password" }}
{{- end }}
{{- with .Values.database.connectionPool -}}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"

View file

@ -82,23 +82,39 @@ tests:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- it: pool with IAM auth fails at render time
- it: pool with IAM auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: pool with Entra auth fails at render time
- it: pool with Entra auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
- it: IAM auth without the pool still renders
template: gateway/deployment.yaml

View file

@ -235,9 +235,9 @@ database:
# network hop. The chart emits LITELLM_PGBOUNCER_ENABLED /
# LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on
# the gateway container only: the backend runs a single worker and the
# migrations Job must keep a direct connection. The pool holds a static
# password, so it cannot be combined with `database.writer.useIAMAuth` or
# `useAzureEntraAuth` (rendering fails). Starting profile for
# migrations Job must keep a direct connection. With
# `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and
# renews the database token itself, so the workers never see it. Starting profile for
# `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
# 5000-connection ceiling fits roughly 200 gateway replicas.
connectionPool:

View file

@ -52,6 +52,7 @@ from typing import Annotated, Final, Protocol, TypeAlias, cast
from pydantic import AliasChoices, BeforeValidator, Field
from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm.proxy.db.pgbouncer import database_url_is_pooled
from litellm.proxy.db.token_auth import (
AZURE_POSTGRESQL_AUTH_ENV_VAR,
DEFAULT_POSTGRES_PORT,
@ -358,8 +359,12 @@ class DatabaseURLSettings(BaseSettings):
Raises ``RuntimeError`` (naming the offending vars) when token auth is
enabled but a required field is missing the proxy cannot recover
from this and a clear startup error beats a Prisma connect failure.
A ``DATABASE_URL`` the supervisor pointed at the in-container PgBouncer
is kept even under token auth: the pooler renews the token upstream.
"""
auth: Final = self.token_auth()
if auth is not None and database_url_is_pooled():
return None
if auth is not None:
missing: Final = tuple(
env

View file

@ -18,19 +18,28 @@ Migrations and the schema diff run in the supervisor before the pooler is
started, so they always go straight to Postgres. ``DATABASE_URL_READ_REPLICA``
is left untouched.
The pooler holds the database password from startup, so it cannot be combined
with ``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH``: those rotate the
password inside every worker on their own schedule, and PgBouncer would keep
authenticating upstream with the expired token.
The workers never hold the upstream credential: they log in to PgBouncer as
``litellm_pgbouncer`` with a random password made at startup, and PgBouncer
takes the database user's password from its auth file. Under
``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH`` that password is a
short-lived token, so the supervisor mints a new one before it expires,
rewrites the auth file and asks PgBouncer to reload; only new upstream
connections authenticate, so live ones are unaffected. The pooled
``DATABASE_URL`` then carries a static password, and the workers must not run
their own token refresh against it: ``LITELLM_PGBOUNCER_POOLED_DATABASE_URL``
tells them so, while a read replica keeps refreshing its own token.
"""
from __future__ import annotations
import atexit
import functools
import os
import re
import secrets
import shlex
import shutil
import signal
import socket
import subprocess
import tempfile
@ -39,6 +48,7 @@ import time
import urllib.parse
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -47,10 +57,18 @@ from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.token_auth import AZURE_POSTGRESQL_AUTH_ENV_VAR, IAM_TOKEN_DB_AUTH_ENV_VAR
from litellm.proxy.db.token_auth import (
DatabaseTokenAuth,
IAMEndpoint,
mint_database_token,
parse_database_token_expiration,
parse_iam_endpoint_from_url,
)
PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_"
PGBOUNCER_POOLED_ENV_VAR: Final = "LITELLM_PGBOUNCER_POOLED_DATABASE_URL"
PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1"
PGBOUNCER_POOL_USER: Final = "litellm_pgbouncer"
PGBOUNCER_INI_NAME: Final = "pgbouncer.ini"
PGBOUNCER_USERLIST_NAME: Final = "userlist.txt"
PGBOUNCER_CA_NAME: Final = "server-ca.pem"
@ -59,13 +77,11 @@ PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0
PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0
PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody"
PGBOUNCER_MIN_VERSION: Final = (1, 19)
PGBOUNCER_MAX_PASSWORD_BYTES: Final = 2048
PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)")
PGBOUNCER_LIST_DELIMITER_PATTERN: Final = re.compile(r"[,\s]")
PGBOUNCER_TOKEN_AUTH_CONFLICT: Final = (
f"the in-container pgbouncer cannot be combined with {IAM_TOKEN_DB_AUTH_ENV_VAR} or "
f"{AZURE_POSTGRESQL_AUTH_ENV_VAR}: each worker rotates the database password on its own schedule and the pooler "
"would keep using the expired token upstream. Disable the pooler or use a static database password"
)
PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS: Final = 180.0
PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS: Final = 600.0
PGBOUNCER_TOKEN_RETRY_SECONDS: Final = 30.0
# Prisma's client-side TLS params describe the hop to Postgres, which becomes
# PgBouncer's server side. They move into ``server_tls_*`` and must not stay on
@ -97,10 +113,18 @@ class PgBouncerSettings(BaseSettings):
@dataclass(frozen=True, slots=True)
class PgBouncerPlan:
ini: str
userlist: str
pooled_url: str
upstream_user: str
upstream_password: str | None
pool_password: str
ca_source: str | None = None
def userlist(self, upstream_password: str) -> str:
return "".join(
f"{_userlist_quote(user)} {_userlist_quote(password)}\n"
for user, password in ((self.upstream_user, upstream_password), (PGBOUNCER_POOL_USER, self.pool_password))
)
@dataclass(frozen=True, slots=True)
class PgBouncerError:
@ -173,7 +197,9 @@ def plan_pgbouncer(
Params describing Prisma's own pool (``connection_limit``, ``pool_timeout``,
...) stay on the pooled URL; the TLS params and ``options`` describe the hop
to Postgres and move into the PgBouncer config. ``run_as_user`` is the
to Postgres and move into the PgBouncer config. The upstream password is
left out of the config on purpose: PgBouncer then takes it from the auth
file, which can be rewritten while it runs. ``run_as_user`` is the
unprivileged user PgBouncer drops to when the proxy runs as root, which
PgBouncer itself refuses to do.
"""
@ -184,14 +210,12 @@ def plan_pgbouncer(
dbname: Final = urllib.parse.unquote(parsed.path.lstrip("/"))
username: Final = urllib.parse.unquote(parsed.username or "")
password: Final = None if parsed.password is None else urllib.parse.unquote(parsed.password)
if not parsed.hostname or not username or password is None or not dbname:
if not parsed.hostname or not username or not dbname:
return PgBouncerError("DATABASE_URL must carry a host, user and database name for the in-container PgBouncer")
if username == PGBOUNCER_POOL_USER:
return PgBouncerError(
"DATABASE_URL must carry a host, user, password and database name for the in-container PgBouncer"
)
if PGBOUNCER_LIST_DELIMITER_PATTERN.search(username):
return PgBouncerError(
f"the database user {username!r} cannot be named in PgBouncer's stats_users list: "
"PgBouncer splits list settings on commas and whitespace and has no quoting for them"
f"the database user cannot be named {PGBOUNCER_POOL_USER!r}: that is the user the workers log in to the "
"in-container PgBouncer as, and PgBouncer keeps one password per user"
)
if "sslidentity" in params:
return PgBouncerError("client certificates (sslidentity) are not supported with the in-container PgBouncer")
@ -212,7 +236,6 @@ def plan_pgbouncer(
f"port={parsed.port or 5432}",
f"dbname={_single_quoted(dbname)}",
f"user={_single_quoted(username)}",
f"password={_single_quoted(password)}",
*((f"connect_query={_single_quoted(connect_query)}",) if connect_query else ()),
)
)
@ -227,7 +250,7 @@ def plan_pgbouncer(
f"unix_socket_dir = {runtime_dir}",
f"auth_file = {runtime_dir / PGBOUNCER_USERLIST_NAME}",
"auth_type = scram-sha-256",
f"stats_users = {username}",
f"stats_users = {PGBOUNCER_POOL_USER}",
"pool_mode = transaction",
f"max_client_conn = {settings.max_client_conn}",
f"default_pool_size = {settings.max_db_connections}",
@ -238,41 +261,175 @@ def plan_pgbouncer(
"",
)
)
userlist: Final = f"{_userlist_quote(username)} {_userlist_quote(password)}\n"
pooled_query: Final = urllib.parse.urlencode(
(*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true"))
)
credentials: Final = f"{urllib.parse.quote(username, safe='')}:{urllib.parse.quote(password, safe='')}"
pool_password: Final = secrets.token_urlsafe(32)
pooled_url: Final = urllib.parse.urlunsplit(
parsed._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query)
parsed._replace(
netloc=f"{PGBOUNCER_POOL_USER}:{pool_password}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query
)
)
return PgBouncerPlan(
ini=ini,
pooled_url=pooled_url,
upstream_user=username,
upstream_password=password,
pool_password=pool_password,
ca_source=params.get("sslcert") or None,
)
return PgBouncerPlan(ini=ini, userlist=userlist, pooled_url=pooled_url, ca_source=params.get("sslcert") or None)
def write_pgbouncer_files(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError:
"""Write the ini, userlist (both hold the password, so mode 0600) and CA copy, and return the ini path.
def _write_private(path: Path, content: str, run_as_user: str | None) -> None:
with open(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w", encoding="utf-8") as handle:
handle.write(content)
if run_as_user is not None:
shutil.chown(path, user=run_as_user)
def write_pgbouncer_ini(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError:
"""Write the ini (mode 0600) and the CA copy, and return the ini path.
``run_as_user`` is the user PgBouncer drops to when started as root; it has
to own the files it re-reads on reload and the socket directory.
"""
ini_path: Final = runtime_dir / PGBOUNCER_INI_NAME
userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME
ca_path: Final = runtime_dir / PGBOUNCER_CA_NAME
if plan.ca_source is not None:
try:
shutil.copyfile(plan.ca_source, ca_path)
except OSError as error:
return PgBouncerError(f"cannot read the CA bundle {plan.ca_source!r} named by sslcert: {error}")
for path, content in ((userlist_path, plan.userlist), (ini_path, plan.ini)):
path.touch(mode=0o600)
path.write_text(content, encoding="utf-8")
_write_private(ini_path, plan.ini, run_as_user)
if run_as_user is not None:
runtime_dir.chmod(0o700)
for path in (runtime_dir, ini_path, userlist_path, *((ca_path,) if plan.ca_source is not None else ())):
for path in (runtime_dir, *((ca_path,) if plan.ca_source is not None else ())):
shutil.chown(path, user=run_as_user)
return ini_path
def write_userlist(userlist: str, runtime_dir: Path, run_as_user: str | None) -> Path:
"""Replace the auth file in one step, so a PgBouncer starting or reloading meanwhile reads the old or the new one whole."""
userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME
staged_path: Final = runtime_dir / f".{PGBOUNCER_USERLIST_NAME}.next"
_write_private(staged_path, userlist, run_as_user)
os.replace(staged_path, userlist_path)
return userlist_path
def export_pooled_database_url(pooled_url: str) -> None:
os.environ["DATABASE_URL"] = pooled_url
os.environ[PGBOUNCER_POOLED_ENV_VAR] = "true"
def database_url_is_pooled(environ: Mapping[str, str] = os.environ) -> bool:
return environ.get(PGBOUNCER_POOLED_ENV_VAR) == "true"
@dataclass(frozen=True, slots=True)
class PgBouncerTokenSource:
auth: DatabaseTokenAuth
endpoint: IAMEndpoint
def mint(self) -> str:
"""The token as Postgres expects it: ``mint_database_token`` returns it percent-encoded for a URL."""
return urllib.parse.unquote(mint_database_token(self.auth, self.endpoint))
def expires_at(self, token: str) -> datetime | None:
return parse_database_token_expiration(self.auth, token)
def _utcnow() -> datetime:
return datetime.now(timezone.utc).replace(tzinfo=None)
class PgBouncerTokenRefresher:
"""Keeps the token in PgBouncer's auth file current from a daemon thread.
``install`` gets each fresh token and is expected to rewrite the auth file
and reload PgBouncer. The next refresh is due ``buffer_seconds`` before the
token expires, or ``fallback_seconds`` later when the expiry cannot be read.
A refresh that fails leaves the previous auth file in place and is retried
after ``retry_seconds``: the old token stays good until it expires, so a
transient credential-provider error costs nothing unless it persists.
"""
def __init__(
self,
source: PgBouncerTokenSource,
install: Callable[[str], None],
*,
buffer_seconds: float = PGBOUNCER_TOKEN_REFRESH_BUFFER_SECONDS,
fallback_seconds: float = PGBOUNCER_TOKEN_FALLBACK_REFRESH_SECONDS,
retry_seconds: float = PGBOUNCER_TOKEN_RETRY_SECONDS,
now: Callable[[], datetime] = _utcnow,
) -> None:
self._source: Final = source
self._install: Final = install
self._buffer_seconds: Final = buffer_seconds
self._fallback_seconds: Final = fallback_seconds
self._retry_seconds: Final = retry_seconds
self._now: Final = now
self._stopping: Final = threading.Event()
self._delay: float = 0.0
self._thread: threading.Thread | None = None
def refresh(self) -> float | PgBouncerError:
label: Final = self._source.auth.label
try:
token: Final = self._source.mint()
except Exception as mint_error:
return PgBouncerError(f"could not mint a {label} for the in-container pgbouncer: {mint_error!r}")
if len(token.encode()) >= PGBOUNCER_MAX_PASSWORD_BYTES:
return PgBouncerError(
f"the {label} is {len(token.encode())} bytes long, but PgBouncer's auth file holds passwords of at "
f"most {PGBOUNCER_MAX_PASSWORD_BYTES - 1} bytes"
)
try:
self._install(token)
except OSError as install_error:
return PgBouncerError(f"could not install the {label} into the pgbouncer auth file: {install_error}")
expires_at: Final = self._source.expires_at(token)
if expires_at is None:
return self._fallback_seconds
return max(self._retry_seconds, (expires_at - self._now()).total_seconds() - self._buffer_seconds)
def start(self) -> PgBouncerError | None:
primed: Final = self.refresh()
if isinstance(primed, PgBouncerError):
return primed
self._delay = primed
self._thread = threading.Thread(target=self._run, daemon=True, name="litellm-pgbouncer-token-refresh")
self._thread.start()
return None
def _run(self) -> None:
while not self._stopping.wait(self._delay):
self._delay = self._refresh_and_report()
def _refresh_and_report(self) -> float:
outcome: Final = self.refresh()
if isinstance(outcome, PgBouncerError):
verbose_proxy_logger.error(
"In-container pgbouncer keeps its current %s (%s); retrying in %.0fs.",
self._source.auth.label,
outcome.reason,
self._retry_seconds,
)
return self._retry_seconds
verbose_proxy_logger.info(
"In-container pgbouncer picked up a fresh %s; the next one is due in %.0fs.",
self._source.auth.label,
outcome,
)
return outcome
def stop(self) -> None:
self._stopping.set()
if self._thread is not None:
self._thread.join()
def _port_open(port: int) -> bool:
try:
with socket.create_connection((PGBOUNCER_LISTEN_ADDR, port), timeout=0.5):
@ -453,6 +610,11 @@ class PgBouncerProcess:
)
threading.Thread(target=self._restart_after_delay, daemon=True, name="litellm-pgbouncer-supervisor").start()
def reload(self) -> None:
with self._lock:
if self._process is not None:
self._process.send_signal(signal.SIGHUP)
def stop(self) -> None:
with self._lock:
self._stopping.set()
@ -461,6 +623,39 @@ class PgBouncerProcess:
_end(process)
def install_pgbouncer_token(
plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None, pooler: PgBouncerProcess, token: str
) -> None:
write_userlist(plan.userlist(token), runtime_dir, run_as_user)
pooler.reload()
def _install_upstream_password(
plan: PgBouncerPlan,
runtime_dir: Path,
run_as_user: str | None,
pooler: PgBouncerProcess,
token_auth: DatabaseTokenAuth | None,
upstream_url: str,
) -> PgBouncerTokenRefresher | None | PgBouncerError:
if token_auth is None:
if plan.upstream_password is None:
return PgBouncerError(
"DATABASE_URL carries no password and neither IAM_TOKEN_DB_AUTH nor AZURE_POSTGRESQL_AUTH is on, "
"so the in-container PgBouncer has nothing to authenticate to Postgres with"
)
write_userlist(plan.userlist(plan.upstream_password), runtime_dir, run_as_user)
return None
refresher: Final = PgBouncerTokenRefresher(
PgBouncerTokenSource(auth=token_auth, endpoint=parse_iam_endpoint_from_url(upstream_url)),
functools.partial(install_pgbouncer_token, plan, runtime_dir, run_as_user, pooler),
)
failed: Final = refresher.start()
if failed is not None:
return failed
return refresher
def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]:
"""An exit hook that does nothing in a forked child, which inherits the parent's ``atexit`` table."""
owner_pid: Final = os.getpid()
@ -475,7 +670,7 @@ def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]:
def start_in_container_pgbouncer(
settings: PgBouncerSettings,
upstream_url: str,
token_auth_enabled: bool = False,
token_auth: DatabaseTokenAuth | None = None,
register_exit_hook: Callable[[Callable[[], None]], object] = atexit.register,
) -> str | PgBouncerError:
"""Start the pooler for ``upstream_url`` and return the loopback URL the workers must use.
@ -484,10 +679,9 @@ def start_in_container_pgbouncer(
once the worker manager has returned, and only by the process that started
it (gunicorn forks its workers, so they carry the hooks too). PgBouncer
refuses to run as root, so a root proxy (the default image) has it drop to
``nobody``.
``nobody``. With ``token_auth`` the password on ``upstream_url`` is ignored:
the pooler mints its own tokens and renews them for as long as it runs.
"""
if token_auth_enabled:
return PgBouncerError(PGBOUNCER_TOKEN_AUTH_CONFLICT)
version: Final = pgbouncer_version(settings.binary)
if isinstance(version, PgBouncerError):
return version
@ -503,7 +697,7 @@ def start_in_container_pgbouncer(
plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir, run_as_user)
if isinstance(plan, PgBouncerError):
return plan
ini_path: Final = write_pgbouncer_files(plan, runtime_dir, run_as_user)
ini_path: Final = write_pgbouncer_ini(plan, runtime_dir, run_as_user)
if isinstance(ini_path, PgBouncerError):
return ini_path
pooler: Final = PgBouncerProcess(
@ -511,15 +705,23 @@ def start_in_container_pgbouncer(
port=settings.port,
socket_path=unix_socket_path(runtime_dir, settings.port),
)
refresher: Final = _install_upstream_password(plan, runtime_dir, run_as_user, pooler, token_auth, upstream_url)
if isinstance(refresher, PgBouncerError):
return refresher
failed: Final = pooler.start()
if failed is not None:
if refresher is not None:
refresher.stop()
return failed
register_exit_hook(_only_in_this_process(pooler.stop))
if refresher is not None:
register_exit_hook(_only_in_this_process(refresher.stop))
verbose_proxy_logger.info(
"In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections.",
"In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections%s.",
pooler.pid,
PGBOUNCER_LISTEN_ADDR,
settings.port,
settings.max_db_connections,
"" if token_auth is None else f" and renewing its {token_auth.label} before each one expires",
)
return plan.pooled_url

View file

@ -18,7 +18,12 @@ from pydantic import BaseModel, ConfigDict
import litellm
from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerSettings,
export_pooled_database_url,
start_in_container_pgbouncer,
)
from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper
if TYPE_CHECKING:
@ -1109,6 +1114,7 @@ def run_server(
from litellm.proxy.db.token_auth import (
AZURE_POSTGRESQL_AUTH_ENV_VAR,
IAM_TOKEN_DB_AUTH_ENV_VAR,
resolve_database_token_auth,
token_auth_flag_enabled,
)
@ -1382,7 +1388,7 @@ def run_server(
upstream_database_url: Final = os.getenv("DATABASE_URL")
if pgbouncer_settings.enabled and upstream_database_url is not None:
pooled_database_url: Final = start_in_container_pgbouncer(
pgbouncer_settings, upstream_database_url, token_auth_enabled=wants_rds_iam or wants_azure_entra
pgbouncer_settings, upstream_database_url, token_auth=resolve_database_token_auth()
)
if isinstance(pooled_database_url, PgBouncerError):
print(
@ -1392,7 +1398,7 @@ def run_server(
flush=True,
)
sys.exit(1)
os.environ["DATABASE_URL"] = pooled_database_url
export_pooled_database_url(pooled_database_url)
if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port):
port = random.randint(1024, 49152)
if prometheus_metrics_port == port:

View file

@ -131,6 +131,7 @@ from litellm.proxy.db.health_check_latest import (
fetch_latest_health_checks_for_models,
)
from litellm.proxy.db.log_db_metrics import log_db_metrics
from litellm.proxy.db.pgbouncer import database_url_is_pooled
from litellm.proxy.db.prisma_client import (
PrismaWrapper,
parse_iam_endpoint_from_url,
@ -4007,6 +4008,7 @@ class PrismaClient:
verbose_proxy_logger.error("Please run 'prisma generate' to generate the Prisma client.")
raise Exception("Unable to find Prisma binaries. Please run 'prisma generate' first.")
token_auth: Final = self.token_auth
writer_token_auth: Final = None if database_url_is_pooled() else token_auth
# When read-replica routing is on, tag log lines with [writer]/[reader]
# so the two wrappers' interleaved token refresh logs can be told apart.
# Single-DB deployments get an empty prefix (logs unchanged).
@ -4015,13 +4017,13 @@ class PrismaClient:
if http_client is not None:
writer_wrapper = PrismaWrapper(
original_prisma=Prisma(http=http_client),
token_auth=token_auth,
token_auth=writer_token_auth,
log_prefix=writer_log_prefix,
)
else:
writer_wrapper = PrismaWrapper(
original_prisma=Prisma(),
token_auth=token_auth,
token_auth=writer_token_auth,
log_prefix=writer_log_prefix,
)

View file

@ -273,18 +273,17 @@ connections the pooler accepts. The module sets
and the migration task keep the direct connection.
```hcl
create_database = false
database_url = "postgresql://litellm:<password>@db.internal:5432/litellm"
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 20
gateway_pool_max_client_conn = 1000
```
The pool needs a static database password, so it is only valid with an
existing database via `database_url`. The module-created Aurora authenticates
with rotating IAM tokens (see [Aurora + IAM auth](#aurora--iam-auth)), which
the pooler cannot follow, and `terraform plan` rejects that combination.
The pool works with the module-created Aurora as well as an existing database
via `database_url`. Against Aurora it authenticates with the same rotating IAM
tokens the workers used to (see [Aurora + IAM auth](#aurora--iam-auth)): the
pooler mints a token from the task role, renews it before it expires and hands
the workers a loopback URL with a static password instead
The componentized `gateway_image` starts through `python -m gateway.launch`,
which reads these variables, starts the pooler once per task and hands the

View file

@ -306,8 +306,8 @@ resource "aws_ecs_task_definition" "gateway" {
}
precondition {
condition = !var.gateway_connection_pool_enabled || local.byo_database
error_message = "gateway_connection_pool_enabled requires an existing database via database_url with create_database = false: the module-created Aurora authenticates with IAM tokens, which the in-container pgbouncer cannot follow because it holds a static database password."
condition = !var.gateway_connection_pool_enabled || local.database_enabled
error_message = "gateway_connection_pool_enabled needs a database: set create_database = true or pass database_url."
}
}

View file

@ -88,16 +88,20 @@ run "gateway_starts_through_the_pool_aware_launcher" {
}
}
run "pool_with_module_created_iam_aurora_fails_at_plan" {
run "pool_with_module_created_iam_aurora_plans_with_both_the_pool_and_iam_auth" {
command = plan
variables {
gateway_connection_pool_enabled = true
}
expect_failures = [
aws_ecs_task_definition.gateway,
]
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
contains(local.managed_db_env, { name = "IAM_TOKEN_DB_AUTH", value = "true" }),
])
error_message = "With the module-created Aurora the gateway must get the pool env alongside IAM token auth."
}
}
run "pool_without_any_database_fails_at_plan" {

View file

@ -208,10 +208,9 @@ variable "gateway_connection_pool_enabled" {
Postgres, so a task's footprint against the database connection ceiling is
workers x connection_limit and grows with every task. Sets
LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS /
LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. Requires
an existing database via `database_url`: the module-created Aurora
authenticates with IAM tokens, which the pooler cannot follow because it
holds one static password for the life of the task.
LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. Works with
the module-created Aurora too: the pooler mints the IAM token itself and
renews it before it expires.
EOT
type = bool
default = false

View file

@ -5,6 +5,7 @@ import textwrap
import urllib.parse
from pathlib import Path
from typing import Final, cast
from unittest.mock import MagicMock, patch
import pytest
from uvicorn.importer import import_from_string
@ -13,7 +14,7 @@ from uvicorn.main import main as uvicorn_main
import gateway.main
from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings
DB_ENV: Final = {
"DATABASE_HOST": "db.internal",
@ -66,7 +67,13 @@ def _query(url: str) -> dict[str, str]:
@pytest.fixture
def password_env(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]:
for var in ("DATABASE_URL", "IAM_TOKEN_DB_AUTH", "AZURE_POSTGRESQL_AUTH", "DATABASE_HOST_READ_REPLICA"):
for var in (
"DATABASE_URL",
"IAM_TOKEN_DB_AUTH",
"AZURE_POSTGRESQL_AUTH",
"DATABASE_HOST_READ_REPLICA",
PGBOUNCER_POOLED_ENV_VAR,
):
monkeypatch.setenv(var, "")
monkeypatch.delenv(var)
for var, value in DB_ENV.items():
@ -74,6 +81,12 @@ def password_env(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]:
return dict(DB_ENV)
def _minted_iam_token(token: str):
rds: Final = MagicMock()
rds.generate_db_auth_token.return_value = token
return patch("boto3.client", return_value=rds)
def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]:
return uvicorn_main.make_context("uvicorn", list(argv)).params
@ -109,16 +122,23 @@ class TestPoolDatabaseUrl:
assert isinstance(outcome, PgBouncerError)
assert "DATABASE_URL" in outcome.reason
def test_token_auth_is_refused(self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
def test_token_auth_hands_the_workers_the_pool_user_not_the_token(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
environ: Final = {"DATABASE_URL": "postgresql://litellm:token@db.internal:5432/litellm"}
outcome: Final = pool_database_url(
DatabaseURLSettings.from_env(),
PgBouncerSettings(enabled=True, port=_free_port(), binary=str(_fake_pooler(tmp_path))),
environ,
)
assert isinstance(outcome, PgBouncerError)
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
port: Final = _free_port()
environ: Final = {"DATABASE_URL": "postgresql://litellm:MINTED_TOKEN@db.internal:5432/litellm"}
with _minted_iam_token("MINTED_TOKEN"):
outcome: Final = pool_database_url(
DatabaseURLSettings.from_env(),
PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path))),
environ,
)
assert isinstance(outcome, str), outcome
pooled: Final = urllib.parse.urlsplit(outcome)
assert (pooled.username, pooled.hostname, pooled.port) == ("litellm_pgbouncer", "127.0.0.1", port)
assert "MINTED_TOKEN" not in outcome
class TestMain:
@ -134,14 +154,39 @@ class TestMain:
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
pooled: Final = os.environ["DATABASE_URL"]
assert urllib.parse.urlsplit(pooled).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}"
assert urllib.parse.urlsplit(pooled).hostname == "127.0.0.1"
assert urllib.parse.urlsplit(pooled).port == port
assert urllib.parse.urlsplit(pooled).username == "litellm_pgbouncer"
assert "p%40ss" not in pooled
assert _query(pooled)["pgbouncer"] == "true"
assert _uvicorn_params(served[0])["timeout_keep_alive"] == 75
DatabaseURLSettings.from_env().apply_to_env()
worker_url: Final = os.environ["DATABASE_URL"]
assert urllib.parse.urlsplit(worker_url).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}"
assert _query(worker_url)["pgbouncer"] == "true"
assert urllib.parse.urlsplit(os.environ["DATABASE_URL"]).netloc == urllib.parse.urlsplit(pooled).netloc
assert _query(os.environ["DATABASE_URL"])["pgbouncer"] == "true"
def test_iam_workers_keep_the_loopback_url_instead_of_minting_their_own(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):
port: Final = _free_port()
monkeypatch.delenv("DATABASE_PASSWORD")
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path)))
served: Final[list[tuple[str, ...]]] = []
with _minted_iam_token("SUPERVISOR_TOKEN"):
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
pooled: Final = os.environ["DATABASE_URL"]
assert urllib.parse.urlsplit(pooled).netloc.endswith(f"@127.0.0.1:{port}")
assert "SUPERVISOR_TOKEN" not in pooled
assert os.environ[PGBOUNCER_POOLED_ENV_VAR] == "true"
assert len(served) == 1
with _minted_iam_token("WORKER_TOKEN"):
DatabaseURLSettings.from_env().apply_to_env()
assert os.environ["DATABASE_URL"] == pooled
def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path

View file

@ -1,4 +1,6 @@
import base64
import configparser
import json
import logging
import os
import signal
@ -9,8 +11,10 @@ import tempfile
import textwrap
import time
import urllib.parse
from collections import deque
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final, cast
@ -18,16 +22,24 @@ import pytest
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.pgbouncer import (
PGBOUNCER_POOLED_ENV_VAR,
PgBouncerError,
PgBouncerPlan,
PgBouncerProcess,
PgBouncerSettings,
PgBouncerTokenRefresher,
PgBouncerTokenSource,
database_url_is_pooled,
export_pooled_database_url,
install_pgbouncer_token,
pgbouncer_version,
plan_pgbouncer,
start_in_container_pgbouncer,
unix_socket_path,
write_pgbouncer_files,
write_pgbouncer_ini,
write_userlist,
)
from litellm.proxy.db.token_auth import AzureEntraTokenAuth, IAMEndpoint
UPSTREAM: Final = (
"postgresql://app:p%40ss%27w@db.internal:5433/litellm"
@ -55,17 +67,24 @@ def _query(url: str) -> dict[str, str]:
class TestPlanPgBouncer:
def test_upstream_credentials_and_timeouts_move_into_the_pgbouncer_config(self):
def test_upstream_route_and_timeouts_move_into_the_pgbouncer_config_without_the_password(self):
ini: Final = _ini(_plan())
assert ini["databases"]["litellm"] == (
"host='db.internal' port=5433 dbname='litellm' user='app' password='p@ss''w' "
"host='db.internal' port=5433 dbname='litellm' user='app' "
"connect_query='SET statement_timeout TO ''7000''; SET lock_timeout TO ''3000'''"
)
assert _plan().userlist == '"app" "p@ss\'w"\n'
def test_the_auth_file_holds_the_upstream_password_and_the_pool_users_own(self):
plan: Final = _plan()
assert plan.upstream_password == "p@ss'w"
assert plan.userlist("p@ss'w") == f'"app" "p@ss\'w"\n"litellm_pgbouncer" "{plan.pool_password}"\n'
def test_a_token_with_quotes_is_escaped_the_way_pgbouncer_reads_it(self):
assert _plan().userlist('to"ken').startswith('"app" "to""ken"\n')
def test_an_upstream_without_a_port_is_reached_on_the_postgres_default(self):
ini: Final = _ini(_plan("postgresql://app:pw@db/litellm"))
assert ini["databases"]["litellm"] == "host='db' port=5432 dbname='litellm' user='app' password='pw'"
assert ini["databases"]["litellm"] == "host='db' port=5432 dbname='litellm' user='app'"
def test_pool_is_sized_from_settings_in_transaction_mode(self):
pgb: Final = _ini(_plan())["pgbouncer"]
@ -79,20 +98,24 @@ class TestPlanPgBouncer:
assert pgb["auth_file"] == "/run/pgb/userlist.txt"
assert pgb["unix_socket_dir"] == "/run/pgb"
def test_the_app_user_can_read_the_pgbouncer_console(self):
assert _ini(_plan())["pgbouncer"]["stats_users"] == "app"
def test_the_pool_user_can_read_the_pgbouncer_console(self):
assert _ini(_plan())["pgbouncer"]["stats_users"] == "litellm_pgbouncer"
@pytest.mark.parametrize("user", ["app,admin", "app%20admin", "app%09admin"])
def test_a_user_pgbouncer_would_split_into_several_console_users_is_refused(self, user: str):
outcome: Final = plan_pgbouncer(f"postgresql://{user}:pw@db/litellm", SETTINGS, Path("/run/pgb"), None)
def test_a_database_user_named_like_the_pool_user_is_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://litellm_pgbouncer:pw@db/litellm", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "stats_users" in outcome.reason
assert "litellm_pgbouncer" in outcome.reason
def test_pooled_url_points_prisma_at_loopback_without_prepared_statements(self):
pooled: Final = urllib.parse.urlsplit(_plan().pooled_url)
def test_pooled_url_points_prisma_at_loopback_as_the_pool_user_without_prepared_statements(self):
plan: Final = _plan()
pooled: Final = urllib.parse.urlsplit(plan.pooled_url)
assert (pooled.hostname, pooled.port, pooled.path) == ("127.0.0.1", 6543, "/litellm")
assert (pooled.username, pooled.password) == ("app", "p%40ss%27w")
assert _query(_plan().pooled_url) == {
assert (pooled.username, pooled.password) == ("litellm_pgbouncer", plan.pool_password)
assert len(plan.pool_password) >= 32
assert "p%40ss" not in plan.pooled_url
assert _query(plan.pooled_url) == {
"schema": "public",
"connection_limit": "10",
"pool_timeout": "20",
@ -118,6 +141,9 @@ class TestPlanPgBouncer:
assert "server_tls_ca_file" not in pgb
assert plan.ca_source is None
def test_every_plan_gets_its_own_pool_password(self):
assert _plan().pool_password != _plan().pool_password
def test_no_tls_params_default_to_prefer(self):
assert _ini(_plan("postgresql://app:pw@db/litellm"))["pgbouncer"]["server_tls_sslmode"] == "prefer"
@ -138,15 +164,20 @@ class TestPlanPgBouncer:
@pytest.mark.parametrize(
"url",
[
"postgresql://app@db/litellm",
"postgresql://app:pw@db",
"postgresql://:pw@db/litellm",
"postgresql://app:pw@/litellm",
],
)
def test_urls_missing_forwardable_credentials_are_refused(self, url: str):
def test_urls_missing_a_route_are_refused(self, url: str):
outcome: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), None)
assert isinstance(outcome, PgBouncerError)
def test_a_url_without_a_password_plans_for_a_token_to_be_installed_later(self):
plan: Final = _plan("postgresql://app@db/litellm")
assert plan.upstream_password is None
assert plan.userlist("minted").startswith('"app" "minted"\n')
def test_every_options_spelling_becomes_a_set_statement(self):
options: Final = urllib.parse.quote("-c a=1 -cb=2 --c=3")
ini: Final = _ini(_plan(f"postgresql://app:pw@db/litellm?options={options}"))
@ -167,12 +198,14 @@ class TestPlanPgBouncer:
class TestWritePgBouncerFiles:
def test_files_hold_the_plan_and_are_private_to_the_owner(self, tmp_path: Path):
plan: Final = _plan("postgresql://app:pw@db/litellm")
ini_path: Final = write_pgbouncer_files(plan, tmp_path, None)
ini_path: Final = write_pgbouncer_ini(plan, tmp_path, None)
assert isinstance(ini_path, Path), ini_path
userlist_path: Final = write_userlist(plan.userlist("pw"), tmp_path, None)
assert ini_path == tmp_path / "pgbouncer.ini"
assert userlist_path == tmp_path / "userlist.txt"
assert ini_path.read_text() == plan.ini
assert (tmp_path / "userlist.txt").read_text() == plan.userlist
for path in (ini_path, tmp_path / "userlist.txt"):
assert userlist_path.read_text() == plan.userlist("pw")
for path in (ini_path, userlist_path):
assert stat.S_IMODE(path.stat().st_mode) == 0o600
assert not (tmp_path / "server-ca.pem").exists()
@ -185,7 +218,7 @@ class TestWritePgBouncerFiles:
f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={bundle}", SETTINGS, runtime_dir, None
)
assert isinstance(plan, PgBouncerPlan), plan
ini_path: Final = write_pgbouncer_files(plan, runtime_dir, None)
ini_path: Final = write_pgbouncer_ini(plan, runtime_dir, None)
assert isinstance(ini_path, Path), ini_path
ca_file: Final = Path(_ini(plan)["pgbouncer"]["server_tls_ca_file"])
assert ca_file.parent == runtime_dir
@ -199,11 +232,32 @@ class TestWritePgBouncerFiles:
None,
)
assert isinstance(plan, PgBouncerPlan), plan
outcome: Final = write_pgbouncer_files(plan, tmp_path, None)
outcome: Final = write_pgbouncer_ini(plan, tmp_path, None)
assert isinstance(outcome, PgBouncerError)
assert "missing.pem" in outcome.reason
assert not (tmp_path / "pgbouncer.ini").exists()
def test_rewriting_the_userlist_replaces_it_whole_and_leaves_nothing_else_behind(self, tmp_path: Path):
write_userlist('"app" "first"\n', tmp_path, None)
with open(tmp_path / "userlist.txt", encoding="utf-8") as before_rewrite:
write_userlist('"app" "second"\n', tmp_path, None)
assert before_rewrite.read() == '"app" "first"\n'
assert (tmp_path / "userlist.txt").read_text() == '"app" "second"\n'
assert stat.S_IMODE((tmp_path / "userlist.txt").stat().st_mode) == 0o600
assert sorted(path.name for path in tmp_path.iterdir()) == ["userlist.txt"]
class TestPooledUrlMarker:
def test_exporting_the_pooled_url_marks_it_for_the_workers(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "")
monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR)
monkeypatch.setenv("DATABASE_URL", "postgresql://app:token@db/litellm")
assert not database_url_is_pooled()
export_pooled_database_url("postgresql://litellm_pgbouncer:pw@127.0.0.1:6432/litellm?pgbouncer=true")
assert os.environ["DATABASE_URL"] == "postgresql://litellm_pgbouncer:pw@127.0.0.1:6432/litellm?pgbouncer=true"
assert database_url_is_pooled()
assert PGBOUNCER_POOLED_ENV_VAR == "LITELLM_PGBOUNCER_POOLED_DATABASE_URL"
def _bound_port(sock: socket.socket) -> int:
return cast(tuple[str, int], sock.getsockname())[1]
@ -222,20 +276,23 @@ def _fake_pooler(
port_file: Path | None = None,
bind_delay_seconds: float = 0.0,
version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable",
auth_log: Path | None = None,
) -> Path:
"""An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.<port>`` in the socket dir.
Port and socket dir come from the ini it is given, else from ``port`` and
``tmp_path``. With ``port_file`` each start reads the port from that file
instead. ``bind_delay_seconds`` holds the bind back, like a slow start.
``--version`` prints ``version_banner``.
``--version`` prints ``version_banner``. With ``auth_log`` it appends the
``auth_file`` it reads at startup and on every SIGHUP, one line per read,
like PgBouncer loading its credentials.
"""
script: Final = tmp_path / "fake-pgbouncer"
script.write_text(
textwrap.dedent(
f"""\
#!{sys.executable}
import configparser, os, pathlib, select, socket, sys, time
import configparser, os, pathlib, select, signal, socket, sys, time
if sys.argv[1:] == ["--version"]:
print({version_banner!r})
sys.exit(0)
@ -243,6 +300,12 @@ def _fake_pooler(
sys.exit(3)
ini = configparser.ConfigParser()
ini.read(sys.argv[1:2])
if not {auth_log is None!r}:
def load_auth_file(*_):
with open({str(auth_log)!r}, "a") as log:
log.write(repr(pathlib.Path(ini.get("pgbouncer", "auth_file")).read_text()) + "\\n")
load_auth_file()
signal.signal(signal.SIGHUP, load_auth_file)
port = ini.getint("pgbouncer", "listen_port", fallback={port})
if not {port_file is None!r}:
port = int(pathlib.Path({str(port_file)!r}).read_text())
@ -511,6 +574,118 @@ class TestPgBouncerProcess:
os.kill(pid, 0)
NOW: Final = datetime(2026, 9, 10, 12, 0, tzinfo=timezone.utc)
ENDPOINT: Final = IAMEndpoint(host="db", port="5432", user="app", name="litellm")
def _entra_jwt(expires_at: datetime) -> str:
payload: Final = base64.urlsafe_b64encode(json.dumps({"exp": int(expires_at.timestamp())}).encode())
return f"aGVhZGVy.{payload.rstrip(b'=').decode()}.c2ln"
def _token_source(*tokens: str | Exception) -> PgBouncerTokenSource:
"""A token source handing out ``tokens`` in order, raising the exceptions among them, then repeating the last."""
pending: Final = deque(tokens)
def provide() -> str:
outcome: Final = pending.popleft() if len(pending) > 1 else pending[0]
if isinstance(outcome, Exception):
raise outcome
return outcome
return PgBouncerTokenSource(auth=AzureEntraTokenAuth(token_provider=provide), endpoint=ENDPOINT)
class TestPgBouncerTokenRefresher:
def _refresher(
self,
source: PgBouncerTokenSource,
installed: list[str],
install: Callable[[str], None] | None = None,
**timing: float,
) -> PgBouncerTokenRefresher:
return PgBouncerTokenRefresher(
source,
install if install is not None else installed.append,
now=lambda: NOW.replace(tzinfo=None),
**timing,
)
def test_the_next_refresh_is_due_a_buffer_before_the_token_expires(self):
installed: Final[list[str]] = []
token: Final = _entra_jwt(NOW + timedelta(hours=1))
refresher: Final = self._refresher(_token_source(token), installed, buffer_seconds=180)
assert refresher.refresh() == 3600 - 180
assert installed == [token]
def test_a_token_whose_expiry_cannot_be_read_is_refreshed_on_the_fallback_interval(self):
installed: Final[list[str]] = []
refresher: Final = self._refresher(_token_source("opaque token"), installed, fallback_seconds=600)
assert refresher.refresh() == 600
assert installed == ["opaque token"]
def test_a_token_already_inside_the_buffer_is_refreshed_after_the_retry_delay(self):
token: Final = _entra_jwt(NOW + timedelta(seconds=100))
refresher: Final = self._refresher(_token_source(token), [], buffer_seconds=180, retry_seconds=30)
assert refresher.refresh() == 30
def test_the_token_reaches_the_auth_file_in_wire_form_not_url_encoded(self):
installed: Final[list[str]] = []
self._refresher(_token_source("to ken/with+odd=chars"), installed).refresh()
assert installed == ["to ken/with+odd=chars"]
def test_a_failed_mint_is_reported_and_installs_nothing(self):
installed: Final[list[str]] = []
outcome: Final = self._refresher(_token_source(RuntimeError("no credential")), installed).refresh()
assert isinstance(outcome, PgBouncerError)
assert "Azure Entra token" in outcome.reason
assert "no credential" in outcome.reason
assert installed == []
def test_a_token_pgbouncer_cannot_hold_is_refused(self):
installed: Final[list[str]] = []
outcome: Final = self._refresher(_token_source("x" * 2048), installed).refresh()
assert isinstance(outcome, PgBouncerError)
assert "2047" in outcome.reason
assert installed == []
def test_an_auth_file_that_cannot_be_written_is_reported_not_raised(self):
def refuse(_: str) -> None:
raise PermissionError("read-only runtime dir")
outcome: Final = self._refresher(_token_source("token"), [], install=refuse).refresh()
assert isinstance(outcome, PgBouncerError)
assert "read-only runtime dir" in outcome.reason
def test_start_fails_when_the_first_token_cannot_be_minted_and_schedules_nothing(self):
installed: Final[list[str]] = []
refresher: Final = self._refresher(
_token_source(RuntimeError("no credential"), "later"), installed, fallback_seconds=0.05
)
assert isinstance(refresher.start(), PgBouncerError)
time.sleep(0.3)
assert installed == []
def test_a_failed_renewal_keeps_the_previous_token_until_the_retry_succeeds(self, caplog: pytest.LogCaptureFixture):
installed: Final[list[str]] = []
refresher: Final = self._refresher(
_token_source("first", RuntimeError("blip"), "third"),
installed,
fallback_seconds=0.05,
retry_seconds=0.05,
)
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
assert refresher.start() is None
assert installed == ["first"]
assert _wait_until(lambda: "third" in installed)
assert installed[:2] == ["first", "third"]
assert any("keeps its current Azure Entra token" in record.message for record in caplog.records)
refresher.stop()
settled: Final = len(installed)
time.sleep(0.3)
assert len(installed) == settled
def _runtime_dir_listening_on(port: int) -> Path:
matches: Final = tuple(
ini.parent
@ -524,10 +699,21 @@ def _runtime_dir_listening_on(port: int) -> Path:
class TestStartInContainerPgBouncer:
def test_returns_the_loopback_url_once_the_pooler_listens(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
auth_log: Final = tmp_path / "auth.log"
binary: Final = _fake_pooler(tmp_path, port, auth_log=auth_log)
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm?connection_limit=5")
assert pooled == f"postgresql://app:pw@127.0.0.1:{port}/litellm?connection_limit=5&pgbouncer=true"
assert isinstance(pooled, str), pooled
parsed: Final = urllib.parse.urlsplit(pooled)
assert (parsed.username, parsed.hostname, parsed.port, parsed.path) == (
"litellm_pgbouncer",
"127.0.0.1",
port,
"/litellm",
)
assert _query(pooled) == {"connection_limit": "5", "pgbouncer": "true"}
assert _listening(port)
assert auth_log.read_text() == repr(f'"app" "pw"\n"litellm_pgbouncer" "{parsed.password}"\n') + "\n"
@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning")
def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, tmp_path: Path):
@ -561,21 +747,79 @@ class TestStartInContainerPgBouncer:
def test_a_bad_upstream_url_is_reported_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app@db/litellm")
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db")
assert isinstance(outcome, PgBouncerError)
assert not _listening(port)
def test_token_auth_is_refused_without_starting_anything(self, tmp_path: Path):
def test_a_passwordless_url_without_token_auth_is_refused_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
outcome: Final = start_in_container_pgbouncer(
settings, "postgresql://app:pw@db/litellm", token_auth_enabled=True
)
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app@db/litellm")
assert isinstance(outcome, PgBouncerError)
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
assert "AZURE_POSTGRESQL_AUTH" in outcome.reason
assert not _listening(port)
def test_token_auth_mints_the_first_token_into_the_auth_file_before_the_pooler_starts(self, tmp_path: Path):
port: Final = _free_port()
auth_log: Final = tmp_path / "auth.log"
binary: Final = _fake_pooler(tmp_path, port, auth_log=auth_log)
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
token: Final = _entra_jwt(datetime.now(tz=timezone.utc) + timedelta(hours=1))
pooled: Final = start_in_container_pgbouncer(
settings,
"postgresql://app:stale-token@db/litellm",
token_auth=AzureEntraTokenAuth(token_provider=lambda: token),
)
assert isinstance(pooled, str), pooled
parsed: Final = urllib.parse.urlsplit(pooled)
assert parsed.username == "litellm_pgbouncer"
assert token not in pooled
assert _listening(port)
assert auth_log.read_text() == repr(f'"app" "{token}"\n"litellm_pgbouncer" "{parsed.password}"\n') + "\n"
def test_a_first_token_that_cannot_be_minted_is_reported_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
def fail() -> str:
raise RuntimeError("no Azure credential")
outcome: Final = start_in_container_pgbouncer(
settings, "postgresql://app@db/litellm", token_auth=AzureEntraTokenAuth(token_provider=fail)
)
assert isinstance(outcome, PgBouncerError)
assert "no Azure credential" in outcome.reason
assert not _listening(port)
def test_a_renewed_token_is_written_and_picked_up_by_the_running_and_by_a_restarted_pooler(self, tmp_path: Path):
port: Final = _free_port()
auth_log: Final = tmp_path / "auth.log"
plan: Final = plan_pgbouncer(
"postgresql://app@db/litellm", PgBouncerSettings(enabled=True, port=port), tmp_path, None
)
assert isinstance(plan, PgBouncerPlan), plan
ini_path: Final = write_pgbouncer_ini(plan, tmp_path, None)
write_userlist(plan.userlist("first"), tmp_path, None)
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, auth_log=auth_log)), str(ini_path)),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
install_pgbouncer_token(plan, tmp_path, None, pooler, "second")
assert _wait_until(lambda: auth_log.read_text().count("\n") == 2)
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert auth_log.read_text().splitlines() == [
repr(plan.userlist("first")),
repr(plan.userlist("second")),
repr(plan.userlist("second")),
]
def test_a_pgbouncer_that_survives_a_failed_tcp_bind_is_refused_without_starting(self, tmp_path: Path):
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.18.1\nlibevent 2.1.12-stable")
@ -590,9 +834,9 @@ class TestStartInContainerPgBouncer:
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0")
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
assert start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") == (
f"postgresql://app:pw@127.0.0.1:{port}/litellm?pgbouncer=true"
)
pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm")
assert isinstance(pooled, str), pooled
assert urllib.parse.urlsplit(pooled).port == port
assert _listening(port)

View file

@ -1,5 +1,6 @@
import datetime as real_datetime
import smtplib
from typing import Final
import pytest
from fastapi import HTTPException
@ -8,7 +9,7 @@ from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
@ -2207,3 +2208,49 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp
assert recorder.received_traceback is not None
assert provider_key not in recorder.received_traceback
assert "REDACTED" in recorder.received_traceback
class TestPrismaClientTokenAuthBehindThePool:
"""Behind the in-container pool the supervisor renews the writer's database
token and hands the workers a loopback URL with a static password, so the
writer wrapper must not run its own refresh loop. The reader is not pooled
and keeps refreshing its own token."""
UPSTREAM: Final = "postgresql://litellm:TOKEN@db.internal:5432/litellm"
READER: Final = "postgresql://litellm:TOKEN@reader.internal:5432/litellm"
def _client(self, monkeypatch: pytest.MonkeyPatch, pooled: bool) -> PrismaClient:
from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR
monkeypatch.delenv("AZURE_POSTGRESQL_AUTH", raising=False)
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
monkeypatch.setenv("DATABASE_URL", self.UPSTREAM)
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", self.READER)
if pooled:
monkeypatch.setenv(PGBOUNCER_POOLED_ENV_VAR, "true")
else:
monkeypatch.delenv(PGBOUNCER_POOLED_ENV_VAR, raising=False)
rds: Final = MagicMock()
rds.generate_db_auth_token.return_value = "TOKEN"
with patch("boto3.client", return_value=rds):
return PrismaClient(database_url=self.UPSTREAM, proxy_logging_obj=MagicMock(spec=ProxyLogging))
def test_a_pooled_writer_leaves_token_refresh_to_the_pooler_while_the_reader_keeps_its_own(
self, monkeypatch: pytest.MonkeyPatch
):
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
client = self._client(monkeypatch, pooled=True)
assert isinstance(client.db, RoutingPrismaWrapper)
assert client.db.writer.iam_token_db_auth is False
assert client.db.reader.iam_token_db_auth is True
assert client.token_auth is not None
def test_an_unpooled_writer_still_refreshes_its_own_token(self, monkeypatch: pytest.MonkeyPatch):
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
client = self._client(monkeypatch, pooled=False)
assert isinstance(client.db, RoutingPrismaWrapper)
assert client.db.writer.iam_token_db_auth is True
assert client.db.reader.iam_token_db_auth is True