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