379 lines
13 KiB
Python
379 lines
13 KiB
Python
"""Async ``Database`` wrapper + the bare-minimum CRUD helpers.
|
|
|
|
Async to match the rest of the stack (aiohttp, uvicorn, claude-code-api).
|
|
psycopg3 has native async support — ``postgresql+psycopg://...`` works
|
|
with ``create_async_engine`` directly. SQLite goes through ``aiosqlite``
|
|
(``sqlite+aiosqlite://...``); user-facing config still uses the plain
|
|
``sqlite:///`` form and we normalise the URL here, so nothing leaks into
|
|
``.env`` / docker-compose.
|
|
|
|
No repository layer (PLAN §4.1 explicitly waives it). Helpers take an
|
|
``AsyncSession`` so callers can batch operations into one transaction
|
|
(e.g. touch ``last_used_at`` + write an audit line on the same request).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from sqlalchemy import inspect, text
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlmodel import SQLModel, col, select
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
from beaver_gateway.storage.models import AuditLog, Token, Usage
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
from sqlalchemy import Connection
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
|
|
def _utcnow() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
def _to_async_url(url: str) -> str:
|
|
"""Promote the user-facing sync URL to its async driver variant.
|
|
|
|
Users write ``sqlite:///gw.db`` or ``postgresql://...`` in ``.env``;
|
|
we translate to ``sqlite+aiosqlite://`` / ``postgresql+psycopg://``
|
|
so they don't have to know which driver we use internally.
|
|
"""
|
|
if url.startswith(("sqlite+aiosqlite://", "postgresql+psycopg://")):
|
|
return url
|
|
if url.startswith("sqlite://"):
|
|
return "sqlite+aiosqlite://" + url[len("sqlite://") :]
|
|
if url.startswith("postgresql://"):
|
|
return "postgresql+psycopg://" + url[len("postgresql://") :]
|
|
if url.startswith("postgres://"):
|
|
return "postgresql+psycopg://" + url[len("postgres://") :]
|
|
return url
|
|
|
|
|
|
class Database:
|
|
"""Owner of the async SQLAlchemy engine.
|
|
|
|
Construct once in ``cli`` from ``settings.database_url``, ``await``
|
|
:meth:`create_all` at startup, hand the instance to
|
|
:class:`GatewayRuntime`. Frontends grab sessions via
|
|
:meth:`session` (an async context manager).
|
|
|
|
``echo=False`` keeps SQL out of INFO logs (admin UI is the
|
|
user-facing view); ``connect_args={"check_same_thread": False}``
|
|
isn't needed under async sqlite — aiosqlite already runs each
|
|
connection on a dedicated thread.
|
|
"""
|
|
|
|
__slots__ = ("_engine",)
|
|
|
|
def __init__(self, url: str) -> None:
|
|
self._engine: AsyncEngine = create_async_engine(_to_async_url(url), echo=False)
|
|
|
|
async def create_all(self) -> None:
|
|
"""Create missing tables, then add columns models gained since.
|
|
|
|
Nullable columns only - that is the whole migration story until
|
|
Alembic is worth it.
|
|
"""
|
|
async with self._engine.begin() as conn:
|
|
await conn.run_sync(SQLModel.metadata.create_all)
|
|
await conn.run_sync(_add_missing_columns)
|
|
|
|
def session(self) -> AsyncSession:
|
|
"""Open a fresh :class:`AsyncSession` (use as ``async with``).
|
|
|
|
Callers commit explicitly; we don't auto-commit on exit so a
|
|
request that fails mid-flight rolls back by simply not
|
|
committing.
|
|
"""
|
|
return AsyncSession(self._engine, expire_on_commit=False)
|
|
|
|
async def dispose(self) -> None:
|
|
"""Close the engine's connection pool. Idempotent."""
|
|
await self._engine.dispose()
|
|
|
|
|
|
def _add_missing_columns(conn: Connection) -> None:
|
|
inspector = inspect(conn)
|
|
for table in SQLModel.metadata.sorted_tables:
|
|
existing = {c["name"] for c in inspector.get_columns(table.name)}
|
|
for column in table.columns:
|
|
if column.name in existing:
|
|
continue
|
|
kind = column.type.compile(conn.dialect)
|
|
ddl = f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}"
|
|
default = _column_default(column)
|
|
if default is not None:
|
|
ddl += f" DEFAULT {default}"
|
|
conn.execute(text(ddl))
|
|
|
|
|
|
def _column_default(column: Any) -> str | None:
|
|
"""Literal for ``ADD COLUMN ... DEFAULT`` so old rows get the model default."""
|
|
server_default = getattr(column, "server_default", None)
|
|
if server_default is not None:
|
|
return str(server_default.arg.text)
|
|
default = getattr(column, "default", None)
|
|
if default is None or default.is_callable:
|
|
return None
|
|
value = default.arg
|
|
if isinstance(value, bool):
|
|
return "TRUE" if value else "FALSE"
|
|
if isinstance(value, int | float):
|
|
return repr(value)
|
|
if isinstance(value, str):
|
|
return "'" + value.replace("'", "''") + "'"
|
|
return None
|
|
|
|
|
|
# ---- Token CRUD ---------------------------------------------------------
|
|
|
|
|
|
async def list_active_tokens(session: AsyncSession) -> Sequence[Token]:
|
|
"""Return every non-revoked token (Phase 4.2 seeds the cache from this)."""
|
|
stmt = select(Token).where(Token.revoked_at.is_(None)) # ty: ignore[unresolved-attribute]
|
|
result = await session.exec(stmt)
|
|
return result.all()
|
|
|
|
|
|
async def list_tokens(
|
|
session: AsyncSession, *, include_revoked: bool = False
|
|
) -> Sequence[Token]:
|
|
"""Return tokens ordered newest-first (Phase 4.3 admin table)."""
|
|
stmt = select(Token).order_by(Token.created_at.desc()) # ty: ignore[unresolved-attribute]
|
|
if not include_revoked:
|
|
stmt = stmt.where(Token.revoked_at.is_(None)) # ty: ignore[unresolved-attribute]
|
|
result = await session.exec(stmt)
|
|
return result.all()
|
|
|
|
|
|
async def create_token(
|
|
session: AsyncSession, *, name: str, scope: str, hashed_value: str
|
|
) -> Token:
|
|
"""Persist a new token. Caller hashes the plaintext before passing it in."""
|
|
row = Token(name=name, scope=scope, hashed_value=hashed_value)
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return row
|
|
|
|
|
|
async def revoke_token(session: AsyncSession, *, token_id: int) -> bool:
|
|
"""Mark a token revoked. Returns ``False`` if no such row."""
|
|
row = await session.get(Token, token_id)
|
|
if row is None or row.revoked_at is not None:
|
|
return False
|
|
row.revoked_at = _utcnow()
|
|
session.add(row)
|
|
await session.commit()
|
|
return True
|
|
|
|
|
|
async def touch_token(session: AsyncSession, *, token_id: int) -> None:
|
|
"""Bump ``last_used_at``. Phase 4.2 batches these — not per-request."""
|
|
row = await session.get(Token, token_id)
|
|
if row is None:
|
|
return
|
|
row.last_used_at = _utcnow()
|
|
session.add(row)
|
|
await session.commit()
|
|
|
|
|
|
# ---- Audit --------------------------------------------------------------
|
|
|
|
|
|
async def append_audit(
|
|
session: AsyncSession,
|
|
*,
|
|
actor: str,
|
|
kind: str,
|
|
agent_name: str | None = None,
|
|
detail: dict[str, Any] | None = None,
|
|
) -> AuditLog:
|
|
"""Append-only insert. ``detail`` JSON-serialised here, not by callers."""
|
|
row = AuditLog(
|
|
actor=actor,
|
|
kind=kind,
|
|
agent_name=agent_name,
|
|
detail_json=json.dumps(detail or {}, separators=(",", ":")),
|
|
)
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return row
|
|
|
|
|
|
async def list_audit_records(
|
|
session: AsyncSession, *, limit: int = 50, before_id: int | None = None
|
|
) -> Sequence[AuditLog]:
|
|
"""Return audit entries newest-first, optionally paginated by id cursor.
|
|
|
|
``before_id`` is a forward-only cursor: pass the smallest id from the
|
|
current page to fetch the next slice. Cheap because ``id`` is the
|
|
primary key (ordered insert).
|
|
"""
|
|
stmt = select(AuditLog).order_by(AuditLog.id.desc()).limit(limit) # ty: ignore[unresolved-attribute]
|
|
if before_id is not None:
|
|
stmt = stmt.where(AuditLog.id < before_id) # ty: ignore[unsupported-operator]
|
|
result = await session.exec(stmt)
|
|
return result.all()
|
|
|
|
|
|
# ---- Usage --------------------------------------------------------------
|
|
|
|
|
|
async def append_usage(session: AsyncSession, row: Usage) -> None:
|
|
"""Persist a turn's usage with ``cost_usd`` / ``model_usage`` as per-turn deltas.
|
|
|
|
``row`` arrives with the SDK's cumulative ``total_cost_usd`` and
|
|
``model_usage`` in ``cost_usd`` / ``model_usage``; they move to the
|
|
``*_total`` columns and the stored values become the difference against
|
|
the previous row of the same session (see :func:`usage_deltas`).
|
|
"""
|
|
prev = await _last_usage(session, row.session_id)
|
|
usage_deltas(row, prev)
|
|
session.add(row)
|
|
await session.commit()
|
|
|
|
|
|
async def _last_usage(session: AsyncSession, session_id: str | None) -> Usage | None:
|
|
if not session_id:
|
|
return None
|
|
stmt = (
|
|
select(Usage)
|
|
.where(col(Usage.session_id) == session_id)
|
|
.order_by(col(Usage.ts).desc(), col(Usage.id).desc())
|
|
.limit(1)
|
|
)
|
|
return (await session.exec(stmt)).first()
|
|
|
|
|
|
# Counters inside ``ResultMessage.model_usage[model]``; everything else there
|
|
# (provider, costBasis, contextWindow, maxOutputTokens, ...) is a constant.
|
|
_MODEL_USAGE_COUNTERS = (
|
|
"inputTokens",
|
|
"outputTokens",
|
|
"cacheReadInputTokens",
|
|
"cacheCreationInputTokens",
|
|
"webSearchRequests",
|
|
"costUSD",
|
|
)
|
|
|
|
|
|
def usage_deltas(row: Usage, prev: Usage | None) -> None:
|
|
"""Turn the SDK's cumulative cost counters on ``row`` into per-turn deltas.
|
|
|
|
``total_cost_usd`` and ``model_usage`` in a ``ResultMessage`` are running
|
|
totals of the claude process: they grow with every turn of the session and
|
|
start over when the process is respawned (rotation, resume after a
|
|
restart). A row fresh from the SDK (``cost_usd_total is None``) first moves
|
|
the raw values into ``cost_usd_total`` / ``model_usage_total``; then
|
|
``cost_usd`` / ``model_usage`` become ``current - previous`` against
|
|
``prev`` (the last row of the same session), or the raw value when there
|
|
is no previous row or the counter went backwards - a fresh process.
|
|
Idempotent on rows that already carry their totals, which is what the
|
|
backfill relies on.
|
|
"""
|
|
if row.cost_usd_total is None and row.model_usage_total is None:
|
|
row.cost_usd_total = row.cost_usd
|
|
row.model_usage_total = row.model_usage
|
|
prev_cost = prev.cost_usd_total if prev is not None else None
|
|
prev_models = prev.model_usage_total if prev is not None else None
|
|
row.cost_usd = _counter_delta(row.cost_usd_total, prev_cost)
|
|
row.model_usage = _model_usage_delta(row.model_usage_total, prev_models)
|
|
|
|
|
|
def _counter_delta(current: float | None, previous: float | None) -> float | None:
|
|
if current is None:
|
|
return None
|
|
if previous is None or current < previous:
|
|
return current
|
|
return round(current - previous, 6)
|
|
|
|
|
|
def _model_usage_delta(
|
|
current: dict[str, Any] | None, previous: dict[str, Any] | None
|
|
) -> dict[str, Any] | None:
|
|
if current is None:
|
|
return None
|
|
out: dict[str, Any] = {}
|
|
for model, fields in current.items():
|
|
if not isinstance(fields, dict):
|
|
out[model] = fields
|
|
continue
|
|
before = (previous or {}).get(model)
|
|
before = before if isinstance(before, dict) else {}
|
|
delta = dict(fields)
|
|
for key in _MODEL_USAGE_COUNTERS:
|
|
if key not in fields:
|
|
continue
|
|
cur, prev = fields.get(key), before.get(key)
|
|
if not isinstance(cur, int | float) or isinstance(cur, bool):
|
|
continue
|
|
prev_num = (
|
|
prev
|
|
if isinstance(prev, int | float) and not isinstance(prev, bool)
|
|
else None
|
|
)
|
|
value = _counter_delta(cur, prev_num)
|
|
delta[key] = (
|
|
round(value or 0, 6) if isinstance(cur, float) else int(value or 0)
|
|
)
|
|
out[model] = delta
|
|
return out
|
|
|
|
|
|
async def backfill_usage_deltas(session: AsyncSession) -> int:
|
|
"""Rewrite ``cost_usd`` / ``model_usage`` of every row as per-turn deltas.
|
|
|
|
Rows written before the ``*_total`` columns existed hold the SDK's
|
|
cumulative values in ``cost_usd`` / ``model_usage``; they are moved to the
|
|
totals first. Rows are replayed per session in time order. Idempotent:
|
|
a second run rewrites nothing. Returns the number of rows changed.
|
|
"""
|
|
stmt = select(Usage).order_by(col(Usage.session_id), col(Usage.ts), col(Usage.id))
|
|
rows = (await session.exec(stmt)).all()
|
|
last_by_session: dict[str, Usage] = {}
|
|
changed = 0
|
|
for row in rows:
|
|
before = (
|
|
row.cost_usd,
|
|
row.model_usage,
|
|
row.cost_usd_total,
|
|
row.model_usage_total,
|
|
)
|
|
prev = last_by_session.get(row.session_id) if row.session_id else None
|
|
usage_deltas(row, prev)
|
|
after = (
|
|
row.cost_usd,
|
|
row.model_usage,
|
|
row.cost_usd_total,
|
|
row.model_usage_total,
|
|
)
|
|
if after != before:
|
|
session.add(row)
|
|
changed += 1
|
|
if row.session_id:
|
|
last_by_session[row.session_id] = row
|
|
await session.commit()
|
|
return changed
|
|
|
|
|
|
__all__ = [
|
|
"Database",
|
|
"append_audit",
|
|
"append_usage",
|
|
"backfill_usage_deltas",
|
|
"create_token",
|
|
"list_active_tokens",
|
|
"list_audit_records",
|
|
"list_tokens",
|
|
"revoke_token",
|
|
"touch_token",
|
|
"usage_deltas",
|
|
]
|