diff --git a/gateway/launch.py b/gateway/launch.py index 6b28fafdcf6..d67432caee1 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -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)) diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index 4ad3cfe0484..dc47bc0b0a1 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -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" diff --git a/helm/litellm/tests/connection_pool_tests.yaml b/helm/litellm/tests/connection_pool_tests.yaml index 31be7575f3c..6dd9c274a8c 100644 --- a/helm/litellm/tests/connection_pool_tests.yaml +++ b/helm/litellm/tests/connection_pool_tests.yaml @@ -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 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 45ab5229f38..677e21a5a51 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -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: diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index d1e4b3e92b8..e93bc96da69 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -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 diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py index 7792867600c..eac80f7a712 100644 --- a/litellm/proxy/db/pgbouncer.py +++ b/litellm/proxy/db/pgbouncer.py @@ -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 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 16d76ff0415..c2d60cd5488 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2cc58ca3110..253494b02f4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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, ) diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 986428151e0..5ca45944c54 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -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:@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 diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 149842c14ff..7176921e8fa 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -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." } } diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl index 38aa7051134..ab0a1cfc076 100644 --- a/terraform/litellm/aws/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -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" { diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index ec8363dbbaf..f57bb1ab7f6 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -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 diff --git a/tests/test_gateway/test_launch.py b/tests/test_gateway/test_launch.py index 964a7021416..6039aaafc2b 100644 --- a/tests/test_gateway/test_launch.py +++ b/tests/test_gateway/test_launch.py @@ -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 diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py index da1d4c91f09..7b5a0bf10f9 100644 --- a/tests/test_litellm/proxy/db/test_pgbouncer.py +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -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.`` 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) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 1287cf23157..646596c5b88 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -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