From d0b3ecf4ad273e1e845cea1ed8b3dbfbc3fcc850 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:25:28 -0700 Subject: [PATCH 1/9] perf(spend-logs): index error rows so the Logs page error filters stop scanning every row --- db_scripts/partition_spend_logs.sql | 2 + .../request_log_indexes.py | 36 +++-- .../spend_management_endpoints.py | 3 + .../migrations/test_request_log_indexes.py | 68 +++++++-- .../test_request_log_indexes_at_boot.py | 37 ++--- .../spend/test_spend_log_error_filters.py | 129 ++++++++++++++++++ .../test_request_log_indexes.py | 4 +- 7 files changed, 236 insertions(+), 43 deletions(-) create mode 100644 tests/integration/spend/test_spend_log_error_filters.py diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql index 750c1a934aa..34dbcabedcf 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_startTime_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_error_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..d34c4d18c95 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 @@ -33,20 +35,27 @@ 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
`, and the WHERE predicate of a + partial index.""" table: str name: str definition: str + predicate: str | 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 +68,12 @@ 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_startTime_idx", + '("startTime")', + "jsonb_typeof(metadata->'error_information') = 'object'", + ), ) _IDENTIFIER_MAX_BYTES: Final = 63 @@ -256,9 +271,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 +389,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 +458,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/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3c65fc0bf06..d0ec0d072db 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -96,6 +96,7 @@ _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')" +_HAS_ERROR_INFORMATION_SQL: Final = "jsonb_typeof(metadata->'error_information') = 'object'" _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -2932,6 +2933,8 @@ async def ui_view_spend_logs( sql_conditions.append(f"metadata->>'user_api_key_alias' LIKE ${p}") sql_params.append(f"%{key_alias}%") p += 1 + if error_code is not None or error_message is not None: + sql_conditions.append(_HAS_ERROR_INFORMATION_SQL) if error_code is not None: sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") sql_params.append(error_code) diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 1de2090b487..805663ff9fa 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -44,6 +44,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_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( { @@ -258,6 +260,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) @@ -287,13 +304,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_startTime_idx") ledger: Final = _assert_every_ledger_row_is_finished(partitioned_database) _assert_no_pending_migrations(partitioned_database) oids: Final = _index_oids(partitioned_database) @@ -303,14 +321,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 @@ -320,7 +335,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: @@ -331,6 +346,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_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) @@ -358,7 +374,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 @@ -539,22 +556,28 @@ 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_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_startTime_idx") @requires_db @@ -882,6 +905,21 @@ 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 + + @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: @@ -921,7 +959,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..eec9819770c 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_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..85ed605c393 --- /dev/null +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -0,0 +1,129 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from itertools import product +from pathlib import Path +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(240) + +ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_startTime_idx" +INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" +SUCCESS_ROWS: Final = 5000 + + +@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"), +) +FILTERS: Final = ( + ("error_code", "500", frozenset({"failed-500-overloaded", "failed-500-timeout"})), + ("error_code", "499", frozenset({"client-disconnected"})), + ("error_message", "upstream", frozenset({"failed-500-overloaded", "failed-500-timeout", "failed-429"})), +) +SESSION_VIEWS: Final = ("false", "true") + + +def _seed(database_url: str, started: datetime) -> 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 _error_index_scans(database_url: 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", (ERROR_INDEX,) + ).fetchone() + assert row is not None, f"{ERROR_INDEX} was not built" + return row.scans + + +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 test_error_filters_read_only_the_rows_that_carry_error_information(gateway: Gateway, tmp_path: Path) -> None: + 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)) + scans_before: Final = _error_index_scans(database_url) + + listed: Final = { + (session, name, value): _listed( + booted.gateway, window, {name: value, "group_by_session": session} + ) + for session, (name, value, _) in product(SESSION_VIEWS, FILTERS) + } + assert listed == { + (session, name, value): rows for session, (name, value, rows) in product(SESSION_VIEWS, FILTERS) + } + eventually( + lambda: _error_index_scans(database_url), + lambda scans: scans - scans_before >= len(listed), + seconds=30, + ) 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): From 54dd6d2cbf8668d5f218775803b4e7e26631e02e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:19:37 -0700 Subject: [PATCH 2/9] perf(spend-logs): index the error code with the time so the error code filter reads only matching rows --- db_scripts/partition_spend_logs.sql | 4 ++-- .../litellm_proxy_extras/request_log_indexes.py | 4 ++-- .../database/migrations/test_request_log_indexes.py | 10 +++++----- .../database/test_request_log_indexes_at_boot.py | 2 +- .../integration/spend/test_spend_log_error_filters.py | 2 +- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql index 34dbcabedcf..28107f28043 100644 --- a/db_scripts/partition_spend_logs.sql +++ b/db_scripts/partition_spend_logs.sql @@ -57,8 +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_startTime_idx" - RENAME TO "LiteLLM_SpendLogs_legacy_error_startTime_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 d34c4d18c95..6a261a22d8f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -70,8 +70,8 @@ REQUEST_LOG_INDEXES: Final = ( RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'), RequestLogIndex( "LiteLLM_SpendLogs", - "LiteLLM_SpendLogs_error_startTime_idx", - '("startTime")', + "LiteLLM_SpendLogs_error_code_startTime_idx", + "((metadata->'error_information'->>'error_code'), \"startTime\")", "jsonb_typeof(metadata->'error_information') = 'object'", ), ) diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 805663ff9fa..5ab0e954890 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -44,7 +44,7 @@ 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_startTime_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( @@ -311,7 +311,7 @@ def test_a_partitioned_spend_logs_upgrade_builds_every_index_per_partition_and_a _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_startTime_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) @@ -346,7 +346,7 @@ def test_a_plain_spend_logs_upgrade_builds_every_index_and_a_second_job_run_rebu 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_startTime_idx") == {ERROR_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) @@ -573,11 +573,11 @@ def test_partitioning_a_table_the_index_job_already_covered_moves_its_indexes_as _indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_litellm_call_id_idx") == "LiteLLM_SpendLogs_legacy" ) assert ( - _indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_error_startTime_idx") == "LiteLLM_SpendLogs_legacy" + _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_startTime_idx") + _assert_index_covers_every_partition(scratch_database, ERROR_INDEX, "error_code_startTime_idx") @requires_db 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 eec9819770c..92c3b3f62ef 100644 --- a/tests/integration/database/test_request_log_indexes_at_boot.py +++ b/tests/integration/database/test_request_log_indexes_at_boot.py @@ -44,7 +44,7 @@ INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" INDEX_BUILD_LINES: Final = ("Building index", "Attached index") BUILD_SECONDS: Final = 60 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_startTime_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"), diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index 85ed605c393..04f20bb115f 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -15,7 +15,7 @@ from psycopg.types.json import Jsonb pytestmark: Final = pytest.mark.timeout(240) -ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_startTime_idx" +ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_code_startTime_idx" INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" SUCCESS_ROWS: Final = 5000 From 3bc7c5f7852ab9a2b3cc585eb688da5e32d97257 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:21:04 -0700 Subject: [PATCH 3/9] fix(spend-logs): index error codes with an implied predicate and cap the build size The error code index now covers rows where the error code is not null, which the unchanged filter query implies, so the query stays the same as main and a proxy without the index runs the plan it runs today. The message filter change is reverted. The index build reads every row's metadata, so above about 5M estimated rows the index job skips it with a warning that gives the CREATE INDEX statement, keeping the Helm migration job inside its deadline. An index built by hand under that name is found in place. --- .../request_log_indexes.py | 56 +++++++++++++-- .../spend_management_endpoints.py | 3 - .../migrations/test_request_log_indexes.py | 49 ++++++++++++- .../spend/test_spend_log_error_filters.py | 71 +++++++++++++------ 4 files changed, 149 insertions(+), 30 deletions(-) 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 6a261a22d8f..c7d9193a32d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -36,13 +36,17 @@ if TYPE_CHECKING: @dataclass(frozen=True, slots=True) class RequestLogIndex: """One index the migration job owns: the table, the exact Prisma index name, the - column list as it would be written after `ON
`, and the WHERE predicate of a - partial index.""" + column list as it would be written after `ON
`, the WHERE predicate of a + partial index, and the most rows the build takes on by itself. An index whose build + has to read each row's metadata costs far more per row than a column index, so above + that estimate 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_rows: int | None = None @property def columns(self) -> tuple[str, ...]: @@ -72,7 +76,8 @@ REQUEST_LOG_INDEXES: Final = ( "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_error_code_startTime_idx", "((metadata->'error_information'->>'error_code'), \"startTime\")", - "jsonb_typeof(metadata->'error_information') = 'object'", + "(metadata->'error_information'->>'error_code') IS NOT NULL", + max_rows=5_000_000, ), ) @@ -88,6 +93,10 @@ _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)" +_ESTIMATED_ROWS_SQL: Final = ( + "SELECT COALESCE(SUM(GREATEST(c.reltuples, 0)), 0)::bigint AS rows FROM pg_class c WHERE c.relkind = 'r' " + "AND (c.oid = to_regclass(%(table)s) OR c.oid IN (SELECT relid FROM pg_partition_tree(to_regclass(%(table)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 " @@ -143,6 +152,11 @@ class _TableKind: partitioned: bool +@dataclass(frozen=True, slots=True) +class _EstimatedRows: + rows: int + + 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.""" @@ -176,8 +190,9 @@ 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 over its + `max_rows` is logged and left for an operator; the result says whether every + other index ended up valid. Never raises.""" import psycopg try: @@ -238,11 +253,42 @@ def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: if table is None: logger.info("Table %s does not exist yet, skipping index %s", index.table, index.name) return True + if _too_large_to_build(connection, schema, index): + 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 _too_large_to_build( + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex +) -> bool: + from psycopg.rows import class_row + + if index.max_rows is None: + return False + existing: Final = _index_state(connection, schema, index.name) + if existing is not None and existing.valid: + return False + with connection.cursor(row_factory=class_row(_EstimatedRows)) as cursor: + rows: Final = sum( + estimate.rows + for estimate in cursor.execute(_ESTIMATED_ROWS_SQL, {"table": _regclass_name(connection, schema, index.table)}) + ) + if rows <= index.max_rows: + return False + logger.warning( + "Not building index %s on %s: it reads about %s rows, more than the %s it builds by itself. Build it off-peak " + "under that name as %s, and the next index build finds it in place", + index.name, + index.table, + rows, + index.max_rows, + index.body, + ) + return True + + def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str: from psycopg import sql diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index d0ec0d072db..3c65fc0bf06 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -96,7 +96,6 @@ _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')" -_HAS_ERROR_INFORMATION_SQL: Final = "jsonb_typeof(metadata->'error_information') = 'object'" _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -2933,8 +2932,6 @@ async def ui_view_spend_logs( sql_conditions.append(f"metadata->>'user_api_key_alias' LIKE ${p}") sql_params.append(f"%{key_alias}%") p += 1 - if error_code is not None or error_message is not None: - sql_conditions.append(_HAS_ERROR_INFORMATION_SQL) if error_code is not None: sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") sql_params.append(error_code) diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 5ab0e954890..6d28ca0193b 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -8,7 +8,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 @@ -62,6 +62,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") ) @@ -73,6 +74,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")) @@ -920,6 +922,51 @@ def test_the_partial_error_index_is_built_beside_the_plain_start_time_index_neve assert _index_predicate(scratch_database, ERROR_INDEX) is not None +def _seed_analyzed_spend_logs(database_url: str, partitioned: bool) -> int: + """Seed the spend logs, partitioned or not, ANALYZE them, and return how many rows they hold.""" + 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") + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + return SEEDED_ROWS + (len(PARTITIONS) + 1) * ROWS_PER_PARTITION if partitioned else SEEDED_ROWS + + +@requires_db +@pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned")) +def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_build( + scratch_database: str, partitioned: bool, caplog: pytest.LogCaptureFixture +) -> None: + """A metadata index costs far more per row to build than a column index, so on a table over its cap the + migration job must finish without it instead of running past its deadline and failing the upgrade.""" + rows: Final = _seed_analyzed_spend_logs(scratch_database, partitioned) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=rows - 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 _indexed_table(scratch_database, ERROR_INDEX) is None + + assert ensure_request_log_indexes(scratch_database, "public", (replace(capped, max_rows=rows),)) is True + assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True + + +@requires_db +def test_an_index_an_operator_built_by_hand_on_a_table_over_the_cap_is_found_in_place( + scratch_database: str, caplog: pytest.LogCaptureFixture +) -> None: + _seed_analyzed_spend_logs(scratch_database, partitioned=False) + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute(f'CREATE INDEX CONCURRENTLY "{ERROR_INDEX}" ON "LiteLLM_SpendLogs" {ERROR_INDEX_DEFINITION.body}') + hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX] + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (replace(ERROR_INDEX_DEFINITION, max_rows=1),)) + assert "Not building index" not in caplog.text + assert _index_oids(scratch_database)[ERROR_INDEX] == 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: diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index 04f20bb115f..13bd1405fd1 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -1,8 +1,10 @@ +import time from collections.abc import 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 @@ -13,11 +15,12 @@ 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(240) +pytestmark: Final = pytest.mark.timeout(300) ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_code_startTime_idx" INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" SUCCESS_ROWS: Final = 5000 +STATS_FLUSH_WINDOW_SECONDS: Final = 11.0 @dataclass(frozen=True, slots=True) @@ -39,12 +42,14 @@ ERROR_ROWS: Final = ( ErrorRow("failed-429", "failure", "429", "rate limited by upstream"), ErrorRow("client-disconnected", "success", "499", "Client disconnected the request"), ) -FILTERS: Final = ( - ("error_code", "500", frozenset({"failed-500-overloaded", "failed-500-timeout"})), - ("error_code", "499", frozenset({"client-disconnected"})), - ("error_message", "upstream", frozenset({"failed-500-overloaded", "failed-500-timeout", "failed-429"})), +ERROR_CODES: Final = MappingProxyType( + { + "500": frozenset({"failed-500-overloaded", "failed-500-timeout"}), + "499": frozenset({"client-disconnected"}), + } ) SESSION_VIEWS: Final = ("false", "true") +FILTERS: Final = tuple(product(SESSION_VIEWS, ERROR_CODES)) def _seed(database_url: str, started: datetime) -> None: @@ -92,12 +97,42 @@ def _error_index_scans(database_url: str) -> int: return row.scans +def _settled_error_index_scans(database_url: str, previous: int, unchanged_since: float) -> int: + changed: Final = eventually( + lambda: _error_index_scans(database_url) != 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_error_index_scans(database_url, _error_index_scans(database_url), time.monotonic()) + + 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 test_error_filters_read_only_the_rows_that_carry_error_information(gateway: Gateway, tmp_path: Path) -> None: +def _list_each_filter( + gateway: Gateway, + database_url: 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_error_index_scans(database_url, scans, time.monotonic()) + return ( + ((session, code), listed, settled - scans), + *_list_each_filter(gateway, database_url, window, filters[1:], settled), + ) + + +def test_error_code_filters_read_only_the_rows_that_carry_an_error_code(gateway: Gateway, tmp_path: Path) -> None: now: Final = datetime.now(timezone.utc) window: Final = { "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), @@ -111,19 +146,13 @@ def test_error_filters_read_only_the_rows_that_carry_error_information(gateway: lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=60 ) _seed(database_url, now - timedelta(hours=1)) - scans_before: Final = _error_index_scans(database_url) - - listed: Final = { - (session, name, value): _listed( - booted.gateway, window, {name: value, "group_by_session": session} - ) - for session, (name, value, _) in product(SESSION_VIEWS, FILTERS) - } - assert listed == { - (session, name, value): rows for session, (name, value, rows) in product(SESSION_VIEWS, FILTERS) - } - eventually( - lambda: _error_index_scans(database_url), - lambda scans: scans - scans_before >= len(listed), - seconds=30, + scans: Final = _settled_error_index_scans( + database_url, _error_index_scans(database_url), time.monotonic() ) + + results: Final = _list_each_filter(booted.gateway, database_url, window, FILTERS, scans) + + 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} From f4d31a5401840c4bc514b86c767fc144fe6617dd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:35:21 -0700 Subject: [PATCH 4/9] fix(spend-logs): index an error code prefix, cap the auto build at 1M rows, and leave interrupted capped builds to the operator The error code index keys on left(code, 128) so a long error code can no longer fail a spend log write, and the filter compares that prefix next to the exact code. The index job builds it only up to 1M rows, analyzing a never-analyzed table first, logs a runnable CREATE INDEX CONCURRENTLY or REINDEX statement above that, and names the index it left in the summary. The invalid-index repair at boot skips row-capped indexes. --- .../request_log_indexes.py | 107 ++++++++++++++---- .../litellm_proxy_extras/utils.py | 5 +- .../spend_management_endpoints.py | 5 +- .../migrations/test_request_log_indexes.py | 73 ++++++++++-- 4 files changed, 157 insertions(+), 33 deletions(-) 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 c7d9193a32d..0a57018d2f6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -17,7 +17,7 @@ import re import time from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from litellm_proxy_extras._logging import logger from litellm_proxy_extras.migration_data_checks import ( @@ -75,9 +75,9 @@ REQUEST_LOG_INDEXES: Final = ( RequestLogIndex( "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_error_code_startTime_idx", - "((metadata->'error_information'->>'error_code'), \"startTime\")", + "((left(metadata->'error_information'->>'error_code', 128)), \"startTime\")", "(metadata->'error_information'->>'error_code') IS NOT NULL", - max_rows=5_000_000, + max_rows=1_000_000, ), ) @@ -94,7 +94,8 @@ _CREATE_INDEX_STATEMENT: Final = re.compile( ) _TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)" _ESTIMATED_ROWS_SQL: Final = ( - "SELECT COALESCE(SUM(GREATEST(c.reltuples, 0)), 0)::bigint AS rows FROM pg_class c WHERE c.relkind = 'r' " + "SELECT COALESCE(SUM(GREATEST(c.reltuples, 0)), 0)::bigint AS rows, COALESCE(BOOL_OR(c.reltuples < 0), false) " + "AS unanalyzed FROM pg_class c WHERE c.relkind = 'r' " "AND (c.oid = to_regclass(%(table)s) OR c.oid IN (SELECT relid FROM pg_partition_tree(to_regclass(%(table)s))))" ) _CHILDREN_WITHOUT_THE_INDEX_SQL: Final = ( @@ -155,6 +156,10 @@ class _TableKind: @dataclass(frozen=True, slots=True) class _EstimatedRows: rows: int + unanalyzed: bool + + +_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: @@ -198,13 +203,21 @@ def ensure_request_log_indexes( 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 @@ -245,50 +258,98 @@ 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 _too_large_to_build(connection, schema, index): - 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 + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex, partitioned: bool ) -> bool: - from psycopg.rows import class_row - if index.max_rows is None: return False existing: Final = _index_state(connection, schema, index.name) if existing is not None and existing.valid: return False - with connection.cursor(row_factory=class_row(_EstimatedRows)) as cursor: - rows: Final = sum( - estimate.rows - for estimate in cursor.execute(_ESTIMATED_ROWS_SQL, {"table": _regclass_name(connection, schema, index.table)}) - ) + rows: Final = _estimated_rows(connection, schema, index.table) if rows <= index.max_rows: return False logger.warning( "Not building index %s on %s: it reads about %s rows, more than the %s it builds by itself. Build it off-peak " - "under that name as %s, and the next index build finds it in place", + "and the next index build finds it in place: %s", index.name, index.table, rows, index.max_rows, - index.body, + _operator_build_statement( + connection, + schema, + index, + partitioned, + interrupted=existing is not None and existing.table == index.table, + ), ) return True +def _estimated_rows(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> int: + """The planner's row count for the table and its partitions. A table never analyzed counts zero rows however + many it holds (a freshly restored database), so it is analyzed first.""" + from psycopg import sql + + estimate: Final = _planner_estimate(connection, schema, table) + if not estimate.unanalyzed: + return estimate.rows + connection.execute(sql.SQL("ANALYZE {}").format(sql.Identifier(schema, table))) + return _planner_estimate(connection, schema, table).rows + + +def _planner_estimate(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> _EstimatedRows: + from psycopg.rows import class_row + + with connection.cursor(row_factory=class_row(_EstimatedRows)) as cursor: + estimate: Final = cursor.execute( + _ESTIMATED_ROWS_SQL, {"table": _regclass_name(connection, schema, table)} + ).fetchone() + return estimate if estimate is not None else _EstimatedRows(rows=0, unanalyzed=False) + + +def _operator_build_statement( + connection: "psycopg.Connection[tuple[object, ...]]", + schema: str, + index: RequestLogIndex, + partitioned: bool, + interrupted: bool, +) -> str: + """What an operator runs by hand: CONCURRENTLY wherever Postgres allows it, and a REINDEX of the invalid index an + interrupted build left, which a CREATE ... IF NOT EXISTS would silently skip.""" + from psycopg import sql + + if interrupted: + reindex: Final = "REINDEX INDEX {}" if partitioned else "REINDEX INDEX CONCURRENTLY {}" + return sql.SQL(reindex).format(sql.Identifier(schema, index.name)).as_string(connection) + build: Final = ( + "CREATE INDEX IF NOT EXISTS {} ON {} " if partitioned else "CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} " + ) + prefix: Final = sql.SQL(build).format(sql.Identifier(index.name), sql.Identifier(schema, index.table)) + return _create_index_statement(connection, prefix, index.body).decode() + + def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str: from psycopg import sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 528bf7def20..2adf8b22b16 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -74,6 +74,9 @@ _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*$") +_ROW_CAPPED_REQUEST_LOG_INDEXES: Final = frozenset( + index.name for index in REQUEST_LOG_INDEXES if index.max_rows is not None +) _INVALID_LITELLM_INDEXES_SQL: Final = ( "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " "FROM pg_index i " @@ -851,7 +854,7 @@ class ProxyExtrasDBManager: 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) + return tuple(_InvalidIndex(*row) for row in rows if row[1] not in _ROW_CAPPED_REQUEST_LOG_INDEXES) @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 3c65fc0bf06..b98624ff998 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2933,7 +2933,10 @@ 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( + f"left(metadata->'error_information'->>'error_code', 128) = left(${p}, 128) " + f"AND metadata->'error_information'->>'error_code' = ${p}" + ) sql_params.append(error_code) p += 1 if error_message is not None: diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 6d28ca0193b..6dafe67624a 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -575,7 +575,8 @@ def test_partitioning_a_table_the_index_job_already_covered_moves_its_indexes_as _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" + _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") @@ -946,6 +947,7 @@ def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_b 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", (replace(capped, max_rows=rows),)) is True @@ -953,18 +955,73 @@ def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_b @requires_db -def test_an_index_an_operator_built_by_hand_on_a_table_over_the_cap_is_found_in_place( +@pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned")) +def test_the_statement_logged_for_an_index_over_its_cap_builds_it_and_the_next_build_finds_it_in_place( + scratch_database: str, partitioned: bool, caplog: pytest.LogCaptureFixture +) -> None: + _seed_analyzed_spend_logs(scratch_database, partitioned) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=1) + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + logged: Final = next(message for message in caplog.messages if "Not building index" in message) + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute(logged.split("finds it in place: ", 1)[1].encode()) + 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 + assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True + + +def _leave_an_invalid_error_index(database_url: str) -> None: + with _pin_a_snapshot_on(database_url, "LiteLLM_SpendLogs"): + 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(ERROR_INDEX), sql.Identifier("LiteLLM_SpendLogs") + ) + 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_analyzed_spend_logs(scratch_database, partitioned=False) - with psycopg.connect(scratch_database, autocommit=True) as conn: - conn.execute(f'CREATE INDEX CONCURRENTLY "{ERROR_INDEX}" ON "LiteLLM_SpendLogs" {ERROR_INDEX_DEFINITION.body}') - hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX] + _leave_an_invalid_error_index(scratch_database) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=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", (replace(ERROR_INDEX_DEFINITION, max_rows=1),)) - assert "Not building index" not in caplog.text - assert _index_oids(scratch_database)[ERROR_INDEX] == hand_built + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + logged: Final = next(message for message in caplog.messages if "Not building index" in message) + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute(logged.split("finds it in place: ", 1)[1].encode()) + + 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_table_never_analyzed_is_counted_before_deciding_whether_it_is_over_the_cap( + scratch_database: str, caplog: pytest.LogCaptureFixture +) -> None: + with psycopg.connect(scratch_database, autocommit=True) as conn: + for n in range(SEEDED_ROWS): + _insert_spend_log(conn, f"unanalyzed-{n}", "2026-10-01") + capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=SEEDED_ROWS - 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 _indexed_table(scratch_database, ERROR_INDEX) is None @requires_db From 29691d6dddce8d9d4cbe53370b0b00be1975da6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:39:15 -0700 Subject: [PATCH 5/9] fix(spend-logs): match short error codes on the indexed prefix alone A code under 128 characters is its own prefix, so the prefix comparison is already exact and Postgres no longer reads each matching row's metadata to recheck the full code. Codes of 128 characters or more keep the full-code check. The integration test now seeds long codes that share a prefix, and error_code is declared with Annotated so direct calls default to None. --- .../spend_management_endpoints.py | 21 +++++++--- .../spend/test_spend_log_error_filters.py | 10 +++++ .../test_spend_management_endpoints.py | 42 ++++++++++++++++++- 3 files changed, 67 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index b98624ff998..09fc2285f51 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -96,6 +96,9 @@ _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})" _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -226,6 +229,15 @@ class _SessionSpendRow(TypedDict): session_models: ReadOnly[Sequence[str]] +def _error_code_condition(placeholder: str, error_code: str) -> str: + 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 @@ -2577,7 +2589,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)" ), @@ -2933,10 +2947,7 @@ async def ui_view_spend_logs( sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append( - f"left(metadata->'error_information'->>'error_code', 128) = left(${p}, 128) " - f"AND metadata->'error_information'->>'error_code' = ${p}" - ) + sql_conditions.append(_error_code_condition(f"${p}", error_code)) sql_params.append(error_code) p += 1 if error_message is not None: diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index 13bd1405fd1..a1e3b97e5db 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -1,3 +1,4 @@ +import secrets import time from collections.abc import Mapping from dataclasses import dataclass @@ -22,6 +23,10 @@ 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: @@ -41,11 +46,16 @@ ERROR_ROWS: Final = ( 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(), } ) SESSION_VIEWS: Final = ("false", "true") 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 8d53817f51e..ca08863330c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -124,7 +124,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) @@ -4770,6 +4770,46 @@ async def test_ui_view_spend_logs_with_error_code(client): app.dependency_overrides.pop(ps.user_api_key_auth, None) + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client): + prefix = "E" * 128 + mock_spend_logs = [ + { + "id": f"log-{suffix}", + "request_id": f"req-{suffix}", + "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": prefix + suffix}}), + } + for suffix in ("first", "second") + ] + + def filter_by_exact_error_code(where): + wanted = str(where.get("metadata", {}).get("equals", "")).strip('"') + return [log for log in mock_spend_logs if json.loads(log["metadata"])["error_information"]["error_code"] == 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) + ): + 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""" From 2616c71da8c882b2e6d5913d5fe2fa716eb21c38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:33:07 -0700 Subject: [PATCH 6/9] fix(spend-logs): cap the error index build by table size and log non-blocking partitioned steps The index job now skips the error code index when the spend log table and its partitions take more than 4 GiB on disk, since build time follows size, not row count. On a partitioned table the logged statements build each partition concurrently and attach it, and boot-time repair leaves the index and its partition copies to those statements. --- .../request_log_indexes.py | 132 ++++++++++-------- .../litellm_proxy_extras/utils.py | 19 ++- .../migrations/test_request_log_indexes.py | 96 ++++++++----- .../test_spend_management_endpoints.py | 22 +-- 4 files changed, 165 insertions(+), 104 deletions(-) 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 0a57018d2f6..2a928db50ec 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -17,6 +17,7 @@ import re import time from collections.abc import Callable from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING, Final, Literal from litellm_proxy_extras._logging import logger @@ -37,16 +38,16 @@ if TYPE_CHECKING: class RequestLogIndex: """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 most rows the build takes on by itself. An index whose build - has to read each row's metadata costs far more per row than a column index, so above - that estimate it is left for an operator to build off-peak instead of holding the - migration job past its deadline.""" + 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_rows: int | None = None + max_bytes: int | None = None @property def columns(self) -> tuple[str, ...]: @@ -77,7 +78,7 @@ REQUEST_LOG_INDEXES: Final = ( "LiteLLM_SpendLogs_error_code_startTime_idx", "((left(metadata->'error_information'->>'error_code', 128)), \"startTime\")", "(metadata->'error_information'->>'error_code') IS NOT NULL", - max_rows=1_000_000, + max_bytes=4 * 1024**3, ), ) @@ -93,9 +94,8 @@ _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)" -_ESTIMATED_ROWS_SQL: Final = ( - "SELECT COALESCE(SUM(GREATEST(c.reltuples, 0)), 0)::bigint AS rows, COALESCE(BOOL_OR(c.reltuples < 0), false) " - "AS unanalyzed FROM pg_class c WHERE c.relkind = 'r' " +_TABLE_BYTES_SQL: Final = ( + "SELECT COALESCE(SUM(pg_table_size(c.oid)), 0)::bigint AS total FROM pg_class c WHERE c.relkind = 'r' " "AND (c.oid = to_regclass(%(table)s) OR c.oid IN (SELECT relid FROM pg_partition_tree(to_regclass(%(table)s))))" ) _CHILDREN_WITHOUT_THE_INDEX_SQL: Final = ( @@ -154,9 +154,8 @@ class _TableKind: @dataclass(frozen=True, slots=True) -class _EstimatedRows: - rows: int - unanalyzed: bool +class _TableBytes: + total: int _IndexOutcome = Literal["in_place", "left_for_operator", "not_in_place"] @@ -195,8 +194,8 @@ 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, and an index over its - `max_rows` is logged and left for an operator; the result says whether every + 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 @@ -281,73 +280,94 @@ def _ensure_index( def _too_large_to_build( connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex, partitioned: bool ) -> bool: - if index.max_rows is None: + 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 - rows: Final = _estimated_rows(connection, schema, index.table) - if rows <= index.max_rows: + size: Final = _table_bytes(connection, schema, index.table) + if size <= index.max_bytes: return False logger.warning( - "Not building index %s on %s: it reads about %s rows, more than the %s it builds by itself. Build it off-peak " - "and the next index build finds it in place: %s", + "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, and the next index build finds it in place: %s", index.name, index.table, - rows, - index.max_rows, - _operator_build_statement( - connection, - schema, - index, - partitioned, - interrupted=existing is not None and existing.table == index.table, - ), + _gib(size), + _gib(index.max_bytes), + "; ".join(_operator_build_statements(connection, schema, index, index.table, index.name, partitioned)), ) return True -def _estimated_rows(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> int: - """The planner's row count for the table and its partitions. A table never analyzed counts zero rows however - many it holds (a freshly restored database), so it is analyzed first.""" - from psycopg import sql - - estimate: Final = _planner_estimate(connection, schema, table) - if not estimate.unanalyzed: - return estimate.rows - connection.execute(sql.SQL("ANALYZE {}").format(sql.Identifier(schema, table))) - return _planner_estimate(connection, schema, table).rows +def _gib(size: int) -> str: + return f"{size / 1024**3:.1f} GiB" -def _planner_estimate(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> _EstimatedRows: +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(_EstimatedRows)) as cursor: - estimate: Final = cursor.execute( - _ESTIMATED_ROWS_SQL, {"table": _regclass_name(connection, schema, table)} - ).fetchone() - return estimate if estimate is not None else _EstimatedRows(rows=0, unanalyzed=False) + 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_statement( +def _operator_build_statements( connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex, + table: str, + name: str, partitioned: bool, - interrupted: bool, -) -> str: - """What an operator runs by hand: CONCURRENTLY wherever Postgres allows it, and a REINDEX of the invalid index an - interrupted build left, which a CREATE ... IF NOT EXISTS would silently skip.""" +) -> 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 - if interrupted: - reindex: Final = "REINDEX INDEX {}" if partitioned else "REINDEX INDEX CONCURRENTLY {}" - return sql.SQL(reindex).format(sql.Identifier(schema, index.name)).as_string(connection) - build: Final = ( - "CREATE INDEX IF NOT EXISTS {} ON {} " if partitioned else "CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} " + 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) ) - prefix: Final = sql.SQL(build).format(sql.Identifier(index.name), sql.Identifier(schema, index.table)) - return _create_index_statement(connection, prefix, index.body).decode() + 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: diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2adf8b22b16..4ecda35cd07 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -74,11 +74,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*$") -_ROW_CAPPED_REQUEST_LOG_INDEXES: Final = frozenset( - index.name for index in REQUEST_LOG_INDEXES if index.max_rows is not None -) +_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 " @@ -97,6 +96,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 @@ -853,8 +857,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 if row[1] not in _ROW_CAPPED_REQUEST_LOG_INDEXES) + 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/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 6dafe67624a..7e489a35706 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -923,26 +923,31 @@ def test_the_partial_error_index_is_built_beside_the_plain_start_time_index_neve assert _index_predicate(scratch_database, ERROR_INDEX) is not None -def _seed_analyzed_spend_logs(database_url: str, partitioned: bool) -> int: - """Seed the spend logs, partitioned or not, ANALYZE them, and return how many rows they hold.""" +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") - conn.execute('ANALYZE "LiteLLM_SpendLogs"') - return SEEDED_ROWS + (len(PARTITIONS) + 1) * ROWS_PER_PARTITION if partitioned else SEEDED_ROWS + + +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: ", 1)[1].split("; "): + conn.execute(statement.encode()) @requires_db @pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned")) -def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_build( +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 costs far more per row to build than a column index, so on a table over its cap the - migration job must finish without it instead of running past its deadline and failing the upgrade.""" - rows: Final = _seed_analyzed_spend_logs(scratch_database, partitioned) - capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=rows - 1) + """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 @@ -950,22 +955,20 @@ def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_b 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", (replace(capped, max_rows=rows),)) is True + 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_statement_logged_for_an_index_over_its_cap_builds_it_and_the_next_build_finds_it_in_place( +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_analyzed_spend_logs(scratch_database, partitioned) - capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=1) + _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 - logged: Final = next(message for message in caplog.messages if "Not building index" in message) - with psycopg.connect(scratch_database, autocommit=True) as conn: - conn.execute(logged.split("finds it in place: ", 1)[1].encode()) + _run_the_logged_statements(scratch_database, caplog.messages) hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX] caplog.clear() @@ -973,15 +976,18 @@ def test_the_statement_logged_for_an_index_over_its_cap_builds_it_and_the_next_b 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 - assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True + 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) -> None: - with _pin_a_snapshot_on(database_url, "LiteLLM_SpendLogs"): +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(ERROR_INDEX), sql.Identifier("LiteLLM_SpendLogs") + sql.Identifier(name), sql.Identifier(table) ) with pytest.raises(psycopg.errors.QueryCanceled): builder.execute(prefix.as_bytes(builder) + ERROR_INDEX_DEFINITION.body.encode()) @@ -991,37 +997,63 @@ def _leave_an_invalid_error_index(database_url: str) -> None: 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_analyzed_spend_logs(scratch_database, partitioned=False) - _leave_an_invalid_error_index(scratch_database) - capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=1) + _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 - logged: Final = next(message for message in caplog.messages if "Not building index" in message) - with psycopg.connect(scratch_database, autocommit=True) as conn: - conn.execute(logged.split("finds it in place: ", 1)[1].encode()) + _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_table_never_analyzed_is_counted_before_deciding_whether_it_is_over_the_cap( +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: - for n in range(SEEDED_ROWS): - _insert_spend_log(conn, f"unanalyzed-{n}", "2026-10-01") - capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=SEEDED_ROWS - 1) + 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 f"Not building index {ERROR_INDEX}" in caplog.text - assert _indexed_table(scratch_database, ERROR_INDEX) is None + _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 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 ca08863330c..985dadb9557 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 @@ -4772,24 +4773,25 @@ async def test_ui_view_spend_logs_with_error_code(client): @pytest.mark.asyncio -async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client): - prefix = "E" * 128 - mock_spend_logs = [ +async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client: TestClient) -> None: + prefix: Final = "E" * 128 + codes: Final = MappingProxyType({f"log-{suffix}": prefix + suffix for suffix in ("first", "second")}) + mock_spend_logs: Final = [ { - "id": f"log-{suffix}", - "request_id": f"req-{suffix}", + "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": prefix + suffix}}), + "metadata": json.dumps({"error_information": {"error_code": code}}), } - for suffix in ("first", "second") + for log_id, code in codes.items() ] - def filter_by_exact_error_code(where): - wanted = str(where.get("metadata", {}).get("equals", "")).strip('"') - return [log for log in mock_spend_logs if json.loads(log["metadata"])["error_information"]["error_code"] == wanted] + 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" From b0bf5a004947e19b76ce8b9229dc4ed8ab57d0ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:16:53 -0700 Subject: [PATCH 7/9] fix(spend): keep main's Error Code query until the prefix index is valid, and print operator statements one per line Each filtered Logs request checks pg_index for a valid prefix index first. With it, the filter matches left(code, 128), adding code = $1 only for codes of 128 characters or more. Without it, the filter sends main's exact code = $1, so an index built by hand on the full code still serves it and nothing is slower than main. Always sending both conditions was rejected because it reads the metadata of every matching row. The over-cap warning prints each statement on its own line ending in a semicolon and asks for them to run one at a time, since psql -c runs several in one transaction and CONCURRENTLY refuses that. The table size check walks pg_inherits so it works on Postgres 11. --- .../request_log_indexes.py | 14 ++- .../spend_management_endpoints.py | 21 ++++- .../migrations/test_request_log_indexes.py | 2 +- .../spend/test_spend_log_error_filters.py | 88 ++++++++++++++----- .../test_spend_management_endpoints.py | 24 +++-- 5 files changed, 116 insertions(+), 33 deletions(-) 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 2a928db50ec..c6f58ab0b2a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -95,8 +95,10 @@ _CREATE_INDEX_STATEMENT: Final = re.compile( ) _TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)" _TABLE_BYTES_SQL: Final = ( - "SELECT COALESCE(SUM(pg_table_size(c.oid)), 0)::bigint AS total FROM pg_class c WHERE c.relkind = 'r' " - "AND (c.oid = to_regclass(%(table)s) OR c.oid IN (SELECT relid FROM pg_partition_tree(to_regclass(%(table)s))))" + "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 " @@ -290,12 +292,16 @@ def _too_large_to_build( 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, and the next index build finds it in place: %s", + "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(_operator_build_statements(connection, schema, index, index.table, index.name, partitioned)), + "".join( + f"{statement};\n" + for statement in _operator_build_statements(connection, schema, index, index.table, index.name, partitioned) + ), ) return True diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 09fc2285f51..59ccfe6ea12 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -99,6 +99,10 @@ _BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_bat _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}", @@ -229,7 +233,18 @@ class _SessionSpendRow(TypedDict): session_models: ReadOnly[Sequence[str]] -def _error_code_condition(placeholder: str, error_code: str) -> 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 ( @@ -2947,7 +2962,9 @@ async def ui_view_spend_logs( sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append(_error_code_condition(f"${p}", error_code)) + 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/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 7e489a35706..b21873fbb52 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -935,7 +935,7 @@ def _seed_spend_logs(database_url: str, partitioned: bool) -> None: 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: ", 1)[1].split("; "): + for statement in logged.split("finds it in place:\n", 1)[1].splitlines(): conn.execute(statement.encode()) diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index a1e3b97e5db..8a8a7377240 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -1,6 +1,6 @@ import secrets import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import product @@ -19,6 +19,7 @@ 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 @@ -58,11 +59,11 @@ ERROR_CODES: Final = MappingProxyType( LONG_CODE_PREFIX[:-1]: frozenset(), } ) +SHORT_ERROR_CODES: Final = ("500", "499") SESSION_VIEWS: Final = ("false", "true") -FILTERS: Final = tuple(product(SESSION_VIEWS, ERROR_CODES)) -def _seed(database_url: str, started: datetime) -> None: +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( @@ -92,31 +93,41 @@ def _seed(database_url: str, started: datetime) -> None: ), row.status, ) - for row in ERROR_ROWS + for row in error_rows ), ) connection.execute('ANALYZE "LiteLLM_SpendLogs"') -def _error_index_scans(database_url: str) -> int: +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", (ERROR_INDEX,) + "SELECT idx_scan AS scans FROM pg_stat_user_indexes WHERE indexrelname = %s", (index,) ).fetchone() - assert row is not None, f"{ERROR_INDEX} was not built" + assert row is not None, f"{index} was not built" return row.scans -def _settled_error_index_scans(database_url: str, previous: int, unchanged_since: float) -> int: +def _settled_index_scans(database_url: str, index: str, previous: int, unchanged_since: float) -> int: changed: Final = eventually( - lambda: _error_index_scans(database_url) != previous, + 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_error_index_scans(database_url, _error_index_scans(database_url), time.monotonic()) + 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]: @@ -127,6 +138,7 @@ def _listed(gateway: Gateway, window: Mapping[str, str], query: Mapping[str, str def _list_each_filter( gateway: Gateway, database_url: str, + index: str, window: Mapping[str, str], filters: tuple[tuple[str, str], ...], scans: int, @@ -135,14 +147,39 @@ def _list_each_filter( return () session, code = filters[0] listed: Final = _listed(gateway, window, {"error_code": code, "group_by_session": session}) - settled: Final = _settled_error_index_scans(database_url, scans, time.monotonic()) + settled: Final = _settled_index_scans(database_url, index, scans, time.monotonic()) return ( ((session, code), listed, settled - scans), - *_list_each_filter(gateway, database_url, window, filters[1:], settled), + *_list_each_filter(gateway, database_url, index, window, filters[1:], settled), ) -def test_error_code_filters_read_only_the_rows_that_carry_an_error_code(gateway: Gateway, tmp_path: Path) -> None: +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"), @@ -155,14 +192,23 @@ def test_error_code_filters_read_only_the_rows_that_carry_an_error_code(gateway: eventually( lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=60 ) - _seed(database_url, now - timedelta(hours=1)) - scans: Final = _settled_error_index_scans( - database_url, _error_index_scans(database_url), time.monotonic() + _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) - results: Final = _list_each_filter(booted.gateway, database_url, window, FILTERS, scans) - 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} +@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/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 985dadb9557..49a30019b4c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -208,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=False +): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. @@ -239,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)) @@ -4698,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, error_code_index_ready): """Test filtering spend logs by error code""" mock_spend_logs = [ { @@ -4744,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() @@ -4773,7 +4780,10 @@ async def test_ui_view_spend_logs_with_error_code(client): @pytest.mark.asyncio -async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client: TestClient) -> None: +@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 = [ @@ -4798,7 +4808,11 @@ async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact ) try: with patch.object( - ps, "prisma_client", make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_exact_error_code) + 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( From 4471416c1a3035175999b48a1e3217a75a5dafea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:57:05 -0700 Subject: [PATCH 8/9] test(spend): type the error code index parameter in the Logs filter tests --- .../proxy/spend_tracking/test_spend_management_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 49a30019b4c..c738ff1fb91 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -209,7 +209,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): def make_ui_spend_logs_mock_prisma( - mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None, error_code_index_ready=False + 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. @@ -4703,7 +4703,7 @@ async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client @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_error_code(client, error_code_index_ready): +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 = [ { From 3c24e51e0477c82dc7592064eac5134f1ef0e4af Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:22:11 -0700 Subject: [PATCH 9/9] test(perf): require only the always-built request-log indexes in the nightly database upgrade test The error code index is size-capped and left for an operator on tables over 4 GiB, so the full-size upgrade snapshot never gets it and the lane's index assertion would fail every night. --- tests/e2e/perf/setup/upgrade_worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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",