Merge 36627a61a3 into e768ad55ce

.circleci/config.yml

@@ -3191,6 +3191,13 @@ jobs:
           image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
           server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
       - start_redis
+      - when:
+          condition:
+            equal: [database, << parameters.suite >>]
+          steps:
+            - run:
+                name: Install the system libpq the pure Python psycopg loads
+                command: sudo apt-get update -qq && sudo apt-get install -y --no-install-recommends libpq5
       - run:
           name: Run owned integration contracts
           command: bash .circleci/scripts/run_integration.sh << parameters.suite >> << parameters.mode >>

Dockerfile

@@ -126,7 +126,7 @@ USER root
 RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
 
 # node (without npm) is required by the prisma CLI at runtime
-RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
+RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent libpq-18
 COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 WORKDIR /app

backend/Dockerfile

@@ -75,7 +75,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
 USER root
 
 RUN for i in 1 2 3; do \
-      apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
+      apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libpq-18 && break; \
       [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
       sleep 5; \
     done

docker/Dockerfile.database

@@ -117,7 +117,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
 USER root
 
 # node (without npm) is required by the prisma CLI at runtime
-RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
+RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent libpq-18
 COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 
 WORKDIR /app

docker/Dockerfile.non_root

@@ -129,7 +129,7 @@ RUN for i in 1 2 3; do \
       apk upgrade --no-cache && break || sleep 5; \
     done && \
     for i in 1 2 3; do \
-      apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
+      apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent libpq-18 && break || sleep 5; \
     done
 COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
 

gateway/Dockerfile

@@ -93,7 +93,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
 USER root
 
 RUN for i in 1 2 3; do \
-      apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
+      apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent libpq-18 && break; \
       [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
       sleep 5; \
     done

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -670,13 +670,15 @@ class ProxyExtrasDBManager:
 
         try:
             import psycopg
-        except ImportError:
+        except ImportError as exc:
             logger.warning(
-                "psycopg is not installed; skipping the LiteLLM_SpendLogs "
-                "partition check. If this table is partitioned (see "
+                "psycopg is not installed or found no libpq (%s); skipping the "
+                "LiteLLM_SpendLogs partition check. If this table is partitioned (see "
                 "db_scripts/partition_spend_logs.sql), schema reconciliation "
                 "will try to rewrite its primary key and fail. Install the "
-                "litellm[extra_proxy] extra, which now includes psycopg."
+                "litellm[extra_proxy] extra, which includes psycopg, plus the "
+                "system libpq it loads (the litellm images ship it).",
+                exc,
             )
             return False
 
@@ -882,10 +884,11 @@ class ProxyExtrasDBManager:
         try:
             import psycopg
             from psycopg import sql
-        except ImportError:
+        except ImportError as exc:
             logger.warning(
-                "psycopg is not installed; skipping the invalid index check. "
-                "Install the litellm[extra_proxy] extra, which includes psycopg."
+                "psycopg is not installed or found no libpq (%s); skipping the invalid index check. "
+                "Install the litellm[extra_proxy] extra, which includes psycopg, plus the system libpq.",
+                exc,
             )
             return False
 

migrations/Dockerfile

@@ -89,7 +89,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
 USER root
 
 RUN for i in 1 2 3; do \
-      apk add --no-cache bash openssl tzdata python-3.13 nodejs libsndfile libatomic && break; \
+      apk add --no-cache bash openssl tzdata python-3.13 nodejs libsndfile libatomic libpq-18 && break; \
       [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
       sleep 5; \
     done

pyproject.toml

@@ -103,9 +103,9 @@ extra_proxy = [
     "prisma>=0.11.0,<1.0",
     # Used by ProxyExtrasDBManager.spend_logs_is_partitioned() to detect a
     # partitioned LiteLLM_SpendLogs and keep schema reconciliation from
-    # fighting its composite primary key.
+    # fighting its composite primary key. No psycopg-binary: that wheel bundles
+    # its own OpenSSL; the image ships libpq for the pure Python driver.
     "psycopg>=3.2,<4.0",
-    "psycopg-binary>=3.2,<4.0",
     "azure-identity>=1.25.2,<2.0",
     "azure-keyvault-secrets>=4.10.0,<5.0",
     # Not in PyPI proxy extra.

tests/integration/database/test_spend_logs_partition_driver.py

@@ -0,0 +1,144 @@
+import os
+import socket
+import subprocess
+import sys
+import uuid
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Final
+from urllib.parse import urlsplit, urlunsplit
+
+import psycopg
+from integration._support.client import Gateway
+from integration._support.process import owned_proxy_process
+from packaging.requirements import Requirement
+from psycopg import sql
+
+REPO_ROOT: Final = Path(__file__).resolve().parents[3]
+PRISMA_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras"
+PARTITION_SCRIPT: Final = REPO_ROOT / "db_scripts" / "partition_spend_logs.sql"
+IMAGE_DOCKERFILES: Final = (
+    REPO_ROOT / "Dockerfile",
+    REPO_ROOT / "docker" / "Dockerfile.database",
+    REPO_ROOT / "docker" / "Dockerfile.non_root",
+    REPO_ROOT / "backend" / "Dockerfile",
+    REPO_ROOT / "gateway" / "Dockerfile",
+    REPO_ROOT / "migrations" / "Dockerfile",
+)
+PURE_PYTHON_DRIVER: Final = {"PSYCOPG_IMPL": "python"}
+
+
+def locked_extra_proxy_packages() -> frozenset[str]:
+    export: Final = subprocess.run(
+        ["uv", "export", "--frozen", "--no-hashes", "--no-dev", "--no-emit-project", "--extra", "extra_proxy"],
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=120,
+        cwd=REPO_ROOT,
+    )
+    return frozenset(
+        Requirement(line.split(" ;")[0]).name.lower().replace("_", "-")
+        for line in export.stdout.splitlines()
+        if line and not line.startswith(("#", "-", " "))
+    )
+
+
+def free_port() -> int:
+    with socket.socket() as reserve:
+        reserve.bind(("127.0.0.1", 0))
+        return reserve.getsockname()[1]
+
+
+def runtime_stage(dockerfile: Path) -> str:
+    return dockerfile.read_text().rsplit("FROM $LITELLM_RUNTIME_IMAGE", maxsplit=1)[1]
+
+
+@contextmanager
+def partitioned_database() -> Iterator[str]:
+    name: Final = f"integration_partitioned_{uuid.uuid4().hex}"
+    admin_url: Final = os.environ["DATABASE_URL"]
+    database_url: Final = urlunsplit(urlsplit(admin_url)._replace(path=f"/{name}"))
+    with psycopg.connect(admin_url, autocommit=True) as admin:
+        admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
+        try:
+            subprocess.run(
+                [
+                    sys.executable,
+                    "-I",
+                    "-m",
+                    "prisma",
+                    "migrate",
+                    "deploy",
+                    "--schema",
+                    str(PRISMA_DIR / "schema.prisma"),
+                ],
+                check=True,
+                capture_output=True,
+                text=True,
+                timeout=300,
+                env={**os.environ, "DATABASE_URL": database_url},
+            )
+            with psycopg.connect(database_url, autocommit=True) as connection:
+                connection.execute(PARTITION_SCRIPT.read_text())
+            yield database_url
+        finally:
+            admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
+
+
+def test_proxy_image_extra_ships_psycopg_without_the_binary_wheel_and_uses_the_image_libpq() -> None:
+    packages: Final = locked_extra_proxy_packages()
+    assert "psycopg" in packages, sorted(packages)
+    assert "psycopg-binary" not in packages, sorted(packages)
+    for dockerfile in IMAGE_DOCKERFILES:
+        assert "libpq" in runtime_stage(dockerfile), f"{dockerfile.name} runtime stage installs no libpq for psycopg"
+
+
+def test_pure_python_psycopg_detects_partitioned_spend_logs_and_refuses_db_push(gateway: Gateway) -> None:
+    with partitioned_database() as database_url:
+        proxy: Final = subprocess.run(
+            [
+                sys.executable,
+                "-P",
+                "-m",
+                "integration._support.proxy",
+                "--config",
+                "tests/integration/proxy_config.yaml",
+                "--host",
+                "127.0.0.1",
+                "--port",
+                str(free_port()),
+                "--use_prisma_db_push",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=300,
+            cwd=REPO_ROOT,
+            env={
+                **os.environ,
+                **PURE_PYTHON_DRIVER,
+                "DATABASE_URL": database_url,
+                "LITELLM_MASTER_KEY": gateway.key,
+                "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"),
+            },
+        )
+        output: Final = proxy.stdout + proxy.stderr
+        assert proxy.returncode != 0, output
+        assert "LiteLLM_SpendLogs is a partitioned table" in output, output
+        assert "psycopg is not installed" not in output, output
+        with psycopg.connect(database_url) as connection:
+            partitioned: Final = connection.execute(
+                "SELECT 1 FROM pg_partitioned_table pt JOIN pg_class c ON c.oid = pt.partrelid "
+                "WHERE c.relname = 'LiteLLM_SpendLogs'"
+            ).fetchone()
+        assert partitioned is not None, "db push rewrote the partitioned LiteLLM_SpendLogs table"
+
+
+def test_pure_python_psycopg_lets_an_unpartitioned_proxy_boot_and_serve(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, PURE_PYTHON_DRIVER) as proxy:
+        health: Final = proxy.gateway.request("GET", "/health/readiness")
+        assert health.status_code == 200, health.text
+        assert health.json()["db"] == "connected", health.text
+        log: Final = proxy.log.read_text()
+        assert "psycopg is not installed" not in log, log

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -708,6 +708,23 @@ class TestSpendLogsPartitionDetectionMissingPsycopg:
             "psycopg is not installed" in record.message for record in caplog.records
         )
 
+    def test_missing_libpq_warning_names_the_libpq_error(self, monkeypatch, caplog):
+        class NoLibpq:
+            @staticmethod
+            def find_spec(name, path=None, target=None):
+                if name == "psycopg":
+                    raise ImportError("no pq wrapper available.\nAttempts made:\n- couldn't import psycopg 'python'")
+
+        for name in [m for m in sys.modules if m == "psycopg" or m.startswith("psycopg.")]:
+            monkeypatch.delitem(sys.modules, name)
+        monkeypatch.setattr(sys, "meta_path", [NoLibpq(), *sys.meta_path])
+        monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db")
+        with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
+            assert ProxyExtrasDBManager.spend_logs_is_partitioned() is False
+        warning: Final = next(r.getMessage() for r in caplog.records if "partition check" in r.getMessage())
+        assert "no pq wrapper available" in warning
+        assert "libpq" in warning
+
 
 _ATTEMPT_BUDGET = 4
 _P3009_MIGRATION_NAME = "20260415120000_health_check_latest_per_model_index"

uv.lock

@@ -4546,7 +4546,6 @@ extra-proxy = [
     { name = "google-cloud-kms" },
     { name = "prisma" },
     { name = "psycopg" },
-    { name = "psycopg-binary" },
     { name = "redisvl" },
     { name = "resend" },
 ]
@@ -4817,7 +4816,6 @@ requires-dist = [
     { name = "prisma", marker = "extra == 'extra-proxy'", specifier = ">=0.11.0,<1.0" },
     { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" },
     { name = "psycopg", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" },
-    { name = "psycopg-binary", marker = "extra == 'extra-proxy'", specifier = ">=3.2,<4.0" },
     { name = "pydantic", marker = "python_full_version < '3.14'", specifier = ">=2.11.0,<3.0.0" },
     { name = "pydantic", marker = "python_full_version >= '3.14'", specifier = ">=2.12.0,<3.0.0" },
     { name = "pydantic", marker = "extra == 'mcp'", specifier = ">=2.12.0,<3" },