mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
fix(proxy-extras): build the SpendLogs indexes in the migration job instead of in migrations (#43948)
The two SpendLogs index migrations shipped in v1.103.0 each break one table shape: the plain CREATE INDEX holds a SHARE lock on a large unpartitioned table and the CONCURRENTLY one fails with 0A000 on a partitioned parent. Both files are now inert and the indexes are built by a table-driven, shape-aware step after migrate deploy: CONCURRENTLY on a plain table, ON ONLY the parent plus per-partition CONCURRENTLY and ATTACH PARTITION on a partitioned one. The migration job builds them synchronously and exits non-zero on failure; a serving proxy that ran migrate deploy itself builds them in the background off the readiness path. A valid index of the same definition under another name is renamed and reused, an invalid one is rebuilt, and extra copies are reported with their DROP INDEX statement instead of being dropped. The migration checker rejects any CREATE INDEX on LiteLLM_SpendLogs or LiteLLM_ErrorLogs in future migrations Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
a14af739a4
commit
c66c8288c3
21 changed files with 2495 additions and 120 deletions
|
|
@ -87,3 +87,21 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
|
|||
f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
|
||||
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
|
||||
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit
|
||||
connection, for DDL that cannot run inside a transaction (`CREATE INDEX CONCURRENTLY`).
|
||||
Yields whether the lock was acquired; a v2 resolver or another migration job's index build
|
||||
holding it yields False. Released on exit."""
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
|
||||
row: Final = cursor.execute("SELECT pg_try_advisory_lock(%s) AS acquired", (MIGRATION_LOCK_KEY,)).fetchone()
|
||||
acquired: Final = row is not None and row.acquired
|
||||
try:
|
||||
yield acquired
|
||||
finally:
|
||||
if acquired:
|
||||
connection.execute("SELECT pg_advisory_unlock(%s)", (MIGRATION_LOCK_KEY,))
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import hashlib
|
||||
import re
|
||||
import subprocess
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -156,3 +157,48 @@ def baseline_current_schema(
|
|||
"review any feature-specific backfill requirements.",
|
||||
len(migrations),
|
||||
)
|
||||
|
||||
|
||||
_LINE_COMMENT_RE: Final = re.compile(r"--[^\n]*")
|
||||
_BLOCK_COMMENT_RE: Final = re.compile(r"/\*.*?\*/", re.DOTALL)
|
||||
_NO_OP_STATEMENT_RE: Final = re.compile(r"^\s*SELECT\s+1\s*$", re.IGNORECASE)
|
||||
|
||||
|
||||
def is_inert_migration(script: str) -> bool:
|
||||
"""Whether a migration file changes nothing: only comments and `SELECT 1`, so
|
||||
applying it can neither repeat nor skip a database change."""
|
||||
stripped: Final = _LINE_COMMENT_RE.sub("", _BLOCK_COMMENT_RE.sub("", script))
|
||||
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.
|
||||
"""
|
||||
from psycopg import sql
|
||||
|
||||
if not is_inert_migration(migration.read_text(encoding="utf-8")):
|
||||
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:
|
||||
return False
|
||||
result: Final = coordinator.connection.execute(
|
||||
sql.SQL(
|
||||
"UPDATE {} SET rolled_back_at = current_timestamp "
|
||||
"WHERE id = %s AND finished_at IS NULL AND rolled_back_at IS NULL"
|
||||
).format(sql.Identifier(schema, "_prisma_migrations")),
|
||||
(unfinished[0].id,),
|
||||
)
|
||||
if result.rowcount != 1:
|
||||
raise RuntimeError("Could not roll back the failed inert 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, "
|
||||
"its index is built by the migration job",
|
||||
migration.parent.name,
|
||||
)
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -1,2 +1,6 @@
|
|||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");
|
||||
-- The (api_key, startTime) index on LiteLLM_SpendLogs is built after migrate deploy,
|
||||
-- through litellm_proxy_extras/request_log_indexes.py: concurrently on a plain table and
|
||||
-- per partition on a partitioned one. The migration job builds it; a serving proxy that
|
||||
-- ran the migrations itself builds it in the background once it serves. A migration
|
||||
-- cannot do either without blocking spend-log writes or failing on a partitioned table.
|
||||
SELECT 1;
|
||||
|
|
|
|||
|
|
@ -1,12 +1,6 @@
|
|||
-- CreateIndex (CONCURRENTLY)
|
||||
--
|
||||
-- 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.
|
||||
-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
|
||||
-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
|
||||
-- - Do not edit this file after it has been applied to any database: Prisma checksums
|
||||
-- migrations; add a new migration instead.
|
||||
-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
|
||||
-- without IF NOT EXISTS if you must support older versions).
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
|
||||
-- The litellm_call_id index on LiteLLM_SpendLogs is built after migrate deploy, through
|
||||
-- litellm_proxy_extras/request_log_indexes.py: concurrently on a plain table and per
|
||||
-- partition on a partitioned one. The migration job builds it; a serving proxy that ran
|
||||
-- the migrations itself builds it in the background once it serves. Postgres refuses
|
||||
-- CREATE INDEX CONCURRENTLY on a partitioned parent, so this migration no longer runs it.
|
||||
SELECT 1;
|
||||
|
|
|
|||
448
litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py
Normal file
448
litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py
Normal file
|
|
@ -0,0 +1,448 @@
|
|||
"""The request-log indexes built after `prisma migrate deploy` instead of by a migration:
|
||||
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
|
||||
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
|
||||
index a database already has is recognized and never rebuilt.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import random
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
from psycopg import sql
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RequestLogIndex:
|
||||
"""One index the migration job owns: the table, the exact Prisma index name and the
|
||||
column list as it would be written after `ON <table>`."""
|
||||
|
||||
table: str
|
||||
name: str
|
||||
definition: str
|
||||
|
||||
@property
|
||||
def columns(self) -> tuple[str, ...]:
|
||||
return tuple(re.findall(r'"([^"]+)"', self.definition))
|
||||
|
||||
def partition_index_name(self, partition: str) -> str:
|
||||
"""The child index name for one partition, built the way Postgres names the
|
||||
children of a partitioned index, and kept within the 63 byte identifier limit."""
|
||||
name: Final = f"{partition}_{self.name.removeprefix(f'{self.table}_')}"
|
||||
if len(name.encode()) <= _IDENTIFIER_MAX_BYTES:
|
||||
return name
|
||||
digest: Final = hashlib.sha256(name.encode()).hexdigest()[:_DIGEST_LENGTH]
|
||||
budget: Final = _IDENTIFIER_MAX_BYTES - _DIGEST_LENGTH - 1
|
||||
kept: Final = next(name[:length] for length in range(len(name), 0, -1) if len(name[:length].encode()) <= budget)
|
||||
return f"{kept}_{digest}"
|
||||
|
||||
|
||||
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")'),
|
||||
)
|
||||
|
||||
_IDENTIFIER_MAX_BYTES: Final = 63
|
||||
_PARENT_LOCK_TIMEOUT: Final = "2s"
|
||||
_PARENT_LOCK_ATTEMPTS: Final = 30
|
||||
_LOCK_HANDOVER_SECONDS: Final = 2.0
|
||||
_DIGEST_LENGTH: Final = 8
|
||||
_CREATE_INDEX_STATEMENT: Final = re.compile(
|
||||
r'^\s*CREATE\s+(?:UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?"(?P<index>[^"]+)"\s+ON\b',
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)"
|
||||
_CHILDREN_WITHOUT_THE_INDEX_SQL: Final = (
|
||||
"SELECT child.relname AS name, n.nspname AS schema, child.relkind = 'p' AS partitioned "
|
||||
"FROM pg_inherits i JOIN pg_class child ON child.oid = i.inhrelid "
|
||||
"JOIN pg_namespace n ON n.oid = child.relnamespace "
|
||||
"WHERE i.inhparent = to_regclass(%s) AND NOT EXISTS ("
|
||||
"SELECT 1 FROM pg_inherits attached JOIN pg_index x ON x.indexrelid = attached.inhrelid "
|
||||
"WHERE attached.inhparent = to_regclass(%s) AND x.indrelid = child.oid) "
|
||||
"ORDER BY child.relname"
|
||||
)
|
||||
_EQUIVALENT_INDEXES_SQL: Final = (
|
||||
"SELECT i.relname AS name, x.indisvalid AS valid "
|
||||
"FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid JOIN pg_am am ON am.oid = i.relam "
|
||||
"WHERE x.indrelid = to_regclass(%s) AND i.relname <> %s AND am.amname = 'btree' AND NOT x.indisunique "
|
||||
"AND x.indexprs IS NULL AND x.indpred IS NULL AND x.indnkeyatts = x.indnatts "
|
||||
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indoption::int2[]) o WHERE o <> 0) "
|
||||
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indclass::oid[]) c JOIN pg_opclass oc ON oc.oid = c WHERE NOT oc.opcdefault) "
|
||||
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indcollation::oid[]) WITH ORDINALITY c(coll, ord) "
|
||||
"JOIN unnest(x.indkey::int2[]) WITH ORDINALITY k(attnum, ord) ON k.ord = c.ord "
|
||||
"JOIN pg_attribute a ON a.attrelid = x.indrelid AND a.attnum = k.attnum "
|
||||
"WHERE c.coll <> 0 AND c.coll <> a.attcollation) "
|
||||
"AND (SELECT array_agg(a.attname::text ORDER BY k.ord) FROM unnest(x.indkey::int2[]) WITH ORDINALITY k(attnum, ord) "
|
||||
"JOIN pg_attribute a ON a.attrelid = x.indrelid AND a.attnum = k.attnum) = %s::text[] "
|
||||
"AND NOT EXISTS (SELECT 1 FROM pg_inherits WHERE inhrelid = x.indexrelid) "
|
||||
"ORDER BY x.indisvalid DESC, i.relname"
|
||||
)
|
||||
_INDEX_STATE_SQL: Final = (
|
||||
'SELECT x.indisvalid AS valid, t.relname AS "table" '
|
||||
"FROM pg_index x JOIN pg_class t ON t.oid = x.indrelid WHERE x.indexrelid = to_regclass(%s)"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Relation:
|
||||
name: str
|
||||
schema: str
|
||||
partitioned: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _IndexState:
|
||||
valid: bool
|
||||
table: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _EquivalentIndex:
|
||||
name: str
|
||||
valid: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TableKind:
|
||||
partitioned: bool
|
||||
|
||||
|
||||
def filter_request_log_index_diff(diff_sql: str, indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES) -> str:
|
||||
"""The `prisma migrate diff` script without the statements that create a migration-job-owned
|
||||
index, which the schema declares and the migrations deliberately do not build."""
|
||||
names: Final = frozenset(index.name for index in indexes)
|
||||
statements: Final = diff_sql.split(";")
|
||||
kept: Final = tuple(statement for statement in statements if not _creates_one_of(statement, names))
|
||||
return ";".join(kept) if any(part.strip() for part in kept) else ""
|
||||
|
||||
|
||||
def _creates_one_of(statement: str, names: frozenset[str]) -> bool:
|
||||
match: Final = _CREATE_INDEX_STATEMENT.match(_without_comments(statement))
|
||||
return match is not None and match["index"] in names
|
||||
|
||||
|
||||
def _without_comments(statement: str) -> str:
|
||||
return "\n".join(line for line in statement.splitlines() if not line.lstrip().startswith("--"))
|
||||
|
||||
|
||||
def _connect(database_url: str) -> "psycopg.Connection[tuple[object, ...]]":
|
||||
import psycopg
|
||||
|
||||
return psycopg.connect(database_url, connect_timeout=10, autocommit=True)
|
||||
|
||||
|
||||
def ensure_request_log_indexes(
|
||||
database_url: str,
|
||||
schema: str,
|
||||
indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES,
|
||||
connect: "Callable[[str], psycopg.Connection[tuple[object, ...]]]" = _connect,
|
||||
) -> bool:
|
||||
"""Build every listed index that is missing or invalid. Each build step runs under
|
||||
the migration coordinator lock, held per statement so a resolver booting on another
|
||||
replica gets in between partitions rather than waiting for the whole table. Any
|
||||
failure is logged and left for the next index build; the result says whether
|
||||
every index ended up valid. Never raises."""
|
||||
import psycopg
|
||||
|
||||
try:
|
||||
with connect(database_url) as connection:
|
||||
connection.execute("SET statement_timeout = 0")
|
||||
results: Final = tuple(_ensure_index(connection, schema, index) for index in indexes)
|
||||
except psycopg.Error as exc:
|
||||
logger.warning("Could not build the request-log indexes, leaving them for the next index build: %s", exc)
|
||||
return False
|
||||
if not all(results):
|
||||
logger.warning("Some request-log indexes are not in place yet, leaving them for the next index build")
|
||||
return False
|
||||
logger.info("Request-log indexes are all in place")
|
||||
return True
|
||||
|
||||
|
||||
def _under_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]", step: Callable[[], bool]) -> bool:
|
||||
with held_migration_lock(connection) as held:
|
||||
if not held:
|
||||
logger.info(
|
||||
"Another process holds the migration lock, leaving the request-log indexes to the next index build"
|
||||
)
|
||||
return False
|
||||
return step()
|
||||
|
||||
|
||||
def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool:
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_TableKind)) as cursor:
|
||||
table: Final = cursor.execute(_TABLE_KIND_SQL, (_regclass_name(connection, schema, index.table),)).fetchone()
|
||||
if table is None:
|
||||
logger.info("Table %s does not exist yet, skipping index %s", index.table, index.name)
|
||||
return True
|
||||
if table.partitioned:
|
||||
return build_index_on_partitioned_table(connection, schema, index)
|
||||
return _build_leaf_index(connection, schema, index.table, index.name, index)
|
||||
|
||||
|
||||
def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str:
|
||||
from psycopg import sql
|
||||
|
||||
return sql.Identifier(schema, name).as_string(connection)
|
||||
|
||||
|
||||
def _create_index_statement(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]", prefix: "sql.Composed", definition: str
|
||||
) -> bytes:
|
||||
return (prefix.as_string(connection) + definition).encode()
|
||||
|
||||
|
||||
def _index_state(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: str) -> "_IndexState | None":
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_IndexState)) as cursor:
|
||||
return cursor.execute(_INDEX_STATE_SQL, (_regclass_name(connection, schema, index),)).fetchone()
|
||||
|
||||
|
||||
def _equivalent_indexes(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
table: str,
|
||||
name: str,
|
||||
index: RequestLogIndex,
|
||||
) -> tuple[_EquivalentIndex, ...]:
|
||||
"""The indexes on `table` other than `name` with the same definition: default btree
|
||||
over the same columns in the same order, no expression, predicate, DESC or custom
|
||||
opclass or collation, and not attached under a partitioned index. Valid ones first."""
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_EquivalentIndex)) as cursor:
|
||||
return tuple(
|
||||
cursor.execute(
|
||||
_EQUIVALENT_INDEXES_SQL, (_regclass_name(connection, schema, table), name, list(index.columns))
|
||||
).fetchall()
|
||||
)
|
||||
|
||||
|
||||
def _adopt_equivalent_index(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
table: str,
|
||||
name: str,
|
||||
index: RequestLogIndex,
|
||||
) -> bool:
|
||||
"""Rename a valid index of the same definition under another name (an operator's
|
||||
hand-built copy, say) to the name this code expects, instead of building a second
|
||||
one. RENAME on an index is a catalog change that lets writes through."""
|
||||
from psycopg import sql
|
||||
|
||||
equivalent: Final = next(
|
||||
(found for found in _equivalent_indexes(connection, schema, table, name, index) if found.valid), None
|
||||
)
|
||||
if equivalent is None:
|
||||
return False
|
||||
logger.info(
|
||||
"Renaming the equivalent index %s on %s to %s instead of building a second one", equivalent.name, table, name
|
||||
)
|
||||
connection.execute(
|
||||
sql.SQL("ALTER INDEX {} RENAME TO {}").format(sql.Identifier(schema, equivalent.name), sql.Identifier(name))
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def _report_second_copies(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
table: str,
|
||||
name: str,
|
||||
index: RequestLogIndex,
|
||||
concurrently: bool,
|
||||
) -> None:
|
||||
"""Log every other index of the same definition with the statement that removes it.
|
||||
Dropping is the operator's call: a second copy costs writes and disk, never results."""
|
||||
from psycopg import sql
|
||||
|
||||
drop: Final = "DROP INDEX CONCURRENTLY" if concurrently else "DROP INDEX"
|
||||
for copy in _equivalent_indexes(connection, schema, table, name, index):
|
||||
logger.warning(
|
||||
"Index %s on %s is a second copy of %s and only costs writes and disk; remove it with: %s %s",
|
||||
copy.name,
|
||||
table,
|
||||
name,
|
||||
drop,
|
||||
sql.Identifier(schema, copy.name).as_string(connection),
|
||||
)
|
||||
|
||||
|
||||
def _children_without_the_index(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str, index: str
|
||||
) -> tuple[_Relation, ...]:
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_Relation)) as cursor:
|
||||
return tuple(
|
||||
cursor.execute(
|
||||
_CHILDREN_WITHOUT_THE_INDEX_SQL,
|
||||
(_regclass_name(connection, schema, table), _regclass_name(connection, schema, index)),
|
||||
).fetchall()
|
||||
)
|
||||
|
||||
|
||||
def _build_leaf_index(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
table: str,
|
||||
name: str,
|
||||
index: RequestLogIndex,
|
||||
) -> bool:
|
||||
"""Build one plain table's or partition's index with CONCURRENTLY so writes keep
|
||||
flowing. The catalog is read under the migration lock, so a replica that saw an
|
||||
invalid index before the lock finds the valid one another replica just built and
|
||||
leaves it. An invalid index left by an interrupted build is dropped and rebuilt; a
|
||||
valid index of the same definition under another name is renamed rather than
|
||||
duplicated; an index of that name on another table is a collision this code will
|
||||
not touch."""
|
||||
from psycopg import sql
|
||||
|
||||
def build() -> bool:
|
||||
existing: Final = _index_state(connection, schema, name)
|
||||
if existing is not None and existing.table != table:
|
||||
logger.warning(
|
||||
"Index %s already exists on %s rather than %s, leaving it alone", name, existing.table, table
|
||||
)
|
||||
return False
|
||||
if existing is not None and existing.valid:
|
||||
return True
|
||||
if existing is not None:
|
||||
logger.info("Dropping the invalid index %s left by an interrupted build on %s", name, table)
|
||||
connection.execute(sql.SQL("DROP INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)))
|
||||
elif _adopt_equivalent_index(connection, schema, table, name, index):
|
||||
return True
|
||||
logger.info("Building index %s on %s concurrently", name, table)
|
||||
prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
|
||||
sql.Identifier(name), sql.Identifier(schema, table)
|
||||
)
|
||||
connection.execute(_create_index_statement(connection, prefix, index.definition))
|
||||
built: Final = _index_state(connection, schema, name)
|
||||
return built is not None and built.valid
|
||||
|
||||
current: Final = _index_state(connection, schema, name)
|
||||
if current is None or not current.valid or current.table != table:
|
||||
if not _under_migration_lock(connection, build):
|
||||
return False
|
||||
time.sleep(_LOCK_HANDOVER_SECONDS)
|
||||
_report_second_copies(connection, schema, table, name, index, concurrently=True)
|
||||
return True
|
||||
|
||||
|
||||
def build_index_on_partitioned_table(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
index: RequestLogIndex,
|
||||
table: "str | None" = None,
|
||||
name: "str | None" = None,
|
||||
) -> bool:
|
||||
"""Build the index the way Postgres allows on a partitioned parent: a metadata-only
|
||||
parent index ON ONLY the parent, one CONCURRENTLY build per partition, and ATTACH
|
||||
PARTITION for each child. Partitions that are themselves partitioned get the same
|
||||
treatment one level down. Every step checks the catalog before acting, so an
|
||||
interrupted run resumes where it stopped and a second run finds nothing to do; a
|
||||
parent or child index of the same definition under another name is renamed and
|
||||
used rather than duplicated. The connection must be in autocommit mode. True when
|
||||
the parent index ends up valid."""
|
||||
|
||||
parent_table: Final = index.table if table is None else table
|
||||
parent_index: Final = index.name if name is None else name
|
||||
existing: Final = _index_state(connection, schema, parent_index)
|
||||
if existing is not None and existing.table != parent_table:
|
||||
logger.warning(
|
||||
"Index %s already exists on %s rather than %s, leaving it alone", parent_index, existing.table, parent_table
|
||||
)
|
||||
return False
|
||||
if existing is None and not _under_migration_lock(
|
||||
connection,
|
||||
lambda: (
|
||||
_adopt_equivalent_index(connection, schema, parent_table, parent_index, index)
|
||||
or _create_parent_index(connection, schema, parent_index, parent_table, index)
|
||||
),
|
||||
):
|
||||
return False
|
||||
children: Final = _children_without_the_index(connection, schema, parent_table, parent_index)
|
||||
if not all(_attach_child_index(connection, schema, parent_index, child, index) for child in children):
|
||||
return False
|
||||
final: Final = _index_state(connection, schema, parent_index)
|
||||
if final is None or not final.valid:
|
||||
return False
|
||||
_report_second_copies(connection, schema, parent_table, parent_index, index, concurrently=False)
|
||||
return True
|
||||
|
||||
|
||||
def _create_parent_index(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
name: str,
|
||||
table: str,
|
||||
index: RequestLogIndex,
|
||||
) -> bool:
|
||||
"""Create the metadata-only parent index. Postgres takes a SHARE lock on the
|
||||
parent for that statement, so it waits for in-flight writes and queues new ones
|
||||
behind it; a short lock_timeout with retries keeps every such pause bounded."""
|
||||
import psycopg
|
||||
from psycopg import sql
|
||||
|
||||
prefix: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
|
||||
sql.Identifier(name), sql.Identifier(schema, table)
|
||||
)
|
||||
statement: Final = _create_index_statement(connection, prefix, index.definition)
|
||||
connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_PARENT_LOCK_TIMEOUT)))
|
||||
try:
|
||||
for _ in range(_PARENT_LOCK_ATTEMPTS):
|
||||
try:
|
||||
connection.execute(statement)
|
||||
return True
|
||||
except psycopg.errors.LockNotAvailable:
|
||||
logger.info("Waiting for in-flight writes to %s before creating the parent index %s", table, name)
|
||||
time.sleep(random.uniform(0.1, 0.5))
|
||||
finally:
|
||||
connection.execute("SET lock_timeout = 0")
|
||||
logger.warning("Could not get the parent lock on %s to create %s, leaving it for the next index build", table, name)
|
||||
return False
|
||||
|
||||
|
||||
def _attach_child_index(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
parent_index: str,
|
||||
child: _Relation,
|
||||
index: RequestLogIndex,
|
||||
) -> bool:
|
||||
from psycopg import sql
|
||||
|
||||
child_index: Final = index.partition_index_name(child.name)
|
||||
built: Final = (
|
||||
build_index_on_partitioned_table(connection, child.schema, index, child.name, child_index)
|
||||
if child.partitioned
|
||||
else _build_leaf_index(connection, child.schema, child.name, child_index, index)
|
||||
)
|
||||
if not built:
|
||||
return False
|
||||
|
||||
def attach() -> bool:
|
||||
connection.execute(
|
||||
sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format(
|
||||
sql.Identifier(schema, parent_index), sql.Identifier(child.schema, child_index)
|
||||
)
|
||||
)
|
||||
logger.info("Attached index %s on partition %s to %s", child_index, child.name, parent_index)
|
||||
return True
|
||||
|
||||
return _under_migration_lock(connection, attach)
|
||||
|
|
@ -5,6 +5,7 @@ import re
|
|||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, replace
|
||||
|
|
@ -13,6 +14,7 @@ from typing import TYPE_CHECKING, Final, Optional
|
|||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -24,6 +26,7 @@ from litellm_proxy_extras.replica_identity import (
|
|||
REPLICA_IDENTITY_FULL_ENV_VAR,
|
||||
apply_replica_identity_full,
|
||||
)
|
||||
from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes, filter_request_log_index_diff
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
|
@ -433,6 +436,21 @@ class ProxyExtrasDBManager:
|
|||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _filter_migration_job_owned_drift(diff_sql: str, partitioned: bool | None = None) -> str:
|
||||
"""The drift script without the indexes the migration job builds (the schema
|
||||
declares them, the migrations deliberately do not) and, when LiteLLM_SpendLogs
|
||||
is partitioned, without its primary-key rewrite and partitioning artifacts."""
|
||||
without_indexes: Final = filter_request_log_index_diff(diff_sql)
|
||||
is_partitioned: Final = ProxyExtrasDBManager.spend_logs_is_partitioned() if partitioned is None else partitioned
|
||||
if not is_partitioned:
|
||||
return without_indexes
|
||||
logger.info(
|
||||
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
|
||||
"rewrite and partitioning artifacts from the drift script"
|
||||
)
|
||||
return filter_partitioned_spend_logs_diff(without_indexes)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_all_migrations(
|
||||
migrations_dir: str, schema_path: str, mark_all_applied: bool = True
|
||||
|
|
@ -513,21 +531,14 @@ class ProxyExtrasDBManager:
|
|||
return
|
||||
logger.info(f"Migration diff created at {diff_sql_path}")
|
||||
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
filtered_sql = filter_partitioned_spend_logs_diff(
|
||||
diff_sql_path.read_text()
|
||||
)
|
||||
diff_sql_path.write_text(filtered_sql)
|
||||
logger.info(
|
||||
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
|
||||
"rewrite and partitioning artifacts from the drift script"
|
||||
)
|
||||
if not filtered_sql.strip():
|
||||
logger.info("Drift script is empty after filtering; nothing to apply")
|
||||
if not mark_all_applied:
|
||||
return
|
||||
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
|
||||
filtered_sql: Final = ProxyExtrasDBManager._filter_migration_job_owned_drift(diff_sql_path.read_text())
|
||||
diff_sql_path.write_text(filtered_sql)
|
||||
if not filtered_sql.strip():
|
||||
logger.info("Drift script is empty after filtering; nothing to apply")
|
||||
if not mark_all_applied:
|
||||
return
|
||||
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
|
||||
return
|
||||
|
||||
# 2. Run prisma db execute to apply the migration
|
||||
applied_ok = False
|
||||
|
|
@ -800,7 +811,7 @@ class ProxyExtrasDBManager:
|
|||
conn.execute(statement)
|
||||
except psycopg.Error as e:
|
||||
logger.warning(
|
||||
"Could not repair invalid index %s.%s, will retry on the next startup. "
|
||||
"Could not repair invalid index %s.%s, will retry on the next database setup run. "
|
||||
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
|
||||
index.schema,
|
||||
index.name,
|
||||
|
|
@ -811,16 +822,21 @@ class ProxyExtrasDBManager:
|
|||
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
|
||||
|
||||
@staticmethod
|
||||
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
|
||||
def repair_invalid_indexes(
|
||||
lock_timeout: str = "30s",
|
||||
repair: "Callable[[psycopg.Connection[tuple[str, str, str]], _InvalidIndex], None] | None" = None,
|
||||
) -> bool:
|
||||
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
|
||||
INVALID (a migration deadlock between replicas is the usual cause; the
|
||||
retried migration skips them because of IF NOT EXISTS). Never raises:
|
||||
returns True when no invalid index remains, False when the repair was
|
||||
skipped or failed and will be retried on the next startup. Looks in the
|
||||
skipped or failed and will be retried on the next database setup run. Looks in the
|
||||
schema DATABASE_URL names, the only URL Prisma migrates through, but
|
||||
connects over DIRECT_URL when set: the session settings, the advisory
|
||||
lock and REINDEX CONCURRENTLY all need one server session, which a
|
||||
transaction pooler does not give."""
|
||||
transaction pooler does not give. Each rebuild holds the migration
|
||||
coordinator lock on its own, like the migration job's index build, so a resolver
|
||||
booting on another replica waits for one index at most."""
|
||||
prisma_url: Final = os.getenv("DATABASE_URL")
|
||||
if not prisma_url:
|
||||
return False
|
||||
|
|
@ -856,20 +872,53 @@ class ProxyExtrasDBManager:
|
|||
if lock_row is None or not lock_row[0]:
|
||||
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
|
||||
return False
|
||||
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
|
||||
ProxyExtrasDBManager._repair_index(conn, index)
|
||||
repair_one: Final = repair or ProxyExtrasDBManager._repair_index
|
||||
repaired: Final = all(
|
||||
ProxyExtrasDBManager._repair_under_migration_lock(conn, schema, index, repair_one)
|
||||
for index in found
|
||||
)
|
||||
if not repaired:
|
||||
return False
|
||||
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
|
||||
except psycopg.Error as e:
|
||||
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
|
||||
logger.warning(
|
||||
"Could not check for invalid indexes, will retry on the next database setup run. Error: %s", e
|
||||
)
|
||||
return False
|
||||
return not remaining
|
||||
|
||||
@staticmethod
|
||||
def _repair_under_migration_lock(
|
||||
conn: "psycopg.Connection[tuple[str, str, str]]",
|
||||
schema: str,
|
||||
index: _InvalidIndex,
|
||||
repair: "Callable[[psycopg.Connection[tuple[str, str, str]], _InvalidIndex], None]",
|
||||
) -> bool:
|
||||
"""Rebuild one index under the migration coordinator lock, skipping it when a
|
||||
migration job finished or dropped it in the meantime. False when another process
|
||||
holds the lock, so the check waits for the next database setup run."""
|
||||
with held_migration_lock(conn) as held:
|
||||
if not held:
|
||||
logger.info(
|
||||
"Another process is building indexes under the migration lock, leaving the "
|
||||
"invalid index check to the next database setup run"
|
||||
)
|
||||
return False
|
||||
still_invalid: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
|
||||
if any(found.schema == index.schema and found.name == index.name for found in still_invalid):
|
||||
repair(conn, index)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
if not use_migrate:
|
||||
return ProxyExtrasDBManager._run_database_v2(False)
|
||||
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
|
||||
from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration
|
||||
from litellm_proxy_extras.migration_recovery import (
|
||||
baseline_current_schema,
|
||||
recover_completed_migration,
|
||||
roll_back_failed_inert_migration,
|
||||
)
|
||||
|
||||
database_url: Final = os.environ.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
|
|
@ -884,7 +933,9 @@ class ProxyExtrasDBManager:
|
|||
if not migration.is_file():
|
||||
return False
|
||||
with migration_lock(lock_url) as coordinator:
|
||||
return recover_completed_migration(coordinator, schema, migration)
|
||||
return recover_completed_migration(coordinator, schema, migration) or roll_back_failed_inert_migration(
|
||||
coordinator, schema, migration
|
||||
)
|
||||
|
||||
def baseline_existing(migrations_dir: str) -> None:
|
||||
with migration_lock(lock_url) as coordinator:
|
||||
|
|
@ -1177,13 +1228,16 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(
|
||||
use_migrate: bool = False, use_v2_resolver: bool = False
|
||||
) -> bool:
|
||||
def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool:
|
||||
"""
|
||||
Set up the database using either prisma migrate or prisma db push
|
||||
Uses migrations from litellm-proxy-extras package
|
||||
|
||||
The request-log indexes in `REQUEST_LOG_INDEXES` are not built here: the
|
||||
migration job builds them through `run_migration_job`, and a serving proxy that
|
||||
ran the migrations itself starts them through `start_request_log_index_build`
|
||||
once it is ready to serve.
|
||||
|
||||
Args:
|
||||
use_migrate: Whether to use prisma migrate instead of db push
|
||||
use_v2_resolver: Opt into the v2 migration resolver (safer during
|
||||
|
|
@ -1200,10 +1254,48 @@ class ProxyExtrasDBManager:
|
|||
migrated = ProxyExtrasDBManager._run_migrations(
|
||||
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
|
||||
)
|
||||
if migrated:
|
||||
ProxyExtrasDBManager.repair_invalid_indexes()
|
||||
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
|
||||
return migrated
|
||||
if not migrated:
|
||||
return False
|
||||
ProxyExtrasDBManager.repair_invalid_indexes()
|
||||
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
|
||||
"""Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the
|
||||
migrations target. Idempotent and never raises; False when an index is still
|
||||
missing or invalid, so the migration job reports it and gets rerun instead of
|
||||
leaving the table unindexed until the next deploy."""
|
||||
database_url: Final = os.environ.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
return True
|
||||
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
|
||||
os.environ.get("DIRECT_URL") or database_url
|
||||
)
|
||||
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
|
||||
return build(direct_url, schema)
|
||||
|
||||
@staticmethod
|
||||
def run_migration_job(
|
||||
use_migrate: bool = False,
|
||||
use_v2_resolver: bool = False,
|
||||
setup: Callable[[bool, bool], bool] = setup_database,
|
||||
build: Callable[[], bool] = build_request_log_indexes,
|
||||
) -> bool:
|
||||
"""The migration job's whole run: `setup_database`, then the request-log indexes,
|
||||
built synchronously so the job exits only once they are in place. False when the
|
||||
migrations failed or an index could not be built, so the Job is rerun."""
|
||||
return setup(use_migrate, use_v2_resolver) and build()
|
||||
|
||||
@staticmethod
|
||||
def start_request_log_index_build(build: Callable[[], bool] = build_request_log_indexes) -> threading.Thread:
|
||||
"""A serving proxy that ran the migrations itself (schema updates not disabled)
|
||||
builds the request-log indexes on a daemon thread, so a long build never delays
|
||||
readiness. A build that could not finish is logged and picked up by the next boot
|
||||
or the migration job."""
|
||||
thread: Final = threading.Thread(target=build, name="litellm-request-log-indexes", daemon=True)
|
||||
thread.start()
|
||||
return thread
|
||||
|
||||
@staticmethod
|
||||
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
|
|
@ -1247,15 +1339,16 @@ class ProxyExtrasDBManager:
|
|||
logger.info("✅ Post-migration sanity check completed")
|
||||
return True
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}")
|
||||
if "P3009" in e.stderr:
|
||||
stderr: Final = str(e.stderr or "")
|
||||
logger.info(f"prisma db error: {stderr}, e: {e.stdout}")
|
||||
if "P3009" in stderr:
|
||||
# Extract the failed migration name from the error message
|
||||
migration_match = re.search(
|
||||
r"`(\d+_.*)` migration", e.stderr
|
||||
r"`(\d+_.*)` migration", stderr
|
||||
)
|
||||
if migration_match:
|
||||
failed_migration = migration_match.group(1)
|
||||
if ProxyExtrasDBManager._is_idempotent_error(e.stderr):
|
||||
if ProxyExtrasDBManager._is_idempotent_error(stderr):
|
||||
logger.info(
|
||||
f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied"
|
||||
)
|
||||
|
|
@ -1311,8 +1404,8 @@ class ProxyExtrasDBManager:
|
|||
f"✅ Migration {failed_migration} marked as rolled back... retrying"
|
||||
)
|
||||
elif (
|
||||
"P3005" in e.stderr
|
||||
and "database schema is not empty" in e.stderr
|
||||
"P3005" in stderr
|
||||
and "database schema is not empty" in stderr
|
||||
):
|
||||
logger.info(
|
||||
"Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
|
||||
|
|
@ -1326,13 +1419,13 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
logger.info("✅ All migrations resolved.")
|
||||
return True
|
||||
elif "P3018" in e.stderr:
|
||||
elif "P3018" in stderr:
|
||||
# Check if this is a permission error or idempotent error
|
||||
if ProxyExtrasDBManager._is_permission_error(e.stderr):
|
||||
if ProxyExtrasDBManager._is_permission_error(stderr):
|
||||
# Permission errors should NOT be marked as applied
|
||||
# Extract migration name for logging
|
||||
migration_match = re.search(
|
||||
r"Migration name: (\d+_.*)", e.stderr
|
||||
r"Migration name: (\d+_.*)", stderr
|
||||
)
|
||||
migration_name = (
|
||||
migration_match.group(1)
|
||||
|
|
@ -1342,7 +1435,7 @@ class ProxyExtrasDBManager:
|
|||
|
||||
logger.error(
|
||||
f"❌ Migration {migration_name} failed due to insufficient permissions. "
|
||||
f"Please check database user privileges. Error: {e.stderr}"
|
||||
f"Please check database user privileges. Error: {stderr}"
|
||||
)
|
||||
|
||||
# Mark as rolled back and exit with error
|
||||
|
|
@ -1365,7 +1458,7 @@ class ProxyExtrasDBManager:
|
|||
f"was NOT applied. Please grant necessary database permissions and retry."
|
||||
) from e
|
||||
|
||||
elif ProxyExtrasDBManager._is_idempotent_error(e.stderr):
|
||||
elif ProxyExtrasDBManager._is_idempotent_error(stderr):
|
||||
# Idempotent errors mean the migration has effectively been applied
|
||||
logger.info(
|
||||
"Migration failed due to idempotent error (e.g., column already exists), "
|
||||
|
|
@ -1373,7 +1466,7 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
# Extract the migration name from the error message
|
||||
migration_match = re.search(
|
||||
r"Migration name: (\d+_.*)", e.stderr
|
||||
r"Migration name: (\d+_.*)", stderr
|
||||
)
|
||||
if migration_match:
|
||||
migration_name = migration_match.group(1)
|
||||
|
|
@ -1422,7 +1515,7 @@ class ProxyExtrasDBManager:
|
|||
logger.warning(
|
||||
f"P3018 error encountered but could not classify "
|
||||
f"as permission or idempotent error. "
|
||||
f"Error: {e.stderr}"
|
||||
f"Error: {stderr}"
|
||||
)
|
||||
raise
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -984,6 +984,30 @@ class PrismaManager:
|
|||
os.chdir(original_dir)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def build_request_log_indexes() -> bool:
|
||||
"""Build the request-log indexes the migrations leave out and wait for them, for the
|
||||
migration job (`--skip_server_startup`) after `setup_database` succeeds. False when
|
||||
an index could not be built, so the job exits non-zero and is rerun."""
|
||||
try:
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
except ImportError as e:
|
||||
verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
|
||||
return False
|
||||
return ProxyExtrasDBManager.build_request_log_indexes()
|
||||
|
||||
@staticmethod
|
||||
def start_request_log_index_build() -> None:
|
||||
"""Build the request-log indexes on a daemon thread, for a serving proxy that ran the
|
||||
migrations itself (`DISABLE_SCHEMA_UPDATE` unset), so a long build never delays
|
||||
readiness. A build that could not finish is logged and retried on the next boot."""
|
||||
try:
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
except ImportError as e:
|
||||
verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
|
||||
return
|
||||
ProxyExtrasDBManager.start_request_log_index_build()
|
||||
|
||||
|
||||
def should_update_prisma_schema(
|
||||
disable_updates: bool | str | None = None,
|
||||
|
|
|
|||
|
|
@ -1410,10 +1410,15 @@ def run_server(
|
|||
"LiteLLM versions contend for the same DB.\033[0m"
|
||||
)
|
||||
try:
|
||||
setup_ok: Final = PrismaManager.setup_database(
|
||||
migrated: Final = PrismaManager.setup_database(
|
||||
use_migrate=not use_prisma_db_push,
|
||||
use_v2_resolver=use_v2_resolver,
|
||||
)
|
||||
setup_ok: Final = migrated and (
|
||||
not skip_server_startup or PrismaManager.build_request_log_indexes()
|
||||
)
|
||||
if migrated and not skip_server_startup:
|
||||
PrismaManager.start_request_log_index_build()
|
||||
except RuntimeError as e:
|
||||
# Raised on unrecoverable migration errors: the v2
|
||||
# resolver's non-idempotent failures and permission
|
||||
|
|
|
|||
|
|
@ -2,7 +2,12 @@
|
|||
|
||||
Runs `prisma migrate deploy` against the LiteLLM writer database using the
|
||||
recovery logic in `litellm_proxy_extras.ProxyExtrasDBManager.setup_database`
|
||||
(P3005 baseline + P3009/P3018 idempotent-error handling, retries, etc.).
|
||||
(P3005 baseline + P3009/P3018 idempotent-error handling, retries, etc.), then
|
||||
builds the request-log indexes the migrations leave out
|
||||
(`litellm_proxy_extras.request_log_indexes`), waiting for them. The job exits
|
||||
non-zero when an index could not be built so that it is rerun. A serving proxy
|
||||
that runs the migrations itself builds the same indexes in the background once
|
||||
it serves.
|
||||
|
||||
Env vars:
|
||||
DATABASE_URL required unless it can be assembled at
|
||||
|
|
@ -23,10 +28,11 @@ Env vars:
|
|||
import os
|
||||
import sys
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager, str_to_bool
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
|
||||
|
||||
def main() -> int:
|
||||
# Assemble DATABASE_URL from the discrete DATABASE_* env vars, matching
|
||||
|
|
@ -52,7 +58,7 @@ def main() -> int:
|
|||
not use_db_push,
|
||||
use_v2,
|
||||
)
|
||||
ok = ProxyExtrasDBManager.setup_database(
|
||||
ok = ProxyExtrasDBManager.run_migration_job(
|
||||
use_migrate=not use_db_push,
|
||||
use_v2_resolver=use_v2,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,15 +7,16 @@ anything whose cost scales with existing table size turns into downtime. A singl
|
|||
plus a doubled heap that plain autovacuum will not give back.
|
||||
|
||||
What is banned is the row-rewriting DML behind that, not everything whose cost
|
||||
scales that way. A non-concurrent `CREATE INDEX`, an `ALTER COLUMN ... TYPE` that is
|
||||
not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE ... AS
|
||||
SELECT` or `SELECT ... INTO` filling a new table from an existing one, the rename
|
||||
that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED VIEW`
|
||||
all read the whole table and all pass. That is deliberate: a rule wide enough 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.
|
||||
scales that way. A non-concurrent `CREATE INDEX` passes except on a request-log
|
||||
table, where it blocks writes until the build finishes. An `ALTER COLUMN ... TYPE`
|
||||
that is not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE
|
||||
... AS SELECT` or `SELECT ... INTO` filling a new table from an existing one, the
|
||||
rename that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED
|
||||
VIEW` all read the whole table and all pass. That is deliberate: a rule wide enough
|
||||
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.
|
||||
|
||||
The one schema change banned outright is `ADD COLUMN ... DEFAULT` on a table in
|
||||
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
|
||||
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,
|
||||
|
|
@ -23,6 +24,13 @@ which on a spend-log-sized table is the same outage as a backfill. Every other t
|
|||
is small enough that the rewrite is not worth a rule, and a column added to a log
|
||||
table without a default is still free on every version.
|
||||
|
||||
An index on a request-log table cannot ship as a migration at all. A plain `CREATE
|
||||
INDEX` blocks writes to the table until the build finishes, and `CREATE INDEX
|
||||
CONCURRENTLY` is refused by Postgres on a partitioned parent, which LiteLLM_SpendLogs
|
||||
is wherever the operator ran db_scripts/partition_spend_logs.sql. The migration job builds
|
||||
those indexes after `migrate deploy`, concurrently and per partition, from the list in
|
||||
litellm_proxy_extras/request_log_indexes.py, so that list is where a new one goes.
|
||||
|
||||
Flagged, per statement, by its leading keyword:
|
||||
|
||||
UPDATE rewrites every matching row, and `WHERE` does not bound the scan
|
||||
|
|
@ -44,6 +52,8 @@ Flagged, per statement, by its leading keyword:
|
|||
actions adds a column with a `DEFAULT`. An `ALTER COLUMN ... SET
|
||||
DEFAULT` written after the column exists changes metadata alone, so it
|
||||
passes, as does an `ADD CONSTRAINT`
|
||||
CREATE only a `CREATE [UNIQUE] INDEX` on a request-log table, concurrent or
|
||||
not; the migration job builds those
|
||||
|
||||
Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a
|
||||
statement's leading keyword, so they pass.
|
||||
|
|
@ -85,7 +95,9 @@ below line up with the statements they exempt.
|
|||
Add a column and let the application populate it, or run the rewrite as an opt-in
|
||||
batched job outside boot. When a rewrite is genuinely bounded and must ship inside
|
||||
the migration, put `-- data-migration-ok: <reason>` on the statement or on the line
|
||||
above it, naming what bounds it. The reason is required. A marker sharing a line
|
||||
above it, naming what bounds it. The reason is required. A marker never exempts a
|
||||
`CREATE INDEX` on a request-log table, since no bound makes that statement safe:
|
||||
the migration job is the only place such an index is built. A marker sharing a line
|
||||
with the statement it follows exempts that statement alone, so the next statement
|
||||
down is still checked rather than picking the marker up as its own. A marker on an
|
||||
`EXECUTE` or on the assignment feeding one covers the single-quoted SQL that
|
||||
|
|
@ -108,6 +120,7 @@ import sys
|
|||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
|
||||
|
|
@ -115,6 +128,9 @@ 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",
|
||||
"20260818224500_add_shadow_eval_stopped_by",
|
||||
|
|
@ -139,9 +155,7 @@ WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?<![<>!:=])=(?![=>])")
|
|||
PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$")
|
||||
QUALIFIER_GAP = re.compile(r"[\s.]*")
|
||||
EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE)
|
||||
DEFINES_A_ROUTINE = re.compile(
|
||||
r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE
|
||||
)
|
||||
DEFINES_A_ROUTINE = re.compile(r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE)
|
||||
QUALIFIED_NAME = r"(?:\"[^\"]*\"|[A-Za-z_][A-Za-z0-9_$]*)"
|
||||
ROUTINE_NAME = re.compile(rf"\s*(?:{QUALIFIED_NAME}\s*\.\s*)?({QUALIFIED_NAME})")
|
||||
TABLE_NAME = ROUTINE_NAME
|
||||
|
|
@ -204,6 +218,12 @@ statement with the bound spelled out:
|
|||
-- data-migration-ok: <what bounds this>
|
||||
UPDATE ...
|
||||
|
||||
An index on a request-log table is not a migration, and no marker exempts one. Declare it with `@@index` in
|
||||
schema.prisma and add it to REQUEST_LOG_INDEXES in
|
||||
litellm_proxy_extras/request_log_indexes.py under the name Prisma derives for it; the
|
||||
migration job builds it after `migrate deploy`, concurrently on a plain table and per
|
||||
partition on a partitioned one, which no single migration statement can do.
|
||||
|
||||
On Postgres 10 an `ADD COLUMN ... DEFAULT` on a request-log table rewrites the table
|
||||
too. Add the column nullable with no default, then set the default in a separate
|
||||
`ALTER COLUMN ... SET DEFAULT`, which never touches existing rows.
|
||||
|
|
@ -215,10 +235,11 @@ class Violation:
|
|||
migration: str
|
||||
line: int
|
||||
keyword: str
|
||||
consequence: str = "rewrites existing rows at boot"
|
||||
|
||||
def render(self) -> str:
|
||||
location = f"{MIGRATIONS_DIR.relative_to(REPO_ROOT)}/{self.migration}/migration.sql"
|
||||
return f"{location}:{self.line}: {self.keyword} rewrites existing rows at boot"
|
||||
return f"{location}:{self.line}: {self.keyword} {self.consequence}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -578,6 +599,35 @@ def rewrites_a_log_table(clause: str, region: str, base: int) -> str | None:
|
|||
return f"ADD COLUMN ... DEFAULT on {named.group(1)}"
|
||||
|
||||
|
||||
def builds_a_log_index(clause: str, region: str, base: int) -> str | None:
|
||||
"""The keyword to report when a `CREATE INDEX` targets a request-log table, concurrent or
|
||||
not: a plain build blocks writes for its whole duration, and a concurrent one fails with
|
||||
P3018 on a partitioned parent, so the migration job builds those instead."""
|
||||
created: Final[re.Match[str] | None] = re.match(
|
||||
r"\s*CREATE\s+(?:UNIQUE\s+)?INDEX\b(?:\s+CONCURRENTLY\b)?", clause, re.IGNORECASE
|
||||
)
|
||||
if created is None:
|
||||
return None
|
||||
on: Final[re.Match[str] | None] = re.search(r"\bON\b(?:\s+ONLY\b)?", clause[created.end() :], re.IGNORECASE)
|
||||
if on is None:
|
||||
return None
|
||||
named: Final[re.Match[str] | None] = TABLE_NAME.match(
|
||||
region, skip_comments(region, base + created.end() + on.end())
|
||||
)
|
||||
if named is None or named.group(1).strip('"') not in REQUEST_LOG_TABLES:
|
||||
return None
|
||||
return f"CREATE INDEX on {named.group(1)}"
|
||||
|
||||
|
||||
def consequence_of(found: str) -> str:
|
||||
if found.startswith("CREATE INDEX"):
|
||||
return (
|
||||
"blocks writes until the build finishes, or fails on a partitioned table; "
|
||||
"add it to REQUEST_LOG_INDEXES in litellm_proxy_extras/request_log_indexes.py instead"
|
||||
)
|
||||
return "rewrites existing rows at boot"
|
||||
|
||||
|
||||
def skip_comments(sql: str, start: int) -> int:
|
||||
index = start
|
||||
while index < len(sql):
|
||||
|
|
@ -746,8 +796,7 @@ def read_markers(sql: str) -> Markers:
|
|||
return Markers(
|
||||
sql,
|
||||
tuple(
|
||||
Marker(match.start(), match.end(), alone_on_its_line(sql, match.start()))
|
||||
for match in MARKER.finditer(sql)
|
||||
Marker(match.start(), match.end(), alone_on_its_line(sql, match.start())) for match in MARKER.finditer(sql)
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -760,9 +809,7 @@ def scan(sql: str, migration: str, markers: Markers) -> Iterator[Violation]:
|
|||
yield from scan_region(sql, sql, migration, markers, 0)
|
||||
|
||||
|
||||
def scan_region(
|
||||
document: str, region: str, migration: str, markers: Markers, offset: int
|
||||
) -> Iterator[Violation]:
|
||||
def scan_region(document: str, region: str, migration: str, markers: Markers, offset: int) -> Iterator[Violation]:
|
||||
"""Violations in one region of `document`, whose text begins at `offset`. Positions are
|
||||
always counted against the whole document, so a statement nested in a dollar-quoted body
|
||||
reports its real file line and lines up with the markers read from that file. A single-quoted
|
||||
|
|
@ -790,13 +837,18 @@ def scan_region(
|
|||
offset + start,
|
||||
)
|
||||
|
||||
keyword = offending_keyword(clause)
|
||||
if exempt:
|
||||
continue
|
||||
found = keyword or rewrites_a_log_table(clause, region, base)
|
||||
index = builds_a_log_index(clause, region, base)
|
||||
found = (
|
||||
index if exempt else offending_keyword(clause) or rewrites_a_log_table(clause, region, base) or index
|
||||
)
|
||||
if found is None:
|
||||
continue
|
||||
yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), found)
|
||||
yield Violation(
|
||||
migration,
|
||||
line_of(document, offset + keyword_start(clause, base)),
|
||||
found,
|
||||
consequence_of(found),
|
||||
)
|
||||
|
||||
for body in bodies:
|
||||
if not runs_when_applied(masked, region, bodies, runnable, identifiers, body):
|
||||
|
|
|
|||
|
|
@ -16,6 +16,10 @@ import httpx
|
|||
import psutil
|
||||
from integration._support.client import GATEWAY_LIMITS, Gateway
|
||||
|
||||
DB_PUSH: Final = ("--use_prisma_db_push",)
|
||||
MIGRATE_DEPLOY: Final = ()
|
||||
LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",)
|
||||
|
||||
|
||||
def proxy_database_environment() -> Mapping[str, str]:
|
||||
writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "")
|
||||
|
|
@ -73,9 +77,16 @@ def owned_proxy(
|
|||
config: Path | None = None,
|
||||
remove_environment: tuple[str, ...] = (),
|
||||
workers: int = 1,
|
||||
database_setup: tuple[str, ...] = DB_PUSH,
|
||||
) -> Iterator[Gateway]:
|
||||
with owned_proxy_process(
|
||||
gateway, directory, overrides, config=config, remove_environment=remove_environment, workers=workers
|
||||
gateway,
|
||||
directory,
|
||||
overrides,
|
||||
config=config,
|
||||
remove_environment=remove_environment,
|
||||
workers=workers,
|
||||
database_setup=database_setup,
|
||||
) as owned:
|
||||
yield owned.gateway
|
||||
|
||||
|
|
@ -171,6 +182,7 @@ def owned_proxy_process(
|
|||
config: Path | None = None,
|
||||
remove_environment: tuple[str, ...] = (),
|
||||
workers: int = 1,
|
||||
database_setup: tuple[str, ...] = DB_PUSH,
|
||||
) -> Iterator[OwnedProxy]:
|
||||
root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
|
||||
environment: Final = {
|
||||
|
|
@ -196,7 +208,7 @@ def owned_proxy_process(
|
|||
"127.0.0.1",
|
||||
"--num_workers",
|
||||
str(workers),
|
||||
"--use_prisma_db_push",
|
||||
*database_setup,
|
||||
"--enforce_prisma_migration_check",
|
||||
)
|
||||
launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS)
|
||||
|
|
|
|||
364
tests/integration/database/test_request_log_indexes_at_boot.py
Normal file
364
tests/integration/database/test_request_log_indexes_at_boot.py
Normal file
|
|
@ -0,0 +1,364 @@
|
|||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import scratch_database
|
||||
from integration._support.process import LEGACY_MIGRATE_DEPLOY, MIGRATE_DEPLOY, owned_proxy_process
|
||||
from psycopg import sql
|
||||
from psycopg.rows import class_row
|
||||
|
||||
REPO_ROOT: Final = Path(__file__).resolve().parents[3]
|
||||
PRISMA_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras"
|
||||
PARTITION_SCRIPT: Final = REPO_ROOT / "db_scripts" / "partition_spend_logs.sql"
|
||||
SHIPPED_MIGRATIONS: Final = tuple(sorted(path.name for path in (PRISMA_DIR / "migrations").iterdir() if path.is_dir()))
|
||||
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"
|
||||
ORIGINAL_MIGRATION_SQL: Final = MappingProxyType(
|
||||
{
|
||||
API_KEY_INDEX_MIGRATION: (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
|
||||
),
|
||||
CALL_ID_INDEX_MIGRATION: (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
|
||||
),
|
||||
}
|
||||
)
|
||||
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")
|
||||
POPULATED_PARTITIONS: Final = MappingProxyType(
|
||||
{
|
||||
"LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"),
|
||||
"LiteLLM_SpendLogs_p2026_09": ("2026-09-01", "2026-10-01"),
|
||||
}
|
||||
)
|
||||
DEFAULT_PARTITION: Final = "LiteLLM_SpendLogs_pdefault"
|
||||
ROWS_PER_PARTITION: Final = 500
|
||||
PARTITIONED_PARENT_ERROR: Final = 'cannot create index on partitioned table "LiteLLM_SpendLogs" concurrently'
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Resolver:
|
||||
"""One migration resolver as the serving proxy selects it (CLI flags) and as the
|
||||
migration job selects it (environment)."""
|
||||
|
||||
proxy_flags: tuple[str, ...]
|
||||
job_environment: Mapping[str, str]
|
||||
|
||||
|
||||
V2: Final = Resolver(MIGRATE_DEPLOY, MappingProxyType({"USE_V2_MIGRATION_RESOLVER": "true"}))
|
||||
LEGACY: Final = Resolver(LEGACY_MIGRATE_DEPLOY, MappingProxyType({"USE_V2_MIGRATION_RESOLVER": "false"}))
|
||||
RESOLVERS: Final = pytest.mark.parametrize("resolver", (V2, LEGACY), ids=("v2", "legacy"))
|
||||
|
||||
|
||||
def release_layout(directory: Path, migrations: tuple[str, ...]) -> Path:
|
||||
"""The Prisma layout of the release that shipped `migrations`: the two index migrations
|
||||
carry the SQL they shipped with, not the inert files of this build."""
|
||||
(directory / "migrations").mkdir(parents=True)
|
||||
shutil.copy(PRISMA_DIR / "schema.prisma", directory / "schema.prisma")
|
||||
shutil.copy(PRISMA_DIR / "migrations" / "migration_lock.toml", directory / "migrations" / "migration_lock.toml")
|
||||
for name in migrations:
|
||||
shutil.copytree(PRISMA_DIR / "migrations" / name, directory / "migrations" / name)
|
||||
for name, original in ORIGINAL_MIGRATION_SQL.items():
|
||||
if name in migrations:
|
||||
(directory / "migrations" / name / "migration.sql").write_text(original)
|
||||
return directory / "schema.prisma"
|
||||
|
||||
|
||||
def migrate_deploy(database_url: str, schema: Path) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
[sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(schema)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300,
|
||||
env={**os.environ, "DATABASE_URL": database_url},
|
||||
)
|
||||
|
||||
|
||||
def migration_job(database_url: str, resolver: Resolver) -> subprocess.CompletedProcess[str]:
|
||||
"""The migrations image entrypoint, as the helm migration Job runs it."""
|
||||
return subprocess.run(
|
||||
[sys.executable, "-I", str(REPO_ROOT / "migrations" / "run.py")],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=MIGRATION_JOB_SECONDS,
|
||||
cwd=REPO_ROOT,
|
||||
env={**os.environ, "DATABASE_URL": database_url, **resolver.job_environment},
|
||||
)
|
||||
|
||||
|
||||
def migration_cli(database_url: str, gateway: Gateway, resolver: Resolver) -> subprocess.CompletedProcess[str]:
|
||||
"""The proxy CLI as a migration job: `--skip_server_startup` migrates, builds the
|
||||
indexes and exits by their result."""
|
||||
return subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-P",
|
||||
"-m",
|
||||
"integration._support.proxy",
|
||||
"--config",
|
||||
"tests/integration/proxy_config.yaml",
|
||||
*resolver.proxy_flags,
|
||||
"--skip_server_startup",
|
||||
"--enforce_prisma_migration_check",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=MIGRATION_JOB_SECONDS,
|
||||
cwd=REPO_ROOT,
|
||||
env={**os.environ, "DATABASE_URL": database_url, "LITELLM_MASTER_KEY": gateway.key},
|
||||
)
|
||||
|
||||
|
||||
def deploy_schema_before_the_index_migrations(database_url: str, directory: Path) -> None:
|
||||
older: Final = tuple(name for name in SHIPPED_MIGRATIONS if name < API_KEY_INDEX_MIGRATION)
|
||||
deployed: Final = migrate_deploy(database_url, release_layout(directory / "older-release", older))
|
||||
assert deployed.returncode == 0, deployed.stdout + deployed.stderr
|
||||
|
||||
|
||||
def deploy_the_original_index_migrations(database_url: str, directory: Path) -> subprocess.CompletedProcess[str]:
|
||||
"""Boot the v1.103.0 layout once: on a partitioned table its CONCURRENTLY call_id index
|
||||
fails with P3018 and leaves the ledger row unfinished, on a plain table both apply."""
|
||||
return migrate_deploy(database_url, release_layout(directory / "v1.103.0", SHIPPED_MIGRATIONS))
|
||||
|
||||
|
||||
def fail_the_call_id_index_migration_like_the_shipped_release(database_url: str, directory: Path) -> None:
|
||||
deployed: Final = deploy_the_original_index_migrations(database_url, directory)
|
||||
assert deployed.returncode != 0, deployed.stdout
|
||||
assert "P3018" in deployed.stderr and PARTITIONED_PARENT_ERROR in deployed.stderr, deployed.stderr
|
||||
assert ledger(database_url)[CALL_ID_INDEX_MIGRATION] is False
|
||||
|
||||
|
||||
def partition_spend_logs(database_url: str) -> None:
|
||||
with psycopg.connect(database_url, autocommit=True) as connection:
|
||||
connection.execute(PARTITION_SCRIPT.read_bytes())
|
||||
for partition, (start, stop) in POPULATED_PARTITIONS.items():
|
||||
add_partition(connection, partition, start, stop)
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_SpendLogs" ("request_id", "call_type", "api_key", "startTime", "endTime") '
|
||||
"SELECT %s || n, 'acompletion', 'sk-' || (n %% 7), %s::timestamp + (n * interval '1 minute'), "
|
||||
"%s::timestamp + (n * interval '1 minute') + interval '1 second' FROM generate_series(1, %s) AS n",
|
||||
(partition, start, start, ROWS_PER_PARTITION),
|
||||
)
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_SpendLogs" ("request_id", "call_type", "startTime", "endTime") '
|
||||
"SELECT 'default-' || n, 'acompletion', '2026-07-01'::timestamp + (n * interval '1 minute'), "
|
||||
"'2026-07-01'::timestamp + (n * interval '1 minute') FROM generate_series(1, %s) AS n",
|
||||
(ROWS_PER_PARTITION,),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LedgerRow:
|
||||
name: str
|
||||
finished: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _IndexRow:
|
||||
index: str
|
||||
valid: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OidRow:
|
||||
index: str
|
||||
oid: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AttachedRow:
|
||||
partition: str
|
||||
parent_index: str
|
||||
valid: bool
|
||||
|
||||
|
||||
def ledger(database_url: str) -> Mapping[str, bool]:
|
||||
"""Every migration in the ledger that was not rolled back, mapped to whether it finished."""
|
||||
with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_LedgerRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
'SELECT migration_name AS name, finished_at IS NOT NULL AS finished FROM "_prisma_migrations" '
|
||||
"WHERE rolled_back_at IS NULL ORDER BY migration_name"
|
||||
).fetchall()
|
||||
return MappingProxyType({row.name: row.finished for row in rows})
|
||||
|
||||
|
||||
def parent_index_validity(database_url: str) -> Mapping[str, bool]:
|
||||
with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_IndexRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
"SELECT c.relname AS index, x.indisvalid AS valid FROM pg_index x JOIN pg_class c ON c.oid = x.indexrelid "
|
||||
"WHERE x.indrelid = '\"LiteLLM_SpendLogs\"'::regclass AND c.relname = ANY(%s)",
|
||||
(list(SPEND_LOGS_INDEXES),),
|
||||
).fetchall()
|
||||
return MappingProxyType({row.index: row.valid for row in rows})
|
||||
|
||||
|
||||
def index_oids(database_url: str) -> Mapping[str, int]:
|
||||
"""index name -> oid for the SpendLogs indexes on the parent or any partition; a rebuild changes the oid."""
|
||||
with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_OidRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
"SELECT c.relname AS index, c.oid::int AS oid FROM pg_index x JOIN pg_class c ON c.oid = x.indexrelid "
|
||||
"WHERE x.indrelid = '\"LiteLLM_SpendLogs\"'::regclass OR x.indrelid IN "
|
||||
"(SELECT inhrelid FROM pg_inherits WHERE inhparent = '\"LiteLLM_SpendLogs\"'::regclass)"
|
||||
).fetchall()
|
||||
return MappingProxyType({row.index: row.oid for row in rows})
|
||||
|
||||
|
||||
def attached_partition_indexes(database_url: str) -> frozenset[tuple[str, str, bool]]:
|
||||
"""(partition, parent index, child is valid) for every child index attached under a SpendLogs parent index."""
|
||||
with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_AttachedRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
"SELECT part.relname AS partition, parent_index.relname AS parent_index, child.indisvalid AS valid "
|
||||
"FROM pg_inherits attached "
|
||||
"JOIN pg_class parent_index ON parent_index.oid = attached.inhparent "
|
||||
"JOIN pg_index child ON child.indexrelid = attached.inhrelid "
|
||||
"JOIN pg_class part ON part.oid = child.indrelid "
|
||||
"WHERE parent_index.relname = ANY(%s)",
|
||||
(list(SPEND_LOGS_INDEXES),),
|
||||
).fetchall()
|
||||
return frozenset((row.partition, row.parent_index, row.valid) for row in rows)
|
||||
|
||||
|
||||
def expected_attachments(partitions: tuple[str, ...]) -> frozenset[tuple[str, str, bool]]:
|
||||
return frozenset((partition, index, True) for partition, index in product(partitions, SPEND_LOGS_INDEXES))
|
||||
|
||||
|
||||
def add_partition(connection: psycopg.Connection[tuple[object, ...]], partition: str, start: str, stop: str) -> None:
|
||||
connection.execute(
|
||||
sql.SQL('CREATE TABLE {} PARTITION OF "LiteLLM_SpendLogs" FOR VALUES FROM ({}) TO ({})').format(
|
||||
sql.Identifier(partition), sql.Literal(start), sql.Literal(stop)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def assert_ready(booted_gateway: Gateway) -> None:
|
||||
readiness: Final = booted_gateway.request("GET", "/health/readiness")
|
||||
assert readiness.status_code == 200, readiness.text
|
||||
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."""
|
||||
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)
|
||||
assert parent_index_validity(database_url) == {index: True for index in SPEND_LOGS_INDEXES}
|
||||
with psycopg.connect(database_url, autocommit=True) as connection:
|
||||
add_partition(connection, "LiteLLM_SpendLogs_p2026_10", "2026-10-01", "2026-11-01")
|
||||
assert attached_partition_indexes(database_url) == expected_attachments((*populated, "LiteLLM_SpendLogs_p2026_10"))
|
||||
|
||||
|
||||
def assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(
|
||||
gateway: Gateway, directory: Path, database_url: str, resolver: Resolver
|
||||
) -> None:
|
||||
"""The serving proxy applies the inert files, reports ready, and its background build
|
||||
finds every index already there, so it builds nothing and the catalog is untouched."""
|
||||
oids: Final = index_oids(database_url)
|
||||
with owned_proxy_process(
|
||||
gateway, directory, {"DATABASE_URL": database_url}, database_setup=resolver.proxy_flags
|
||||
) as booted:
|
||||
assert_ready(booted.gateway)
|
||||
log: Final = eventually(
|
||||
lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=BUILD_SECONDS
|
||||
)
|
||||
assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
|
||||
assert not any(line in log for line in INDEX_BUILD_LINES), log[-4000:]
|
||||
assert index_oids(database_url) == oids
|
||||
|
||||
|
||||
@RESOLVERS
|
||||
def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_both_indexes_per_partition(
|
||||
gateway: Gateway, tmp_path: Path, resolver: Resolver
|
||||
) -> 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_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_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
|
||||
|
||||
|
||||
@RESOLVERS
|
||||
def test_the_migration_job_heals_a_partitioned_table_left_with_the_failed_call_id_ledger_row(
|
||||
gateway: Gateway, tmp_path: Path, resolver: Resolver
|
||||
) -> None:
|
||||
with scratch_database() as database_url:
|
||||
deploy_schema_before_the_index_migrations(database_url, tmp_path)
|
||||
partition_spend_logs(database_url)
|
||||
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_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
|
||||
|
||||
|
||||
@RESOLVERS
|
||||
def test_the_migration_job_leaves_a_plain_table_that_applied_the_original_index_migrations_alone(
|
||||
gateway: Gateway, tmp_path: Path, resolver: Resolver
|
||||
) -> None:
|
||||
with scratch_database() as database_url:
|
||||
deploy_schema_before_the_index_migrations(database_url, tmp_path)
|
||||
deployed: Final = deploy_the_original_index_migrations(database_url, tmp_path)
|
||||
assert deployed.returncode == 0, deployed.stdout + deployed.stderr
|
||||
before: Final = index_oids(database_url)
|
||||
assert set(SPEND_LOGS_INDEXES) <= set(before), before
|
||||
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 "Building index" not in job.stderr + job.stdout, job.stdout + job.stderr
|
||||
assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
|
||||
assert index_oids(database_url) == before
|
||||
assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
|
||||
|
||||
|
||||
@RESOLVERS
|
||||
def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_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."""
|
||||
with scratch_database() as database_url:
|
||||
deploy_schema_before_the_index_migrations(database_url, tmp_path)
|
||||
partition_spend_logs(database_url)
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in parent_index_validity(database_url)
|
||||
with owned_proxy_process(
|
||||
gateway, tmp_path, {"DATABASE_URL": database_url}, database_setup=resolver.proxy_flags
|
||||
) as booted:
|
||||
assert_ready(booted.gateway)
|
||||
assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
|
||||
log: Final = eventually(
|
||||
lambda: booted.log.read_text(errors="replace"),
|
||||
lambda text: INDEXES_IN_PLACE in text,
|
||||
seconds=BUILD_SECONDS,
|
||||
)
|
||||
assert "Building index" in log and "Attached index" in log, log[-4000:]
|
||||
assert_both_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:
|
||||
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)
|
||||
|
|
@ -1,12 +1,11 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import litellm_proxy_extras
|
||||
import psycopg
|
||||
import pytest
|
||||
from litellm_proxy_extras.request_log_indexes import REQUEST_LOG_INDEXES
|
||||
from psycopg.types.json import Jsonb
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
|
@ -29,11 +28,8 @@ _SPEND_LOGS_DDL: Final = """
|
|||
)
|
||||
"""
|
||||
|
||||
_API_KEY_START_TIME_INDEX_MIGRATION: Final = (
|
||||
Path(litellm_proxy_extras.__file__).parent
|
||||
/ "migrations"
|
||||
/ "20260823000000_add_spend_logs_api_key_starttime_index"
|
||||
/ "migration.sql"
|
||||
_API_KEY_START_TIME_INDEX: Final = next(
|
||||
index for index in REQUEST_LOG_INDEXES if index.name == "LiteLLM_SpendLogs_api_key_startTime_idx"
|
||||
)
|
||||
|
||||
_STATS_SQL: Final = """
|
||||
|
|
@ -56,7 +52,12 @@ class _Settle:
|
|||
|
||||
def _create_spend_logs_table(database_url: str) -> None:
|
||||
write_rows(_SPEND_LOGS_DDL, (), database_url=database_url)
|
||||
write_rows(_API_KEY_START_TIME_INDEX_MIGRATION.read_text(), (), database_url=database_url)
|
||||
write_rows(
|
||||
f'CREATE INDEX "{_API_KEY_START_TIME_INDEX.name}" ON "{_API_KEY_START_TIME_INDEX.table}" ' # pyright: ignore[reportArgumentType] # DDL from the migration job index list
|
||||
f"{_API_KEY_START_TIME_INDEX.definition}",
|
||||
(),
|
||||
database_url=database_url,
|
||||
)
|
||||
|
||||
|
||||
def _spend_log_stats(database_url: str) -> dict[str, int]:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import tempfile
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from litellm_proxy_extras.request_log_indexes import filter_request_log_index_diff
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
|
|
@ -16,7 +17,9 @@ def test_schema_migration_in_sync():
|
|||
|
||||
Applies every committed migration to an empty database, then diffs the result
|
||||
against schema.prisma. A non-empty diff means the schema was changed without a
|
||||
matching migration being generated.
|
||||
matching migration being generated. The request-log indexes the migration job
|
||||
builds are declared in the schema and deliberately absent from the migrations,
|
||||
so those statements are filtered out before the diff is judged.
|
||||
"""
|
||||
db_url = os.environ["DATABASE_URL"]
|
||||
source_migrations_dir = Path(
|
||||
|
|
@ -60,11 +63,14 @@ def test_schema_migration_in_sync():
|
|||
)
|
||||
|
||||
if diff.returncode == 2:
|
||||
pytest.fail(
|
||||
"Schema changes detected that no migration captures. Run "
|
||||
"`python litellm/ci_cd/run_migration.py <migration_name>`.\n\n"
|
||||
+ diff.stdout
|
||||
)
|
||||
assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}"
|
||||
drift = filter_request_log_index_diff(diff.stdout)
|
||||
if drift.strip():
|
||||
pytest.fail(
|
||||
"Schema changes detected that no migration captures. Run "
|
||||
"`python litellm/ci_cd/run_migration.py <migration_name>`.\n\n"
|
||||
+ drift
|
||||
)
|
||||
else:
|
||||
assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}"
|
||||
finally:
|
||||
shutil.rmtree(temp_base, ignore_errors=True)
|
||||
|
|
|
|||
|
|
@ -1,12 +1,14 @@
|
|||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager
|
||||
from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY
|
||||
from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager, _InvalidIndex
|
||||
|
||||
psycopg = pytest.importorskip("psycopg")
|
||||
|
||||
|
|
@ -20,6 +22,7 @@ requires_db: Final = pytest.mark.skipif(
|
|||
HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable"
|
||||
HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx"
|
||||
HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC'
|
||||
SECOND_HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_name_idx"
|
||||
LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable"
|
||||
LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx"
|
||||
PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable"
|
||||
|
|
@ -167,7 +170,74 @@ def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: st
|
|||
|
||||
|
||||
@requires_db
|
||||
def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None:
|
||||
def test_repair_yields_to_the_migration_job_building_indexes_under_the_migration_lock(scratch_schema: str) -> None:
|
||||
"""A migration job's index build holds the migration lock while its CREATE INDEX CONCURRENTLY
|
||||
is cataloged as invalid; the repair must not rebuild that in-flight index."""
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
|
||||
with psycopg.connect(_base_url(), autocommit=True) as index_builder:
|
||||
index_builder.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False}
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
||||
|
||||
|
||||
def _hold_migration_lock_once_free(release: threading.Event) -> None:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as resolver:
|
||||
resolver.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
|
||||
release.wait(timeout=60)
|
||||
|
||||
|
||||
def _wait_until_a_session_queues_for_the_migration_lock() -> None:
|
||||
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
||||
for _ in range(200):
|
||||
queued: Final = conn.execute(
|
||||
"SELECT count(*) FROM pg_locks WHERE locktype = 'advisory' AND NOT granted "
|
||||
"AND classid = %s AND objid = %s",
|
||||
(MIGRATION_LOCK_KEY >> 32, MIGRATION_LOCK_KEY & 0xFFFFFFFF),
|
||||
).fetchone()
|
||||
if queued is not None and queued[0]:
|
||||
return
|
||||
time.sleep(0.05)
|
||||
pytest.fail("no session queued for the migration lock")
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_releases_the_migration_lock_between_indexes_so_a_booting_resolver_gets_in(
|
||||
scratch_schema: str,
|
||||
) -> None:
|
||||
"""A v2 resolver on another replica waits for the migration lock; with two invalid
|
||||
indexes to rebuild it must get the lock after the first REINDEX, not after both."""
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, SECOND_HEALTH_INDEX, '"model_name"')
|
||||
release: Final = threading.Event()
|
||||
resolver: Final = threading.Thread(target=_hold_migration_lock_once_free, args=(release,))
|
||||
|
||||
def repair_then_let_a_resolver_queue_for_the_lock(
|
||||
conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex
|
||||
) -> None:
|
||||
ProxyExtrasDBManager._repair_index(conn, index)
|
||||
if not resolver.is_alive():
|
||||
resolver.start()
|
||||
_wait_until_a_session_queues_for_the_migration_lock()
|
||||
|
||||
try:
|
||||
assert (
|
||||
ProxyExtrasDBManager.repair_invalid_indexes(repair=repair_then_let_a_resolver_queue_for_the_lock) is False
|
||||
)
|
||||
assert sorted(_index_validity(scratch_schema).values()) == [False, True]
|
||||
finally:
|
||||
release.set()
|
||||
resolver.join()
|
||||
|
||||
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
||||
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, SECOND_HEALTH_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_boot(scratch_schema: str) -> None:
|
||||
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
||||
|
||||
with psycopg.connect(_base_url()) as pin:
|
||||
|
|
|
|||
763
tests/proxy_migration_tests/test_request_log_indexes.py
Normal file
763
tests/proxy_migration_tests/test_request_log_indexes.py
Normal file
|
|
@ -0,0 +1,763 @@
|
|||
import os
|
||||
import queue
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
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.request_log_indexes import (
|
||||
REQUEST_LOG_INDEXES,
|
||||
RequestLogIndex,
|
||||
build_index_on_partitioned_table,
|
||||
ensure_request_log_indexes,
|
||||
)
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
from psycopg import sql
|
||||
from psycopg.abc import Params, QueryNoTemplate
|
||||
from psycopg.rows import class_row
|
||||
|
||||
pytestmark = pytest.mark.timeout(900)
|
||||
|
||||
requires_db: Final = pytest.mark.skipif(
|
||||
"DATABASE_URL" not in os.environ,
|
||||
reason="requires a postgres database (DATABASE_URL)",
|
||||
)
|
||||
|
||||
REPO: Final = Path(__file__).resolve().parents[2]
|
||||
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"
|
||||
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(
|
||||
{
|
||||
API_KEY_INDEX_MIGRATION: (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
|
||||
),
|
||||
CALL_ID_INDEX_MIGRATION: (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
|
||||
),
|
||||
}
|
||||
)
|
||||
CALL_ID_INDEX_DEFINITION: Final = next(index for index in REQUEST_LOG_INDEXES if index.name == CALL_ID_INDEX)
|
||||
RELEASES: Final = pytest.mark.parametrize(
|
||||
"release", (API_KEY_INDEX_MIGRATION, CALL_ID_INDEX_MIGRATION), ids=("v1.102.1", "v1.103.0")
|
||||
)
|
||||
PARTITIONS: Final = MappingProxyType(
|
||||
{
|
||||
"LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"),
|
||||
"LiteLLM_SpendLogs_p2026_09": ("2026-09-01", "2026-10-01"),
|
||||
}
|
||||
)
|
||||
DEFAULT_PARTITION: Final = "LiteLLM_SpendLogs_pdefault"
|
||||
ROWS_PER_PARTITION: Final = 200
|
||||
RESOLVERS: Final = pytest.mark.parametrize("use_v2_resolver", (True, False), ids=("v2", "v1"))
|
||||
|
||||
|
||||
def _base_url() -> str:
|
||||
return os.environ["DATABASE_URL"].split("?")[0]
|
||||
|
||||
|
||||
def _migrate_deploy(database_url: str, schema: Path) -> "subprocess.CompletedProcess[str]":
|
||||
return subprocess.run(
|
||||
[sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(schema)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={**os.environ, "DATABASE_URL": database_url},
|
||||
)
|
||||
|
||||
|
||||
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."""
|
||||
(prisma_dir / "migrations").mkdir(parents=True)
|
||||
shutil.copy(PACKAGE / "schema.prisma", prisma_dir / "schema.prisma")
|
||||
for migration in sorted((PACKAGE / "migrations").iterdir()):
|
||||
if migration.is_dir() and migration.name < before:
|
||||
shutil.copytree(migration, prisma_dir / "migrations" / migration.name)
|
||||
for name, original in ORIGINAL_MIGRATION_SQL.items():
|
||||
if (prisma_dir / "migrations" / name).is_dir():
|
||||
(prisma_dir / "migrations" / name / "migration.sql").write_text(original)
|
||||
return prisma_dir / "schema.prisma"
|
||||
|
||||
|
||||
def _deploy_release(database_url: str, prisma_dir: Path, before: str) -> None:
|
||||
deployed: Final = _migrate_deploy(database_url, _release_layout(prisma_dir, before))
|
||||
assert deployed.returncode == 0, deployed.stderr
|
||||
|
||||
|
||||
def _insert_spend_log(
|
||||
conn: "psycopg.Connection[tuple[object, ...]]", request_id: str, day: str, table: str = "LiteLLM_SpendLogs"
|
||||
) -> None:
|
||||
conn.execute(
|
||||
sql.SQL(
|
||||
'INSERT INTO {} ("request_id", "call_type", "startTime", "endTime", "api_key") VALUES (%s, %s, %s, %s, %s)'
|
||||
).format(sql.Identifier(table)),
|
||||
(request_id, "acompletion", day, day, f"key-{request_id[-1]}"),
|
||||
)
|
||||
|
||||
|
||||
def _partition_spend_logs(database_url: str) -> None:
|
||||
with psycopg.connect(database_url, autocommit=True) as conn:
|
||||
conn.execute(PARTITION_SCRIPT.read_bytes())
|
||||
for partition, (start, stop) in PARTITIONS.items():
|
||||
conn.execute(
|
||||
sql.SQL('CREATE TABLE {} PARTITION OF "LiteLLM_SpendLogs" FOR VALUES FROM ({}) TO ({})').format(
|
||||
sql.Identifier(partition), sql.Literal(start), sql.Literal(stop)
|
||||
)
|
||||
)
|
||||
for row in range(ROWS_PER_PARTITION):
|
||||
_insert_spend_log(conn, f"{partition}-{row}", start)
|
||||
for row in range(ROWS_PER_PARTITION):
|
||||
_insert_spend_log(conn, f"default-{row}", "2020-01-01")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def release() -> str:
|
||||
"""The first migration a database has not applied yet; the v1.103.0 shape unless a test parametrizes it."""
|
||||
return CALL_ID_INDEX_MIGRATION
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scratch_database(release: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[str]:
|
||||
"""A deployment stopped before `release`, with DATABASE_URL pointed at it so
|
||||
ProxyExtrasDBManager upgrades it like a booting proxy."""
|
||||
admin_url: Final = _base_url()
|
||||
name: Final = f"spend_logs_index_{uuid.uuid4().hex[:8]}"
|
||||
with psycopg.connect(admin_url, autocommit=True) as conn:
|
||||
conn.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
|
||||
database_url: Final = f"{admin_url.rsplit('/', 1)[0]}/{name}"
|
||||
try:
|
||||
_deploy_release(database_url, tmp_path / "prisma", release)
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
monkeypatch.setenv("DATABASE_URL", database_url)
|
||||
yield database_url
|
||||
finally:
|
||||
with psycopg.connect(admin_url, autocommit=True) as conn:
|
||||
conn.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def partitioned_database(scratch_database: str) -> str:
|
||||
_partition_spend_logs(scratch_database)
|
||||
return scratch_database
|
||||
|
||||
|
||||
def _fail_the_call_id_migration_like_the_shipped_release(database_url: str, tmp_path: Path) -> None:
|
||||
"""Boot the original v1.103.0 layout once: its CONCURRENTLY statement fails on the
|
||||
partitioned parent and leaves the call_id ledger row unfinished."""
|
||||
failed: Final = _migrate_deploy(database_url, _release_layout(tmp_path / "v1.103.0", "99999999999999"))
|
||||
assert failed.returncode != 0 and PARTITIONED_PARENT_ERROR in failed.stderr, failed.stderr
|
||||
assert _ledger(database_url)[CALL_ID_INDEX_MIGRATION] == (False, False)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _IndexRow:
|
||||
name: str
|
||||
valid: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AttachedRow:
|
||||
table: str
|
||||
index: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LedgerRow:
|
||||
name: str
|
||||
finished: bool
|
||||
rolled_back: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _OidRow:
|
||||
name: str
|
||||
oid: int
|
||||
|
||||
|
||||
def _index_validity(database_url: str, suffix: str) -> Mapping[str, bool]:
|
||||
"""index name -> indisvalid for every index ending in `suffix` on the SpendLogs parent or one of its partitions."""
|
||||
with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_IndexRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
"SELECT c.relname AS name, i.indisvalid AS valid FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid "
|
||||
"WHERE c.relname LIKE %s AND (i.indrelid = to_regclass('\"LiteLLM_SpendLogs\"') OR i.indrelid IN "
|
||||
"(SELECT inhrelid FROM pg_inherits WHERE inhparent = to_regclass('\"LiteLLM_SpendLogs\"'))) "
|
||||
"ORDER BY c.relname",
|
||||
(f"%{suffix}",),
|
||||
).fetchall()
|
||||
return MappingProxyType({row.name: row.valid for row in rows})
|
||||
|
||||
|
||||
def _attached_children(database_url: str, parent_index: str) -> frozenset[tuple[str, str]]:
|
||||
"""(partition, child index) pairs attached under the parent index."""
|
||||
with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_AttachedRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
'SELECT t.relname AS "table", c.relname AS index FROM pg_inherits i '
|
||||
"JOIN pg_class c ON c.oid = i.inhrelid JOIN pg_index x ON x.indexrelid = c.oid "
|
||||
"JOIN pg_class t ON t.oid = x.indrelid "
|
||||
"WHERE i.inhparent = to_regclass(%s)",
|
||||
(f'"{parent_index}"',),
|
||||
).fetchall()
|
||||
return frozenset((row.table, row.index) for row in rows)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TableRow:
|
||||
name: str
|
||||
|
||||
|
||||
def _indexed_table(database_url: str, index: str) -> "str | None":
|
||||
with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_TableRow)) as cursor:
|
||||
row: Final = cursor.execute(
|
||||
"SELECT t.relname AS name FROM pg_index x JOIN pg_class t ON t.oid = x.indrelid "
|
||||
"WHERE x.indexrelid = to_regclass(%s)",
|
||||
(f'"{index}"',),
|
||||
).fetchone()
|
||||
return None if row is None else row.name
|
||||
|
||||
|
||||
def _ledger(database_url: str) -> Mapping[str, tuple[bool, bool]]:
|
||||
"""migration name -> (finished, rolled back) for the newest ledger row of each migration."""
|
||||
with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_LedgerRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
"SELECT DISTINCT ON (migration_name) migration_name AS name, finished_at IS NOT NULL AS finished, "
|
||||
"rolled_back_at IS NOT NULL AS rolled_back FROM _prisma_migrations ORDER BY migration_name, started_at DESC"
|
||||
).fetchall()
|
||||
return MappingProxyType({row.name: (row.finished, row.rolled_back) for row in rows})
|
||||
|
||||
|
||||
def _index_oids(database_url: str) -> Mapping[str, int]:
|
||||
"""index name -> oid for every index on the SpendLogs parent or one of its partitions; a rebuild changes the oid."""
|
||||
with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_OidRow)) as cursor:
|
||||
rows: Final = cursor.execute(
|
||||
"SELECT c.relname AS name, c.oid::int AS oid FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid "
|
||||
"WHERE i.indrelid = to_regclass('\"LiteLLM_SpendLogs\"') OR i.indrelid IN "
|
||||
"(SELECT inhrelid FROM pg_inherits WHERE inhparent = to_regclass('\"LiteLLM_SpendLogs\"'))"
|
||||
).fetchall()
|
||||
return MappingProxyType({row.name: row.oid for row in rows})
|
||||
|
||||
|
||||
def _migration_job(use_v2_resolver: bool) -> bool:
|
||||
return ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver)
|
||||
|
||||
|
||||
def _assert_no_pending_migrations(database_url: str) -> None:
|
||||
status: Final = _migrate_deploy(database_url, PACKAGE / "schema.prisma")
|
||||
assert status.returncode == 0 and "No pending migrations" in status.stdout, status.stdout + status.stderr
|
||||
|
||||
|
||||
def _assert_every_ledger_row_is_finished(database_url: str) -> Mapping[str, tuple[bool, bool]]:
|
||||
ledger: Final = _ledger(database_url)
|
||||
assert ledger[API_KEY_INDEX_MIGRATION] == (True, False) and ledger[CALL_ID_INDEX_MIGRATION] == (True, False)
|
||||
assert all(finished and not rolled_back for finished, rolled_back in ledger.values()), ledger
|
||||
return ledger
|
||||
|
||||
|
||||
def _expected_children(partitions: tuple[str, ...], suffix: str) -> frozenset[tuple[str, str]]:
|
||||
return frozenset((partition, f"{partition}_{suffix}") for partition in partitions)
|
||||
|
||||
|
||||
def _assert_index_covers_every_partition(database_url: str, parent_index: str, suffix: str) -> None:
|
||||
partitions: Final = (*PARTITIONS, DEFAULT_PARTITION)
|
||||
assert _index_validity(database_url, suffix) == {parent_index: True} | {f"{p}_{suffix}": True for p in partitions}
|
||||
assert _attached_children(database_url, parent_index) == _expected_children(partitions, suffix)
|
||||
|
||||
|
||||
@requires_db
|
||||
@RESOLVERS
|
||||
@RELEASES
|
||||
def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_a_rerun_is_idempotent(
|
||||
partitioned_database: str, use_v2_resolver: bool
|
||||
) -> None:
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
|
||||
_assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
ledger: Final = _assert_every_ledger_row_is_finished(partitioned_database)
|
||||
_assert_no_pending_migrations(partitioned_database)
|
||||
oids: Final = _index_oids(partitioned_database)
|
||||
|
||||
with psycopg.connect(partitioned_database, autocommit=True) as conn:
|
||||
conn.execute(
|
||||
'CREATE TABLE "LiteLLM_SpendLogs_p2026_10" PARTITION OF "LiteLLM_SpendLogs" '
|
||||
"FOR VALUES FROM ('2026-10-01') TO ('2026-11-01')"
|
||||
)
|
||||
inherited: Final = frozenset(
|
||||
("LiteLLM_SpendLogs_p2026_10", f"LiteLLM_SpendLogs_p2026_10_{suffix}")
|
||||
for suffix in ("api_key_startTime_idx", "litellm_call_id_idx")
|
||||
)
|
||||
attached: Final = _attached_children(partitioned_database, API_KEY_INDEX) | _attached_children(
|
||||
partitioned_database, CALL_ID_INDEX
|
||||
)
|
||||
assert inherited <= attached, attached
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
assert _ledger(partitioned_database) == ledger
|
||||
assert {name: oid for name, oid in _index_oids(partitioned_database).items() if name in oids} == oids
|
||||
|
||||
|
||||
@requires_db
|
||||
@RESOLVERS
|
||||
@RELEASES
|
||||
def test_a_plain_spend_logs_upgrade_builds_both_indexes_and_a_second_job_run_rebuilds_nothing(
|
||||
scratch_database: str, use_v2_resolver: bool
|
||||
) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
for row in range(ROWS_PER_PARTITION):
|
||||
_insert_spend_log(conn, f"flat-{row}", "2026-09-01")
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
|
||||
assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
_assert_every_ledger_row_is_finished(scratch_database)
|
||||
_assert_no_pending_migrations(scratch_database)
|
||||
oids: Final = _index_oids(scratch_database)
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
assert _index_oids(scratch_database) == oids
|
||||
|
||||
|
||||
@requires_db
|
||||
@RESOLVERS
|
||||
def test_a_database_that_applied_the_original_migration_files_sees_no_pending_migrations_and_no_rebuild(
|
||||
scratch_database: str, use_v2_resolver: bool, tmp_path: Path
|
||||
) -> None:
|
||||
"""A plain table upgraded on v1.103.0 applied both original files. The inert files in
|
||||
this build must neither re-run nor fail those rows, and the migration job must keep the
|
||||
indexes the migrations built."""
|
||||
deployed: Final = _migrate_deploy(scratch_database, _release_layout(tmp_path / "v1.103.0", "99999999999999"))
|
||||
assert deployed.returncode == 0, deployed.stderr
|
||||
before: Final = _ledger(scratch_database)
|
||||
assert before[API_KEY_INDEX_MIGRATION] == (True, False) and before[CALL_ID_INDEX_MIGRATION] == (True, False)
|
||||
oids: Final = _index_oids(scratch_database)
|
||||
assert {API_KEY_INDEX, CALL_ID_INDEX} <= set(oids)
|
||||
|
||||
_assert_no_pending_migrations(scratch_database)
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
|
||||
assert _ledger(scratch_database) == before
|
||||
assert _index_oids(scratch_database) == oids
|
||||
|
||||
|
||||
@requires_db
|
||||
@RESOLVERS
|
||||
def test_a_failed_call_id_ledger_row_from_a_v1_103_boot_is_rolled_back_and_the_inert_file_applied(
|
||||
partitioned_database: str, use_v2_resolver: bool, tmp_path: Path
|
||||
) -> None:
|
||||
_fail_the_call_id_migration_like_the_shipped_release(partitioned_database, tmp_path)
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
|
||||
_assert_every_ledger_row_is_finished(partitioned_database)
|
||||
_assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
_assert_no_pending_migrations(partitioned_database)
|
||||
with psycopg.connect(partitioned_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",
|
||||
(CALL_ID_INDEX_MIGRATION,),
|
||||
).fetchall()
|
||||
assert rows == [(False, True), (True, False)], rows
|
||||
|
||||
|
||||
@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
|
||||
) -> None:
|
||||
_fail_the_call_id_migration_like_the_shipped_release(partitioned_database, tmp_path)
|
||||
still_building: Final = tmp_path / "edited" / CALL_ID_INDEX_MIGRATION / "migration.sql"
|
||||
still_building.parent.mkdir(parents=True)
|
||||
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 _ledger(partitioned_database)[CALL_ID_INDEX_MIGRATION] == (False, False)
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_migration_without_a_failed_row_is_not_touched(partitioned_database: str) -> None:
|
||||
inert: Final = PACKAGE / "migrations" / CALL_ID_INDEX_MIGRATION / "migration.sql"
|
||||
before: Final = _ledger(partitioned_database)
|
||||
|
||||
with migration_lock(partitioned_database) as coordinator:
|
||||
assert roll_back_failed_inert_migration(coordinator, "public", inert) is False
|
||||
|
||||
assert _ledger(partitioned_database) == before
|
||||
|
||||
|
||||
def _pin_a_snapshot_on(database_url: str, table: str) -> "psycopg.Connection[tuple[object, ...]]":
|
||||
pin: Final = psycopg.connect(database_url)
|
||||
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
||||
pin.execute(sql.SQL("SELECT count(*) FROM {}").format(sql.Identifier(table)))
|
||||
return pin
|
||||
|
||||
|
||||
def _leave_an_invalid_index(database_url: str, name: str, table: str, column: str) -> None:
|
||||
with _pin_a_snapshot_on(database_url, table):
|
||||
with psycopg.connect(database_url, autocommit=True) as builder:
|
||||
builder.execute("SET statement_timeout = '1s'")
|
||||
with pytest.raises(psycopg.errors.QueryCanceled):
|
||||
builder.execute(
|
||||
sql.SQL("CREATE INDEX CONCURRENTLY {} ON {} ({})").format(
|
||||
sql.Identifier(name), sql.Identifier(table), sql.Identifier(column)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_an_invalid_index_of_the_managed_name_on_a_plain_table_is_rebuilt(scratch_database: str) -> None:
|
||||
_leave_an_invalid_index(scratch_database, CALL_ID_INDEX, "LiteLLM_SpendLogs", "litellm_call_id")
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: False}
|
||||
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
|
||||
|
||||
def _rebuild_as_another_replica(database_url: str, name: str, table: str, column: str) -> int:
|
||||
"""Drop and rebuild the index from a second connection, as a replica that won the
|
||||
race would, and return the oid of the index it built."""
|
||||
with psycopg.connect(database_url, autocommit=True) as other_replica:
|
||||
other_replica.execute(sql.SQL("DROP INDEX {}").format(sql.Identifier(name)))
|
||||
other_replica.execute(
|
||||
sql.SQL("CREATE INDEX {} ON {} ({})").format(
|
||||
sql.Identifier(name), sql.Identifier(table), sql.Identifier(column)
|
||||
)
|
||||
)
|
||||
return _index_oids(database_url)[name]
|
||||
|
||||
|
||||
def _connecting_with_another_replica_acting_first(
|
||||
statement: str, other_replica: Callable[[QueryNoTemplate], None]
|
||||
) -> Callable[[str], "psycopg.Connection[tuple[object, ...]]"]:
|
||||
"""A connect function whose cursors let `other_replica` act, once, right before the
|
||||
first statement containing `statement` runs: the interleaving two replicas booting
|
||||
together can produce, made deterministic."""
|
||||
raced: Final = threading.Event()
|
||||
|
||||
class _RacedCursor(psycopg.Cursor[tuple[object, ...]]):
|
||||
def execute( # pyright: ignore[reportIncompatibleMethodOverride] # the builder never runs a Template query
|
||||
self,
|
||||
query: QueryNoTemplate,
|
||||
params: "Params | None" = None,
|
||||
*,
|
||||
prepare: "bool | None" = None,
|
||||
binary: "bool | None" = None,
|
||||
) -> "_RacedCursor":
|
||||
text: Final = query.as_string(self.connection) if isinstance(query, sql.Composable) else query
|
||||
if isinstance(text, str) and statement in text and not raced.is_set():
|
||||
raced.set()
|
||||
other_replica(query)
|
||||
return super().execute(query, params, prepare=prepare, binary=binary)
|
||||
|
||||
def connect(database_url: str) -> "psycopg.Connection[tuple[object, ...]]":
|
||||
return psycopg.connect(database_url, autocommit=True, cursor_factory=_RacedCursor)
|
||||
|
||||
return connect
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_an_index_another_replica_made_valid_before_the_lock_was_taken_is_kept(scratch_database: str) -> None:
|
||||
"""Two replicas boot against the same invalid index. The one that takes the lock
|
||||
second must read the catalog again under it, or it drops the valid index the first
|
||||
one just finished and starts the whole build over."""
|
||||
_leave_an_invalid_index(scratch_database, CALL_ID_INDEX, "LiteLLM_SpendLogs", "litellm_call_id")
|
||||
theirs: Final[queue.SimpleQueue[int]] = queue.SimpleQueue()
|
||||
connect: Final = _connecting_with_another_replica_acting_first(
|
||||
"pg_try_advisory_lock",
|
||||
lambda _: theirs.put(
|
||||
_rebuild_as_another_replica(scratch_database, CALL_ID_INDEX, "LiteLLM_SpendLogs", "litellm_call_id")
|
||||
),
|
||||
)
|
||||
|
||||
assert ensure_request_log_indexes(scratch_database, "public", (CALL_ID_INDEX_DEFINITION,), connect) is True
|
||||
|
||||
assert _index_oids(scratch_database)[CALL_ID_INDEX] == theirs.get_nowait()
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_child_index_another_replica_attached_first_is_not_attached_twice(partitioned_database: str) -> None:
|
||||
"""A replica that reaches the attach step after another one attached the same child
|
||||
relies on ATTACH PARTITION being a no-op for an index already under that parent
|
||||
(PostgreSQL 14 ALTER INDEX, ATExecAttachPartitionIdx, checked 2026-10-01); this test
|
||||
is where that would surface if a future version or a code change made it an error."""
|
||||
|
||||
def attach_as_another_replica(statement: QueryNoTemplate) -> None:
|
||||
with psycopg.connect(partitioned_database, autocommit=True) as other_replica:
|
||||
other_replica.execute(statement)
|
||||
|
||||
connect: Final = _connecting_with_another_replica_acting_first("ATTACH PARTITION", attach_as_another_replica)
|
||||
|
||||
assert ensure_request_log_indexes(partitioned_database, "public", (CALL_ID_INDEX_DEFINITION,), connect) is True
|
||||
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_an_invalid_child_index_left_by_an_interrupted_build_is_rebuilt_and_attached(
|
||||
partitioned_database: str,
|
||||
) -> None:
|
||||
partition: Final = "LiteLLM_SpendLogs_p2026_08"
|
||||
child: Final = f"{partition}_litellm_call_id_idx"
|
||||
_leave_an_invalid_index(partitioned_database, child, partition, "litellm_call_id")
|
||||
assert _index_validity(partitioned_database, "litellm_call_id_idx") == {child: False}
|
||||
|
||||
assert ensure_request_log_indexes(partitioned_database, "public") is True
|
||||
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_an_index_of_that_name_on_another_table_is_left_alone_and_reported(partitioned_database: str) -> None:
|
||||
with psycopg.connect(partitioned_database, autocommit=True) as conn:
|
||||
conn.execute(f'CREATE INDEX "{CALL_ID_INDEX}" ON "LiteLLM_ErrorLogs" ("request_id")')
|
||||
|
||||
assert ensure_request_log_indexes(partitioned_database, "public") is False
|
||||
|
||||
assert _index_validity(partitioned_database, "litellm_call_id_idx") == {}
|
||||
assert _indexed_table(partitioned_database, CALL_ID_INDEX) == "LiteLLM_ErrorLogs"
|
||||
_assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_an_invalid_index_of_a_child_name_on_another_table_is_not_dropped(partitioned_database: str) -> None:
|
||||
child: Final = "LiteLLM_SpendLogs_p2026_08_litellm_call_id_idx"
|
||||
_leave_an_invalid_index(partitioned_database, child, "LiteLLM_ErrorLogs", "request_id")
|
||||
|
||||
assert ensure_request_log_indexes(partitioned_database, "public") is False
|
||||
|
||||
assert _indexed_table(partitioned_database, child) == "LiteLLM_ErrorLogs"
|
||||
assert _attached_children(partitioned_database, CALL_ID_INDEX) == frozenset()
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_process_holding_the_migration_lock_makes_the_build_wait_for_the_next_job_run(scratch_database: str) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as other_replica:
|
||||
other_replica.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is False
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {}
|
||||
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
@RESOLVERS
|
||||
def test_a_migration_job_that_could_not_build_the_indexes_reports_failure_and_succeeds_when_rerun(
|
||||
scratch_database: str, use_v2_resolver: bool
|
||||
) -> None:
|
||||
"""The migration job waits for the build and exits by run_migration_job's result; a job
|
||||
that exits 0 with the indexes missing would leave the table unindexed until the next
|
||||
deploy or until a serving proxy's background build gets to them."""
|
||||
with psycopg.connect(scratch_database, autocommit=True) as other_replica:
|
||||
other_replica.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
|
||||
assert _migration_job(use_v2_resolver) is False
|
||||
_assert_every_ledger_row_is_finished(scratch_database)
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {}
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
@RESOLVERS
|
||||
def test_the_serving_proxy_setup_applies_the_inert_migrations_and_builds_no_index(
|
||||
partitioned_database: str, use_v2_resolver: bool
|
||||
) -> None:
|
||||
"""setup_database alone applies the inert files and builds nothing, so a serving proxy's
|
||||
readiness is never held up by an index build; the build it starts afterwards, or the
|
||||
migration job, is what puts the indexes in place."""
|
||||
api_key_index_before: Final = _index_validity(partitioned_database, "api_key_startTime_idx")
|
||||
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
|
||||
|
||||
_assert_every_ledger_row_is_finished(partitioned_database)
|
||||
_assert_no_pending_migrations(partitioned_database)
|
||||
assert _index_validity(partitioned_database, "litellm_call_id_idx") == {}
|
||||
assert _index_validity(partitioned_database, "api_key_startTime_idx") == api_key_index_before
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
_assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_role_that_may_not_create_indexes_is_logged_and_left_for_the_next_job_run(
|
||||
scratch_database: str, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC")
|
||||
conn.execute("CREATE ROLE spend_logs_reader LOGIN PASSWORD 'reader'")
|
||||
conn.execute("GRANT USAGE ON SCHEMA public TO spend_logs_reader")
|
||||
conn.execute('GRANT SELECT ON "LiteLLM_SpendLogs" TO spend_logs_reader')
|
||||
reader_url: Final = scratch_database.replace("postgres:postgres@", "spend_logs_reader:reader@", 1)
|
||||
try:
|
||||
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
|
||||
assert ensure_request_log_indexes(reader_url, "public") is False
|
||||
finally:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute("DROP OWNED BY spend_logs_reader")
|
||||
conn.execute("DROP ROLE spend_logs_reader")
|
||||
assert "leaving them for the next index build" in caplog.text
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_inserts_keep_flowing_while_the_partition_indexes_build(partitioned_database: str) -> None:
|
||||
"""With a write open on one partition, the parent index goes on ONLY the parent and
|
||||
the CONCURRENTLY child build waits for that write without blocking new INSERTs. A
|
||||
plain CREATE INDEX on the parent would wait for the same write while holding SHARE
|
||||
on the parent, queueing every new INSERT behind it."""
|
||||
outcome: Final[list[bool]] = [] # mutable-ok: the builder thread hands its result back through it
|
||||
with psycopg.connect(partitioned_database) as writer:
|
||||
_insert_spend_log(writer, "LiteLLM_SpendLogs_p2026_08-open", "2026-08-15", table="LiteLLM_SpendLogs_p2026_08")
|
||||
builder_thread: Final = threading.Thread(
|
||||
target=lambda: outcome.append(_build_in_its_own_session(partitioned_database, CALL_ID_INDEX_DEFINITION))
|
||||
)
|
||||
builder_thread.start()
|
||||
try:
|
||||
_wait_until_the_build_is_waiting(partitioned_database)
|
||||
with psycopg.connect(partitioned_database, autocommit=True) as late_writer:
|
||||
late_writer.execute("SET lock_timeout = '1s'")
|
||||
_insert_spend_log(late_writer, "LiteLLM_SpendLogs_p2026_08-late", "2026-08-16")
|
||||
finally:
|
||||
writer.commit()
|
||||
builder_thread.join()
|
||||
assert outcome == [True]
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
|
||||
|
||||
def _build_in_its_own_session(database_url: str, index: RequestLogIndex) -> bool:
|
||||
with psycopg.connect(database_url, autocommit=True) as builder:
|
||||
return build_index_on_partitioned_table(builder, "public", index)
|
||||
|
||||
|
||||
def _wait_until_the_build_is_waiting(database_url: str) -> None:
|
||||
deadline: Final = time.monotonic() + 30
|
||||
with psycopg.connect(database_url, autocommit=True) as conn:
|
||||
while time.monotonic() < deadline:
|
||||
waiting = conn.execute(
|
||||
"SELECT 1 FROM pg_stat_activity WHERE query LIKE 'CREATE INDEX%' AND wait_event_type IS NOT NULL"
|
||||
).fetchone()
|
||||
if waiting is not None:
|
||||
return
|
||||
time.sleep(0.05)
|
||||
pytest.fail("the partition index build never started waiting on the open write")
|
||||
|
||||
|
||||
def _create_index(database_url: str, name: str, table: str, columns: str) -> int:
|
||||
"""Create a plain index by hand, the way an operator's workaround would, and return its oid."""
|
||||
with psycopg.connect(database_url, autocommit=True) as conn:
|
||||
conn.execute(
|
||||
sql.SQL("CREATE INDEX {} ON {} {}").format(sql.Identifier(name), sql.Identifier(table), sql.SQL(columns))
|
||||
)
|
||||
return _index_oids(database_url)[name]
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_valid_index_of_the_same_definition_under_another_name_is_renamed_instead_of_rebuilt(
|
||||
scratch_database: str,
|
||||
) -> None:
|
||||
hand_built: Final = _create_index(scratch_database, "call_id_by_hand", "LiteLLM_SpendLogs", '("litellm_call_id")')
|
||||
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
|
||||
oids: Final = _index_oids(scratch_database)
|
||||
assert "call_id_by_hand" not in oids and oids[CALL_ID_INDEX] == hand_built
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_hand_built_child_index_under_another_name_is_renamed_and_attached(partitioned_database: str) -> None:
|
||||
partition: Final = "LiteLLM_SpendLogs_p2026_08"
|
||||
hand_built: Final = _create_index(
|
||||
partitioned_database, "p2026_08_call_id_by_hand", partition, '("litellm_call_id")'
|
||||
)
|
||||
|
||||
assert ensure_request_log_indexes(partitioned_database, "public") is True
|
||||
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
oids: Final = _index_oids(partitioned_database)
|
||||
assert "p2026_08_call_id_by_hand" not in oids and oids[f"{partition}_litellm_call_id_idx"] == hand_built
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_an_index_with_another_definition_is_not_taken_for_the_managed_one(scratch_database: str) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
conn.execute(sql.SQL("DROP INDEX {}").format(sql.Identifier(API_KEY_INDEX)))
|
||||
others: Final = {
|
||||
"time_then_key": _create_index(
|
||||
scratch_database, "time_then_key", "LiteLLM_SpendLogs", '("startTime", "api_key")'
|
||||
),
|
||||
"call_id_desc": _create_index(
|
||||
scratch_database, "call_id_desc", "LiteLLM_SpendLogs", '("litellm_call_id" DESC)'
|
||||
),
|
||||
"call_id_then_key": _create_index(
|
||||
scratch_database, "call_id_then_key", "LiteLLM_SpendLogs", '("litellm_call_id", "api_key")'
|
||||
),
|
||||
"call_id_pattern": _create_index(
|
||||
scratch_database, "call_id_pattern", "LiteLLM_SpendLogs", '("litellm_call_id" text_pattern_ops)'
|
||||
),
|
||||
}
|
||||
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
|
||||
oids: Final = _index_oids(scratch_database)
|
||||
assert {name: oids[name] for name in others} == others
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_valid_partitioned_parent_index_under_another_name_is_renamed_with_its_children_kept(
|
||||
partitioned_database: str, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
hand_built: Final = _create_index(
|
||||
partitioned_database, "call_id_parent_by_hand", "LiteLLM_SpendLogs", '("litellm_call_id")'
|
||||
)
|
||||
children_before: Final = _attached_children(partitioned_database, "call_id_parent_by_hand")
|
||||
|
||||
with caplog.at_level("INFO", logger="litellm_proxy_extras"):
|
||||
assert ensure_request_log_indexes(partitioned_database, "public") is True
|
||||
|
||||
assert "Building index" not in caplog.text
|
||||
oids: Final = _index_oids(partitioned_database)
|
||||
assert "call_id_parent_by_hand" not in oids and oids[CALL_ID_INDEX] == hand_built
|
||||
assert _attached_children(partitioned_database, CALL_ID_INDEX) == children_before
|
||||
assert _index_validity(partitioned_database, "litellm_call_id_idx")[CALL_ID_INDEX] is True
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_a_second_copy_of_a_managed_index_is_reported_with_its_drop_statement_and_left_in_place(
|
||||
scratch_database: str, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
copy: Final = _create_index(scratch_database, "call_id_copy", "LiteLLM_SpendLogs", '("litellm_call_id")')
|
||||
|
||||
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
|
||||
assert 'remove it with: DROP INDEX CONCURRENTLY "public"."call_id_copy"' in caplog.text
|
||||
assert _index_oids(scratch_database)["call_id_copy"] == copy
|
||||
|
|
@ -2,7 +2,9 @@ import glob
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -1024,3 +1026,143 @@ class TestJWTKeyMappingCascade:
|
|||
f"{path} must declare onDelete: Cascade on the JWT key mapping "
|
||||
"relation (issue #33702)"
|
||||
)
|
||||
|
||||
|
||||
|
||||
class TestBuildRequestLogIndexes:
|
||||
"""The migration job hands the index build the direct database URL and the schema
|
||||
the migrations target, waits for it, and reports its result."""
|
||||
|
||||
@pytest.fixture
|
||||
def builds(self):
|
||||
return []
|
||||
|
||||
@pytest.fixture
|
||||
def build(self, builds):
|
||||
def record(database_url: str, schema: str) -> bool:
|
||||
builds.append((database_url, schema))
|
||||
return True
|
||||
|
||||
return record
|
||||
|
||||
def test_the_build_gets_the_direct_url_without_prisma_params_and_the_prisma_schema(self, monkeypatch, builds, build):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@pooler:6543/db?schema=tenant&pgbouncer=true")
|
||||
monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@primary:5432/db?connection_limit=1")
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
|
||||
|
||||
assert builds == [("postgresql://u:p@primary:5432/db", "tenant")]
|
||||
|
||||
def test_the_build_defaults_to_the_database_url_and_the_public_schema(self, monkeypatch, builds, build):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
|
||||
monkeypatch.delenv("DIRECT_URL", raising=False)
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
|
||||
|
||||
assert builds == [("postgresql://u:p@primary:5432/db", "public")]
|
||||
|
||||
def test_a_build_that_leaves_indexes_missing_is_reported_so_the_job_reruns(self, monkeypatch):
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=lambda url, schema: False) is False
|
||||
|
||||
def test_without_a_database_url_nothing_is_built(self, monkeypatch, builds, build):
|
||||
monkeypatch.delenv("DATABASE_URL", raising=False)
|
||||
|
||||
assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
|
||||
|
||||
assert builds == []
|
||||
|
||||
|
||||
class TestStartRequestLogIndexBuild:
|
||||
"""A serving proxy that ran the migrations starts the index build on a daemon thread
|
||||
and goes on to serve while it runs."""
|
||||
|
||||
def test_the_build_runs_on_a_daemon_thread_that_does_not_hold_up_the_caller(self):
|
||||
release: Final = threading.Event()
|
||||
builds: Final[list[str]] = [] # mutable-ok: the builder thread hands back the thread it ran on
|
||||
|
||||
def build() -> bool:
|
||||
assert release.wait(5), "the caller never came back from start_request_log_index_build"
|
||||
builds.append(threading.current_thread().name)
|
||||
return True
|
||||
|
||||
thread: Final = ProxyExtrasDBManager.start_request_log_index_build(build=build)
|
||||
|
||||
assert builds == [], "the build ran before start_request_log_index_build returned"
|
||||
assert thread.daemon is True
|
||||
release.set()
|
||||
thread.join(5)
|
||||
assert builds == ["litellm-request-log-indexes"]
|
||||
|
||||
|
||||
class TestRunMigrationJob:
|
||||
"""`run_migration_job` is `setup_database` followed by the index build, each step's
|
||||
result deciding whether the job reports success."""
|
||||
|
||||
@pytest.fixture
|
||||
def calls(self):
|
||||
return []
|
||||
|
||||
@pytest.fixture
|
||||
def setup(self, calls):
|
||||
def record(result: bool):
|
||||
def setup_database(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
calls.append(("setup", use_migrate, use_v2_resolver))
|
||||
return result
|
||||
|
||||
return setup_database
|
||||
|
||||
return record
|
||||
|
||||
@pytest.fixture
|
||||
def build(self, calls):
|
||||
def record(result: bool):
|
||||
def build_request_log_indexes() -> bool:
|
||||
calls.append(("build",))
|
||||
return result
|
||||
|
||||
return build_request_log_indexes
|
||||
|
||||
return record
|
||||
|
||||
def test_the_job_builds_the_indexes_after_the_migrations_succeed(self, calls, setup, build):
|
||||
assert ProxyExtrasDBManager.run_migration_job(True, False, setup=setup(True), build=build(True)) is True
|
||||
|
||||
assert calls == [("setup", True, False), ("build",)]
|
||||
|
||||
def test_the_job_fails_without_building_when_the_migrations_fail(self, calls, setup, build):
|
||||
assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(False), build=build(True)) is False
|
||||
|
||||
assert calls == [("setup", True, True)]
|
||||
|
||||
def test_the_job_fails_when_an_index_could_not_be_built(self, calls, setup, build):
|
||||
assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is False
|
||||
|
||||
assert calls == [("setup", True, True), ("build",)]
|
||||
|
||||
|
||||
class TestMigrationJobOwnedDrift:
|
||||
JOB_INDEXES = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
|
||||
"\n-- CreateIndex\n"
|
||||
'CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
|
||||
)
|
||||
|
||||
def test_a_plain_spend_logs_table_only_loses_the_migration_job_indexes(self):
|
||||
filtered = ProxyExtrasDBManager._filter_migration_job_owned_drift(
|
||||
_PARTITIONED_DRIFT_SQL + self.JOB_INDEXES, partitioned=False
|
||||
)
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
|
||||
assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in filtered
|
||||
assert 'PRIMARY KEY ("request_id")' in filtered
|
||||
|
||||
def test_a_partitioned_spend_logs_table_also_loses_its_partitioning_artifacts(self):
|
||||
filtered = ProxyExtrasDBManager._filter_migration_job_owned_drift(
|
||||
_PARTITIONED_DRIFT_SQL + self.JOB_INDEXES, partitioned=True
|
||||
)
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
|
||||
assert 'PRIMARY KEY ("request_id")' not in filtered
|
||||
assert "LiteLLM_SpendLogs_legacy" not in filtered
|
||||
assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in filtered
|
||||
|
|
|
|||
149
tests/unit/litellm_proxy_extras/test_request_log_indexes.py
Normal file
149
tests/unit/litellm_proxy_extras/test_request_log_indexes.py
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
import re
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from litellm_proxy_extras.migration_recovery import is_inert_migration
|
||||
from litellm_proxy_extras.request_log_indexes import (
|
||||
REQUEST_LOG_INDEXES,
|
||||
RequestLogIndex,
|
||||
filter_request_log_index_diff,
|
||||
)
|
||||
|
||||
PACKAGE: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras"
|
||||
SCHEMA: Final = PACKAGE / "schema.prisma"
|
||||
INERT_MIGRATIONS: Final = (
|
||||
"20260823000000_add_spend_logs_api_key_starttime_index",
|
||||
"20260831120001_spend_logs_litellm_call_id_index",
|
||||
)
|
||||
CALL_ID_INDEX: Final = RequestLogIndex(
|
||||
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'
|
||||
)
|
||||
|
||||
|
||||
def _prisma_indexes_of(schema: str, model: str) -> frozenset[str]:
|
||||
"""The index names Prisma derives for a model's @@index declarations: <model>_<columns joined by _>_idx."""
|
||||
body: Final = re.search(rf"model {model} \{{(.*?)\n\}}", schema, re.DOTALL)
|
||||
assert body is not None, model
|
||||
declarations: Final[tuple[str, ...]] = tuple(
|
||||
match.group(1) for match in re.finditer(r"@@index\(\[([^\]]+)\]\)", body.group(1))
|
||||
)
|
||||
return frozenset(
|
||||
f"{model}_{'_'.join(column.strip() for column in columns.split(','))}_idx" for columns in declarations
|
||||
)
|
||||
|
||||
|
||||
class TestTheIndexList:
|
||||
def test_every_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self):
|
||||
schema: Final = SCHEMA.read_text()
|
||||
for index in REQUEST_LOG_INDEXES:
|
||||
assert index.name in _prisma_indexes_of(schema, index.table), index
|
||||
|
||||
@pytest.mark.parametrize("name", INERT_MIGRATIONS)
|
||||
def test_the_migrations_that_used_to_build_these_indexes_run_no_sql(self, name: str):
|
||||
assert is_inert_migration((PACKAGE / "migrations" / name / "migration.sql").read_text())
|
||||
|
||||
|
||||
class TestIsInertMigration:
|
||||
@pytest.mark.parametrize(
|
||||
"script",
|
||||
(
|
||||
"",
|
||||
"-- only a comment\n",
|
||||
"/* block */\n-- line\n",
|
||||
"-- a semicolon; in a comment\n",
|
||||
";\n;",
|
||||
"-- why\nSELECT 1;\n",
|
||||
"select 1",
|
||||
),
|
||||
ids=(
|
||||
"empty",
|
||||
"line-comment",
|
||||
"both-comments",
|
||||
"semicolon-in-comment",
|
||||
"bare-separators",
|
||||
"select-1",
|
||||
"lowercase",
|
||||
),
|
||||
)
|
||||
def test_comments_and_a_select_1_alone_are_inert(self, script: str):
|
||||
assert is_inert_migration(script) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"script",
|
||||
(
|
||||
"SELECT 2;",
|
||||
'SELECT 1 FROM "LiteLLM_SpendLogs";',
|
||||
'-- comment\nCREATE INDEX "ix" ON "t" ("a");',
|
||||
"/* c */ ALTER TABLE t ADD COLUMN a TEXT",
|
||||
'SELECT 1; DROP INDEX "ix";',
|
||||
),
|
||||
ids=("select-2", "select-from", "index-after-comment", "alter-after-block-comment", "drop-after-select-1"),
|
||||
)
|
||||
def test_any_statement_is_not_inert(self, script: str):
|
||||
assert is_inert_migration(script) is False
|
||||
|
||||
|
||||
class TestPartitionIndexName:
|
||||
def test_a_partition_gets_the_name_postgres_would_give_an_inherited_index(self):
|
||||
assert CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p2026_09") == (
|
||||
"LiteLLM_SpendLogs_p2026_09_litellm_call_id_idx"
|
||||
)
|
||||
|
||||
def test_an_index_not_prefixed_by_its_table_keeps_its_whole_name(self):
|
||||
index = RequestLogIndex("LiteLLM_SpendLogs", "call_id_lookup", '("litellm_call_id")')
|
||||
assert index.partition_index_name("LiteLLM_SpendLogs_pdefault") == "LiteLLM_SpendLogs_pdefault_call_id_lookup"
|
||||
|
||||
def test_a_long_name_is_cut_to_63_bytes_with_a_digest_that_keeps_partitions_apart(self):
|
||||
first = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_09")
|
||||
second = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_10")
|
||||
assert len(first.encode()) == 63 and len(second.encode()) == 63
|
||||
assert first != second
|
||||
assert first.startswith("LiteLLM_SpendLogs_p") and first[-9] == "_"
|
||||
|
||||
def test_the_byte_limit_counts_multibyte_characters(self):
|
||||
name = CALL_ID_INDEX.partition_index_name("é" * 40)
|
||||
assert len(name.encode()) <= 63 and len(name) < 63
|
||||
|
||||
|
||||
class TestColumns:
|
||||
def test_the_columns_are_the_quoted_names_of_the_definition_in_order(self):
|
||||
index = RequestLogIndex(
|
||||
"LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'
|
||||
)
|
||||
assert index.columns == ("api_key", "startTime")
|
||||
|
||||
def test_every_migration_job_index_names_at_least_one_column(self):
|
||||
assert all(index.columns for index in REQUEST_LOG_INDEXES)
|
||||
|
||||
|
||||
DRIFT_WITH_BOTH_INDEXES: Final = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
|
||||
"\n"
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
|
||||
)
|
||||
|
||||
|
||||
class TestFilterRequestLogIndexDiff:
|
||||
def test_a_drift_script_that_only_creates_the_migration_job_indexes_becomes_empty(self):
|
||||
assert filter_request_log_index_diff(DRIFT_WITH_BOTH_INDEXES) == ""
|
||||
|
||||
def test_other_statements_survive_with_the_migration_job_indexes_removed(self):
|
||||
other: Final = '-- AlterTable\nALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;\n'
|
||||
filtered = filter_request_log_index_diff(other + DRIFT_WITH_BOTH_INDEXES)
|
||||
assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN "updated_by" TEXT;' in filtered
|
||||
assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
|
||||
assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in filtered
|
||||
|
||||
def test_an_index_of_another_name_on_spend_logs_is_kept(self):
|
||||
sql: Final = 'CREATE INDEX "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user");\n'
|
||||
assert filter_request_log_index_diff(sql) == sql
|
||||
|
||||
def test_a_drop_of_a_migration_job_index_is_kept_for_the_operator_to_see(self):
|
||||
sql: Final = 'DROP INDEX "LiteLLM_SpendLogs_litellm_call_id_idx";\n'
|
||||
assert filter_request_log_index_diff(sql) == sql
|
||||
|
||||
def test_an_empty_script_stays_empty(self):
|
||||
assert filter_request_log_index_diff("") == ""
|
||||
|
|
@ -454,3 +454,19 @@ def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashi
|
|||
|
||||
assert PrismaManager.setup_database(use_migrate=False) is False
|
||||
assert fake_prisma_cli.calls == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("run", "outcome"),
|
||||
(
|
||||
(PrismaManager.build_request_log_indexes, False),
|
||||
(PrismaManager.start_request_log_index_build, None),
|
||||
),
|
||||
ids=("wait-for-the-build", "start-the-build"),
|
||||
)
|
||||
def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising(monkeypatch, run, outcome):
|
||||
"""The migration job exits non-zero and a serving proxy keeps booting when the extras
|
||||
package that owns the index build is not installed."""
|
||||
monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None)
|
||||
|
||||
assert run() is outcome
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import inspect
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -2128,12 +2129,14 @@ class TestRunServerDbSetup:
|
|||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_use_prisma_db_push_flag_behavior(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
|
|
@ -2187,9 +2190,7 @@ class TestRunServerDbSetup:
|
|||
# Test 1: Without --use_prisma_db_push flag (default behavior)
|
||||
# use_prisma_db_push should be False (default), so use_migrate should be True
|
||||
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
|
||||
mock_setup_database.assert_called_with(
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
mock_setup_database.assert_called_with(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
# Reset mocks
|
||||
mock_setup_database.reset_mock()
|
||||
|
|
@ -2202,18 +2203,18 @@ class TestRunServerDbSetup:
|
|||
["--local", "--skip_server_startup", "--use_prisma_db_push"],
|
||||
standalone_mode=False,
|
||||
)
|
||||
mock_setup_database.assert_called_with(
|
||||
use_migrate=False, use_v2_resolver=True
|
||||
)
|
||||
mock_setup_database.assert_called_with(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
def test_migrations_run_when_the_prisma_cli_is_not_on_path(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
tmp_path,
|
||||
|
|
@ -2262,19 +2263,19 @@ class TestRunServerDbSetup:
|
|||
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
|
||||
|
||||
assert "prisma CLI is neither on PATH" not in capsys.readouterr().out
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_startup_fails_when_db_setup_fails(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
|
|
@ -2329,19 +2330,19 @@ class TestRunServerDbSetup:
|
|||
standalone_mode=False,
|
||||
)
|
||||
assert exc_info.value.code == 1
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_startup_exits_on_non_postgres_database_url(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
|
|
@ -2387,12 +2388,14 @@ class TestRunServerDbSetup:
|
|||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_v2_migration_resolver_opts_in_via_env_var(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
|
|
@ -2439,9 +2442,7 @@ class TestRunServerDbSetup:
|
|||
["--local", "--skip_server_startup"], standalone_mode=False
|
||||
)
|
||||
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=True
|
||||
)
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
|
||||
assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -2479,12 +2480,14 @@ class TestRunServerDbSetup:
|
|||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
|
||||
def test_legacy_resolver_flag_reaches_database_setup(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
|
|
@ -2533,9 +2536,76 @@ class TestRunServerDbSetup:
|
|||
standalone_mode=False,
|
||||
)
|
||||
|
||||
mock_setup_database.assert_called_once_with(
|
||||
use_migrate=True, use_v2_resolver=False
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("arguments", "migrated", "exits", "waits_for_the_build"),
|
||||
(
|
||||
(("--local", "--skip_server_startup"), True, True, True),
|
||||
(("--local",), True, False, False),
|
||||
(("--local",), False, True, False),
|
||||
),
|
||||
ids=("migration-job", "serving-proxy", "serving-proxy-whose-migrations-failed"),
|
||||
)
|
||||
@patch("uvicorn.run")
|
||||
@patch("subprocess.run")
|
||||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True)
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes", return_value=False)
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.start_request_log_index_build")
|
||||
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=True)
|
||||
def test_the_migration_job_waits_for_the_index_build_and_a_serving_proxy_starts_it_in_the_background(
|
||||
self,
|
||||
mock_should_update_schema,
|
||||
mock_check_schema_diff,
|
||||
mock_start_build,
|
||||
mock_build_indexes,
|
||||
mock_setup_database,
|
||||
mock_atexit_register,
|
||||
mock_subprocess_run,
|
||||
mock_uvicorn_run,
|
||||
arguments,
|
||||
migrated,
|
||||
exits,
|
||||
waits_for_the_build,
|
||||
):
|
||||
"""`--skip_server_startup` is the migration job: it waits for the index build after the
|
||||
migrations and exits 1 when one could not be built. A serving proxy that ran the
|
||||
migrations starts the build in the background and serves whatever the build does; one
|
||||
whose migrations failed exits 1 under `--enforce_prisma_migration_check` and starts no build."""
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
mock_setup_database.return_value = migrated
|
||||
mock_subprocess_run.return_value = MagicMock(returncode=0)
|
||||
mock_proxy_module = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
|
||||
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
|
||||
outcome = pytest.raises(SystemExit) if exits else nullcontext()
|
||||
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module},
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
) as mock_get_args,
|
||||
outcome as exc_info,
|
||||
):
|
||||
mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000}
|
||||
run_server.main([*arguments, "--enforce_prisma_migration_check"], standalone_mode=False)
|
||||
|
||||
assert (exc_info is not None and exc_info.value.code == 1) is exits
|
||||
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
|
||||
assert mock_build_indexes.call_count == int(migrated and waits_for_the_build)
|
||||
assert mock_start_build.call_count == int(migrated and not waits_for_the_build)
|
||||
|
||||
|
||||
# --- Module-level helpers for worker startup hook tests ---
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ import importlib.util
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_CHECKER_PATH = Path(__file__).resolve().parents[1] / "code_coverage_tests" / "check_migrations_no_data_rewrites.py"
|
||||
_SPEC = importlib.util.spec_from_file_location("check_migrations_no_data_rewrites", _CHECKER_PATH)
|
||||
assert _SPEC is not None and _SPEC.loader is not None
|
||||
|
|
@ -205,6 +207,96 @@ class TestDefaultedColumnsOnRequestLogTables:
|
|||
assert 'ADD COLUMN ... DEFAULT on "LiteLLM_SpendLogs" rewrites existing rows at boot' in rendered
|
||||
|
||||
|
||||
class TestIndexesOnLogTables:
|
||||
"""Every CREATE INDEX on a request-log table is rejected: a plain one blocks writes for
|
||||
the whole build and a concurrent one fails on a partitioned parent, so the migration job
|
||||
(litellm_proxy_extras/request_log_indexes.py) builds those instead."""
|
||||
|
||||
def test_the_original_spend_log_index_statement_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("api_key", "startTime");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_the_original_concurrent_call_id_index_statement_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" '
|
||||
'ON "LiteLLM_SpendLogs"("litellm_call_id");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_unique_index_with_if_not_exists_on_error_logs_is_flagged(self, tmp_path):
|
||||
sql = 'CREATE UNIQUE INDEX IF NOT EXISTS "ix" ON "LiteLLM_ErrorLogs" ("request_id");'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_ErrorLogs"',)
|
||||
|
||||
def test_a_unique_concurrent_index_on_error_logs_is_flagged(self, tmp_path):
|
||||
sql = 'CREATE UNIQUE INDEX CONCURRENTLY "ix" ON "LiteLLM_ErrorLogs" ("request_id");'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_ErrorLogs"',)
|
||||
|
||||
def test_lowercase_schema_qualified_and_only_forms_are_flagged(self, tmp_path):
|
||||
sql = (
|
||||
'create index on "public"."LiteLLM_SpendLogs" ("api_key");\n'
|
||||
'CREATE INDEX "ix" ON ONLY "LiteLLM_SpendLogs" ("api_key");\n'
|
||||
'CREATE INDEX CONCURRENTLY "iy" ON "public"."LiteLLM_SpendLogs" ("api_key");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) * 3
|
||||
|
||||
def test_a_comment_between_on_and_the_table_is_flagged(self, tmp_path):
|
||||
sql = 'CREATE INDEX CONCURRENTLY "ix" ON /* table */ "LiteLLM_SpendLogs" ("api_key");'
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_a_concurrent_index_with_comments_and_line_breaks_is_flagged(self, tmp_path):
|
||||
sql = (
|
||||
"-- CreateIndex\n"
|
||||
'CREATE INDEX CONCURRENTLY IF NOT EXISTS "ix"\n'
|
||||
' ON "LiteLLM_SpendLogs" /* partitioned in some deployments */\n'
|
||||
' ("api_key", "startTime");\n'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_indexes_on_a_non_log_table_pass_concurrent_or_not(self, tmp_path):
|
||||
sql = (
|
||||
'CREATE INDEX "ix" ON "LiteLLM_VerificationToken" ("token");\n'
|
||||
'CREATE INDEX CONCURRENTLY "iy" ON "LiteLLM_VerificationToken" ("token");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ()
|
||||
|
||||
def test_an_index_run_by_execute_is_flagged(self, tmp_path):
|
||||
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"',)
|
||||
|
||||
def test_a_marker_on_a_rewrite_still_leaves_the_index_below_it_flagged(self, tmp_path):
|
||||
sql = (
|
||||
"-- data-migration-ok: one row\n"
|
||||
'UPDATE "LiteLLM_SpendLogs" SET "api_key" = \'k\' WHERE "request_id" = \'r\';\n'
|
||||
'CREATE INDEX CONCURRENTLY "ix" ON "LiteLLM_SpendLogs" ("api_key");'
|
||||
)
|
||||
assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',)
|
||||
|
||||
def test_render_points_at_the_migration_job_index_list(self, tmp_path):
|
||||
sql = 'CREATE INDEX CONCURRENTLY "ix" ON "public"."LiteLLM_SpendLogs" ("api_key");'
|
||||
rendered = _scan(tmp_path, sql)[0].render()
|
||||
assert "20260101000000_fixture/migration.sql:1" in rendered
|
||||
assert "blocks writes until the build finishes, or fails on a partitioned table" in rendered
|
||||
assert "REQUEST_LOG_INDEXES in litellm_proxy_extras/request_log_indexes.py" in rendered
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name",
|
||||
("20260823000000_add_spend_logs_api_key_starttime_index", "20260831120001_spend_logs_litellm_call_id_index"),
|
||||
)
|
||||
def test_the_inert_index_migrations_scan_clean_without_a_grandfather(self, name):
|
||||
assert checker.scan_migration(checker.MIGRATIONS_DIR / name) == ()
|
||||
assert name not in checker.GRANDFATHERED
|
||||
|
||||
|
||||
class TestInsert:
|
||||
def test_insert_values_is_bounded_and_passes(self, tmp_path):
|
||||
assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue