mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
ef1a37795c
commit
3b93da414e
15 changed files with 703 additions and 131 deletions
|
|
@ -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))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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" {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue