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/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)

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.

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
+        )

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" },