PostgresBudgetStore — durable BudgetStore impl (Sprint B, v5.6.1).
Backed by the token_usage_daily Alembic migration (003). Replaces
the in-memory-only TokenBudgetTracker with a write-through ledger that
multiple app replicas converge on.
Design
Write-through. Every TokenBudgetTracker.record() call fans out to
:meth:upsert_daily, merging the snapshot into the day's row. The
tracker wraps the call in contextlib.suppress(Exception) so a dead
database never takes the LLM pipeline down — telemetry failures never
break user requests.
Multi-replica reads. :meth:read_today / :meth:read_month hydrate
the tracker's in-memory buckets before a budget check so every replica
sees the same running total. One Postgres round-trip per
check_budget is cheap compared to the LLM call it gates.
Framework-agnostic session factory. The store accepts any async
callable that yields a SQLAlchemy AsyncSession — no coupling to
FastAPI's Depends(get_db) or to a specific app scaffold. The
scaffolded app.setup.budget passes the project's SessionLocal
directly.
Graceful degradation
The tracker catches all exceptions from this store. We still log at
WARNING with exc_info=True so ops dashboards can alert on a sick
ledger — silent failure was an issue flagged in Sprint 4's Codex review.
PostgresBudgetStore
PostgresBudgetStore(session_factory: SessionFactory)
Persistent :class:BudgetStore backed by token_usage_daily.
Parameters
session_factory
Zero-arg callable producing either an AsyncContextManager or
an AsyncIterator that yields a SQLAlchemy AsyncSession.
Typical wiring: PostgresBudgetStore(session_factory=SessionLocal).
Source code in src/symfonic/core/observability/postgres_budget_store.py
| def __init__(self, session_factory: SessionFactory) -> None:
self._session_factory = session_factory
|
read_month
async
read_month(
tenant_id: str, year_month: str | None = None
) -> BudgetUsage | None
Return the sum-of-days for tenant_id in year_month.
year_month defaults to the current month (YYYY-MM).
Returns None when no rows exist (so callers can distinguish
"no history" from "zero spend").
Uses an inclusive date range (BETWEEN start AND end) rather
than EXTRACT so the same query works on both Postgres and
SQLite (test harness).
Source code in src/symfonic/core/observability/postgres_budget_store.py
| async def read_month(
self, tenant_id: str, year_month: str | None = None,
) -> BudgetUsage | None:
"""Return the sum-of-days for ``tenant_id`` in ``year_month``.
``year_month`` defaults to the current month (``YYYY-MM``).
Returns ``None`` when no rows exist (so callers can distinguish
"no history" from "zero spend").
Uses an inclusive date range (``BETWEEN start AND end``) rather
than ``EXTRACT`` so the same query works on both Postgres and
SQLite (test harness).
"""
key = year_month or month_key()
start, end = _month_bounds(key)
sql = _text(
"""
SELECT COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens,
COALESCE(SUM(cost_usd), 0) AS cost_usd,
COALESCE(SUM(request_count), 0) AS request_count,
MAX(updated_at) AS updated_at,
COUNT(*) AS row_count
FROM token_usage_daily
WHERE tenant_id = :tenant_id
AND date >= :start_date
AND date <= :end_date
""",
)
async with self._open() as session:
result = await session.execute(sql, {
"tenant_id": tenant_id,
"start_date": start,
"end_date": end,
})
row = result.first()
if row is None or int(row.row_count or 0) == 0:
return None
return _row_to_usage(tenant_id, "month", key, row)
|
read_today
async
read_today(tenant_id: str) -> BudgetUsage | None
Return today's row for tenant_id, or None if none exists.
Source code in src/symfonic/core/observability/postgres_budget_store.py
| async def read_today(self, tenant_id: str) -> BudgetUsage | None:
"""Return today's row for ``tenant_id``, or ``None`` if none exists."""
today = today_key()
sql = _text(
"""
SELECT input_tokens, output_tokens, cache_read_tokens,
cache_write_tokens, cost_usd, request_count, updated_at
FROM token_usage_daily
WHERE tenant_id = :tenant_id AND date = :date
""",
)
async with self._open() as session:
result = await session.execute(sql, {
"tenant_id": tenant_id,
"date": _parse_day_key(today),
})
row = result.first()
if row is None:
return None
return _row_to_usage(tenant_id, "day", today, row)
|
upsert_daily
async
upsert_daily(usage: BudgetUsage) -> None
Merge the tenant's day row with the supplied cumulative snapshot.
The tracker passes the CUMULATIVE in-memory totals, so this is an
assignment not an increment — idempotent on retry. Uses
INSERT ... ON CONFLICT DO UPDATE for atomic upsert.
Source code in src/symfonic/core/observability/postgres_budget_store.py
| async def upsert_daily(self, usage: BudgetUsage) -> None:
"""Merge the tenant's day row with the supplied cumulative snapshot.
The tracker passes the CUMULATIVE in-memory totals, so this is an
assignment not an increment — idempotent on retry. Uses
``INSERT ... ON CONFLICT DO UPDATE`` for atomic upsert.
"""
if usage.window != "day":
# Defensive: a month snapshot must never collide with daily
# rows. Callers should never do this; log + skip if they do.
logger.warning(
"PostgresBudgetStore.upsert_daily skipping non-day window %r",
usage.window,
)
return
sql = _text(
"""
INSERT INTO token_usage_daily (
tenant_id, date, input_tokens, output_tokens,
cache_read_tokens, cache_write_tokens,
cost_usd, request_count, updated_at
) VALUES (
:tenant_id, :date, :input_tokens, :output_tokens,
:cache_read_tokens, :cache_write_tokens,
:cost_usd, :request_count, :updated_at
)
ON CONFLICT (tenant_id, date) DO UPDATE SET
input_tokens = EXCLUDED.input_tokens,
output_tokens = EXCLUDED.output_tokens,
cache_read_tokens = EXCLUDED.cache_read_tokens,
cache_write_tokens = EXCLUDED.cache_write_tokens,
cost_usd = EXCLUDED.cost_usd,
request_count = EXCLUDED.request_count,
updated_at = EXCLUDED.updated_at
""",
)
async with self._open() as session:
await session.execute(sql, {
"tenant_id": usage.tenant_id,
"date": _parse_day_key(usage.window_key),
"input_tokens": usage.input_tokens,
"output_tokens": usage.output_tokens,
"cache_read_tokens": usage.cached_tokens,
"cache_write_tokens": usage.cache_write_tokens,
"cost_usd": usage.cost_usd,
"request_count": usage.request_count,
"updated_at": usage.updated_at,
})
await session.commit()
|