fix(db): build the spend-logs api_key index without CONCURRENTLY so partitioned tables can migrate

Postgres rejects a concurrent index build on a partitioned parent, so the migration aborted on deployments that ran db_scripts/partition_spend_logs.sql. Also carry the index through the partition and unpartition runbooks, and cover the partitioned shape with a migration test that applies the shipped statement for real.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

db_scripts/partition_spend_logs.sql

@@ -48,6 +48,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
     RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
 ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
     RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
+ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
+    RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
 
 CREATE TABLE "LiteLLM_SpendLogs" (
     LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
@@ -73,6 +75,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
 CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
     ON "LiteLLM_SpendLogs" ("session_id");
 
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
+    ON "LiteLLM_SpendLogs" ("api_key", "startTime" DESC);
+
 -- Safety net: any row whose startTime has no explicit partition lands here so
 -- writes never fail. The cleanup job never drops the DEFAULT partition.
 CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"

db_scripts/unpartition_spend_logs.sql

@@ -40,6 +40,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
     RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
 ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
     RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
+ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
+    RENAME TO "LiteLLM_SpendLogs_partitioned_api_key_startTime_idx";
 
 CREATE TABLE "LiteLLM_SpendLogs" (
     LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
@@ -60,6 +62,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
 CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
     ON "LiteLLM_SpendLogs" ("session_id");
 
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
+    ON "LiteLLM_SpendLogs" ("api_key", "startTime" DESC);
+
 INSERT INTO "LiteLLM_SpendLogs"
 SELECT * FROM "LiteLLM_SpendLogs_partitioned"
 ON CONFLICT ("request_id") DO NOTHING;

litellm-proxy-extras/litellm_proxy_extras/migrations/20260818000000_add_spendlogs_api_key_start_time_idx/migration.sql

@@ -1,11 +1,14 @@
--- CreateIndex (CONCURRENTLY)
+-- CreateIndex
 --
 -- Disclaimer:
--- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
---   single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
--- - "LiteLLM_SpendLogs" is the largest table on most deployments, so a blocking build would
---   stall spend-log inserts for the whole build; if the build is interrupted, Postgres may
---   leave an INVALID index that must be dropped and recreated.
+-- - Not CONCURRENTLY on purpose: Postgres refuses a concurrent index build on a partitioned
+--   parent, and "LiteLLM_SpendLogs" is a partitioned parent on every deployment that ran
+--   db_scripts/partition_spend_logs.sql, so a concurrent build would abort the migration run
+--   there and block the rollout.
+-- - The build takes a write lock on "LiteLLM_SpendLogs" for its duration, and that is the largest
+--   table on most deployments. Operators who cannot pause spend-log inserts can build the index
+--   themselves before upgrading, CONCURRENTLY when the table is plain or per partition plus
+--   ATTACH when it is partitioned, and this statement then no-ops on the index name.
 -- - Do not edit this file after it has been applied to any database: Prisma checksums
 --   migrations; add a new migration instead.
-CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC);
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC);

tests/proxy_migration_tests/test_db_schema_migration.py

@@ -95,16 +95,79 @@ def test_spend_logs_declares_api_key_index(schema_path: Path):
     )
 
 
-def test_spend_logs_api_key_index_has_a_migration():
-    """A schema-only index never reaches deployments that run ``prisma migrate deploy``."""
-    migrations = Path("./litellm-proxy-extras/litellm_proxy_extras/migrations")
+MIGRATIONS_DIR = Path("./litellm-proxy-extras/litellm_proxy_extras/migrations")
+
+PARTITION_PROBE_SCHEMA = "litellm_spendlogs_partition_probe"
+
+
+def _api_key_index_migration() -> str:
     creating = [
         sql
-        for sql in (path.read_text() for path in migrations.glob("*/migration.sql"))
+        for sql in (path.read_text() for path in MIGRATIONS_DIR.glob("*/migration.sql"))
         if SPEND_LOGS_API_KEY_INDEX in sql
     ]
     assert len(creating) == 1, f"expected exactly one migration creating {SPEND_LOGS_API_KEY_INDEX}, found {len(creating)}"
-    assert 'ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC)' in creating[0]
-    assert "CONCURRENTLY" in creating[0], (
-        "LiteLLM_SpendLogs is the largest table on most deployments; a blocking build stalls spend-log inserts"
-    )
+    return creating[0]
+
+
+def test_spend_logs_api_key_index_has_a_migration():
+    """A schema-only index never reaches deployments that run ``prisma migrate deploy``."""
+    assert 'ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC)' in _api_key_index_migration()
+
+
+@pytest.mark.skipif(
+    "DATABASE_URL" not in os.environ,
+    reason="requires a postgres database (DATABASE_URL)",
+)
+def test_spend_logs_api_key_index_migration_applies_to_a_partitioned_table():
+    """``db_scripts/partition_spend_logs.sql`` leaves ``LiteLLM_SpendLogs`` a partitioned parent,
+    and Postgres rejects both a concurrent index build on a partitioned parent and any concurrent
+    build inside a transaction, so a ``CONCURRENTLY`` migration would abort the rollout of every
+    partitioned deployment. This applies the shipped statement against that shape for real.
+    """
+    probe_sql = f"""
+DROP SCHEMA IF EXISTS "{PARTITION_PROBE_SCHEMA}" CASCADE;
+CREATE SCHEMA "{PARTITION_PROBE_SCHEMA}";
+SET search_path TO "{PARTITION_PROBE_SCHEMA}";
+
+CREATE TABLE "LiteLLM_SpendLogs" (
+    request_id TEXT NOT NULL,
+    api_key TEXT,
+    "startTime" TIMESTAMP NOT NULL,
+    PRIMARY KEY (request_id, "startTime")
+) PARTITION BY RANGE ("startTime");
+CREATE TABLE "LiteLLM_SpendLogs_pdefault" PARTITION OF "LiteLLM_SpendLogs" DEFAULT;
+
+{_api_key_index_migration()}
+
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_indexes
+        WHERE schemaname = '{PARTITION_PROBE_SCHEMA}'
+          AND tablename = 'LiteLLM_SpendLogs'
+          AND indexname = '{SPEND_LOGS_API_KEY_INDEX}'
+    ) THEN
+        RAISE EXCEPTION 'migration did not create {SPEND_LOGS_API_KEY_INDEX} on the partitioned parent';
+    END IF;
+END
+$$;
+
+DROP SCHEMA "{PARTITION_PROBE_SCHEMA}" CASCADE;
+"""
+
+    temp_base = Path(tempfile.mkdtemp(prefix="litellm_partition_probe_"))
+    try:
+        sql_path = temp_base / "probe.sql"
+        sql_path.write_text(probe_sql)
+        applied = subprocess.run(
+            ["prisma", "db", "execute", "--url", os.environ["DATABASE_URL"], "--file", str(sql_path)],
+            capture_output=True,
+            text=True,
+        )
+        assert applied.returncode == 0, (
+            f"migration for {SPEND_LOGS_API_KEY_INDEX} failed on a partitioned LiteLLM_SpendLogs:\n"
+            f"{applied.stdout}\n{applied.stderr}"
+        )
+    finally:
+        shutil.rmtree(temp_base, ignore_errors=True)