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/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 d4c07d56d90..3f5dff0a537 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -144,7 +144,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/pyproject.toml b/pyproject.toml index ba72378989a..899e1bb5bf8 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. Pure Python psycopg over the image's + # libpq: the psycopg-binary wheel bundles its own EOL OpenSSL 1.1.1k. "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..599d461aa2b --- /dev/null +++ b/tests/integration/database/test_spend_logs_partition_driver.py @@ -0,0 +1,140 @@ +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", +) +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, + "-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/uv.lock b/uv.lock index 85e2b6d4e52..a0bb18d7151 100644 --- a/uv.lock +++ b/uv.lock @@ -4545,7 +4545,6 @@ extra-proxy = [ { name = "google-cloud-kms" }, { name = "prisma" }, { name = "psycopg" }, - { name = "psycopg-binary" }, { name = "redisvl" }, { name = "resend" }, ] @@ -4816,7 +4815,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" },