diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql index 750c1a934aa..28107f28043 100644 --- a/db_scripts/partition_spend_logs.sql +++ b/db_scripts/partition_spend_logs.sql @@ -57,6 +57,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx"; ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" RENAME TO "LiteLLM_SpendLogs_legacy_litellm_call_id_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_error_code_startTime_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_error_code_startTime_idx"; CREATE TABLE "LiteLLM_SpendLogs" ( LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 4c856c741d3..c6f58ab0b2a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -5,8 +5,10 @@ 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. +names of plain indexes match what Prisma derives from the `@@index` declarations in +schema.prisma, so an index a database already has is recognized and never rebuilt. +Prisma cannot declare a partial index, and its schema diff ignores one, so a partial +index lives only here. """ import hashlib @@ -15,7 +17,8 @@ import re import time from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Final +from itertools import chain +from typing import TYPE_CHECKING, Final, Literal from litellm_proxy_extras._logging import logger from litellm_proxy_extras.migration_data_checks import ( @@ -33,20 +36,31 @@ if TYPE_CHECKING: @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 `.""" + """One index the migration job owns: the table, the exact Prisma index name, the + column list as it would be written after `ON
`, the WHERE predicate of a + partial index, and the largest table the build takes on by itself. An index whose + build has to read each row's metadata takes time in proportion to the table's size on + disk, so above that size it is left for an operator to build off-peak instead of + holding the migration job past its deadline.""" table: str name: str definition: str + predicate: str | None = None + max_bytes: int | None = None @property def columns(self) -> tuple[str, ...]: return tuple(re.findall(r'"([^"]+)"', self.definition)) + @property + def body(self) -> str: + return self.definition if self.predicate is None else f"{self.definition} WHERE {self.predicate}" + 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.""" + """The child index name this job builds on one partition, kept within the 63 byte + identifier limit. Postgres names the children it creates for a new partition after + their columns instead, so a partition's child is found by attachment, never by name.""" name: Final = f"{partition}_{self.name.removeprefix(f'{self.table}_')}" if len(name.encode()) <= _IDENTIFIER_MAX_BYTES: return name @@ -59,6 +73,13 @@ class RequestLogIndex: REQUEST_LOG_INDEXES: Final = ( RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'), RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'), + RequestLogIndex( + "LiteLLM_SpendLogs", + "LiteLLM_SpendLogs_error_code_startTime_idx", + "((left(metadata->'error_information'->>'error_code', 128)), \"startTime\")", + "(metadata->'error_information'->>'error_code') IS NOT NULL", + max_bytes=4 * 1024**3, + ), ) _IDENTIFIER_MAX_BYTES: Final = 63 @@ -73,6 +94,12 @@ _CREATE_INDEX_STATEMENT: Final = re.compile( re.IGNORECASE, ) _TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)" +_TABLE_BYTES_SQL: Final = ( + "WITH RECURSIVE tree(oid) AS (SELECT to_regclass(%(table)s)::oid " + "UNION ALL SELECT i.inhrelid FROM pg_inherits i JOIN tree ON i.inhparent = tree.oid) " + "SELECT COALESCE(SUM(pg_table_size(c.oid)), 0)::bigint AS total " + "FROM tree JOIN pg_class c ON c.oid = tree.oid WHERE c.relkind = 'r'" +) _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 " @@ -128,6 +155,14 @@ class _TableKind: partitioned: bool +@dataclass(frozen=True, slots=True) +class _TableBytes: + total: int + + +_IndexOutcome = Literal["in_place", "left_for_operator", "not_in_place"] + + 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.""" @@ -161,20 +196,29 @@ def ensure_request_log_indexes( """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.""" + failure is logged and left for the next index build, and an index on a table over its + `max_bytes` is logged and left for an operator; the result says whether every + other 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) + outcomes: 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): + if "not_in_place" in outcomes: logger.warning("Some request-log indexes are not in place yet, leaving them for the next index build") return False + left_for_operator: Final = tuple( + index.name for index, outcome in zip(indexes, outcomes) if outcome == "left_for_operator" + ) + if left_for_operator: + logger.warning( + "Request-log indexes are in place except %s, left for an operator to build", ", ".join(left_for_operator) + ) + return True logger.info("Request-log indexes are all in place") return True @@ -215,17 +259,121 @@ def _with_bounded_lock( return False -def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool: +def _ensure_index( + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex +) -> _IndexOutcome: 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) + return "in_place" + if _too_large_to_build(connection, schema, index, partitioned=table.partitioned): + return "left_for_operator" + built: Final = ( + build_index_on_partitioned_table(connection, schema, index) + if table.partitioned + else _build_leaf_index(connection, schema, index.table, index.name, index) + ) + return "in_place" if built else "not_in_place" + + +def _too_large_to_build( + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex, partitioned: bool +) -> bool: + if index.max_bytes is None: + return False + existing: Final = _index_state(connection, schema, index.name) + if existing is not None and existing.valid: + return False + size: Final = _table_bytes(connection, schema, index.table) + if size <= index.max_bytes: + return False + logger.warning( + "Not building index %s on %s: the table takes %s on disk, more than the %s it builds by itself. Build it " + "off-peak by running each of these statements on its own, for example with psql -f, and the next index build " + "finds it in place:\n%s", + index.name, + index.table, + _gib(size), + _gib(index.max_bytes), + "".join( + f"{statement};\n" + for statement in _operator_build_statements(connection, schema, index, index.table, index.name, partitioned) + ), + ) + return True + + +def _gib(size: int) -> str: + return f"{size / 1024**3:.1f} GiB" + + +def _table_bytes(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> int: + """The table's size on disk across its partitions, TOAST and dead rows included: an upper bound on what a build + reads, known without ANALYZE, where a row count misses how large each row's metadata is.""" + from psycopg.rows import class_row + + with connection.cursor(row_factory=class_row(_TableBytes)) as cursor: + size: Final = cursor.execute(_TABLE_BYTES_SQL, {"table": _regclass_name(connection, schema, table)}).fetchone() + return size.total if size is not None else 0 + + +def _operator_build_statements( + connection: "psycopg.Connection[tuple[object, ...]]", + schema: str, + index: RequestLogIndex, + table: str, + name: str, + partitioned: bool, +) -> tuple[str, ...]: + """What an operator runs by hand to finish the build from where it stands: the steps this job takes, so no + statement holds spend-log writes back for the length of a build. A partitioned table gets the metadata-only + parent index, then each partition still missing its child gets a CONCURRENTLY build and ATTACH PARTITION. An + invalid index an interrupted build left gets REINDEX CONCURRENTLY, which CREATE ... IF NOT EXISTS would skip.""" + from psycopg import sql + + found: Final = _index_state(connection, schema, name) + existing: Final = found if found is not None and found.table == table else None + if existing is not None and existing.valid: + return () + if not partitioned and existing is not None: + return (sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)).as_string(connection),) + if not partitioned: + leaf: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format( + sql.Identifier(name), sql.Identifier(schema, table) + ) + return (_create_index_statement(connection, leaf, index.body).decode(),) + parent: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format( + sql.Identifier(name), sql.Identifier(schema, table) + ) + parent_statement: Final = ( + () if existing is not None else (_create_index_statement(connection, parent, index.body).decode(),) + ) + children: Final = _children_without_the_index(connection, schema, table, name) + return parent_statement + tuple( + chain.from_iterable(_operator_child_statements(connection, schema, name, child, index) for child in children) + ) + + +def _operator_child_statements( + connection: "psycopg.Connection[tuple[object, ...]]", + schema: str, + parent_index: str, + child: _Relation, + index: RequestLogIndex, +) -> tuple[str, ...]: + from psycopg import sql + + child_index: Final = index.partition_index_name(child.name) + attach: Final = sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format( + sql.Identifier(schema, parent_index), sql.Identifier(child.schema, child_index) + ) + built: Final = _operator_build_statements( + connection, child.schema, index, child.name, child_index, child.partitioned + ) + return (*built, attach.as_string(connection)) def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str: @@ -256,9 +404,12 @@ def _equivalent_indexes( ) -> 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.""" + opclass or collation, and not attached under a partitioned index. Valid ones first. + Empty for a partial index, whose predicate this lookup cannot compare.""" from psycopg.rows import class_row + if index.predicate is not None: + return () with connection.cursor(row_factory=class_row(_EquivalentIndex)) as cursor: return tuple( cursor.execute( @@ -371,7 +522,7 @@ def _build_leaf_index( 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)) + connection.execute(_create_index_statement(connection, prefix, index.body)) built: Final = _index_state(connection, schema, name) return built is not None and built.valid @@ -440,7 +591,7 @@ def _create_parent_index( 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) + statement: Final = _create_index_statement(connection, prefix, index.body) connection.execute(statement) return True diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 8b19576d86e..19649b921c5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -68,8 +68,10 @@ _MIGRATION_DEADLOCK_MARKER = "deadlock detected" _MIGRATION_LOCK_TIMEOUT_MARKER = "canceling statement due to lock timeout" INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") _TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") +_SIZE_CAPPED_REQUEST_LOG_INDEXES: Final = tuple(index for index in REQUEST_LOG_INDEXES if index.max_bytes is not None) _INVALID_LITELLM_INDEXES_SQL: Final = ( - "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " + "SELECT n.nspname AS schema, c.relname AS name, pg_size_pretty(pg_table_size(t.oid)) AS table_size, " + 't.relname AS "table" ' "FROM pg_index i " "JOIN pg_class c ON c.oid = i.indexrelid " "JOIN pg_class t ON t.oid = i.indrelid " @@ -88,6 +90,11 @@ class _InvalidIndex: schema: str name: str table_size: str + table: str + + @property + def left_to_the_request_log_index_build(self) -> bool: + return any(self.name == index.partition_index_name(self.table) for index in _SIZE_CAPPED_REQUEST_LOG_INDEXES) MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 @@ -844,8 +851,11 @@ class ProxyExtrasDBManager: def _invalid_litellm_indexes( conn: "psycopg.Connection[tuple[str, str, str]]", schema: str ) -> tuple[_InvalidIndex, ...]: - rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall() - return tuple(_InvalidIndex(*row) for row in rows) + from psycopg.rows import class_row + + with conn.cursor(row_factory=class_row(_InvalidIndex)) as cursor: + invalid: Final = cursor.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall() + return tuple(index for index in invalid if not index.left_to_the_request_log_index_build) @staticmethod def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 89200792c0b..c74fb6d4edd 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -97,6 +97,13 @@ _SESSION_REPRESENTATIVE_ORDER_SQL: Final = ( f'call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC, request_id' ) _BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')" +_ERROR_CODE_SQL: Final = "metadata->'error_information'->>'error_code'" +_INDEXED_ERROR_CODE_PREFIX_LENGTH: Final = 128 +_INDEXED_ERROR_CODE_PREFIX_SQL: Final = f"left({_ERROR_CODE_SQL}, {_INDEXED_ERROR_CODE_PREFIX_LENGTH})" +_ERROR_CODE_INDEX_READY_SQL: Final = ( + "SELECT EXISTS (SELECT 1 FROM pg_index " + "WHERE indexrelid = to_regclass('\"LiteLLM_SpendLogs_error_code_startTime_idx\"') AND indisvalid) AS ready" +) _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -227,6 +234,26 @@ class _SessionSpendRow(TypedDict): session_models: ReadOnly[Sequence[str]] +class _IndexReadyRow(TypedDict): + ready: ReadOnly[bool] + + +async def _error_code_index_is_ready(prisma_client: PrismaClient) -> bool: + rows: Final[Sequence[_IndexReadyRow]] = await _query_raw(prisma_client, _ERROR_CODE_INDEX_READY_SQL) + return bool(rows) and rows[0]["ready"] is True + + +def _error_code_condition(placeholder: str, error_code: str, index_ready: bool) -> str: + if not index_ready: + return f"{_ERROR_CODE_SQL} = {placeholder}" + if len(error_code) < _INDEXED_ERROR_CODE_PREFIX_LENGTH: + return f"{_INDEXED_ERROR_CODE_PREFIX_SQL} = {placeholder}" + return ( + f"{_INDEXED_ERROR_CODE_PREFIX_SQL} = left({placeholder}, {_INDEXED_ERROR_CODE_PREFIX_LENGTH}) " + f"AND {_ERROR_CODE_SQL} = {placeholder}" + ) + + _SESSION_MODELS_LIMIT: Final = 10 _SESSION_MODEL_NAME_MAX_LEN: Final = 256 @@ -2580,7 +2607,9 @@ async def ui_view_spend_logs( model_group: str | None = fastapi.Query(default=None, description="Filter logs by model group"), key_alias: str | None = fastapi.Query(default=None, description="Filter logs by key alias"), end_user: str | None = fastapi.Query(default=None, description="Filter logs by end user"), - error_code: str | None = fastapi.Query(default=None, description="Filter logs by error code (e.g., '404', '500')"), + error_code: Annotated[ + str | None, fastapi.Query(description="Filter logs by error code (e.g., '404', '500')") + ] = None, error_message: str | None = fastapi.Query( default=None, description="Filter logs by error message (partial string match)" ), @@ -2936,7 +2965,9 @@ async def ui_view_spend_logs( sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") + sql_conditions.append( + _error_code_condition(f"${p}", error_code, await _error_code_index_is_ready(prisma_client)) + ) sql_params.append(error_code) p += 1 if error_message is not None: diff --git a/tests/e2e/perf/setup/upgrade_worker.py b/tests/e2e/perf/setup/upgrade_worker.py index e69e841814f..52ce14e222e 100644 --- a/tests/e2e/perf/setup/upgrade_worker.py +++ b/tests/e2e/perf/setup/upgrade_worker.py @@ -102,7 +102,7 @@ def main() -> int: UNION (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id DESC LIMIT 32)""").fetchall() ids: Final = tuple(row[0] for row in _STRINGS.validate_python(sample)) assert ids, "Restored baseline has no spend logs" - names: Final = tuple(index.name for index in request_log_indexes.REQUEST_LOG_INDEXES) + names: Final = tuple(index.name for index in request_log_indexes.REQUEST_LOG_INDEXES if index.max_bytes is None) migrations: Final = Path(request_log_indexes.__file__).parent / "migrations" initial: Final = UpgradeState( phase="schema", diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index d06f0b13cff..2fafad97160 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -9,7 +9,7 @@ import threading import time import uuid from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from types import MappingProxyType from typing import Final @@ -45,6 +45,8 @@ API_KEY_INDEX_MIGRATION: Final = "20260823000000_add_spend_logs_api_key_starttim 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" +ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_code_startTime_idx" +START_TIME_INDEX: Final = "LiteLLM_SpendLogs_startTime_idx" PARTITIONED_PARENT_ERROR: Final = 'cannot create index on partitioned table "LiteLLM_SpendLogs" concurrently' ORIGINAL_MIGRATION_SQL: Final = MappingProxyType( { @@ -61,6 +63,7 @@ ORIGINAL_MIGRATION_SQL: Final = MappingProxyType( } ) CALL_ID_INDEX_DEFINITION: Final = next(index for index in REQUEST_LOG_INDEXES if index.name == CALL_ID_INDEX) +ERROR_INDEX_DEFINITION: Final = next(index for index in REQUEST_LOG_INDEXES if index.name == ERROR_INDEX) RELEASES: Final = pytest.mark.parametrize( "release", (API_KEY_INDEX_MIGRATION, CALL_ID_INDEX_MIGRATION), ids=("v1.102.1", "v1.103.0") ) @@ -72,6 +75,7 @@ PARTITIONS: Final = MappingProxyType( ) DEFAULT_PARTITION: Final = "LiteLLM_SpendLogs_pdefault" ROWS_PER_PARTITION: Final = 200 +SEEDED_ROWS: Final = 50 RESOLVERS: Final = pytest.mark.parametrize("use_v2_resolver", (True, False), ids=("v2", "v1")) @@ -259,6 +263,21 @@ def _index_oids(database_url: str) -> Mapping[str, int]: return MappingProxyType({row.name: row.oid for row in rows}) +@dataclass(frozen=True, slots=True) +class _PredicateRow: + predicate: "str | None" + + +def _index_predicate(database_url: str, index: str) -> "str | None": + with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_PredicateRow)) as cursor: + row: Final = cursor.execute( + "SELECT pg_get_expr(indpred, indrelid) AS predicate FROM pg_index WHERE indexrelid = to_regclass(%s)", + (f'"{index}"',), + ).fetchone() + assert row is not None, index + return row.predicate + + def _migration_job(use_v2_resolver: bool) -> bool: return ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) @@ -288,13 +307,14 @@ def _assert_index_covers_every_partition(database_url: str, parent_index: str, s @requires_db @RESOLVERS @RELEASES -def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_a_rerun_is_idempotent( +def test_a_partitioned_spend_logs_upgrade_builds_every_index_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") + _assert_index_covers_every_partition(partitioned_database, ERROR_INDEX, "error_code_startTime_idx") ledger: Final = _assert_every_ledger_row_is_finished(partitioned_database) _assert_no_pending_migrations(partitioned_database) oids: Final = _index_oids(partitioned_database) @@ -304,14 +324,11 @@ def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_ '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 + inherited: Final = { + index: "LiteLLM_SpendLogs_p2026_10" in {table for table, _ in _attached_children(partitioned_database, index)} + for index in (API_KEY_INDEX, CALL_ID_INDEX, ERROR_INDEX) + } + assert all(inherited.values()), inherited assert _migration_job(use_v2_resolver) is True assert _ledger(partitioned_database) == ledger @@ -321,7 +338,7 @@ def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_ @requires_db @RESOLVERS @RELEASES -def test_a_plain_spend_logs_upgrade_builds_both_indexes_and_a_second_job_run_rebuilds_nothing( +def test_a_plain_spend_logs_upgrade_builds_every_index_and_a_second_job_run_rebuilds_nothing( scratch_database: str, use_v2_resolver: bool ) -> None: with psycopg.connect(scratch_database, autocommit=True) as conn: @@ -332,6 +349,7 @@ def test_a_plain_spend_logs_upgrade_builds_both_indexes_and_a_second_job_run_reb 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 _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: True} _assert_every_ledger_row_is_finished(scratch_database) _assert_no_pending_migrations(scratch_database) oids: Final = _index_oids(scratch_database) @@ -359,7 +377,8 @@ def test_a_database_that_applied_the_original_migration_files_sees_no_pending_mi assert _migration_job(use_v2_resolver) is True assert _ledger(scratch_database) == before - assert _index_oids(scratch_database) == oids + after: Final = _index_oids(scratch_database) + assert {name: after[name] for name in oids} == oids @requires_db @@ -540,22 +559,29 @@ def test_an_invalid_child_index_left_by_an_interrupted_build_is_rebuilt_and_atta @requires_db -def test_partitioning_a_table_the_index_job_already_covered_moves_its_call_id_index_aside( +def test_partitioning_a_table_the_index_job_already_covered_moves_its_indexes_aside( scratch_database: str, ) -> None: """The runbook renames the Prisma indexes off the legacy table so the partitioned table can take their - names. The call-id index is built by the index job, not Prisma, and a deployment that partitions after - that job ran kept the name on the legacy table, so every later job run reported it as on another table.""" + names. The call-id and error indexes are built by the index job, not Prisma, and a deployment that partitions + after that job ran kept their names on the legacy table, so every later job run reported them as on another + table.""" assert ensure_request_log_indexes(scratch_database, "public") is True assert _indexed_table(scratch_database, CALL_ID_INDEX) == "LiteLLM_SpendLogs" + assert _indexed_table(scratch_database, ERROR_INDEX) == "LiteLLM_SpendLogs" _partition_spend_logs(scratch_database) assert ( _indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_litellm_call_id_idx") == "LiteLLM_SpendLogs_legacy" ) + assert ( + _indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_error_code_startTime_idx") + == "LiteLLM_SpendLogs_legacy" + ) assert ensure_request_log_indexes(scratch_database, "public") is True _assert_index_covers_every_partition(scratch_database, CALL_ID_INDEX, "litellm_call_id_idx") + _assert_index_covers_every_partition(scratch_database, ERROR_INDEX, "error_code_startTime_idx") @requires_db @@ -883,6 +909,154 @@ def test_a_hand_built_child_index_under_another_name_is_renamed_and_attached(par assert "p2026_08_call_id_by_hand" not in oids and oids[f"{partition}_litellm_call_id_idx"] == hand_built +@requires_db +def test_the_partial_error_index_is_built_beside_the_plain_start_time_index_never_renamed_from_it( + scratch_database: str, +) -> None: + start_time_index: Final = _index_oids(scratch_database)[START_TIME_INDEX] + + assert ensure_request_log_indexes(scratch_database, "public") is True + + oids: Final = _index_oids(scratch_database) + assert oids[START_TIME_INDEX] == start_time_index + assert oids[ERROR_INDEX] != start_time_index + assert _index_predicate(scratch_database, START_TIME_INDEX) is None + assert _index_predicate(scratch_database, ERROR_INDEX) is not None + + +def _seed_spend_logs(database_url: str, partitioned: bool) -> None: + """Seed the spend logs, partitioned or not, and never ANALYZE them, as in a freshly restored database.""" + if partitioned: + _partition_spend_logs(database_url) + with psycopg.connect(database_url, autocommit=True) as conn: + for n in range(SEEDED_ROWS): + _insert_spend_log(conn, f"capped-{n}", "2026-10-01") + + +def _run_the_logged_statements(database_url: str, messages: list[str]) -> None: + logged: Final = next(message for message in messages if "Not building index" in message) + with psycopg.connect(database_url, autocommit=True) as conn: + for statement in logged.split("finds it in place:\n", 1)[1].splitlines(): + conn.execute(statement.encode()) + + +@requires_db +@pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned")) +def test_an_index_on_a_table_over_its_size_cap_is_left_for_an_operator_without_failing_the_build( + scratch_database: str, partitioned: bool, caplog: pytest.LogCaptureFixture +) -> None: + """A metadata index takes time to build in proportion to the table's size on disk, so on a table over its cap + the migration job must finish without it instead of running past its deadline and failing the upgrade.""" + _seed_spend_logs(scratch_database, partitioned) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1) + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + assert f"Not building index {ERROR_INDEX}" in caplog.text + assert f"in place except {ERROR_INDEX}" in caplog.text + assert _indexed_table(scratch_database, ERROR_INDEX) is None + + assert ensure_request_log_indexes(scratch_database, "public", (ERROR_INDEX_DEFINITION,)) is True + assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True + + +@requires_db +@pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned")) +def test_the_statements_logged_for_an_index_over_its_cap_build_it_and_the_next_build_finds_it_in_place( + scratch_database: str, partitioned: bool, caplog: pytest.LogCaptureFixture +) -> None: + _seed_spend_logs(scratch_database, partitioned) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1) + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + _run_the_logged_statements(scratch_database, caplog.messages) + hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX] + caplog.clear() + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + assert "Not building index" not in caplog.text + assert _index_oids(scratch_database)[ERROR_INDEX] == hand_built + if partitioned: + _assert_index_covers_every_partition(scratch_database, ERROR_INDEX, "error_code_startTime_idx") + else: + assert _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: True} + + +def _leave_an_invalid_error_index(database_url: str, name: str, table: 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'") + prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY {} ON {} ").format( + sql.Identifier(name), sql.Identifier(table) + ) + with pytest.raises(psycopg.errors.QueryCanceled): + builder.execute(prefix.as_bytes(builder) + ERROR_INDEX_DEFINITION.body.encode()) + + +@requires_db +def test_an_interrupted_build_of_an_index_over_its_cap_is_left_to_the_logged_reindex_instead_of_repaired_at_boot( + scratch_database: str, caplog: pytest.LogCaptureFixture +) -> None: + _seed_spend_logs(scratch_database, partitioned=False) + _leave_an_invalid_error_index(scratch_database, ERROR_INDEX, "LiteLLM_SpendLogs") + capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1) + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: False} + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + _run_the_logged_statements(scratch_database, caplog.messages) + + assert _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: True} + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + + +@requires_db +def test_a_partitioned_build_interrupted_between_partitions_over_its_cap_is_finished_by_the_logged_statements( + scratch_database: str, caplog: pytest.LogCaptureFixture +) -> None: + """A partitioned build that stopped with one child attached, one left invalid and one never started must be + finished by the logged statements, and the boot-time repair must leave the invalid child to them.""" + _seed_spend_logs(scratch_database, partitioned=True) + suffix: Final = "error_code_startTime_idx" + attached, interrupted = PARTITIONS + with psycopg.connect(scratch_database, autocommit=True) as conn: + parent: Final = sql.SQL("CREATE INDEX {} ON ONLY {} ").format( + sql.Identifier(ERROR_INDEX), sql.Identifier("LiteLLM_SpendLogs") + ) + conn.execute(parent.as_bytes(conn) + ERROR_INDEX_DEFINITION.body.encode()) + child: Final = sql.SQL("CREATE INDEX CONCURRENTLY {} ON {} ").format( + sql.Identifier(f"{attached}_{suffix}"), sql.Identifier(attached) + ) + conn.execute(child.as_bytes(conn) + ERROR_INDEX_DEFINITION.body.encode()) + conn.execute( + sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format( + sql.Identifier(ERROR_INDEX), sql.Identifier(f"{attached}_{suffix}") + ) + ) + _leave_an_invalid_error_index(scratch_database, f"{interrupted}_{suffix}", interrupted) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1) + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_database, suffix) == { + ERROR_INDEX: False, + f"{attached}_{suffix}": True, + f"{interrupted}_{suffix}": False, + } + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + _run_the_logged_statements(scratch_database, caplog.messages) + caplog.clear() + + _assert_index_covers_every_partition(scratch_database, ERROR_INDEX, suffix) + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + assert "Not building index" not in caplog.text + + @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: @@ -922,7 +1096,7 @@ def test_a_valid_partitioned_parent_index_under_another_name_is_renamed_with_its 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 + assert not [line for line in caplog.messages if line.startswith("Building index") and "litellm_call_id" in line] 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 diff --git a/tests/integration/database/test_request_log_indexes_at_boot.py b/tests/integration/database/test_request_log_indexes_at_boot.py index e72477de901..92c3b3f62ef 100644 --- a/tests/integration/database/test_request_log_indexes_at_boot.py +++ b/tests/integration/database/test_request_log_indexes_at_boot.py @@ -43,7 +43,8 @@ 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") +ORIGINAL_SPEND_LOGS_INDEXES: Final = ("LiteLLM_SpendLogs_api_key_startTime_idx", "LiteLLM_SpendLogs_litellm_call_id_idx") +SPEND_LOGS_INDEXES: Final = (*ORIGINAL_SPEND_LOGS_INDEXES, "LiteLLM_SpendLogs_error_code_startTime_idx") POPULATED_PARTITIONS: Final = MappingProxyType( { "LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"), @@ -254,9 +255,9 @@ def assert_ready(booted_gateway: Gateway) -> None: assert readiness.json()["db"] == "connected", readiness.text -def assert_both_indexes_cover_every_partition(database_url: str) -> None: - """Every populated partition's index is attached and valid, both parents are valid, and - a partition created afterwards inherits both indexes.""" +def assert_every_index_covers_every_partition(database_url: str) -> None: + """Every populated partition's index is attached and valid, every parent is valid, and + a partition created afterwards inherits every index.""" populated: Final = (*POPULATED_PARTITIONS, DEFAULT_PARTITION) assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS} assert attached_partition_indexes(database_url) == expected_attachments(populated) @@ -285,7 +286,7 @@ def assert_the_serving_proxy_boots_and_finds_the_indexes_in_place( @RESOLVERS -def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_both_indexes_per_partition( +def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_every_index_per_partition( gateway: Gateway, tmp_path: Path, resolver: Resolver ) -> None: with scratch_database() as database_url: @@ -294,7 +295,7 @@ def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_bot job: Final = migration_job(database_url, resolver) assert job.returncode == 0, job.stdout + job.stderr assert INDEXES_IN_PLACE in job.stderr + job.stdout, job.stdout + job.stderr - assert_both_indexes_cover_every_partition(database_url) + assert_every_index_covers_every_partition(database_url) assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver) @@ -308,7 +309,7 @@ def test_the_migration_job_heals_a_partitioned_table_left_with_the_failed_call_i fail_the_call_id_index_migration_like_the_shipped_release(database_url, tmp_path) job: Final = migration_job(database_url, resolver) assert job.returncode == 0, job.stdout + job.stderr - assert_both_indexes_cover_every_partition(database_url) + assert_every_index_covers_every_partition(database_url) assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver) @@ -321,23 +322,25 @@ def test_the_migration_job_leaves_a_plain_table_that_applied_the_original_index_ 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 + assert set(ORIGINAL_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 + output: Final = job.stdout + job.stderr + assert job.returncode == 0, output + assert INDEXES_IN_PLACE in output, output + assert not any(f"Building index {name} " in output for name in ORIGINAL_SPEND_LOGS_INDEXES), output assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS} - assert index_oids(database_url) == before + after: Final = index_oids(database_url) + assert {name: after[name] for name in before} == 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( +def test_a_serving_proxy_that_runs_the_migrations_itself_builds_every_index_after_it_is_ready( gateway: Gateway, tmp_path: Path, resolver: Resolver ) -> None: """A deployment that runs migrate deploy from the serving proxy and never runs the migration job answers readiness with the inert files applied, then its background build - puts both indexes on every partition.""" + puts every index on every partition.""" with scratch_database() as database_url: deploy_schema_before_the_index_migrations(database_url, tmp_path) partition_spend_logs(database_url) @@ -353,13 +356,13 @@ def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_indexes_aft seconds=BUILD_SECONDS, ) assert "Building index" in log and "Attached index" in log, log[-4000:] - assert_both_indexes_cover_every_partition(database_url) + assert_every_index_covers_every_partition(database_url) -def test_the_cli_run_as_a_migration_job_builds_both_indexes_before_it_exits(gateway: Gateway, tmp_path: Path) -> None: +def test_the_cli_run_as_a_migration_job_builds_every_index_before_it_exits(gateway: Gateway, tmp_path: Path) -> None: with scratch_database() as database_url: deploy_schema_before_the_index_migrations(database_url, tmp_path) partition_spend_logs(database_url) job: Final = migration_cli(database_url, gateway, V2) assert job.returncode == 0, job.stdout + job.stderr - assert_both_indexes_cover_every_partition(database_url) + assert_every_index_covers_every_partition(database_url) diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py new file mode 100644 index 00000000000..8a8a7377240 --- /dev/null +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -0,0 +1,214 @@ +import secrets +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +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, list_value, object_value, string_value +from integration._support.database import scratch_database +from integration._support.process import MIGRATE_DEPLOY, owned_proxy_process +from psycopg.rows import class_row +from psycopg.types.json import Jsonb + +pytestmark: Final = pytest.mark.timeout(300) + +ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_code_startTime_idx" +HAND_BUILT_INDEX: Final = "spend_logs_error_code_by_hand_idx" +INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" +SUCCESS_ROWS: Final = 5000 +STATS_FLUSH_WINDOW_SECONDS: Final = 11.0 + +LONG_CODE_PREFIX: Final = secrets.token_hex(64) +LONG_CODE: Final = LONG_CODE_PREFIX + secrets.token_hex(1500) +LONG_CODE_SIBLING: Final = LONG_CODE_PREFIX + secrets.token_hex(1500) + + +@dataclass(frozen=True, slots=True) +class ErrorRow: + request_id: str + status: str + error_code: str + error_message: str + + +@dataclass(frozen=True, slots=True) +class _ScanRow: + scans: int + + +ERROR_ROWS: Final = ( + ErrorRow("failed-500-overloaded", "failure", "500", "upstream overloaded"), + ErrorRow("failed-500-timeout", "failure", "500", "upstream timed out"), + ErrorRow("failed-429", "failure", "429", "rate limited by upstream"), + ErrorRow("client-disconnected", "success", "499", "Client disconnected the request"), + ErrorRow("failed-long-code", "failure", LONG_CODE, "provider error with a long code"), + ErrorRow("failed-long-code-sibling", "failure", LONG_CODE_SIBLING, "provider error with a long code"), +) +ERROR_CODES: Final = MappingProxyType( + { + "500": frozenset({"failed-500-overloaded", "failed-500-timeout"}), + "499": frozenset({"client-disconnected"}), + LONG_CODE: frozenset({"failed-long-code"}), + LONG_CODE_PREFIX: frozenset(), + LONG_CODE_PREFIX[:-1]: frozenset(), + } +) +SHORT_ERROR_CODES: Final = ("500", "499") +SESSION_VIEWS: Final = ("false", "true") + + +def _seed(database_url: str, started: datetime, error_rows: tuple[ErrorRow, ...]) -> None: + utc: Final = started.astimezone(timezone.utc).replace(tzinfo=None) + with psycopg.connect(database_url, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime", metadata, status) ' + "SELECT 'ok-' || n, 'acompletion', %(started)s + n * interval '1 millisecond', " + "%(started)s + n * interval '1 millisecond', %(metadata)s, 'success' " + "FROM generate_series(1, %(rows)s) AS n", + {"started": utc, "metadata": Jsonb({"error_information": None}), "rows": SUCCESS_ROWS}, + ) + with connection.cursor() as cursor: + cursor.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime", metadata, status) ' + "VALUES (%s, 'acompletion', %s, %s, %s, %s)", + tuple( + ( + row.request_id, + utc, + utc, + Jsonb( + { + "error_information": { + "error_code": row.error_code, + "error_message": row.error_message, + "error_class": "APIError", + } + } + ), + row.status, + ) + for row in error_rows + ), + ) + connection.execute('ANALYZE "LiteLLM_SpendLogs"') + + +def _index_scans(database_url: str, index: str) -> int: + with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_ScanRow)) as cursor: + row: Final = cursor.execute( + "SELECT idx_scan AS scans FROM pg_stat_user_indexes WHERE indexrelname = %s", (index,) + ).fetchone() + assert row is not None, f"{index} was not built" + return row.scans + + +def _settled_index_scans(database_url: str, index: str, previous: int, unchanged_since: float) -> int: + changed: Final = eventually( + lambda: _index_scans(database_url, index) != previous, + lambda drifted: drifted, + seconds=STATS_FLUSH_WINDOW_SECONDS - (time.monotonic() - unchanged_since), + return_last_on_timeout=True, + ) + if not changed: + return previous + return _settled_index_scans(database_url, index, _index_scans(database_url, index), time.monotonic()) + + +def _replace_the_error_index_with_one_on_the_full_error_code(database_url: str) -> None: + with psycopg.connect(database_url, autocommit=True) as connection: + connection.execute(f'DROP INDEX "{ERROR_INDEX}"') + connection.execute( + f'CREATE INDEX "{HAND_BUILT_INDEX}" ON "LiteLLM_SpendLogs" ' + "((metadata->'error_information'->>'error_code'), \"startTime\")" + ) + connection.execute('ANALYZE "LiteLLM_SpendLogs"') + + +def _listed(gateway: Gateway, window: Mapping[str, str], query: Mapping[str, str]) -> frozenset[str]: + page: Final = gateway.get("/spend/logs/ui", params={**window, **query}) + return frozenset(string_value(object_value(row)["request_id"]) for row in list_value(page["data"])) + + +def _list_each_filter( + gateway: Gateway, + database_url: str, + index: str, + window: Mapping[str, str], + filters: tuple[tuple[str, str], ...], + scans: int, +) -> tuple[tuple[tuple[str, str], frozenset[str], int], ...]: + if not filters: + return () + session, code = filters[0] + listed: Final = _listed(gateway, window, {"error_code": code, "group_by_session": session}) + settled: Final = _settled_index_scans(database_url, index, scans, time.monotonic()) + return ( + ((session, code), listed, settled - scans), + *_list_each_filter(gateway, database_url, index, window, filters[1:], settled), + ) + + +def _no_change(_: str) -> None: + return None + + +@dataclass(frozen=True, slots=True) +class IndexCase: + index: str + error_rows: tuple[ErrorRow, ...] + codes: tuple[str, ...] + prepare: Callable[[str], None] + + +INDEX_CASES: Final = ( + IndexCase(ERROR_INDEX, ERROR_ROWS, tuple(ERROR_CODES), _no_change), + IndexCase( + HAND_BUILT_INDEX, + tuple(row for row in ERROR_ROWS if row.error_code in ("500", "429", "499")), + SHORT_ERROR_CODES, + _replace_the_error_index_with_one_on_the_full_error_code, + ), +) + + +def _filter_through_index( + gateway: Gateway, tmp_path: Path, case: IndexCase, filters: tuple[tuple[str, str], ...] +) -> tuple[tuple[tuple[str, str], frozenset[str], int], ...]: + now: Final = datetime.now(timezone.utc) + window: Final = { + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + } + with scratch_database() as database_url: + with owned_proxy_process( + gateway, tmp_path, {"DATABASE_URL": database_url}, database_setup=MIGRATE_DEPLOY + ) as booted: + eventually( + lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=60 + ) + _seed(database_url, now - timedelta(hours=1), case.error_rows) + case.prepare(database_url) + scans: Final = _settled_index_scans( + database_url, case.index, _index_scans(database_url, case.index), time.monotonic() + ) + return _list_each_filter(booted.gateway, database_url, case.index, window, filters, scans) + + +@pytest.mark.parametrize("case", INDEX_CASES, ids=("built_by_the_index_job", "built_by_hand_on_the_full_error_code")) +def test_error_code_filters_read_only_the_rows_that_carry_an_error_code( + gateway: Gateway, tmp_path: Path, case: IndexCase +) -> None: + filters: Final = tuple(product(SESSION_VIEWS, case.codes)) + + results: Final = _filter_through_index(gateway, tmp_path, case, filters) + + assert {combo: listed for combo, listed, _ in results} == { + (session, code): ERROR_CODES[code] for session, code in filters + } + assert {combo: index_scans > 0 for combo, _, index_scans in results} == {combo: True for combo in filters} diff --git a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py index 5cf7593c5cd..89236d79625 100644 --- a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py +++ b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py @@ -34,10 +34,10 @@ def _prisma_indexes_of(schema: str, model: str) -> frozenset[str]: class TestTheIndexList: - def test_every_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self): + def test_every_plain_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 + assert (index.predicate is None) == (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): diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 6598b94e60f..ec324b51b5b 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7,6 +7,7 @@ import re import sqlite3 from collections.abc import Mapping from datetime import timezone +from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -124,7 +125,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): gte = re.search(r'"startTime" >= \(\$(\d+)', cond) lte = re.search(r'"startTime" <= \(\$(\d+)', cond) alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) - code = re.search(r"error_code' = \$(\d+)", cond) + code = re.search(r"error_code'(?:, 128\))? = \$(\d+)", cond) msg = re.search(r"error_message' LIKE \$(\d+)", cond) credential = re.fullmatch(r"metadata->>'used_client_oauth_token' = \$(\d+)", cond) sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond) @@ -207,7 +208,9 @@ def _reconstruct_ui_where_from_sql(sql_query, params): return where -def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None): +def make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None, error_code_index_ready: bool = False +): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. @@ -238,6 +241,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No query_observer(sql_query, params) if "mcp_tool_call_count" in sql_query: return [] + if "FROM pg_index" in sql_query: + return [{"ready": error_code_index_ready}] if 'SELECT DISTINCT "user", team_id' in sql_query: return _emulate_spend_log_owner_lookup(mock_spend_logs, sql_query, params) filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) @@ -4697,7 +4702,8 @@ async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client @pytest.mark.asyncio -async def test_ui_view_spend_logs_with_error_code(client): +@pytest.mark.parametrize("error_code_index_ready", [True, False], ids=["with_the_index", "without_the_index"]) +async def test_ui_view_spend_logs_with_error_code(client: TestClient, error_code_index_ready: bool): """Test filtering spend logs by error code""" mock_spend_logs = [ { @@ -4743,7 +4749,9 @@ async def test_ui_view_spend_logs_with_error_code(client): with patch.object( ps, "prisma_client", - make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code), + make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_by_error_code, error_code_index_ready=error_code_index_ready + ), ): start_date, end_date = _default_date_range() @@ -4770,6 +4778,54 @@ async def test_ui_view_spend_logs_with_error_code(client): app.dependency_overrides.pop(ps.user_api_key_auth, None) + +@pytest.mark.asyncio +@pytest.mark.parametrize("error_code_index_ready", [True, False], ids=["with_the_index", "without_the_index"]) +async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code( + client: TestClient, error_code_index_ready: bool +) -> None: + prefix: Final = "E" * 128 + codes: Final = MappingProxyType({f"log-{suffix}": prefix + suffix for suffix in ("first", "second")}) + mock_spend_logs: Final = [ + { + "id": log_id, + "request_id": f"req-{log_id}", + "api_key": "sk-test-key", + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + "metadata": json.dumps({"error_information": {"error_code": code}}), + } + for log_id, code in codes.items() + ] + + def filter_by_exact_error_code(where: Mapping[str, Mapping[str, object]]) -> list[dict[str, str | float]]: + wanted: Final = str(where.get("metadata", {}).get("equals", "")).strip('"') + return [log for log in mock_spend_logs if codes[str(log["id"])] == wanted] + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + with patch.object( + ps, + "prisma_client", + make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_by_exact_error_code, error_code_index_ready=error_code_index_ready + ), + ): + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"error_code": prefix + "first", "start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [row["id"] for row in response.json()["data"]] == ["log-first"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + @pytest.mark.asyncio async def test_ui_view_spend_logs_with_error_message(client): """Test filtering spend logs by error message"""