diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233000_add_spend_log_api_key_window_covering_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233000_add_spend_log_api_key_window_covering_index/migration.sql new file mode 100644 index 00000000000..44331db245e --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233000_add_spend_log_api_key_window_covering_index/migration.sql @@ -0,0 +1,6 @@ +-- Keep budget-window reseeds off the wide spend-log heap, including for keys +-- that own most rows. CONCURRENTLY avoids blocking inserts on production tables. +-- This migration must remain a single statement because CREATE INDEX +-- CONCURRENTLY cannot run inside a transaction. +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" +ON "LiteLLM_SpendLogs" ("api_key", "startTime") INCLUDE ("spend"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233100_add_spend_log_team_window_covering_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233100_add_spend_log_team_window_covering_index/migration.sql new file mode 100644 index 00000000000..947124e7353 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233100_add_spend_log_team_window_covering_index/migration.sql @@ -0,0 +1,5 @@ +-- Match the team budget-window aggregate with the same index-only access path. +-- This migration is separate from the api-key index because each concurrent +-- index build must be applied outside a transaction as a single statement. +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_team_id_startTime_idx" +ON "LiteLLM_SpendLogs" ("team_id", "startTime") INCLUDE ("spend"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d9959677116..f71eeac3232 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -645,6 +645,8 @@ model LiteLLM_SpendLogs { updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) + @@index([api_key, startTime]) + @@index([team_id, startTime]) @@index([end_user]) @@index([session_id]) } diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d9959677116..f71eeac3232 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -645,6 +645,8 @@ model LiteLLM_SpendLogs { updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) + @@index([api_key, startTime]) + @@index([team_id, startTime]) @@index([end_user]) @@index([session_id]) } diff --git a/schema.prisma b/schema.prisma index d9959677116..f71eeac3232 100644 --- a/schema.prisma +++ b/schema.prisma @@ -645,6 +645,8 @@ model LiteLLM_SpendLogs { updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) + @@index([api_key, startTime]) + @@index([team_id, startTime]) @@index([end_user]) @@index([session_id]) } diff --git a/tests/test_litellm/proxy/db/test_spend_log_budget_window_indexes.py b/tests/test_litellm/proxy/db/test_spend_log_budget_window_indexes.py new file mode 100644 index 00000000000..29f4d03e9e9 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_log_budget_window_indexes.py @@ -0,0 +1,37 @@ +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[4] +EXTRAS_ROOT = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" +SCHEMAS = ( + REPO_ROOT / "schema.prisma", + REPO_ROOT / "litellm" / "proxy" / "schema.prisma", + EXTRAS_ROOT / "schema.prisma", +) + +INDEXES = { + "api_key": ( + EXTRAS_ROOT / "migrations" / "20260804233000_add_spend_log_api_key_window_covering_index" / "migration.sql" + ), + "team_id": ( + EXTRAS_ROOT / "migrations" / "20260804233100_add_spend_log_team_window_covering_index" / "migration.sql" + ), +} + + +def test_spend_log_budget_window_indexes_are_declared_in_all_schemas(): + schemas = [schema.read_text(encoding="utf-8") for schema in SCHEMAS] + + assert all(schema == schemas[0] for schema in schemas[1:]) + for schema in schemas: + for field in INDEXES: + assert f"@@index([{field}, startTime])" in schema + + +def test_spend_log_budget_window_migrations_create_covering_indexes_concurrently(): + for field, migration_path in INDEXES.items(): + migration = migration_path.read_text(encoding="utf-8") + + assert migration.count("CREATE INDEX CONCURRENTLY IF NOT EXISTS") == 1 + assert "CREATE INDEX CONCURRENTLY IF NOT EXISTS" in migration + assert f'("{field}", "startTime") INCLUDE ("spend")' in migration