"""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.ext.asyncio import create_async_engine from sqlmodel import SQLModel, 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.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: """Idempotent ``CREATE TABLE IF NOT EXISTS`` for every model.""" async with self._engine.begin() as conn: await conn.run_sync(SQLModel.metadata.create_all) 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() # ---- 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: session.add(row) await session.commit() __all__ = [ "Database", "append_audit", "append_usage", "create_token", "list_active_tokens", "list_audit_records", "list_tokens", "revoke_token", "touch_token", ]