From 94b6d0642bf535be558ee98e4b8a8a117b20fa57 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 2 Sep 2026 11:57:36 +0200 Subject: [PATCH] perf(db): index SpendLogs on (user, startTime), (api_key, startTime) --- db_scripts/partition_spend_logs.sql | 10 ++ db_scripts/unpartition_spend_logs.sql | 10 ++ .../migration.sql | 6 + .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/proxy/schema.prisma | 2 + schema.prisma | 2 + .../test_db_schema_migration.py | 128 ++++++++++++++++++ 8 files changed, 165 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_spend_logs_user_start_time_idx/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000100_add_spend_logs_api_key_start_time_idx/migration.sql diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql index 4e4a93539d7..5ee6bcccd5d 100644 --- a/db_scripts/partition_spend_logs.sql +++ b/db_scripts/partition_spend_logs.sql @@ -53,6 +53,10 @@ 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_user_startTime_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_user_startTime_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 @@ -78,6 +82,12 @@ 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_user_startTime_idx" + ON "LiteLLM_SpendLogs" ("user", "startTime" DESC); + +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" diff --git a/db_scripts/unpartition_spend_logs.sql b/db_scripts/unpartition_spend_logs.sql index 0bd82513e4a..0b3c5f20e03 100644 --- a/db_scripts/unpartition_spend_logs.sql +++ b/db_scripts/unpartition_spend_logs.sql @@ -40,6 +40,10 @@ 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_user_startTime_idx" + RENAME TO "LiteLLM_SpendLogs_partitioned_user_startTime_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 +64,12 @@ 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_user_startTime_idx" + ON "LiteLLM_SpendLogs" ("user", "startTime" DESC); + +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; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_spend_logs_user_start_time_idx/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_spend_logs_user_start_time_idx/migration.sql new file mode 100644 index 00000000000..ebec1354fd9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_spend_logs_user_start_time_idx/migration.sql @@ -0,0 +1,6 @@ +-- CreateIndex +-- Plain build on purpose: Postgres refuses CONCURRENTLY on a partitioned parent (what +-- db_scripts/partition_spend_logs.sql leaves behind) and inside the transaction prisma migrate +-- deploy runs this in. Inserts to "LiteLLM_SpendLogs" block for the build; operators can +-- pre-create the index under this exact name and this statement then no-ops. +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_user_startTime_idx" ON "LiteLLM_SpendLogs"("user", "startTime" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000100_add_spend_logs_api_key_start_time_idx/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000100_add_spend_logs_api_key_start_time_idx/migration.sql new file mode 100644 index 00000000000..c4937d0e0c5 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000100_add_spend_logs_api_key_start_time_idx/migration.sql @@ -0,0 +1,5 @@ +-- CreateIndex +-- Same constraints as 20260902000000_add_spend_logs_user_start_time_idx: plain build so +-- partitioned parents can take it, inserts block for the duration, pre-creating under this +-- name makes it a no-op. +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7604ceadf7a..1caa7813313 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -662,6 +662,8 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([user, startTime(sort: Desc)]) + @@index([api_key, startTime(sort: Desc)]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7604ceadf7a..1caa7813313 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -662,6 +662,8 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([user, startTime(sort: Desc)]) + @@index([api_key, startTime(sort: Desc)]) } model LiteLLM_BudgetWindowSpend { diff --git a/schema.prisma b/schema.prisma index 7604ceadf7a..1caa7813313 100644 --- a/schema.prisma +++ b/schema.prisma @@ -662,6 +662,8 @@ model LiteLLM_SpendLogs { @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([user, startTime(sort: Desc)]) + @@index([api_key, startTime(sort: Desc)]) } model LiteLLM_BudgetWindowSpend { diff --git a/tests/proxy_migration_tests/test_db_schema_migration.py b/tests/proxy_migration_tests/test_db_schema_migration.py index b0d44cd3e1c..ef9eb315327 100644 --- a/tests/proxy_migration_tests/test_db_schema_migration.py +++ b/tests/proxy_migration_tests/test_db_schema_migration.py @@ -1,8 +1,12 @@ +import json import os import shutil import subprocess import tempfile +import uuid +from collections.abc import Iterator from pathlib import Path +from typing import Final import pytest @@ -68,3 +72,127 @@ def test_schema_migration_in_sync(): assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}" finally: shutil.rmtree(temp_base, ignore_errors=True) + + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + +MIGRATIONS_DIR: Final = Path("./litellm-proxy-extras/litellm_proxy_extras/migrations") + +SPEND_LOGS_FILTER_INDEXES: Final = ( + ('"user"', "user-7", "LiteLLM_SpendLogs_user_startTime_idx"), + ("api_key", "key-7", "LiteLLM_SpendLogs_api_key_startTime_idx"), +) + +SPEND_LOGS_PROBE_SQL: Final = """ +CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT PRIMARY KEY, + api_key TEXT NOT NULL DEFAULT '', + "startTime" TIMESTAMP(3) NOT NULL, + "user" TEXT DEFAULT '' +); +CREATE INDEX "LiteLLM_SpendLogs_startTime_idx" ON "LiteLLM_SpendLogs"("startTime"); +CREATE INDEX "LiteLLM_SpendLogs_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("startTime", "request_id"); +INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user") +SELECT 'req-' || g, 'key-' || (g % 200), TIMESTAMP '2026-08-01' + g * INTERVAL '1 minute', 'user-' || (g % 200) +FROM generate_series(1, 60000) AS g; +ANALYZE "LiteLLM_SpendLogs"; +""" + +SPEND_LOGS_PARTITIONED_PROBE_SQL: Final = """ +CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT NOT NULL, + api_key TEXT NOT NULL DEFAULT '', + "startTime" TIMESTAMP(3) NOT NULL, + "user" TEXT DEFAULT '', + PRIMARY KEY (request_id, "startTime") +) PARTITION BY RANGE ("startTime"); +CREATE TABLE "LiteLLM_SpendLogs_pdefault" PARTITION OF "LiteLLM_SpendLogs" DEFAULT; +""" + + +def _in_schema(schema: str, sql: str) -> list[tuple[object, ...]]: + psycopg: Final = pytest.importorskip("psycopg") + with psycopg.connect(os.environ["DATABASE_URL"].split("?")[0], autocommit=True) as conn: + conn.execute(f'SET search_path TO "{schema}"') + cursor: Final = conn.execute(sql) + return cursor.fetchall() if cursor.description is not None else [] + + +def _migration_creating(index_name: str) -> str: + creating: Final = tuple( + sql + for sql in (path.read_text() for path in sorted(MIGRATIONS_DIR.glob("*/migration.sql"))) + if index_name in sql + ) + assert len(creating) == 1, f"expected exactly one migration creating {index_name}, found {len(creating)}" + return creating[0] + + +def _bounded_count_plan(schema: str, column: str, value: str) -> str: + rows: Final = _in_schema( + schema, + f""" + EXPLAIN (FORMAT JSON) + SELECT COUNT(*) AS total_count + FROM ( + SELECT 1 + FROM "LiteLLM_SpendLogs" + WHERE "startTime" >= ('2026-08-10T00:00:00Z'::timestamptz AT TIME ZONE 'UTC') + AND "startTime" <= ('2026-08-17T00:00:00Z'::timestamptz AT TIME ZONE 'UTC') + AND {column} = '{value}' + LIMIT 10001 + ) AS bounded_matches + """, + ) + return json.dumps(rows[0][0]) + + +@pytest.fixture +def spend_logs_scratch_schema() -> Iterator[str]: + schema: Final = f"spend_logs_idx_{uuid.uuid4().hex[:8]}" + _in_schema("public", f'CREATE SCHEMA "{schema}"') + yield schema + _in_schema("public", f'DROP SCHEMA "{schema}" CASCADE') + + +@requires_db +@pytest.mark.parametrize(("column", "value", "index_name"), SPEND_LOGS_FILTER_INDEXES, ids=("user", "api_key")) +def test_spend_logs_bounded_count_is_served_by_the_filter_column_index( + spend_logs_scratch_schema: str, column: str, value: str, index_name: str +) -> None: + """The /spend/logs/v2 count for one user or key inside a startTime window must read that + user's rows only. Without a filter-column-led index it range-scans the whole window and + post-filters every row, which times out on multi-day windows at production volume. + """ + _in_schema(spend_logs_scratch_schema, SPEND_LOGS_PROBE_SQL) + _in_schema(spend_logs_scratch_schema, _migration_creating(index_name)) + _in_schema(spend_logs_scratch_schema, 'ANALYZE "LiteLLM_SpendLogs"') + + plan: Final = _bounded_count_plan(spend_logs_scratch_schema, column, value) + + assert f'"Index Name": "{index_name}"' in plan, plan + assert '"Filter"' not in plan, plan + + +@requires_db +def test_spend_logs_filter_index_migrations_apply_to_a_partitioned_table(spend_logs_scratch_schema: str) -> None: + """db_scripts/partition_spend_logs.sql leaves LiteLLM_SpendLogs a partitioned parent, and + Postgres rejects CONCURRENTLY there, so a concurrent build would abort every partitioned + deployment's rollout. Apply the shipped statements against that shape for real. + """ + _in_schema(spend_logs_scratch_schema, SPEND_LOGS_PARTITIONED_PROBE_SQL) + for _column, _value, index_name in SPEND_LOGS_FILTER_INDEXES: + _in_schema(spend_logs_scratch_schema, _migration_creating(index_name)) + + parent_indexes: Final = frozenset( + str(row[0]) + for row in _in_schema( + spend_logs_scratch_schema, + "SELECT indexname FROM pg_indexes WHERE schemaname = current_schema() AND tablename = 'LiteLLM_SpendLogs'", + ) + ) + + assert {index_name for _column, _value, index_name in SPEND_LOGS_FILTER_INDEXES} <= parent_indexes, parent_indexes