From bd75b5bf3cc451c04d96611a4062f46fb00568e6 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 10 Oct 2026 00:49:13 -0700 Subject: [PATCH] fix(spend): run the MonthlyGlobalSpend refresh on the writer when a read replica is configured (#45699) --- .../spend_management_endpoints.py | 2 +- .../test_global_spend_refresh_writer.py | 182 ++++++++++++++++++ .../test_spend_management_endpoints.py | 60 ++++++ 3 files changed, 243 insertions(+), 1 deletion(-) create mode 100644 tests/integration/database/test_global_spend_refresh_writer.py diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f40500028a2..5482661c3de 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3767,7 +3767,7 @@ async def global_spend_refresh(): }, ) await new_client.db.connect() - await _query_raw(new_client, sql_query) + await new_client.db.execute_raw(sql_query) verbose_proxy_logger.info("MonthlyGlobalSpend view refreshed") return { "message": "MonthlyGlobalSpend view refreshed", diff --git a/tests/integration/database/test_global_spend_refresh_writer.py b/tests/integration/database/test_global_spend_refresh_writer.py new file mode 100644 index 00000000000..7fec85196b2 --- /dev/null +++ b/tests/integration/database/test_global_spend_refresh_writer.py @@ -0,0 +1,182 @@ +import os +import uuid +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import psycopg +import pytest +from psycopg import sql + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows, scratch_database +from integration._support.process import owned_proxy + +pytestmark: Final = pytest.mark.timeout(240) + +READER_PASSWORD: Final = "integration-reader-password" +SPEND: Final = 1.25 +REFRESHED: Final = {"message": "MonthlyGlobalSpend view refreshed", "status": "success"} +FAILED: Final = {"message": "Failed to refresh materialized view", "status": "failure"} + + +@dataclass(frozen=True, slots=True) +class SpendDatabase: + candidate: Gateway + writer_url: str + read_url: str + + +@contextmanager +def readonly_role() -> Generator[str]: + role: Final = f"integration_reader_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD {} NOSUPERUSER NOINHERIT").format( + sql.Identifier(role), sql.Literal(READER_PASSWORD) + ) + ) + try: + admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role))) + yield role + finally: + admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) + + +@contextmanager +def spend_database(gateway: Gateway, tmp_path: Path, *, with_reader: bool) -> Generator[SpendDatabase]: + """A proxy on its own database with its spend views created, reading through a real read-only role when + `with_reader`, and with no read replica at all otherwise""" + with readonly_role() as role, scratch_database() as writer_url: + parsed: Final = urlsplit(writer_url) + host: Final = parsed.netloc.rpartition("@")[2] + reader_url: Final = urlunsplit(parsed._replace(netloc=f"{role}:{READER_PASSWORD}@{host}")) + with psycopg.connect(writer_url, autocommit=True) as owner: + owner.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role))) + owner.execute( + sql.SQL("ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO {}").format( + sql.Identifier(role) + ) + ) + with owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": writer_url, **({"DATABASE_URL_READ_REPLICA": reader_url} if with_reader else {})}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as candidate: + eventually( + lambda: read_rows( + "SELECT relname FROM pg_class WHERE relname=%s AND relkind='v'", + ("MonthlyGlobalSpend",), + database_url=writer_url, + ), + bool, + seconds=60, + ) + connected: Final = bool(read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,))) + assert connected == with_reader, f"read-only reader connected={connected}, expected {with_reader}" + yield SpendDatabase(candidate, writer_url, reader_url if with_reader else writer_url) + + +def materialize_stock_view(writer_url: str) -> None: + with psycopg.connect(writer_url, autocommit=True) as owner: + owner.execute(""" + DO $$ + DECLARE definition text := pg_get_viewdef('"MonthlyGlobalSpend"'::regclass); + BEGIN + EXECUTE 'DROP VIEW "MonthlyGlobalSpend"'; + EXECUTE format('CREATE MATERIALIZED VIEW "MonthlyGlobalSpend" AS %s WITH NO DATA', rtrim(definition, '; ')); + END $$ + """) + + +def insert_spend(writer_url: str) -> None: + with psycopg.connect(writer_url, autocommit=True) as owner: + owner.execute( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime", spend) ' + "VALUES (%s, 'acompletion', now(), now(), %s)", + (f"integration-{uuid.uuid4().hex}", SPEND), + ) + + +def test_global_spend_refresh_runs_on_the_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None: + with spend_database(gateway, tmp_path, with_reader=True) as database: + materialize_stock_view(database.writer_url) + insert_spend(database.writer_url) + with psycopg.connect(database.read_url) as reader: + with pytest.raises(psycopg.errors.ReadOnlySqlTransaction): + reader.execute('REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend"') + + response: Final = database.candidate.request("POST", "/global/spend/refresh") + + assert response.status_code == 200 and response.json() == REFRESHED, response.text + assert read_rows('SELECT spend FROM "MonthlyGlobalSpend"', (), database_url=database.read_url) == [ + {"spend": SPEND} + ] + + +def test_concurrent_global_spend_refreshes_all_run_on_the_writer(gateway: Gateway, tmp_path: Path) -> None: + with spend_database(gateway, tmp_path, with_reader=True) as database: + materialize_stock_view(database.writer_url) + insert_spend(database.writer_url) + + with ThreadPoolExecutor(max_workers=5) as executor: + pending: Final = [ + executor.submit(database.candidate.request, "POST", "/global/spend/refresh") for _ in range(5) + ] + responses: Final = [future.result() for future in pending] + + assert [(response.status_code, response.json()) for response in responses] == [(200, REFRESHED)] * 5 + assert read_rows('SELECT spend FROM "MonthlyGlobalSpend"', (), database_url=database.read_url) == [ + {"spend": SPEND} + ] + + +def test_global_spend_refresh_without_a_read_replica_still_refreshes(gateway: Gateway, tmp_path: Path) -> None: + with spend_database(gateway, tmp_path, with_reader=False) as database: + materialize_stock_view(database.writer_url) + insert_spend(database.writer_url) + + response: Final = database.candidate.request("POST", "/global/spend/refresh") + + assert response.status_code == 200 and response.json() == REFRESHED, response.text + assert read_rows('SELECT spend FROM "MonthlyGlobalSpend"', (), database_url=database.writer_url) == [ + {"spend": SPEND} + ] + + +def test_global_spend_refresh_leaves_the_stock_plain_view_alone(gateway: Gateway, tmp_path: Path) -> None: + with spend_database(gateway, tmp_path, with_reader=True) as database: + response: Final = database.candidate.request("POST", "/global/spend/refresh") + + assert response.status_code == 200 and response.json() is None, response.text + assert read_rows( + "SELECT relkind FROM pg_class WHERE relname=%s", ("MonthlyGlobalSpend",), database_url=database.writer_url + ) == [{"relkind": "v"}] + + +def test_global_spend_refresh_reports_failure_when_the_writer_refresh_errors(gateway: Gateway, tmp_path: Path) -> None: + with spend_database(gateway, tmp_path, with_reader=True) as database: + with psycopg.connect(database.writer_url, autocommit=True) as owner: + owner.execute("CREATE SEQUENCE refresh_attempts") + owner.execute('DROP VIEW "MonthlyGlobalSpend"') + owner.execute( + 'CREATE MATERIALIZED VIEW "MonthlyGlobalSpend" AS ' + "SELECT nextval('refresh_attempts') / 0 AS spend WITH NO DATA" + ) + + response: Final = database.candidate.request("POST", "/global/spend/refresh") + + assert response.status_code == 200 and response.json() == FAILED, response.text + assert read_rows("SELECT is_called FROM refresh_attempts", (), database_url=database.writer_url) == [ + {"is_called": True} + ], "the refresh never ran on the writer" + assert read_rows( + "SELECT relispopulated FROM pg_class WHERE relname=%s", + ("MonthlyGlobalSpend",), + database_url=database.writer_url, + ) == [{"relispopulated": False}] 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 2c1b5a4dbb5..25180038e2a 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -9,6 +9,7 @@ import sqlite3 from datetime import timezone from unittest.mock import AsyncMock, MagicMock, patch +import prisma import pytest from fastapi import HTTPException from fastapi.testclient import TestClient @@ -277,6 +278,7 @@ from litellm.proxy._types import ( from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.management.teams import authz as team_access from litellm.proxy.proxy_server import app +from litellm.proxy.utils import PrismaClient from litellm.proxy.spend_tracking import spend_management_endpoints from litellm.router import Router from litellm.types.utils import BudgetConfig @@ -7999,3 +8001,61 @@ async def test_management_team_lookup_without_memberships_keeps_own_user_scope() scope = await resolve_owned_read_scope(auth.user_id, lookup) assert scope == OwnedRows("caller") assert team_reads == [] + + +class _PrismaJournal: + def __init__(self) -> None: + self.statements: tuple[tuple[str, str], ...] = () + + def __call__(self, **kwargs: object) -> "_ReadOnlyReplicaPrisma": + return _ReadOnlyReplicaPrisma(self, "reader" if "datasource" in kwargs else "writer") + + def record(self, role: str, query: str) -> None: + self.statements = (*self.statements, (role, query.strip())) + + +class _ReadOnlyReplicaPrisma: + def __init__(self, journal: _PrismaJournal, role: str) -> None: + self.journal = journal + self.role = role + + def is_connected(self) -> bool: + return False + + async def connect(self, timeout: object = None) -> None: + return None + + def _run(self, query: str) -> None: + self.journal.record(self.role, query) + if self.role == "reader" and not query.lstrip().upper().startswith("SELECT"): + raise Exception(f"cannot execute {query.split()[0]} in a read-only transaction") + + async def query_raw(self, query: str, *args: object) -> list[dict[str, str]]: + self._run(query) + return [{"relname": "MonthlyGlobalSpend", "relkind": "m"}] + + async def execute_raw(self, query: str, *args: object) -> int: + self._run(query) + return 0 + + +@pytest.mark.asyncio +async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader_is_configured( + monkeypatch: pytest.MonkeyPatch, +) -> None: + journal: Final = _PrismaJournal() + monkeypatch.setattr(prisma, "Prisma", journal) + monkeypatch.setenv("DATABASE_URL", "postgresql://writer@localhost/litellm") + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader@localhost/litellm") + monkeypatch.setattr( + ps, + "prisma_client", + PrismaClient(database_url="postgresql://writer@localhost/litellm", proxy_logging_obj=ps.proxy_logging_obj), + ) + + response: Final = await spend_management_endpoints.global_spend_refresh() + + assert response == {"message": "MonthlyGlobalSpend view refreshed", "status": "success"}, journal.statements + assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [ + ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') + ]