ci(spend-logs): gate new spend logs and daily aggregate queries behind a CODEOWNERS-reviewed budget

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

.github/CODEOWNERS

@@ -1,3 +1,13 @@
 /ui/ @yuneng-jiang @ryan-crabbe-berri
 /litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
 /ui/litellm-dashboard/src/lib/http/schema.d.ts
+
+/spend-logs-query-budget.json @krrishdholakia @ishaan-jaff @yuneng-jiang
+/tests/code_coverage_tests/check_spend_logs_query_budget.py @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/proxy/schema.prisma @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/models/spend_logs.py @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/proxy/spend_tracking/ @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/proxy/db/db_transaction_queue/ @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/proxy/db/db_spend_update_writer.py @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/proxy/db/create_views.py @krrishdholakia @ishaan-jaff @yuneng-jiang
+/litellm/proxy/management_endpoints/common_daily_activity.py @krrishdholakia @ishaan-jaff @yuneng-jiang

.github/workflows/test-code-quality.yml

@@ -115,6 +115,9 @@ jobs:
       - name: check_fastuuid_usage
         run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
 
+      - name: check_spend_logs_query_budget
+        run: uv run --no-sync python ./tests/code_coverage_tests/check_spend_logs_query_budget.py
+
       - name: check_e2e_no_raw_requests
         run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
 

CLAUDE.md

@@ -45,6 +45,8 @@ Run tests before you commit. Also, run `make pre-commit` right before each commi
 
 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
 
+Don't add new queries against `LiteLLM_SpendLogs`, the spend log index tables, or the daily aggregate tables. They're the biggest tables in a production database, so a new scan there is how a gateway gets taken down. Reuse an existing repository or helper; if you truly need a new query, `make check-spend-logs-query-budget` will fail until you run `make spend-logs-query-budget-update` and commit `spend-logs-query-budget.json`, which is CODEOWNERS-gated so the query gets reviewed by the spend logs owners
+
 If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
 
 If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason

Makefile

@@ -9,7 +9,8 @@
 	lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
 	install-dev install-proxy-dev install-test-deps install-hooks \
 	install-helm-unittest check-circular-imports check-import-safety pre-commit \
-	lint-install lint-fetch-base bootstrap
+	lint-install lint-fetch-base bootstrap \
+	check-spend-logs-query-budget spend-logs-query-budget-update
 
 # Default target
 help:
@@ -35,6 +36,8 @@ help:
 	@echo "  make lint-gate        - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
 	@echo "  make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
 	@echo "  make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + basedpyright)"
+	@echo "  make check-spend-logs-query-budget - Gate new LiteLLM_SpendLogs / daily aggregate queries"
+	@echo "  make spend-logs-query-budget-update - Rewrite spend-logs-query-budget.json from the working tree"
 	@echo "  make check-circular-imports - Check for circular imports"
 	@echo "  make check-import-safety - Check import safety"
 	@echo "  make test               - Run all tests"
@@ -211,6 +214,12 @@ lint-type-discipline-budget-update: install-dev lint-fetch-base
 # Ratchet all budgets in one shot (ruff strict + type-discipline + basedpyright)
 lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-basedpyright-budget-update
 
+check-spend-logs-query-budget: $(LINT_DEP_INSTALL)
+	$(UV_RUN) python tests/code_coverage_tests/check_spend_logs_query_budget.py
+
+spend-logs-query-budget-update: $(LINT_DEP_INSTALL)
+	$(UV_RUN) python tests/code_coverage_tests/check_spend_logs_query_budget.py --update
+
 check-circular-imports: $(LINT_DEP_INSTALL)
 	cd litellm && $(UV_RUN) python ../tests/documentation_tests/test_circular_imports.py && cd ..
 

spend-logs-query-budget.json

@@ -0,0 +1,27 @@
+{
+  "litellm/integrations/cloudzero/database.py": 1,
+  "litellm/integrations/focus/database.py": 2,
+  "litellm/models/__init__.py": 1,
+  "litellm/proxy/_types.py": 1,
+  "litellm/proxy/agent_endpoints/endpoints.py": 1,
+  "litellm/proxy/analytics_endpoints/analytics_endpoints.py": 1,
+  "litellm/proxy/db/create_views.py": 7,
+  "litellm/proxy/db/db_spend_update_writer.py": 12,
+  "litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py": 1,
+  "litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py": 1,
+  "litellm/proxy/management_endpoints/common_daily_activity.py": 12,
+  "litellm/proxy/management_endpoints/customer_endpoints.py": 1,
+  "litellm/proxy/management_endpoints/internal_user_endpoints.py": 2,
+  "litellm/proxy/management_endpoints/organization_endpoints.py": 1,
+  "litellm/proxy/management_endpoints/tag_management_endpoints.py": 1,
+  "litellm/proxy/management_endpoints/team_endpoints.py": 1,
+  "litellm/proxy/management_endpoints/tool_management_endpoints.py": 2,
+  "litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py": 3,
+  "litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py": 5,
+  "litellm/proxy/proxy_server.py": 4,
+  "litellm/proxy/spend_tracking/spend_management_endpoints.py": 25,
+  "litellm/proxy/spend_tracking/spend_tracking_utils.py": 2,
+  "litellm/proxy/utils.py": 1,
+  "litellm/repositories/table_repositories.py": 6,
+  "litellm/responses/litellm_completion_transformation/session_handler.py": 1
+}

tests/code_coverage_tests/check_spend_logs_query_budget.py

@@ -0,0 +1,145 @@
+"""LiteLLM_SpendLogs and the daily aggregate tables are the largest, hottest tables in
+the proxy database, so an unreviewed query against them can take a production gateway
+down. This check counts every query site against those tables in litellm/ and
+enterprise/ and fails when a file exceeds the budget recorded in
+spend-logs-query-budget.json.
+
+Adding a query therefore means editing that budget file, which is owned in
+.github/CODEOWNERS, so it cannot merge without a review from the spend logs owners. Run
+`python tests/code_coverage_tests/check_spend_logs_query_budget.py --update` to rewrite
+the budgets from the working tree once the new query is signed off."""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import json
+import re
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+BUDGET_PATH = REPO_ROOT / "spend-logs-query-budget.json"
+SCANNED_DIRS = ("litellm", "enterprise")
+
+GUARDED_TABLES = frozenset(
+    {
+        "LiteLLM_SpendLogs",
+        "LiteLLM_SpendLogGuardrailIndex",
+        "LiteLLM_SpendLogToolIndex",
+        "LiteLLM_DailyUserSpend",
+        "LiteLLM_DailyTeamSpend",
+        "LiteLLM_DailyTagSpend",
+        "LiteLLM_DailyOrganizationSpend",
+        "LiteLLM_DailyEndUserSpend",
+        "LiteLLM_DailyAgentSpend",
+        "LiteLLM_DailyGuardrailMetrics",
+        "LiteLLM_DailyPolicyMetrics",
+    }
+)
+GUARDED_PRISMA_ATTRS = frozenset(table.lower() for table in GUARDED_TABLES)
+SQL_TABLE_REFERENCE = re.compile(
+    r"\b(?:from|join|into|update|truncate|table|only)\s+\"?("
+    + "|".join(sorted(GUARDED_TABLES, key=len, reverse=True))
+    + r")\"?",
+    re.IGNORECASE,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class QuerySite:
+    path: str
+    lineno: int
+    detail: str
+
+
+def _sql_table_in(text: str) -> str | None:
+    match = SQL_TABLE_REFERENCE.search(text)
+    return match.group(1) if match is not None else None
+
+
+def _site_for(node: ast.AST, relative: str) -> QuerySite | None:
+    if isinstance(node, ast.Attribute) and node.attr in GUARDED_PRISMA_ATTRS:
+        return QuerySite(relative, node.lineno, f"prisma model access '{node.attr}'")
+    if not isinstance(node, ast.Constant) or not isinstance(node.value, str):
+        return None
+    if node.value.lower() in GUARDED_PRISMA_ATTRS:
+        return QuerySite(relative, node.lineno, f"table name reference '{node.value}'")
+    table = _sql_table_in(node.value)
+    return QuerySite(relative, node.lineno, f"raw SQL against '{table}'") if table is not None else None
+
+
+def sites_in_file(path: Path, root: Path) -> tuple[QuerySite, ...]:
+    relative = path.relative_to(root).as_posix()
+    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+    return tuple(site for node in ast.walk(tree) if (site := _site_for(node, relative)) is not None)
+
+
+def scan_sites(root: Path, directories: tuple[str, ...] = SCANNED_DIRS) -> tuple[QuerySite, ...]:
+    return tuple(
+        site
+        for directory in directories
+        for path in sorted((root / directory).rglob("*.py"))
+        for site in sites_in_file(path, root)
+    )
+
+
+def counts_by_file(sites: tuple[QuerySite, ...]) -> dict[str, int]:
+    return {path: sum(1 for site in sites if site.path == path) for path in sorted({site.path for site in sites})}
+
+
+def over_budget(counts: dict[str, int], budget: dict[str, int]) -> dict[str, int]:
+    return {path: count for path, count in counts.items() if count > budget.get(path, 0)}
+
+
+def load_budget(budget_path: Path) -> dict[str, int]:
+    parsed: object = json.loads(budget_path.read_text(encoding="utf-8"))
+    if not isinstance(parsed, dict):
+        raise ValueError(f"{budget_path.name} must be an object mapping file paths to query counts")
+    return {str(path): int(limit) for path, limit in parsed.items()}
+
+
+def write_budget(budget_path: Path, counts: dict[str, int]) -> None:
+    _ = budget_path.write_text(json.dumps(counts, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+
+
+def _report(sites: tuple[QuerySite, ...], violations: dict[str, int], budget: dict[str, int]) -> None:
+    for path, count in violations.items():
+        print(f"{path}: {count} spend logs query site(s), budget is {budget.get(path, 0)}")
+        for site in sites:
+            if site.path == path:
+                print(f"  {site.path}:{site.lineno}: {site.detail}")
+    print(
+        "\nQuerying LiteLLM_SpendLogs or the daily aggregate tables is restricted. Reuse an existing "
+        "repository or helper if one fits; otherwise get sign-off from the spend logs owners in "
+        ".github/CODEOWNERS, run `python tests/code_coverage_tests/check_spend_logs_query_budget.py --update`, "
+        "and commit the updated spend-logs-query-budget.json so the new query is reviewed explicitly"
+    )
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Enforce the spend logs query budget")
+    _ = parser.add_argument("--update", action="store_true", help="rewrite the budget file from the working tree")
+    args = parser.parse_args()
+
+    sites = scan_sites(REPO_ROOT)
+    counts = counts_by_file(sites)
+
+    if bool(args.update):
+        write_budget(BUDGET_PATH, counts)
+        print(f"wrote {BUDGET_PATH.name}: {sum(counts.values())} query site(s) across {len(counts)} file(s)")
+        return 0
+
+    budget = load_budget(BUDGET_PATH)
+    violations = over_budget(counts, budget)
+    if not violations:
+        print(f"spend logs query budget ok: {sum(counts.values())} query site(s) across {len(counts)} file(s)")
+        return 0
+
+    _report(sites, violations, budget)
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/test_litellm/test_check_spend_logs_query_budget.py

@@ -0,0 +1,86 @@
+"""Tests for the spend logs query guard at
+tests/code_coverage_tests/check_spend_logs_query_budget.py.
+
+The guard exists so a new query against LiteLLM_SpendLogs or a daily aggregate table
+cannot land without editing the budget file, which is CODEOWNERS-gated. These tests pin
+the two things that would silently defeat it: detecting every shape of query we use
+(raw SQL, prisma model access, table name passed to a generic helper) while ignoring
+prose mentions, and failing the budget comparison for files that grow or appear.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+_CODE_COVERAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "code_coverage_tests")
+sys.path.insert(0, _CODE_COVERAGE_DIR)
+
+import check_spend_logs_query_budget as guard  # noqa: E402
+
+_REPO_ROOT = Path(_CODE_COVERAGE_DIR).resolve().parents[1]
+
+
+def _scan(tmp_path: Path, source: str) -> tuple[guard.QuerySite, ...]:
+    package = tmp_path / "litellm"
+    package.mkdir(exist_ok=True)
+    module = package / "module.py"
+    module.write_text(source, encoding="utf-8")
+    return guard.scan_sites(tmp_path, ("litellm",))
+
+
+def test_detects_raw_sql_select(tmp_path):
+    sites = _scan(
+        tmp_path,
+        "rows = await prisma.db.query_raw('SELECT spend FROM \"LiteLLM_SpendLogs\" sl WHERE sl.spend > 0')\n",
+    )
+    assert [site.detail for site in sites] == ["raw SQL against 'LiteLLM_SpendLogs'"]
+    assert sites[0].path == "litellm/module.py"
+
+
+def test_detects_raw_sql_join_on_daily_aggregate(tmp_path):
+    sites = _scan(tmp_path, 'sql = """\n  SELECT 1\n  JOIN LiteLLM_DailyTeamSpend d ON d.team_id = t.team_id\n"""\n')
+    assert [site.detail for site in sites] == ["raw SQL against 'LiteLLM_DailyTeamSpend'"]
+    assert sites[0].lineno == 1
+
+
+def test_detects_prisma_model_access(tmp_path):
+    sites = _scan(tmp_path, "rows = await prisma_client.db.litellm_dailyuserspend.find_many(where={})\n")
+    assert [site.detail for site in sites] == ["prisma model access 'litellm_dailyuserspend'"]
+
+
+def test_detects_table_name_passed_to_generic_helper(tmp_path):
+    sites = _scan(tmp_path, 'result = await get_daily_activity(table_name="litellm_dailytagspend")\n')
+    assert [site.detail for site in sites] == ["table name reference 'litellm_dailytagspend'"]
+
+
+def test_ignores_prose_mentions_and_unguarded_tables(tmp_path):
+    sites = _scan(
+        tmp_path,
+        '"""Spend for these calls lands in LiteLLM_SpendLogs with zero tokens."""\n'
+        "rows = await prisma.db.query_raw('SELECT * FROM \"LiteLLM_TeamTable\"')\n"
+        "keys = await prisma_client.db.litellm_verificationtoken.find_many()\n",
+    )
+    assert sites == ()
+
+
+def test_counts_every_query_site_in_a_file(tmp_path):
+    sites = _scan(
+        tmp_path,
+        "a = await prisma.db.query_raw('SELECT 1 FROM \"LiteLLM_SpendLogs\"')\n"
+        "b = await prisma.db.litellm_spendlogs.count()\n"
+        "c = await prisma.db.query_raw('DELETE FROM \"LiteLLM_DailyTagSpend\"')\n",
+    )
+    assert guard.counts_by_file(sites) == {"litellm/module.py": 3}
+
+
+def test_over_budget_flags_new_file_and_growth_but_not_shrinkage():
+    counts = {"a.py": 2, "b.py": 1, "c.py": 1}
+    budget = {"a.py": 1, "c.py": 3}
+    assert guard.over_budget(counts, budget) == {"a.py": 2, "b.py": 1}
+
+
+def test_committed_budget_matches_the_repository():
+    counts = guard.counts_by_file(guard.scan_sites(_REPO_ROOT))
+    budget = guard.load_budget(_REPO_ROOT / "spend-logs-query-budget.json")
+    assert guard.over_budget(counts, budget) == {}
+    assert set(budget) == set(counts), "stale entries in spend-logs-query-budget.json; rerun the guard with --update"