From 33fcd5140758f67790c54a23fdedd2243c097f6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Tue, 4 Aug 2026 23:27:16 +0800 Subject: [PATCH 1/2] fix(db): cover spend window reseed queries --- .../migration.sql | 6 ++++ .../migration.sql | 5 +++ .../litellm_proxy_extras/schema.prisma | 2 ++ schema.prisma | 2 ++ .../test_spend_log_budget_window_indexes.py | 33 +++++++++++++++++++ 5 files changed, 48 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233000_add_spend_log_api_key_window_covering_index/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804233100_add_spend_log_team_window_covering_index/migration.sql create mode 100644 tests/test_litellm/proxy/db/test_spend_log_budget_window_indexes.py 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 0d7fa8692c8..8be4d394eac 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -639,6 +639,8 @@ model LiteLLM_SpendLogs { proxy_server_request Json? @default("{}") @@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 0d7fa8692c8..8be4d394eac 100644 --- a/schema.prisma +++ b/schema.prisma @@ -639,6 +639,8 @@ model LiteLLM_SpendLogs { proxy_server_request Json? @default("{}") @@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..b499dbdc9f1 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_spend_log_budget_window_indexes.py @@ -0,0 +1,33 @@ +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[4] +EXTRAS_ROOT = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" + +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_both_schemas(): + root_schema = (REPO_ROOT / "schema.prisma").read_text(encoding="utf-8") + packaged_schema = (EXTRAS_ROOT / "schema.prisma").read_text(encoding="utf-8") + + for field in INDEXES: + declaration = f"@@index([{field}, startTime])" + assert declaration in root_schema + assert declaration in packaged_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 From 7ed3291502415c4365a74d67805513a77841d057 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Tue, 4 Aug 2026 23:35:35 +0800 Subject: [PATCH 2/2] test(db): enforce spend log schema sync --- litellm/proxy/schema.prisma | 2 ++ .../db/test_spend_log_budget_window_indexes.py | 18 +++++++++++------- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 0d7fa8692c8..8be4d394eac 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -639,6 +639,8 @@ model LiteLLM_SpendLogs { proxy_server_request Json? @default("{}") @@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 index b499dbdc9f1..29f4d03e9e9 100644 --- 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 @@ -3,6 +3,11 @@ 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": ( @@ -14,14 +19,13 @@ INDEXES = { } -def test_spend_log_budget_window_indexes_are_declared_in_both_schemas(): - root_schema = (REPO_ROOT / "schema.prisma").read_text(encoding="utf-8") - packaged_schema = (EXTRAS_ROOT / "schema.prisma").read_text(encoding="utf-8") +def test_spend_log_budget_window_indexes_are_declared_in_all_schemas(): + schemas = [schema.read_text(encoding="utf-8") for schema in SCHEMAS] - for field in INDEXES: - declaration = f"@@index([{field}, startTime])" - assert declaration in root_schema - assert declaration in packaged_schema + 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():