From ad8de0e1927c18d5d14c92939bbd531c54573874 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:42:19 +0000 Subject: [PATCH] fix(proxy): roll up only closed days into LiteLLM_DailyGlobalSpend and split the key-free read at the marker The write path no longer dual-writes the global table. The cron rolls up closed UTC days only, so a pod still flushing the current day can never leave the global table short. The key-free arm reads days through the marker from the global table and later days from LiteLLM_DailyUserSpend in one UNION ALL, and the marker comes from the config cache rather than a per-request database lookup. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/daily_spend_bulk_upsert.py | 98 +++--------- litellm/proxy/db/db_spend_update_writer.py | 7 +- .../common_daily_activity.py | 79 ++++++--- .../daily_global_spend_rollup.py | 43 ++--- .../proxy/db/test_daily_spend_bulk_upsert.py | 150 ------------------ .../proxy/db/test_db_spend_update_writer.py | 101 +----------- .../test_common_daily_activity.py | 90 ++++++----- .../test_daily_global_spend_rollup.py | 92 +++++------ 8 files changed, 204 insertions(+), 456 deletions(-) diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index c83043101eb..a143643577e 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -25,41 +25,29 @@ SpendRow = Mapping[str, object] @dataclass(frozen=True, slots=True) class DailySpendTable: - """A daily rollup table and the unique constraint its upserts arbitrate on.""" + """The physical table behind one entity's daily rollup.""" name: str - key_columns: tuple[str, ...] + entity_id_column: str carries_request_id: bool = False +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), + "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), + "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), + "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), + "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), + "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + } +) + # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. _KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") - -def _entity_table(name: str, entity_id_column: str, carries_request_id: bool = False) -> DailySpendTable: - return DailySpendTable( - name=name, key_columns=(entity_id_column, *_KEY_COLUMNS), carries_request_id=carries_request_id - ) - - -DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( - { - "user": _entity_table("LiteLLM_DailyUserSpend", "user_id"), - "team": _entity_table("LiteLLM_DailyTeamSpend", "team_id"), - "org": _entity_table("LiteLLM_DailyOrganizationSpend", "organization_id"), - "end_user": _entity_table("LiteLLM_DailyEndUserSpend", "end_user_id"), - "agent": _entity_table("LiteLLM_DailyAgentSpend", "agent_id"), - "tag": _entity_table("LiteLLM_DailyTagSpend", "tag", carries_request_id=True), - } -) - -GLOBAL_SPEND_TABLE: Final = DailySpendTable( - name="LiteLLM_DailyGlobalSpend", - key_columns=("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"), -) - _COUNTER_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -104,7 +92,7 @@ def _as_float(value: object) -> float: def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in table.key_columns) + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -142,11 +130,7 @@ def _row_params( return ( str(uuid.uuid4()), *key, - *( - () - if "model_group" in table.key_columns - else (None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),) - ), + None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), @@ -156,25 +140,26 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - *table.key_columns, - *(() if "model_group" in table.key_columns else ("model_group",)), + table.entity_id_column, + *_KEY_COLUMNS, + "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), ) -def _upsert_statement( +def build_bulk_upsert( table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]], - first_param: int, -) -> str: +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" columns: Final = _insert_columns(table) quoted_table: Final = f'"{table.name}"' rows: Final = ", ".join( "(" + ", ".join( - f"${first_param + row_index * len(columns) + offset}::{_CASTS.get(column, 'text')}" + f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" for offset, column in enumerate(columns) ) + ", (NOW() AT TIME ZONE 'UTC'))" @@ -191,44 +176,11 @@ def _upsert_statement( if table.carries_request_id else "" ) - return ( + sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted(table.key_columns)}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) - - -def _params(table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]]) -> tuple[SqlValue, ...]: - return tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) - - -def build_bulk_upsert( - table: DailySpendTable, - batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """The single statement writing one merged batch, plus its positional arguments.""" - return _upsert_statement(table, batch, first_param=1), _params(table, batch) - - -def build_bulk_upsert_with_global_rollup( - table: DailySpendTable, - batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """One statement writing a batch to its table and, atomically, its key-free rollup - to ``LiteLLM_DailyGlobalSpend``. - - A data-modifying CTE runs both inserts in the same snapshot and transaction, so a - batch that lands in one table lands in both and a retried deadlock replays both. - Postgres does not order the CTE against the main statement, so two writers can still - deadlock across the tables; the caller's deadlock retry covers that, and each insert - takes its own rows in key order so same-table lock order stays deterministic. - """ - global_batch: Final = merge_by_conflict_key(GLOBAL_SPEND_TABLE, tuple(row for _, row in batch)) - entity_params: Final = _params(table, batch) - sql: Final = ( - f"WITH entity_rows AS (\n{_upsert_statement(table, batch, first_param=1)}\nRETURNING 1)\n" - f"{_upsert_statement(GLOBAL_SPEND_TABLE, global_batch, first_param=len(entity_params) + 1)}" - ) - return sql, (*entity_params, *_params(GLOBAL_SPEND_TABLE, global_batch)) + return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d5c839a9be8..eaa03c5d7f7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -46,7 +46,6 @@ from litellm.proxy._types import ( from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, - build_bulk_upsert_with_global_rollup, merge_by_conflict_key, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -1940,11 +1939,7 @@ class DBSpendUpdateWriter: merged_batch = merge_by_conflict_key( table=table, transactions=tuple(transactions_to_process.values()) ) - sql, params = ( - build_bulk_upsert_with_global_rollup(table=table, batch=merged_batch) - if entity_type == "user" - else build_bulk_upsert(table=table, batch=merged_batch) - ) + sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f1d78dca201..b90f874c04c 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -11,8 +11,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT from litellm.proxy._types import CommonProxyErrors -from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE -from litellm.proxy.spend_tracking.daily_global_spend_rollup import reconciled_through +from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, @@ -736,28 +735,62 @@ def _rollup_metric_select(table_name: str) -> str: _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" -async def key_free_source_table(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: - """The table the key-free arm reads from, when the global rollup can answer instead of the per-key table. +_KEY_FREE_SOURCE_COLUMNS: Final = ( + "date", + "model", + "model_group", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + "spend", + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", + "compression_savings_spend", + "prompt_caching_savings_spend", + "gateway_injected_caching_savings_spend", + "autorouter_savings_spend", + "api_requests", + "successful_requests", + "failed_requests", +) - Only an unfiltered read of the user table has the same rows as ``LiteLLM_DailyGlobalSpend``, - and only through the day the reconcile marker has reached: the writer keeps that day - current, later days are covered once the next run advances the marker. + +async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: + """The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to + read it all from the per-key table. + + Only an unfiltered read of the user table sums to the same rows as the global table. The + marker read is served from the config cache, so this is not a database round trip per request. """ if query["table_name"] != "litellm_dailyuserspend": return None if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: return None - _, adjusted_end = _adjust_dates_for_timezone( - query["start_date"], query["end_date"], query["timezone_offset_minutes"], query["include_current_utc_day"] - ) try: - marker: Final = await reconciled_through(prisma_client) + return await reconciled_through(prisma_client) except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) return None - if marker is None or adjusted_end > marker: - return None - return GLOBAL_SPEND_TABLE.name + + +def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str: + """The relation the key-free arm aggregates: the per-key table alone, or the global rollup + for days through the marker plus the per-key table for the days still open after it.""" + if marker_param is None: + return f'"{pg_table}"\n WHERE {where_clause}' + columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS) + return f"""( + SELECT {columns} + FROM "{GLOBAL_SPEND_TABLE_NAME}" + WHERE {where_clause} AND date <= {marker_param} + UNION ALL + SELECT {columns} + FROM "{pg_table}" + WHERE {where_clause} AND date > {marker_param} + ) AS key_free_source""" def _build_aggregated_sql_query( @@ -772,15 +805,16 @@ def _build_aggregated_sql_query( exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path timezone_offset_minutes: int | None = None, include_current_utc_day: bool = False, - key_free_table: str | None = None, + global_rollup_through: str | None = None, ) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params """Build the GROUPING SETS query for aggregated daily activity. One statement, two UNION ALL arms over the same WHERE clause. The first arm is key-free: grand total, per-date totals and the (date, model / model_group / provider / mcp / endpoint) rollups, so its row count never grows with the number - of keys; it reads ``key_free_table`` when given (the global rollup, whose row count - never grew with the number of keys to begin with) and the entity table otherwise. + of keys. With ``global_rollup_through`` it reads days through that marker from + ``LiteLLM_DailyGlobalSpend`` (whose row count never grew with the number of keys to + begin with) and only the days after it from the per-key table. The second arm emits the (date, , api_key) rollups for the USAGE_TOP_API_KEYS_LIMIT highest-spend keys only. Both arms share the 7-bit group_level bitmask (date, api_key, model, model_group, provider, mcp, endpoint). @@ -806,8 +840,8 @@ def _build_aggregated_sql_query( exclude_entity_ids=exclude_entity_ids, ) sentinel_param: Final = f"${len(where_params) + 1}" + marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}" metric_select: Final = _rollup_metric_select(table_name) - key_free_source: Final = key_free_table or pg_table # TODO: drop the successful_requests/failed_requests aggregates (and the # total_successful_requests metadata they feed) once the admin UI reads SGR @@ -826,8 +860,7 @@ def _build_aggregated_sql_query( | GROUPING(model, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level,{metric_select} - FROM "{key_free_source}" - WHERE {where_clause} + FROM {_key_free_source(pg_table, where_clause, marker_param)} GROUP BY GROUPING SETS ( (date), (date, model), @@ -869,7 +902,8 @@ def _build_aggregated_sql_query( )) """ - return sql_query, [*where_params, PTU_SENTINEL_API_KEY] + marker_params: Final = () if global_rollup_through is None else (global_rollup_through,) + return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params] def _build_entity_rollup_sql_query( @@ -1418,7 +1452,8 @@ async def get_daily_activity_aggregated( include_current_utc_day=include_current_utc_day, ) sql_query, sql_params = _build_aggregated_sql_query( - **query_kwargs, key_free_table=await key_free_source_table(prisma_client, query_kwargs) + **query_kwargs, + global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs), ) entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 9d344421332..a9fb7669785 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -1,9 +1,10 @@ -"""Reconcile ``LiteLLM_DailyGlobalSpend`` from ``LiteLLM_DailyUserSpend``, one day per transaction. +"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``. -The spend writer keeps both tables in step from the moment it is deployed; this job rolls up -the days before that and records how far it has reached in ``LiteLLM_Config`` so usage reads -know when the global table can answer for a date range. It runs as a background cron, never -in a Prisma migration, since on a large deployment the aggregate is minutes of work. +Only days that are over get rolled up, so a pod still flushing per-key spend for the current +day can never leave the global table short; usage reads serve days through the recorded +marker from the global table and later days live from the per-key table. The marker lives in +``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on a +large deployment the first backfill is minutes of work. """ from collections.abc import Awaitable, Callable @@ -19,7 +20,6 @@ from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ) -from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE from litellm.repositories.config_repository import ConfigRepository if TYPE_CHECKING: @@ -27,8 +27,11 @@ if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient -_DAY_TRANSACTION_TIMEOUT: Final = timedelta(minutes=10) _REPLAY_DAYS: Final = 1 +GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" +# The unique constraint, in constraint order. NULL never matches itself in a unique index, so +# every column is normalized to '' or the same group would be inserted again on every run. +_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") _METRIC_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -51,23 +54,21 @@ def _quoted(columns: tuple[str, ...]) -> str: def _reconcile_day_sql() -> str: - key_columns: Final = GLOBAL_SPEND_TABLE.key_columns - normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in key_columns) + normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS) sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) return ( - f'INSERT INTO "{GLOBAL_SPEND_TABLE.name}" ("id", {_quoted(key_columns)}, {_quoted(_METRIC_COLUMNS)}, ' + f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, ' '"updated_at")\n' f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' f"GROUP BY {normalized_keys}\n" - f"ON CONFLICT ({_quoted(key_columns)}) DO UPDATE SET {overwrite}, " + f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, " "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_LOCK_GLOBAL_TABLE_SQL: Final = f'LOCK TABLE "{GLOBAL_SPEND_TABLE.name}" IN EXCLUSIVE MODE' _PENDING_DAYS_SQL: Final = ( 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' ) @@ -134,19 +135,19 @@ def _first_pending_day(marker: str | None) -> str: async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - """Every UTC day through today still to roll up, oldest first; the marker day and the one - before it are replayed so rows flushed by a pre-writer pod during a rolling deploy are folded in.""" + """Every closed UTC day (strictly before today) still to roll up, oldest first. The marker + day and the one before it are replayed so per-key rows that landed after their day was + rolled up (a flush straddling midnight, a late retry) are folded in.""" marker: Final = await reconciled_through(prisma_client) - rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), today.isoformat()) - return tuple(sorted({*(_DateRow.model_validate(row).date for row in rows), today.isoformat()})) + last_closed_day: Final = (today - timedelta(days=1)).isoformat() + rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), last_closed_day) + return tuple(_DateRow.model_validate(row).date for row in rows) async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: - """Rewrite one day of the global table from the per-key sums; the table lock keeps the - writer's increments out between the aggregate and the overwrite so none are lost.""" - async with prisma_client.db.tx(timeout=_DAY_TRANSACTION_TIMEOUT) as transaction: - await transaction.execute_raw(_LOCK_GLOBAL_TABLE_SQL) - await transaction.execute_raw(RECONCILE_DAY_SQL, day) + """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun + overwrites every group with the same totals.""" + await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) async def run_daily_global_spend_reconcile( diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py index cc443a2cfe5..c1efb3e7220 100644 --- a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -1,19 +1,12 @@ """Tests for the single-statement daily spend upsert (LIT-5291).""" -import pathlib import re -from typing import Final -import psycopg import pytest -from psycopg.rows import dict_row -from pytest_postgresql import factories from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, - GLOBAL_SPEND_TABLE, build_bulk_upsert, - build_bulk_upsert_with_global_rollup, conflict_key, merge_by_conflict_key, ) @@ -192,146 +185,3 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null(): _, params = prisma_client.db.statements[0] assert None not in params[:9] assert transactions == {} - - -def user_txn(**overrides): - txn = {**tag_txn(), "user_id": "u-1", **overrides} - del txn["tag"] - del txn["request_id"] - return txn - - -def _bound_rows(insert_sql: str, params: tuple[object, ...]) -> list[dict[str, object]]: - """Each VALUES row of one INSERT as a column -> bound value mapping, consuming params in order.""" - header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", insert_sql) - assert header is not None, insert_sql - columns = [c.strip('"') for c in header.group(1).split(", ") if c != '"updated_at"'] - row_count = insert_sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") - return [dict(zip(columns, params[i * len(columns) : (i + 1) * len(columns)])) for i in range(row_count)] - - -def test_global_rollup_folds_every_key_and_user_into_one_row_per_dimension_tuple(): - """The global table has no api_key or user_id, so a batch spread over many keys and - users must collapse to one row per (date, model, group, provider, mcp, endpoint).""" - batch = merge_by_conflict_key( - USER_TABLE, - tuple(user_txn(user_id=f"u-{i}", api_key=f"sk-{i}", spend=1.0, api_requests=1) for i in range(5)) - + (user_txn(user_id="u-0", api_key="sk-0", model="claude", spend=10.0, api_requests=3),), - ) - - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - - entity_insert, global_insert = sql.split("RETURNING 1)") - entity_rows = _bound_rows(entity_insert, params) - global_rows = _bound_rows(global_insert, params[len(entity_rows) * len(entity_rows[0]) :]) - assert len(entity_rows) == 6 - assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in global_insert - assert [(r["model"], r["spend"], r["api_requests"]) for r in global_rows] == [ - ("claude", 10.0, 3), - ("gpt-4o-mini", 5.0, 5), - ] - assert all("api_key" not in r and "user_id" not in r for r in global_rows) - conflict = re.search(r"ON CONFLICT \(([^)]*)\)", global_insert) - assert conflict is not None - assert conflict.group(1) == ", ".join(f'"{c}"' for c in GLOBAL_SPEND_TABLE.key_columns) - - -def test_global_rollup_params_follow_the_entity_params_in_one_placeholder_sequence(): - """Both inserts bind from one flat tuple, so the global arm's placeholders must start - exactly where the entity arm's stop or every value lands one column off.""" - batch = merge_by_conflict_key(USER_TABLE, (user_txn(),)) - - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - - placeholders = [int(n) for n in re.findall(r"\$(\d+)::", sql)] - assert placeholders == list(range(1, len(params) + 1)) - - -_bulk_upsert_postgresql_proc: Final = factories.postgresql_proc() -_bulk_upsert_postgresql: Final = factories.postgresql("_bulk_upsert_postgresql_proc") - -_MIGRATIONS_DIR: Final = ( - pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" -) -_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" - -_DAILY_USER_SPEND_DDL: Final = """ - CREATE TABLE "LiteLLM_DailyUserSpend" ( - id TEXT PRIMARY KEY, - user_id TEXT, - date TEXT NOT NULL, - api_key TEXT NOT NULL, - model TEXT, - model_group TEXT, - custom_llm_provider TEXT, - mcp_namespaced_tool_name TEXT, - endpoint TEXT, - prompt_tokens BIGINT DEFAULT 0, - completion_tokens BIGINT DEFAULT 0, - cache_read_input_tokens BIGINT DEFAULT 0, - cache_creation_input_tokens BIGINT DEFAULT 0, - compression_saved_tokens BIGINT DEFAULT 0, - compression_savings_spend DOUBLE PRECISION DEFAULT 0, - prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, - spend DOUBLE PRECISION DEFAULT 0, - api_requests BIGINT DEFAULT 0, - successful_requests BIGINT DEFAULT 0, - failed_requests BIGINT DEFAULT 0, - created_at TIMESTAMP DEFAULT now(), - updated_at TIMESTAMP, - UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) - ) -""" - - -def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: - converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) - conn.execute( - converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query - {f"p{i}": v for i, v in enumerate(params, start=1)}, - ) - conn.commit() - - -def test_global_rollup_equals_the_per_key_sums_after_repeated_flushes(_bulk_upsert_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: two flushes of a mixed batch leave - the global table exactly equal to the per-key table summed over user and key, with the - NULL and '' spellings of a dimension folded into one row.""" - conn: Final = _bulk_upsert_postgresql - conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - - batch = merge_by_conflict_key( - USER_TABLE, - ( - user_txn(user_id="u-1", api_key="sk-1", spend=1.0, prompt_tokens=10), - user_txn(user_id="u-2", api_key="sk-2", spend=2.0, prompt_tokens=20), - user_txn(user_id="u-1", api_key="sk-3", model=None, custom_llm_provider=None, spend=4.0), - user_txn(user_id="u-3", api_key="sk-4", model="", custom_llm_provider="", spend=8.0), - ), - ) - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - _execute_dollar_sql(conn, sql, params) - _execute_dollar_sql(conn, sql, params) - - with conn.cursor(row_factory=dict_row) as cur: - global_rows = cur.execute( - 'SELECT model, spend, prompt_tokens, api_requests FROM "LiteLLM_DailyGlobalSpend" ORDER BY model' - ).fetchall() - per_key = cur.execute( - """ - SELECT COALESCE(model, '') AS model, SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, - SUM(api_requests) AS api_requests - FROM "LiteLLM_DailyUserSpend" GROUP BY COALESCE(model, '') ORDER BY 1 - """ - ).fetchall() - - assert [row["model"] for row in global_rows] == ["", "gpt-4o-mini"] - assert [(r["model"], r["spend"], int(r["prompt_tokens"]), int(r["api_requests"])) for r in global_rows] == [ - (r["model"], float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"])) for r in per_key - ] - assert global_rows[0]["spend"] == pytest.approx(24.0) - assert global_rows[1]["spend"] == pytest.approx(6.0) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index d8a9013398e..5e977712a1e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -254,19 +254,14 @@ class _RecordingPrisma: def _row_values(statement: Statement, column: str) -> list[object]: - """Every row's value for one column of the first INSERT, read out of the flat parameter tuple. - - The user-table statement chains a global rollup INSERT after its own, so the row count - comes from the first INSERT's VALUES rather than from the parameter count. - """ + """Every row's value for one column, read out of the flat parameter tuple.""" sql, params = statement header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) assert header is not None, sql columns = header.group(1).split(", ") stride = len(columns) - 1 # updated_at is inlined, not bound offset = columns.index(f'"{column}"') - rows = sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") - return [params[row * stride + offset] for row in range(rows)] + return [params[row * stride + offset] for row in range(len(params) // stride)] @pytest.mark.asyncio @@ -1468,98 +1463,6 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected -def _entity_txn(entity_field: str, entity_id: str, api_key: str) -> dict[str, object]: - txn = _daily_txn() - del txn["user_id"] - return {**txn, entity_field: entity_id, "api_key": api_key} - - -@pytest.mark.asyncio -async def test_user_flush_writes_the_global_rollup_in_the_same_statement(): - """The user flush is the one place per-key spend becomes key-free spend, so a batch spread - over many keys must land in LiteLLM_DailyGlobalSpend as one row in the same statement. - A separate statement would let a crash between the two leave the tables out of sync.""" - prisma_client = _RecordingPrisma() - txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(4)} - - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=MagicMock(), - daily_spend_transactions=txns, - entity_type="user", - entity_id_field="user_id", - ) - - assert len(prisma_client.db.statements) == 1 - sql, params = prisma_client.db.statements[0] - assert sql.count('INSERT INTO "LiteLLM_DailyUserSpend"') == 1 - assert sql.count('INSERT INTO "LiteLLM_DailyGlobalSpend"') == 1 - assert sql.index('"LiteLLM_DailyUserSpend"') < sql.index('"LiteLLM_DailyGlobalSpend"') - global_insert = sql.split('INSERT INTO "LiteLLM_DailyGlobalSpend"', 1)[1] - assert global_insert.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") == 1 - assert "api_key" not in global_insert - assert params.count(0.4) == 1 - assert txns == {} - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("entity_type", "entity_field"), - [ - ("team", "team_id"), - ("org", "organization_id"), - ("tag", "tag"), - ("end_user", "end_user_id"), - ("agent", "agent_id"), - ], -) -async def test_other_entity_flushes_leave_the_global_table_alone(entity_type, entity_field): - """Every entity table sees the same request, so writing the rollup from more than one of - them would count each request once per entity type.""" - prisma_client = _RecordingPrisma() - txn = _entity_txn(entity_field, "e-1", "sk-1") - if entity_type == "tag": - txn["request_id"] = "req-1" - - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=MagicMock(), - daily_spend_transactions={"k": txn}, - entity_type=entity_type, - entity_id_field=entity_field, - ) - - (sql, _params) = prisma_client.db.statements[0] - assert "LiteLLM_DailyGlobalSpend" not in sql - - -@pytest.mark.asyncio -async def test_a_failed_chained_user_flush_keeps_every_transaction_for_retry(): - def raise_outage(): - raise ValueError("simulated database outage") - - prisma_client = _RecordingPrisma(execute_raw=raise_outage) - txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(3)} - expected = dict(txns) - mock_proxy_logging = MagicMock() - mock_proxy_logging.failure_handler = AsyncMock() - - with pytest.raises(ValueError, match="simulated database outage"): - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=mock_proxy_logging, - daily_spend_transactions=txns, - entity_type="user", - entity_id_field="user_id", - ) - - assert txns == expected - assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in prisma_client.db.statements[0][0] - - @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 5c74facae6a..12e5fe6af4d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,3 +1,4 @@ +import pathlib import re from collections.abc import Sequence from datetime import datetime, timedelta, timezone @@ -10,11 +11,6 @@ import pytest from psycopg.rows import dict_row from pytest_postgresql import factories -from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR - - -import pathlib - from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, PTU_SENTINEL_API_KEY, @@ -29,10 +25,11 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, - key_free_source_table, + global_rollup_reconciled_through, update_metrics, ) from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, @@ -1632,7 +1629,9 @@ def _prisma_with_marker(marker: str | None) -> MagicMock: prisma.db = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - row = None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + row = ( + None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + ) prisma.get_generic_data = AsyncMock(return_value=row) return prisma @@ -1657,9 +1656,9 @@ def _unfiltered_user_query(**overrides): @pytest.mark.parametrize( ("marker", "overrides", "expected"), [ - ("2026-06-02", {}, "LiteLLM_DailyGlobalSpend"), - ("2026-06-02", {"model": "gpt-5"}, "LiteLLM_DailyGlobalSpend"), - ("2026-06-01", {}, None), + ("2026-06-02", {}, "2026-06-02"), + ("2026-06-02", {"model": "gpt-5"}, "2026-06-02"), + ("2026-05-01", {}, "2026-05-01"), (None, {}, None), ("2026-06-02", {"api_key": "sk-1"}, None), ("2026-06-02", {"api_key": []}, None), @@ -1668,33 +1667,50 @@ def _unfiltered_user_query(**overrides): ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), ], ) -async def test_key_free_source_table_routes_only_unfiltered_user_reads_within_the_marker(marker, overrides, expected): - """Anything that filters by key or entity has no counterpart in the global table, and a - range the reconcile has not reached must stay on the per-key table.""" +async def test_global_rollup_marker_is_used_only_for_unfiltered_user_reads(marker, overrides, expected): + """Anything that filters by key or entity has no counterpart in the global table; the + SQL splits the range at the marker itself, so the marker passes through unchanged.""" await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) prisma = _prisma_with_marker(marker) - assert await key_free_source_table(prisma, _unfiltered_user_query(**overrides)) == expected + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query(**overrides)) == expected await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) @pytest.mark.asyncio -async def test_key_free_source_table_judges_the_timezone_extended_end_not_the_requested_one(): - """A caller west of UTC asking through their local today gets today's UTC bucket added to - the range; the marker must cover that extended day, not just the requested end.""" +async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table(): await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - today_utc: Final = datetime.now(timezone.utc).date() - yesterday: Final = (today_utc - timedelta(days=1)).isoformat() - query: Final = _unfiltered_user_query( - start_date=yesterday, end_date=yesterday, timezone_offset_minutes=24 * 60, include_current_utc_day=True - ) + prisma = _prisma_with_marker(None) + prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down")) - assert await key_free_source_table(_prisma_with_marker(yesterday), query) is None - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - assert await key_free_source_table(_prisma_with_marker(today_utc.isoformat()), query) == "LiteLLM_DailyGlobalSpend" + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) +def test_aggregated_sql_splits_the_key_free_arm_at_the_marker_and_keeps_the_key_arm_per_key(): + sql, params = _build_aggregated_sql_query(**_unfiltered_user_query(), global_rollup_through="2026-06-01") + marker_param: Final = f"${len(params)}" + + assert params[-1] == "2026-06-01" + assert ( + f'FROM "LiteLLM_DailyGlobalSpend"\n WHERE date >= $1 AND date <= $2 AND date <= {marker_param}' + in sql + ) + assert ( + f'FROM "LiteLLM_DailyUserSpend"\n WHERE date >= $1 AND date <= $2 AND date > {marker_param}' in sql + ) + key_arm: Final = sql.split("UNION ALL\n (WITH top_api_keys")[1] + assert "LiteLLM_DailyGlobalSpend" not in key_arm + assert marker_param not in key_arm + + +def test_aggregated_sql_without_a_marker_reads_the_per_key_table_only(): + sql, params = _build_aggregated_sql_query(**_unfiltered_user_query()) + + assert "LiteLLM_DailyGlobalSpend" not in sql + assert params[-1] == PTU_SENTINEL_API_KEY + + _GLOBAL_SPEND_MIGRATION: Final = ( pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" @@ -1706,12 +1722,12 @@ _GLOBAL_SPEND_MIGRATION: Final = ( @pytest.mark.asyncio -async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_free_arm( +async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_table_and_open_days_live( _aggregated_postgresql: psycopg.Connection, ): - """With the range reconciled, the key-free arm reads LiteLLM_DailyGlobalSpend while the - per-key arm stays on the user table, and the response is identical to the all-per-key - read: same totals, same rollups, same top keys.""" + """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must + give the same response as reading everything per-key: day 1 from the global table, day 2 + live, one grand total across both. The per-key arm stays on the user table throughout.""" n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 rows: Final = [ ( @@ -1734,11 +1750,10 @@ async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_ _seed_daily_user_spend(_aggregated_postgresql, rows) with _aggregated_postgresql.cursor() as cur: cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - for day in ("2026-06-01", "2026-06-02"): - cur.execute( - re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg - {"p1": day}, - ) + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg + {"p1": "2026-06-01"}, + ) _aggregated_postgresql.commit() async def read(marker: str | None, sql_seen: list[str]): @@ -1760,16 +1775,19 @@ async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_ per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim from_per_key = await read(None, per_key_sql) - from_global = await read("2026-06-02", global_sql) + from_global = await read("2026-06-01", global_sql) await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 - assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 2 + assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} + + def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace( diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 13dc757cbbd..11ca72e7b3d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -2,7 +2,6 @@ import pathlib import re -from contextlib import asynccontextmanager from datetime import date from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -13,11 +12,7 @@ from psycopg.rows import dict_row from pytest_postgresql import factories from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM -from litellm.proxy.db.daily_spend_bulk_upsert import ( - DAILY_SPEND_TABLES, - build_bulk_upsert_with_global_rollup, - merge_by_conflict_key, -) +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( RECONCILE_DAY_SQL, reconciled_through, @@ -45,21 +40,6 @@ class _FakeConfigTable: return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) -class _FakeTransaction: - def __init__(self, prisma: "_FakePrisma") -> None: - self._prisma = prisma - - async def execute_raw(self, sql: str, *params: str) -> int: - if "LOCK TABLE" in sql: - self._prisma.locks_taken += 1 - return 0 - (day,) = params - if day in self._prisma.failing_days: - raise RuntimeError(f"day {day} exploded") - self._prisma.reconciled.append(day) - return 1 - - class _FakeDb: def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma @@ -69,19 +49,21 @@ class _FakeDb: first, last = params return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] - @asynccontextmanager - async def tx(self, timeout: object): - yield _FakeTransaction(self._prisma) + async def execute_raw(self, sql: str, *params: str) -> int: + (day,) = params + if day in self._prisma.failing_days: + raise RuntimeError(f"day {day} exploded") + self._prisma.reconciled.append(day) + return 1 class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and a transaction.""" + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: self.user_days = user_days self.failing_days = failing_days self.reconciled: list[str] = [] - self.locks_taken = 0 self.db = _FakeDb(self) async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: @@ -97,33 +79,45 @@ async def _fresh_marker_cache(): @pytest.mark.asyncio -async def test_first_run_rolls_up_every_historical_day_and_today_then_marks_today(): - """Before any marker exists, every day with per-key rows is rolled up, plus today even - with no rows yet, so reads for ranges ending today can switch to the global table.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14")) +async def test_first_run_rolls_up_every_closed_day_and_never_today(): + """Before any marker exists every closed day with per-key rows is rolled up. Today is left + out: pods are still flushing it, so it is served live from the per-key table until it closes.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15") + assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") assert result.failed_day is None - assert result.reconciled_through == "2026-09-15" - assert await reconciled_through(prisma) == "2026-09-15" - assert prisma.locks_taken == 4 + assert result.reconciled_through == "2026-09-14" + assert await reconciled_through(prisma) == "2026-09-14" + assert "2026-09-15" not in prisma.reconciled @pytest.mark.asyncio async def test_later_run_replays_the_marker_day_and_the_day_before_only(): """Days older than marker-1 are settled; the marker day and its predecessor are replayed so - rows a pre-writer pod flushed around midnight during a rolling deploy get folded in.""" + per-key rows that landed after their day was rolled up get folded in.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.reconciled.clear() result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14", "2026-09-15") + assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14") assert "2026-09-01" not in prisma.reconciled - assert await reconciled_through(prisma) == "2026-09-15" + assert await reconciled_through(prisma) == "2026-09-14" + + +@pytest.mark.asyncio +async def test_a_run_with_no_new_closed_days_keeps_the_marker(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + + assert result.days_reconciled == ("2026-09-13",) + assert result.reconciled_through == "2026-09-13" @pytest.mark.asyncio @@ -149,16 +143,16 @@ async def test_the_next_run_resumes_from_the_failed_day(): result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03", "2026-09-15") - assert await reconciled_through(prisma) == "2026-09-15" + assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") + assert await reconciled_through(prisma) == "2026-09-03" @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """A pre-writer pod flushing rows for the day before the marker is exactly the replay case; - when that replay fails the marker must stay put and the operator must hear about it.""" + """A late flush for the day before the marker is exactly the replay case; when that replay + fails the marker must stay put and the operator must hear about it.""" prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.user_days = ("2026-09-12", "2026-09-13") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() @@ -212,7 +206,7 @@ async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) - assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_awaited_once() @@ -226,7 +220,7 @@ async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read() result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) - assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_not_awaited() @@ -341,9 +335,9 @@ def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: rows the writer never saw (a - pre-writer pod's flush, NULL and '' dimension spellings) end up folded into the global - day, running the day twice changes nothing, and other days are left alone.""" + """Against real Postgres and the shipped migration: writer-shaped rows and legacy rows + (NULL and '' dimension spellings) fold into one global day, running the day twice changes + nothing, and other days are left alone.""" conn: Final = _rollup_postgresql conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal @@ -353,7 +347,7 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p USER_TABLE, (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), ) - _execute_dollar_sql(conn, *build_bulk_upsert_with_global_rollup(USER_TABLE, written_batch)) + _execute_dollar_sql(conn, *build_bulk_upsert(USER_TABLE, written_batch)) conn.execute( """