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] 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(