mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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>
This commit is contained in:
yassin
parent
572832dea1
commit
df915b8779
12 changed files with 159 additions and 40 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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_]*)?\$")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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")'
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue