diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2145f891318..b44cc1d98f9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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[^"]+)"\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. @@ -433,6 +459,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 @@ -513,21 +553,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 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 e5826b18668..ec994c2186c 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -16,6 +16,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, ) @@ -538,7 +539,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 @@ -558,7 +559,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() @@ -599,6 +600,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): diff --git a/tests/proxy_migration_tests/test_db_schema_migration.py b/tests/proxy_migration_tests/test_db_schema_migration.py index b0d44cd3e1c..617ace37290 100644 --- a/tests/proxy_migration_tests/test_db_schema_migration.py +++ b/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 `.\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) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 4d6915914c2..aeac491c959 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -3,12 +3,10 @@ import re import time from collections.abc import Mapping, Sequence from datetime import datetime, timedelta -from pathlib import Path from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock -import litellm_proxy_extras import psycopg import pytest from prisma.errors import PrismaError @@ -620,11 +618,8 @@ _SPEND_LOGS_DDL: Final = """ ) """ -_API_KEY_START_TIME_INDEX_MIGRATION: Final = ( - Path(litellm_proxy_extras.__file__).parent - / "migrations" - / "20260823000000_add_spend_logs_api_key_starttime_index" - / "migration.sql" +_API_KEY_START_TIME_INDEX_DDL: Final = ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime")' ) _SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL: Final = """ @@ -636,7 +631,7 @@ _SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL: Final = """ def _create_spend_logs_table(conn: psycopg.Connection) -> None: conn.execute(_SPEND_LOGS_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.execute(_API_KEY_START_TIME_INDEX_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # migration file + conn.execute(_API_KEY_START_TIME_INDEX_DDL) # pyright: ignore[reportArgumentType] # DDL literal def _psycopg_prisma(conn: psycopg.Connection) -> MagicMock: