diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql index 08fcbddb6f8..c720d705fbe 100644 --- a/db_scripts/partition_spend_logs.sql +++ b/db_scripts/partition_spend_logs.sql @@ -48,6 +48,8 @@ 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_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 @@ -73,6 +75,9 @@ 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_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..fe5b7335241 100644 --- a/db_scripts/unpartition_spend_logs.sql +++ b/db_scripts/unpartition_spend_logs.sql @@ -40,6 +40,8 @@ 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_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 +62,9 @@ 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_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/20260818000000_add_spendlogs_api_key_start_time_idx/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260818000000_add_spendlogs_api_key_start_time_idx/migration.sql index 1d57e456e08..c15c9d65836 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260818000000_add_spendlogs_api_key_start_time_idx/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260818000000_add_spendlogs_api_key_start_time_idx/migration.sql @@ -1,11 +1,14 @@ --- CreateIndex (CONCURRENTLY) +-- CreateIndex -- -- Disclaimer: --- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a --- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. --- - "LiteLLM_SpendLogs" is the largest table on most deployments, so a blocking build would --- stall spend-log inserts for the whole build; if the build is interrupted, Postgres may --- leave an INVALID index that must be dropped and recreated. +-- - Not CONCURRENTLY on purpose: Postgres refuses a concurrent index build on a partitioned +-- parent, and "LiteLLM_SpendLogs" is a partitioned parent on every deployment that ran +-- db_scripts/partition_spend_logs.sql, so a concurrent build would abort the migration run +-- there and block the rollout. +-- - The build takes a write lock on "LiteLLM_SpendLogs" for its duration, and that is the largest +-- table on most deployments. Operators who cannot pause spend-log inserts can build the index +-- themselves before upgrading, CONCURRENTLY when the table is plain or per partition plus +-- ATTACH when it is partitioned, and this statement then no-ops on the index name. -- - Do not edit this file after it has been applied to any database: Prisma checksums -- migrations; add a new migration instead. -CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC); diff --git a/tests/proxy_migration_tests/test_db_schema_migration.py b/tests/proxy_migration_tests/test_db_schema_migration.py index 0d2eced5620..ae60992798f 100644 --- a/tests/proxy_migration_tests/test_db_schema_migration.py +++ b/tests/proxy_migration_tests/test_db_schema_migration.py @@ -95,16 +95,79 @@ def test_spend_logs_declares_api_key_index(schema_path: Path): ) -def test_spend_logs_api_key_index_has_a_migration(): - """A schema-only index never reaches deployments that run ``prisma migrate deploy``.""" - migrations = Path("./litellm-proxy-extras/litellm_proxy_extras/migrations") +MIGRATIONS_DIR = Path("./litellm-proxy-extras/litellm_proxy_extras/migrations") + +PARTITION_PROBE_SCHEMA = "litellm_spendlogs_partition_probe" + + +def _api_key_index_migration() -> str: creating = [ sql - for sql in (path.read_text() for path in migrations.glob("*/migration.sql")) + for sql in (path.read_text() for path in MIGRATIONS_DIR.glob("*/migration.sql")) if SPEND_LOGS_API_KEY_INDEX in sql ] assert len(creating) == 1, f"expected exactly one migration creating {SPEND_LOGS_API_KEY_INDEX}, found {len(creating)}" - assert 'ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC)' in creating[0] - assert "CONCURRENTLY" in creating[0], ( - "LiteLLM_SpendLogs is the largest table on most deployments; a blocking build stalls spend-log inserts" - ) + return creating[0] + + +def test_spend_logs_api_key_index_has_a_migration(): + """A schema-only index never reaches deployments that run ``prisma migrate deploy``.""" + assert 'ON "LiteLLM_SpendLogs"("api_key", "startTime" DESC)' in _api_key_index_migration() + + +@pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) +def test_spend_logs_api_key_index_migration_applies_to_a_partitioned_table(): + """``db_scripts/partition_spend_logs.sql`` leaves ``LiteLLM_SpendLogs`` a partitioned parent, + and Postgres rejects both a concurrent index build on a partitioned parent and any concurrent + build inside a transaction, so a ``CONCURRENTLY`` migration would abort the rollout of every + partitioned deployment. This applies the shipped statement against that shape for real. + """ + probe_sql = f""" +DROP SCHEMA IF EXISTS "{PARTITION_PROBE_SCHEMA}" CASCADE; +CREATE SCHEMA "{PARTITION_PROBE_SCHEMA}"; +SET search_path TO "{PARTITION_PROBE_SCHEMA}"; + +CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT NOT NULL, + api_key TEXT, + "startTime" TIMESTAMP NOT NULL, + PRIMARY KEY (request_id, "startTime") +) PARTITION BY RANGE ("startTime"); +CREATE TABLE "LiteLLM_SpendLogs_pdefault" PARTITION OF "LiteLLM_SpendLogs" DEFAULT; + +{_api_key_index_migration()} + +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_indexes + WHERE schemaname = '{PARTITION_PROBE_SCHEMA}' + AND tablename = 'LiteLLM_SpendLogs' + AND indexname = '{SPEND_LOGS_API_KEY_INDEX}' + ) THEN + RAISE EXCEPTION 'migration did not create {SPEND_LOGS_API_KEY_INDEX} on the partitioned parent'; + END IF; +END +$$; + +DROP SCHEMA "{PARTITION_PROBE_SCHEMA}" CASCADE; +""" + + temp_base = Path(tempfile.mkdtemp(prefix="litellm_partition_probe_")) + try: + sql_path = temp_base / "probe.sql" + sql_path.write_text(probe_sql) + applied = subprocess.run( + ["prisma", "db", "execute", "--url", os.environ["DATABASE_URL"], "--file", str(sql_path)], + capture_output=True, + text=True, + ) + assert applied.returncode == 0, ( + f"migration for {SPEND_LOGS_API_KEY_INDEX} failed on a partitioned LiteLLM_SpendLogs:\n" + f"{applied.stdout}\n{applied.stderr}" + ) + finally: + shutil.rmtree(temp_base, ignore_errors=True)