diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py index 7eebf32cde9..ede47fce0c3 100644 --- a/litellm/proxy/db/pgbouncer.py +++ b/litellm/proxy/db/pgbouncer.py @@ -28,6 +28,7 @@ from __future__ import annotations import atexit import os +import re import shlex import shutil import socket @@ -56,6 +57,8 @@ PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0 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_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)") 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 " @@ -272,6 +275,25 @@ def unix_socket_path(runtime_dir: Path, port: int) -> Path: return runtime_dir / f".s.PGSQL.{port}" +def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError: + """``(major, minor)`` from `` --version``. + + Readiness relies on PgBouncer exiting when it cannot bind its TCP port, + which it does from 1.19 on. Older releases log a warning and serve the unix + socket alone, so their socket would vouch for a port held by someone else. + """ + try: + output: Final = subprocess.run( + (binary, "--version"), capture_output=True, text=True, check=False, timeout=10 + ).stdout + except (OSError, subprocess.TimeoutExpired) as run_error: + return PgBouncerError(f"could not run {binary!r} --version: {run_error}") + found: Final = PGBOUNCER_VERSION_PATTERN.search(output) + if found is None: + return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}") + return int(found[1]), int(found[2]) + + def _end(process: subprocess.Popen[bytes]) -> None: if process.poll() is not None: return @@ -296,7 +318,8 @@ class PgBouncerProcess: A connect probe of ``port`` cannot tell the child from another process that grabbed the port after the availability check, so readiness also needs ``socket_path``: the unix socket PgBouncer creates in the private - runtime directory, which it only does once every TCP listener is bound. + runtime directory, which it only does once every TCP listener is bound + (PgBouncer 1.19 or newer, see ``pgbouncer_version``). """ def __init__( @@ -429,6 +452,15 @@ def start_in_container_pgbouncer( """ if token_auth_enabled: return PgBouncerError(PGBOUNCER_TOKEN_AUTH_CONFLICT) + version: Final = pgbouncer_version(settings.binary) + if isinstance(version, PgBouncerError): + return version + if version < PGBOUNCER_MIN_VERSION: + return PgBouncerError( + f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy " + f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} " + "or newer is required" + ) runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-")) atexit.register(shutil.rmtree, runtime_dir, ignore_errors=True) run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py index 28ca91e7fbe..9c7273a0afa 100644 --- a/tests/test_litellm/proxy/db/test_pgbouncer.py +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -21,6 +21,7 @@ from litellm.proxy.db.pgbouncer import ( PgBouncerPlan, PgBouncerProcess, PgBouncerSettings, + pgbouncer_version, plan_pgbouncer, start_in_container_pgbouncer, unix_socket_path, @@ -170,12 +171,14 @@ def _fake_pooler( exit_immediately: bool = False, port_file: Path | None = None, bind_delay_seconds: float = 0.0, + version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable", ) -> 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``. """ script: Final = tmp_path / "fake-pgbouncer" script.write_text( @@ -183,6 +186,9 @@ def _fake_pooler( f"""\ #!{sys.executable} import configparser, os, pathlib, select, socket, sys, time + if sys.argv[1:] == ["--version"]: + print({version_banner!r}) + sys.exit(0) if {exit_immediately!r}: sys.exit(3) ini = configparser.ConfigParser() @@ -481,6 +487,40 @@ class TestStartInContainerPgBouncer: assert "AZURE_POSTGRESQL_AUTH" in outcome.reason assert not _listening(port) + 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") + settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary)) + outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") + assert isinstance(outcome, PgBouncerError) + assert "PgBouncer 1.18" in outcome.reason + assert "1.19" in outcome.reason + assert not _listening(port) + + def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path): + 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" + ) + assert _listening(port) + + +class TestPgBouncerVersion: + def test_reads_major_and_minor_from_the_banner(self, tmp_path: Path): + assert pgbouncer_version(str(_fake_pooler(tmp_path, _free_port()))) == (1, 25) + + def test_a_binary_that_cannot_run_is_reported(self, tmp_path: Path): + outcome: Final = pgbouncer_version(str(tmp_path / "missing-pgbouncer")) + assert isinstance(outcome, PgBouncerError) + assert "missing-pgbouncer" in outcome.reason + + def test_a_banner_without_a_version_is_reported(self, tmp_path: Path): + outcome: Final = pgbouncer_version(str(_fake_pooler(tmp_path, _free_port(), version_banner="something else"))) + assert isinstance(outcome, PgBouncerError) + assert "something else" in outcome.reason + class TestPgBouncerSettings: def test_reads_the_litellm_pgbouncer_env_vars(self, monkeypatch: pytest.MonkeyPatch):