From e34f43328c8f6e0bd0df10747f2453ac0433b684 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 31 Aug 2026 17:33:12 -0700 Subject: [PATCH] fix(proxy): ship psycopg so partitioned SpendLogs detection actually runs (#38994) ProxyExtrasDBManager.spend_logs_is_partitioned() (#38452) silently returns False when psycopg can't be imported, and psycopg was never added to the extra_proxy install, so every production image lacks it. Schema reconciliation then generates the unfiltered primary-key rewrite against a genuinely partitioned LiteLLM_SpendLogs and Postgres rejects it, exactly the failure the fix was meant to prevent. Ships psycopg via extra_proxy and logs a warning when it's still missing instead of failing silently. --- .../litellm_proxy_extras/utils.py | 7 ++++++ pyproject.toml | 5 +++++ .../test_litellm_proxy_extras_utils.py | 22 +++++++++++++++++++ uv.lock | 4 ++++ 4 files changed, 38 insertions(+) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b2dc0a52c8f..b8032dd0d28 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -512,6 +512,13 @@ class ProxyExtrasDBManager: try: import psycopg except ImportError: + logger.warning( + "psycopg is not installed; 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." + ) return False cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url) diff --git a/pyproject.toml b/pyproject.toml index 34c1fec1c11..96dcf1121bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,11 @@ cli = [ ] 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. + "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/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 498d0cb4723..b3d457707b8 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -681,3 +681,25 @@ class TestSpendLogsPartitionDetectionSchemaScope: def test_only_partitioned_relations_match(self, monkeypatch): query, _ = self._detect(monkeypatch, "postgresql://u:p@localhost:5432/db") assert "pg_partitioned_table" in query + + +class TestSpendLogsPartitionDetectionMissingPsycopg: + """psycopg ships in the `extra_proxy` install, but a stripped-down image + can still lack it. When it does, detection must fail closed to False + (never crash the migration path) and say so loudly, because a silent + False here is what let a genuinely partitioned LiteLLM_SpendLogs hit the + unfiltered primary-key rewrite in production.""" + + def test_missing_psycopg_returns_false(self, monkeypatch): + monkeypatch.setitem(sys.modules, "psycopg", None) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") + assert ProxyExtrasDBManager.spend_logs_is_partitioned() is False + + def test_missing_psycopg_logs_a_warning(self, monkeypatch, caplog): + monkeypatch.setitem(sys.modules, "psycopg", None) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:5432/db") + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + ProxyExtrasDBManager.spend_logs_is_partitioned() + assert any( + "psycopg is not installed" in record.message for record in caplog.records + ) diff --git a/uv.lock b/uv.lock index 8ef72116466..7df1a35b28a 100644 --- a/uv.lock +++ b/uv.lock @@ -4306,6 +4306,8 @@ extra-proxy = [ { name = "google-cloud-iam" }, { name = "google-cloud-kms" }, { name = "prisma" }, + { name = "psycopg" }, + { name = "psycopg-binary" }, { name = "redisvl" }, { name = "resend" }, ] @@ -4544,6 +4546,8 @@ requires-dist = [ { name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" }, { 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", specifier = ">=2.10.0,<3.0.0" }, { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" },