From bfc6cb184c98fb84c2a8fda389961d9da4e41970 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 2 Nov 2026 13:29:13 +0000 Subject: [PATCH] fix(proxy-extras): build the remaining request-log indexes in the migration job instead of in migrations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration_recovery.py | 49 +++--- .../migration.sql | 5 +- .../migration.sql | 3 - .../migration.sql | 4 +- .../migration.sql | 4 +- .../request_log_indexes.py | 10 +- .../litellm_proxy_extras/utils.py | 6 +- .../check_migrations_no_data_rewrites.py | 6 +- .../migrations/test_request_log_indexes.py | 139 +++++++++++++++++- .../test_request_log_indexes_at_boot.py | 52 +++++-- .../test_litellm_proxy_extras_utils.py | 2 +- .../test_request_log_indexes.py | 3 + .../test_check_migrations_no_data_rewrites.py | 9 +- 13 files changed, 233 insertions(+), 59 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py index 1c16eb5e6b7..8fbfd3285a7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -4,6 +4,7 @@ import subprocess from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import TYPE_CHECKING, Final from uuid import uuid4 @@ -15,6 +16,17 @@ if TYPE_CHECKING: import psycopg +# SHA-256 of earlier shipped files whose successful run applied everything the current file does, +# and whose failed run the current file can replay. +SUPERSEDED_CHECKSUMS: Final = MappingProxyType( + { + "20260226120000_add_spend_log_tool_index": frozenset( + {"019ae3e4e71f35763d6d515e1d463d97e75b8c50e7637c842e4e16ce59009a4b"} + ) + } +) + + @dataclass(frozen=True, slots=True) class MigrationProgress: checksum: str @@ -23,12 +35,8 @@ class MigrationProgress: id: str = "" finished: bool = False - def confirms_completion(self, script: bytes) -> bool: - return ( - self.applied_steps_count == 1 - and not self.logs.strip() - and self.checksum == hashlib.sha256(script).hexdigest() - ) + def confirms_completion(self, checksums: frozenset[str]) -> bool: + return self.applied_steps_count == 1 and not self.logs.strip() and self.checksum in checksums def _migration_records( @@ -60,9 +68,12 @@ def recover_completed_migration(coordinator: MigrationCoordinator, schema: str, records: Final = _migration_records(coordinator.connection, schema, migration) unfinished: Final = tuple(record for record in records if not record.finished) script: Final = migration.read_bytes() + checksums: Final = frozenset({hashlib.sha256(script).hexdigest()}) | SUPERSEDED_CHECKSUMS.get( + migration.parent.name, frozenset() + ) if not unfinished: - return any(record.checksum == hashlib.sha256(script).hexdigest() for record in records) - if len(unfinished) != 1 or not unfinished[0].confirms_completion(script): + return any(record.checksum in checksums for record in records) + if len(unfinished) != 1 or not unfinished[0].confirms_completion(checksums): return False progress: Final = unfinished[0] result: Final = coordinator.connection.execute( @@ -178,21 +189,23 @@ def is_inert_migration(script: str) -> bool: return all(not part.strip() or _NO_OP_STATEMENT_RE.match(part) for part in stripped.split(";")) -def roll_back_failed_inert_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool: - """Roll back the failed ledger row of a migration whose file in this build is inert, - so `migrate deploy` applies the inert file on its next pass. The row records an - earlier build's attempt at SQL this build no longer ships (an index now built by the - migration job), so no database change can be repeated or skipped by replaying - the empty file. The caller commits this checkpoint before the next Prisma command. +def roll_back_replayable_failed_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool: + """Roll back the failed ledger row of a migration this build can safely replay, so + `migrate deploy` applies the current file on its next pass. Either the file is now + inert, or the row recorded an earlier shipped file listed in SUPERSEDED_CHECKSUMS + whose index the migration job now builds. A failed run of the current file is left + for the operator. The caller commits this checkpoint before the next Prisma command. """ from psycopg import sql - if not is_inert_migration(migration.read_text(encoding="utf-8")): + inert: Final = is_inert_migration(migration.read_text(encoding="utf-8")) + superseded: Final = SUPERSEDED_CHECKSUMS.get(migration.parent.name, frozenset()) + if not inert and not superseded: return False coordinator.acquire_prisma_lock() records: Final = _migration_records(coordinator.connection, schema, migration) unfinished: Final = tuple(record for record in records if not record.finished) - if len(unfinished) != 1: + if len(unfinished) != 1 or not (inert or unfinished[0].checksum in superseded): return False result: Final = coordinator.connection.execute( sql.SQL( @@ -202,9 +215,9 @@ def roll_back_failed_inert_migration(coordinator: MigrationCoordinator, schema: (unfinished[0].id,), ) if result.rowcount != 1: - raise RuntimeError("Could not roll back the failed inert migration history row; rerun the database setup.") + raise RuntimeError("Could not roll back the failed migration history row; rerun the database setup.") logger.info( - "Rolled back the failed history row of %s: this build ships it as an inert migration, " + "Rolled back the failed history row of %s: this build replays it safely, " "its index is built by the migration job", migration.parent.name, ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql index a647f386d9e..03ef4885a58 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql @@ -1,3 +1,2 @@ --- CreateIndex -CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id"); - +-- Built after migrate deploy by litellm_proxy_extras/request_log_indexes.py, never by a migration. +SELECT 1; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql index d9c234696c8..edda1cc0a7e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql @@ -6,6 +6,3 @@ CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogToolIndex" ( CONSTRAINT "LiteLLM_SpendLogToolIndex_pkey" PRIMARY KEY ("request_id","tool_name") ); - --- CreateIndex -CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql index 66d792fe386..03ef4885a58 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql @@ -1,2 +1,2 @@ --- CreateIndex -CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("startTime", "request_id"); +-- Built after migrate deploy by litellm_proxy_extras/request_log_indexes.py, never by a migration. +SELECT 1; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260724000000_add_spend_log_tool_index_start_time_idx/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260724000000_add_spend_log_tool_index_start_time_idx/migration.sql index 548c3bd5683..03ef4885a58 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260724000000_add_spend_log_tool_index_start_time_idx/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260724000000_add_spend_log_tool_index_start_time_idx/migration.sql @@ -1,2 +1,2 @@ --- CreateIndex -CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("start_time"); +-- Built after migrate deploy by litellm_proxy_extras/request_log_indexes.py, never by a migration. +SELECT 1; diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 4c856c741d3..1c16ef97065 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -2,7 +2,7 @@ by the migration job, or by a serving proxy that ran the migrations itself (in the background, once it serves). -A migration cannot build them: a plain `CREATE INDEX` blocks spend-log inserts for the +A migration cannot build them: a plain `CREATE INDEX` blocks writes to the indexed table for the whole build, and `CREATE INDEX CONCURRENTLY` is refused on a partitioned parent (db_scripts/partition_spend_logs.sql). `REQUEST_LOG_INDEXES` is the one list to extend; names match what Prisma derives from the `@@index` declarations in schema.prisma, so an @@ -59,6 +59,14 @@ class RequestLogIndex: REQUEST_LOG_INDEXES: Final = ( RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'), RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'), + RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_session_id_idx", '("session_id")'), + RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_startTime_request_id_idx", '("startTime", "request_id")'), + RequestLogIndex( + "LiteLLM_SpendLogToolIndex", + "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx", + '("tool_name", "start_time")', + ), + RequestLogIndex("LiteLLM_SpendLogToolIndex", "LiteLLM_SpendLogToolIndex_start_time_idx", '("start_time")'), ) _IDENTIFIER_MAX_BYTES: Final = 63 diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 8b19576d86e..a31ca9c3d1a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -974,7 +974,7 @@ class ProxyExtrasDBManager: from litellm_proxy_extras.migration_recovery import ( baseline_current_schema, recover_completed_migration, - roll_back_failed_inert_migration, + roll_back_replayable_failed_migration, ) database_url: Final = os.environ.get("DATABASE_URL") @@ -990,9 +990,9 @@ class ProxyExtrasDBManager: if not migration.is_file(): return False with migration_lock(lock_url) as coordinator: - return recover_completed_migration(coordinator, schema, migration) or roll_back_failed_inert_migration( + return recover_completed_migration( coordinator, schema, migration - ) + ) or roll_back_replayable_failed_migration(coordinator, schema, migration) def baseline_existing(migrations_dir: str) -> None: with migration_lock(lock_url) as coordinator: diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index 5c694702c6e..b8a2e0ee1c9 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -17,7 +17,7 @@ to reach them fires on most ordinary migrations, and a marker everyone adds by reflex stops carrying information. The outage this was written for was a backfill. One column change banned outright is `ADD COLUMN ... DEFAULT` on a table in -`REQUEST_LOG_TABLES`, the tables that hold a row per request. Postgres 11 stores such +`REQUEST_LOG_TABLES`, the tables that hold request-log rows. Postgres 11 stores such a default as metadata and touches no rows, but Postgres 10, which is supported, rewrites the whole heap and rebuilds every index under an `ACCESS EXCLUSIVE` lock, which on a spend-log-sized table is the same outage as a backfill. Every other table @@ -128,8 +128,6 @@ MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / " GRANDFATHERED = frozenset( { "20250425182129_add_session_id", - "20250510142544_add_session_id_index_spend_logs", - "20260228100000_add_spend_logs_composite_index", "20250326162113_baseline", "20260817000000_shadow_eval_multi_key", "20260818000000_add_spend_log_timestamps", @@ -137,7 +135,7 @@ GRANDFATHERED = frozenset( } ) -REQUEST_LOG_TABLES = frozenset({"LiteLLM_SpendLogs", "LiteLLM_ErrorLogs"}) +REQUEST_LOG_TABLES = frozenset({"LiteLLM_SpendLogs", "LiteLLM_ErrorLogs", "LiteLLM_SpendLogToolIndex"}) MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTILINE) DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 5741c85470b..c4db4a2b019 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -18,7 +18,7 @@ import psycopg import pytest from litellm_proxy_extras import migration_data_checks, migration_dry_run, request_log_indexes from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lock -from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration +from litellm_proxy_extras.migration_recovery import recover_completed_migration, roll_back_replayable_failed_migration from litellm_proxy_extras.request_log_indexes import ( REQUEST_LOG_INDEXES, RequestLogIndex, @@ -43,11 +43,37 @@ PACKAGE: Final = REPO / "litellm-proxy-extras" / "litellm_proxy_extras" PARTITION_SCRIPT: Final = REPO / "db_scripts" / "partition_spend_logs.sql" API_KEY_INDEX_MIGRATION: Final = "20260823000000_add_spend_logs_api_key_starttime_index" CALL_ID_INDEX_MIGRATION: Final = "20260831120001_spend_logs_litellm_call_id_index" +TOOL_INDEX_MIGRATION: Final = "20260226120000_add_spend_log_tool_index" API_KEY_INDEX: Final = "LiteLLM_SpendLogs_api_key_startTime_idx" CALL_ID_INDEX: Final = "LiteLLM_SpendLogs_litellm_call_id_idx" PARTITIONED_PARENT_ERROR: Final = 'cannot create index on partitioned table "LiteLLM_SpendLogs" concurrently' ORIGINAL_MIGRATION_SQL: Final = MappingProxyType( { + "20250510142544_add_session_id_index_spend_logs": ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" ' + 'ON "LiteLLM_SpendLogs"("session_id");\n' + ), + TOOL_INDEX_MIGRATION: ( + "-- CreateTable\n" + 'CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogToolIndex" (\n' + ' "request_id" TEXT NOT NULL,\n' + ' "tool_name" TEXT NOT NULL,\n' + ' "start_time" TIMESTAMP(3) NOT NULL,\n' + '\n' + ' CONSTRAINT "LiteLLM_SpendLogToolIndex_pkey" PRIMARY KEY ("request_id","tool_name")\n' + ');\n\n' + "-- CreateIndex\n" + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ' + 'ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time");\n' + ), + "20260228100000_add_spend_logs_composite_index": ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" ' + 'ON "LiteLLM_SpendLogs"("startTime", "request_id");\n' + ), + "20260724000000_add_spend_log_tool_index_start_time_idx": ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_start_time_idx" ' + 'ON "LiteLLM_SpendLogToolIndex"("start_time");\n' + ), API_KEY_INDEX_MIGRATION: ( "-- CreateIndex\n" 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ' @@ -89,8 +115,8 @@ def _migrate_deploy(database_url: str, schema: Path) -> "subprocess.CompletedPro def _release_layout(prisma_dir: Path, before: str) -> Path: - """The shipped migrations older than `before`, with the two index migrations written - the way the releases that shipped them did: the Prisma layout of a proxy on that release.""" + """The shipped migrations older than `before`, with index migrations written the way + the releases that shipped them did: the Prisma layout of a proxy on that release.""" (prisma_dir / "migrations").mkdir(parents=True) shutil.copy(PACKAGE / "schema.prisma", prisma_dir / "schema.prisma") for migration in sorted((PACKAGE / "migrations").iterdir()): @@ -172,6 +198,26 @@ def _fail_the_call_id_migration_like_the_shipped_release(database_url: str, tmp_ assert _ledger(database_url)[CALL_ID_INDEX_MIGRATION] == (False, False) +def _leave_the_tool_index_migration_unfinished(database_url: str, checksum: str) -> None: + original: Final = ORIGINAL_MIGRATION_SQL[TOOL_INDEX_MIGRATION] + with psycopg.connect(database_url, autocommit=True) as conn: + conn.execute("DELETE FROM _prisma_migrations WHERE migration_name = %s", (TOOL_INDEX_MIGRATION,)) + conn.execute('DROP TABLE "LiteLLM_SpendLogToolIndex"') + conn.execute(original) + conn.execute( + "INSERT INTO _prisma_migrations " + "(id, checksum, migration_name, logs, started_at, applied_steps_count) " + "VALUES (%s, %s, %s, '', current_timestamp, 1)", + (str(uuid.uuid4()), checksum, TOOL_INDEX_MIGRATION), + ) + + +def _recover_tool_index_migration(database_url: str) -> bool: + migration: Final = PACKAGE / "migrations" / TOOL_INDEX_MIGRATION / "migration.sql" + with migration_lock(database_url) as coordinator: + return recover_completed_migration(coordinator, "public", migration) + + @dataclass(frozen=True, slots=True) class _IndexRow: name: str @@ -384,6 +430,89 @@ def test_a_failed_call_id_ledger_row_from_a_v1_103_boot_is_rolled_back_and_the_i assert rows == [(False, True), (True, False)], rows +@requires_db +@pytest.mark.parametrize( + "checksum,use_v2_migration_job,expected", + ( + (hashlib.sha256(ORIGINAL_MIGRATION_SQL[TOOL_INDEX_MIGRATION].encode()).hexdigest(), True, True), + (hashlib.sha256(b"unrelated checksum").hexdigest(), False, False), + ), + ids=("shipped-checksum", "unrelated-checksum"), +) +def test_tool_index_recovery_accepts_only_the_shipped_checksum( + scratch_database: str, checksum: str, use_v2_migration_job: bool, expected: bool +) -> None: + _leave_the_tool_index_migration_unfinished(scratch_database, checksum) + + recovered: Final = ( + _migration_job(True) if use_v2_migration_job else _recover_tool_index_migration(scratch_database) + ) + + assert recovered is expected + with psycopg.connect(scratch_database) as conn: + row: Final = conn.execute( + "SELECT checksum, finished_at IS NOT NULL FROM _prisma_migrations WHERE migration_name = %s", + (TOOL_INDEX_MIGRATION,), + ).fetchone() + assert row == (checksum, expected) + + +def _fail_the_tool_index_migration(database_url: str, checksum: str) -> None: + """The ledger row an interrupted blocking CREATE INDEX leaves: nothing applied, the table rolled back with it.""" + with psycopg.connect(database_url, autocommit=True) as conn: + conn.execute("DELETE FROM _prisma_migrations WHERE migration_name = %s", (TOOL_INDEX_MIGRATION,)) + conn.execute('DROP TABLE "LiteLLM_SpendLogToolIndex"') + conn.execute( + "INSERT INTO _prisma_migrations " + "(id, checksum, migration_name, logs, started_at, applied_steps_count) " + "VALUES (%s, %s, %s, 'canceling statement due to user request', current_timestamp, 0)", + (str(uuid.uuid4()), checksum, TOOL_INDEX_MIGRATION), + ) + + +@requires_db +def test_a_failed_tool_index_run_of_the_shipped_file_is_rolled_back_and_its_indexes_built_by_the_job( + scratch_database: str, +) -> None: + _fail_the_tool_index_migration( + scratch_database, hashlib.sha256(ORIGINAL_MIGRATION_SQL[TOOL_INDEX_MIGRATION].encode()).hexdigest() + ) + + assert _migration_job(True) is True + + _assert_no_pending_migrations(scratch_database) + with psycopg.connect(scratch_database) as conn: + rows: Final = conn.execute( + "SELECT finished_at IS NOT NULL, rolled_back_at IS NOT NULL FROM _prisma_migrations " + "WHERE migration_name = %s ORDER BY started_at", + (TOOL_INDEX_MIGRATION,), + ).fetchall() + indexes: Final = conn.execute( + "SELECT c.relname, i.indisvalid FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid " + "WHERE i.indrelid = '\"LiteLLM_SpendLogToolIndex\"'::regclass AND NOT i.indisprimary ORDER BY 1" + ).fetchall() + assert rows == [(False, True), (True, False)], rows + assert indexes == [ + ("LiteLLM_SpendLogToolIndex_start_time_idx", True), + ("LiteLLM_SpendLogToolIndex_tool_name_start_time_idx", True), + ], indexes + + +@requires_db +def test_a_failed_tool_index_row_with_an_unknown_checksum_is_left_for_the_operator(scratch_database: str) -> None: + _fail_the_tool_index_migration(scratch_database, hashlib.sha256(b"unrelated checksum").hexdigest()) + + with migration_lock(scratch_database) as coordinator: + assert ( + roll_back_replayable_failed_migration( + coordinator, "public", PACKAGE / "migrations" / TOOL_INDEX_MIGRATION / "migration.sql" + ) + is False + ) + + assert _ledger(scratch_database)[TOOL_INDEX_MIGRATION] == (False, False) + + @requires_db def test_a_failed_row_whose_migration_still_runs_sql_in_this_build_is_left_for_the_operator( partitioned_database: str, tmp_path: Path @@ -394,7 +523,7 @@ def test_a_failed_row_whose_migration_still_runs_sql_in_this_build_is_left_for_t still_building.write_text(ORIGINAL_MIGRATION_SQL[CALL_ID_INDEX_MIGRATION]) with migration_lock(partitioned_database) as coordinator: - assert roll_back_failed_inert_migration(coordinator, "public", still_building) is False + assert roll_back_replayable_failed_migration(coordinator, "public", still_building) is False assert _ledger(partitioned_database)[CALL_ID_INDEX_MIGRATION] == (False, False) @@ -405,7 +534,7 @@ def test_a_migration_without_a_failed_row_is_not_touched(partitioned_database: s before: Final = _ledger(partitioned_database) with migration_lock(partitioned_database) as coordinator: - assert roll_back_failed_inert_migration(coordinator, "public", inert) is False + assert roll_back_replayable_failed_migration(coordinator, "public", inert) is False assert _ledger(partitioned_database) == before diff --git a/tests/integration/database/test_request_log_indexes_at_boot.py b/tests/integration/database/test_request_log_indexes_at_boot.py index e72477de901..fcdee6c6ac2 100644 --- a/tests/integration/database/test_request_log_indexes_at_boot.py +++ b/tests/integration/database/test_request_log_indexes_at_boot.py @@ -27,6 +27,29 @@ API_KEY_INDEX_MIGRATION: Final = "20260823000000_add_spend_logs_api_key_starttim CALL_ID_INDEX_MIGRATION: Final = "20260831120001_spend_logs_litellm_call_id_index" ORIGINAL_MIGRATION_SQL: Final = MappingProxyType( { + "20250510142544_add_session_id_index_spend_logs": ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" ' + 'ON "LiteLLM_SpendLogs"("session_id");\n' + ), + "20260226120000_add_spend_log_tool_index": ( + 'CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogToolIndex" (\n' + ' "request_id" TEXT NOT NULL,\n' + ' "tool_name" TEXT NOT NULL,\n' + ' "start_time" TIMESTAMP(3) NOT NULL,\n' + '\n' + ' CONSTRAINT "LiteLLM_SpendLogToolIndex_pkey" PRIMARY KEY ("request_id","tool_name")\n' + ');\n' + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ' + 'ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time");\n' + ), + "20260228100000_add_spend_logs_composite_index": ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" ' + 'ON "LiteLLM_SpendLogs"("startTime", "request_id");\n' + ), + "20260724000000_add_spend_log_tool_index_start_time_idx": ( + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_start_time_idx" ' + 'ON "LiteLLM_SpendLogToolIndex"("start_time");\n' + ), API_KEY_INDEX_MIGRATION: ( "-- CreateIndex\n" 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ' @@ -43,7 +66,12 @@ MIGRATION_JOB_SECONDS: Final = 300 INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" INDEX_BUILD_LINES: Final = ("Building index", "Attached index") BUILD_SECONDS: Final = 60 -SPEND_LOGS_INDEXES: Final = ("LiteLLM_SpendLogs_api_key_startTime_idx", "LiteLLM_SpendLogs_litellm_call_id_idx") +SPEND_LOGS_INDEXES: Final = ( + "LiteLLM_SpendLogs_api_key_startTime_idx", + "LiteLLM_SpendLogs_litellm_call_id_idx", + "LiteLLM_SpendLogs_session_id_idx", + "LiteLLM_SpendLogs_startTime_request_id_idx", +) POPULATED_PARTITIONS: Final = MappingProxyType( { "LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"), @@ -254,9 +282,9 @@ def assert_ready(booted_gateway: Gateway) -> None: assert readiness.json()["db"] == "connected", readiness.text -def assert_both_indexes_cover_every_partition(database_url: str) -> None: - """Every populated partition's index is attached and valid, both parents are valid, and - a partition created afterwards inherits both indexes.""" +def assert_all_indexes_cover_every_partition(database_url: str) -> None: + """Every populated partition's managed index is attached and valid, every parent is valid, + and a partition created afterwards inherits every index.""" populated: Final = (*POPULATED_PARTITIONS, DEFAULT_PARTITION) assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS} assert attached_partition_indexes(database_url) == expected_attachments(populated) @@ -285,7 +313,7 @@ def assert_the_serving_proxy_boots_and_finds_the_indexes_in_place( @RESOLVERS -def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_both_indexes_per_partition( +def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_all_indexes_per_partition( gateway: Gateway, tmp_path: Path, resolver: Resolver ) -> None: with scratch_database() as database_url: @@ -294,7 +322,7 @@ def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_bot job: Final = migration_job(database_url, resolver) assert job.returncode == 0, job.stdout + job.stderr assert INDEXES_IN_PLACE in job.stderr + job.stdout, job.stdout + job.stderr - assert_both_indexes_cover_every_partition(database_url) + assert_all_indexes_cover_every_partition(database_url) assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver) @@ -308,7 +336,7 @@ def test_the_migration_job_heals_a_partitioned_table_left_with_the_failed_call_i fail_the_call_id_index_migration_like_the_shipped_release(database_url, tmp_path) job: Final = migration_job(database_url, resolver) assert job.returncode == 0, job.stdout + job.stderr - assert_both_indexes_cover_every_partition(database_url) + assert_all_indexes_cover_every_partition(database_url) assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver) @@ -332,12 +360,12 @@ def test_the_migration_job_leaves_a_plain_table_that_applied_the_original_index_ @RESOLVERS -def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_indexes_after_it_is_ready( +def test_a_serving_proxy_that_runs_the_migrations_itself_builds_all_indexes_after_it_is_ready( gateway: Gateway, tmp_path: Path, resolver: Resolver ) -> None: """A deployment that runs migrate deploy from the serving proxy and never runs the migration job answers readiness with the inert files applied, then its background build - puts both indexes on every partition.""" + puts every index on each partition.""" with scratch_database() as database_url: deploy_schema_before_the_index_migrations(database_url, tmp_path) partition_spend_logs(database_url) @@ -353,13 +381,13 @@ def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_indexes_aft seconds=BUILD_SECONDS, ) assert "Building index" in log and "Attached index" in log, log[-4000:] - assert_both_indexes_cover_every_partition(database_url) + assert_all_indexes_cover_every_partition(database_url) -def test_the_cli_run_as_a_migration_job_builds_both_indexes_before_it_exits(gateway: Gateway, tmp_path: Path) -> None: +def test_the_cli_run_as_a_migration_job_builds_all_indexes_before_it_exits(gateway: Gateway, tmp_path: Path) -> None: with scratch_database() as database_url: deploy_schema_before_the_index_migrations(database_url, tmp_path) partition_spend_logs(database_url) job: Final = migration_cli(database_url, gateway, V2) assert job.returncode == 0, job.stdout + job.stderr - assert_both_indexes_cover_every_partition(database_url) + assert_all_indexes_cover_every_partition(database_url) diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 02bdf3429ea..3ae29c800af 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -1492,7 +1492,7 @@ def test_migration_completion_requires_a_matching_successful_script(steps, logs, from litellm_proxy_extras.migration_recovery import MigrationProgress progress = MigrationProgress(hashlib.sha256(b"CREATE TABLE item (id int);").hexdigest(), steps, logs) - assert progress.confirms_completion(script) is expected + assert progress.confirms_completion(frozenset({hashlib.sha256(script).hexdigest()})) is expected def test_prisma_lock_waiting_has_its_own_deadline(): diff --git a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py index 5cf7593c5cd..34cb6ef7969 100644 --- a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py +++ b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py @@ -15,6 +15,9 @@ SCHEMA: Final = PACKAGE / "schema.prisma" INERT_MIGRATIONS: Final = ( "20260823000000_add_spend_logs_api_key_starttime_index", "20260831120001_spend_logs_litellm_call_id_index", + "20250510142544_add_session_id_index_spend_logs", + "20260228100000_add_spend_logs_composite_index", + "20260724000000_add_spend_log_tool_index_start_time_idx", ) CALL_ID_INDEX: Final = RequestLogIndex( "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")' diff --git a/tests/unit/test_check_migrations_no_data_rewrites.py b/tests/unit/test_check_migrations_no_data_rewrites.py index fb1053ff887..1de46e8d000 100644 --- a/tests/unit/test_check_migrations_no_data_rewrites.py +++ b/tests/unit/test_check_migrations_no_data_rewrites.py @@ -267,11 +267,10 @@ class TestIndexesOnLogTables: sql = 'DO $$ BEGIN EXECUTE \'CREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key")\'; END $$;' assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) - def test_a_marker_does_not_exempt_the_index(self, tmp_path): - sql = ( - '-- data-migration-ok: table is empty at this point\nCREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key");' - ) - assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + @pytest.mark.parametrize("table", ("LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex")) + def test_a_marker_does_not_exempt_the_index(self, tmp_path: Path, table: str): + sql = f'-- data-migration-ok: table is empty at this point\nCREATE INDEX "ix" ON "{table}" ("api_key");' + assert _keywords(tmp_path, sql) == (f'CREATE INDEX on "{table}"',) def test_a_marker_on_a_rewrite_still_leaves_the_index_below_it_flagged(self, tmp_path): sql = (