diff --git a/.circleci/config.yml b/.circleci/config.yml index 1798abe9de5..7db0c4a066e 100644 --- a/.circleci/config.yml +++ b/.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 >> diff --git a/Dockerfile b/Dockerfile index 4dcecf3ea3d..8c53cbe19bf 100644 --- a/Dockerfile +++ b/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 diff --git a/backend/Dockerfile b/backend/Dockerfile index 59f836b55f8..7f0829223b7 100644 --- a/backend/Dockerfile +++ b/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 diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 61b6faae691..3e952f025a4 100644 --- a/docker/Dockerfile.database +++ b/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 diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index ca526e06834..147d80ed94c 100644 --- a/docker/Dockerfile.non_root +++ b/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 diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 8045a8b64cb..f5667bcbbd0 100644 --- a/gateway/Dockerfile +++ b/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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 245244250ee..2e06ee5cf01 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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 diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 255b94b0ea8..f4b144118f3 100644 --- a/migrations/Dockerfile +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 8f467513079..c55aab3029d 100644 --- a/pyproject.toml +++ b/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. diff --git a/tests/integration/database/test_spend_logs_partition_driver.py b/tests/integration/database/test_spend_logs_partition_driver.py new file mode 100644 index 00000000000..cdcfe9e6c54 --- /dev/null +++ b/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 diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index ea4a25283a1..030aa01291c 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/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" diff --git a/uv.lock b/uv.lock index b227ba183b5..f5d71e7c0de 100644 --- a/uv.lock +++ b/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" },