205 lines
6.7 KiB
Python
205 lines
6.7 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.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",
|
|
]
|