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] 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):