Merge pull request #44397 from BerriAI/litellm_fix_d03a47

fix(proxy-extras): keep the post-migration sanity check from building the hand-built SpendLogs indexes

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -115,6 +115,17 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
     re.IGNORECASE,
 )
 
+HAND_BUILT_SPEND_LOGS_INDEXES: Final = frozenset(
+    {
+        "LiteLLM_SpendLogs_api_key_startTime_idx",
+        "LiteLLM_SpendLogs_litellm_call_id_idx",
+    }
+)
+_CREATE_INDEX_RE: Final = re.compile(
+    r'^\s*CREATE\s+(?:UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?"(?P<index>[^"]+)"\s+ON\b',
+    re.IGNORECASE,
+)
+
 PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
     "LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
     "so its primary key must include the partition key (\"startTime\"). `prisma db push` "
@@ -165,6 +176,21 @@ def filter_partitioned_spend_logs_diff(diff_sql: str) -> str:
     return "".join(f"{statement};\n\n" for statement in kept)
 
 
+def filter_hand_built_spend_logs_index_diff(diff_sql: str) -> str:
+    """The `prisma migrate diff` script without the statements that create the
+    LiteLLM_SpendLogs indexes schema.prisma declares but the migrations deliberately
+    leave for operators to build online."""
+    kept: Final = tuple(
+        statement for statement in diff_sql.split(";") if not _creates_hand_built_spend_logs_index(statement)
+    )
+    return ";".join(kept) if any(part.strip() for part in kept) else ""
+
+
+def _creates_hand_built_spend_logs_index(statement: str) -> bool:
+    match: Final = _CREATE_INDEX_RE.match(_without_sql_comments(statement))
+    return match is not None and match["index"] in HAND_BUILT_SPEND_LOGS_INDEXES
+
+
 def _migration_timestamp(name: str) -> int:
     """Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
 
@@ -441,6 +467,20 @@ class ProxyExtrasDBManager:
                 return True
         return False
 
+    @staticmethod
+    def _filter_drift_script(diff_sql: str) -> str:
+        """The drift script without the hand-built SpendLogs indexes and, when
+        LiteLLM_SpendLogs is partitioned, without its primary-key rewrite and
+        partitioning artifacts."""
+        without_indexes: Final = filter_hand_built_spend_logs_index_diff(diff_sql)
+        if not ProxyExtrasDBManager.spend_logs_is_partitioned():
+            return without_indexes
+        logger.info(
+            "LiteLLM_SpendLogs is partitioned; removed its primary-key "
+            "rewrite and partitioning artifacts from the drift script"
+        )
+        return filter_partitioned_spend_logs_diff(without_indexes)
+
     @staticmethod
     def _resolve_all_migrations(
         migrations_dir: str, schema_path: str, mark_all_applied: bool = True
@@ -521,21 +561,14 @@ class ProxyExtrasDBManager:
             return
         logger.info(f"Migration diff created at {diff_sql_path}")
 
-        if ProxyExtrasDBManager.spend_logs_is_partitioned():
-            filtered_sql = filter_partitioned_spend_logs_diff(
-                diff_sql_path.read_text()
-            )
-            diff_sql_path.write_text(filtered_sql)
-            logger.info(
-                "LiteLLM_SpendLogs is partitioned; removed its primary-key "
-                "rewrite and partitioning artifacts from the drift script"
-            )
-            if not filtered_sql.strip():
-                logger.info("Drift script is empty after filtering; nothing to apply")
-                if not mark_all_applied:
-                    return
-                ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
+        filtered_sql: Final = ProxyExtrasDBManager._filter_drift_script(diff_sql_path.read_text())
+        diff_sql_path.write_text(filtered_sql)
+        if not filtered_sql.strip():
+            logger.info("Drift script is empty after filtering; nothing to apply")
+            if not mark_all_applied:
                 return
+            ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
+            return
 
         # 2. Run prisma db execute to apply the migration
         applied_ok = False

tests/proxy_migration_tests/test_db_schema_migration.py

@@ -5,6 +5,7 @@ import tempfile
 from pathlib import Path
 
 import pytest
+from litellm_proxy_extras.utils import filter_hand_built_spend_logs_index_diff
 
 
 @pytest.mark.skipif(
@@ -59,12 +60,12 @@ def test_schema_migration_in_sync():
             text=True,
         )
 
-        if diff.returncode == 2:
+        if diff.returncode == 2 and filter_hand_built_spend_logs_index_diff(diff.stdout).strip():
             pytest.fail(
                 "Schema changes detected that no migration captures. Run "
                 "`python litellm/ci_cd/run_migration.py <migration_name>`.\n\n"
                 + diff.stdout
             )
-        assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}"
+        assert diff.returncode in (0, 2), f"prisma migrate diff errored: {diff.stderr}"
     finally:
         shutil.rmtree(temp_base, ignore_errors=True)

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -18,6 +18,7 @@ sys.path.insert(
 from litellm_proxy_extras.utils import (
     PARTITIONED_SPEND_LOGS_PUSH_ERROR,
     ProxyExtrasDBManager,
+    filter_hand_built_spend_logs_index_diff,
     filter_partitioned_spend_logs_diff,
 )
 
@@ -540,7 +541,7 @@ class _FakeCompleted:
 
 
 class TestResolveAllMigrationsLedger:
-    def _run(self, monkeypatch, tmp_path, partitioned, execute_fails):
+    def _run(self, monkeypatch, tmp_path, partitioned, execute_fails, drift_sql=_PARTITIONED_DRIFT_SQL):
         import subprocess as subprocess_module
 
         import litellm_proxy_extras.utils as utils_module
@@ -560,7 +561,7 @@ class TestResolveAllMigrationsLedger:
         def fake_run(cmd, **kwargs):
             calls.append(cmd)
             if "diff" in cmd:
-                kwargs["stdout"].write(_PARTITIONED_DRIFT_SQL)
+                kwargs["stdout"].write(drift_sql)
                 return _FakeCompleted()
             if "execute" in cmd:
                 executed_sql = open(cmd[cmd.index("--file") + 1]).read()
@@ -601,6 +602,59 @@ class TestResolveAllMigrationsLedger:
         calls = self._run(monkeypatch, tmp_path, partitioned=False, execute_fails=False)
         assert self._executed_sql(calls) == _PARTITIONED_DRIFT_SQL
 
+    @pytest.mark.parametrize("partitioned", [False, True])
+    def test_hand_built_spend_logs_indexes_never_reach_the_database(self, monkeypatch, tmp_path, partitioned):
+        calls = self._run(
+            monkeypatch, tmp_path, partitioned=partitioned, execute_fails=False, drift_sql=_HAND_BUILT_INDEX_DRIFT_SQL
+        )
+        executed_sql = self._executed_sql(calls)
+        assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in executed_sql
+        assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in executed_sql
+        assert 'CREATE INDEX "LiteLLM_TeamTable_name_idx" ON "LiteLLM_TeamTable"("name");' in executed_sql
+
+    def test_drift_of_only_hand_built_indexes_executes_nothing(self, monkeypatch, tmp_path):
+        calls = self._run(
+            monkeypatch, tmp_path, partitioned=False, execute_fails=False, drift_sql=_ONLY_HAND_BUILT_INDEXES_DRIFT_SQL
+        )
+        assert not any(isinstance(c, list) and "execute" in c for c in calls)
+        assert len(self._resolved(calls)) == 1
+
+
+_ONLY_HAND_BUILT_INDEXES_DRIFT_SQL = """-- CreateIndex
+CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
+
+-- CreateIndex
+CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");
+"""
+
+_HAND_BUILT_INDEX_DRIFT_SQL = (
+    _ONLY_HAND_BUILT_INDEXES_DRIFT_SQL
+    + """
+-- CreateIndex
+CREATE INDEX "LiteLLM_TeamTable_name_idx" ON "LiteLLM_TeamTable"("name");
+"""
+)
+
+
+class TestHandBuiltSpendLogsIndexDriftFilter:
+    @pytest.mark.parametrize(
+        "statement",
+        [
+            'CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id")',
+            'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id")',
+            'create index if not exists "LiteLLM_SpendLogs_api_key_startTime_idx" on "LiteLLM_SpendLogs"("api_key", "startTime")',
+        ],
+    )
+    def test_hand_built_index_creation_is_removed(self, statement):
+        assert filter_hand_built_spend_logs_index_diff(f"-- CreateIndex\n{statement};\n") == ""
+
+    def test_other_spend_logs_indexes_and_statements_are_kept_verbatim(self):
+        kept = (
+            'CREATE INDEX "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id");\n'
+            'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN     "updated_by" TEXT;\n'
+        )
+        assert filter_hand_built_spend_logs_index_diff(kept) == kept
+
 
 class TestPartitionedSpendLogsPushGuard:
     def _forbid_subprocess(self, monkeypatch):