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.

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
 

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:

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

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}

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(