From df915b8779bbba13b38d475efd921117cea1d90f 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 | 22 +++++- .../migration.sql | 5 +- .../migration.sql | 3 - .../migration.sql | 4 +- .../migration.sql | 4 +- .../request_log_indexes.py | 10 ++- .../check_migrations_no_data_rewrites.py | 6 +- .../migrations/test_request_log_indexes.py | 79 ++++++++++++++++++- .../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 +-- 12 files changed, 159 insertions(+), 40 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..5dcf9c8ab82 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,16 @@ if TYPE_CHECKING: import psycopg +# SHA-256 of earlier shipped files whose successful run applied everything the current file does. +SUPERSEDED_CHECKSUMS: Final = MappingProxyType( + { + "20260226120000_add_spend_log_tool_index": frozenset( + {"019ae3e4e71f35763d6d515e1d463d97e75b8c50e7637c842e4e16ce59009a4b"} + ) + } +) + + @dataclass(frozen=True, slots=True) class MigrationProgress: checksum: str @@ -23,11 +34,11 @@ class MigrationProgress: id: str = "" finished: bool = False - def confirms_completion(self, script: bytes) -> bool: + def confirms_completion(self, checksums: frozenset[str]) -> bool: return ( self.applied_steps_count == 1 and not self.logs.strip() - and self.checksum == hashlib.sha256(script).hexdigest() + and self.checksum in checksums ) @@ -60,9 +71,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( 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/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..8231588d03d 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_failed_inert_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,33 @@ 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) + + @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 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..85f4158092c 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, table): + 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 = (