lorebase-store 0.2.3__py3-none-any.whl

lorebase_store/auth/sessions.py

@@ -0,0 +1,153 @@
+"""Browser sessions of the web app (design 13.2, Keeper as backend-for-frontend).
+
+The repository stores opaque values: the session id digest and the encrypted
+provider tokens are produced by Keeper. A session is found by the SHA-256 of its id.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+
+from sqlalchemy import delete, or_, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_store.models import WebSession
+
+
+@dataclass(frozen=True)
+class SessionRecord:
+    """A stored session: a digest and ciphertext only, never the session id."""
+
+    id_hash: bytes
+    principal_id: str
+    subject: str
+    external_groups: tuple[str, ...]
+    tokens: bytes | None
+    created_at: datetime
+    last_seen_at: datetime
+    expires_at: datetime
+    refresh_at: datetime | None
+
+
+def _record(row: WebSession) -> SessionRecord:
+    return SessionRecord(
+        id_hash=row.id_hash,
+        principal_id=row.principal_id,
+        subject=row.subject,
+        external_groups=tuple(str(group) for group in row.external_groups or ()),
+        tokens=row.tokens,
+        created_at=row.created_at,
+        last_seen_at=row.last_seen_at,
+        expires_at=row.expires_at,
+        refresh_at=row.refresh_at,
+    )
+
+
+class SessionRepository:
+    """Web sessions. The table is installation-wide: no Domain scope applies."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def add(
+        self,
+        *,
+        id_hash: bytes,
+        principal_id: str,
+        subject: str,
+        external_groups: Sequence[str],
+        tokens: bytes | None,
+        created_at: datetime,
+        expires_at: datetime,
+        refresh_at: datetime | None,
+    ) -> SessionRecord:
+        """Insert a session first seen at `created_at`."""
+        row = WebSession(
+            id_hash=id_hash,
+            principal_id=principal_id,
+            subject=subject,
+            external_groups=list(external_groups),
+            tokens=tokens,
+            created_at=created_at,
+            last_seen_at=created_at,
+            expires_at=expires_at,
+            refresh_at=refresh_at,
+        )
+        self._session.add(row)
+        await self._session.flush()
+        return _record(row)
+
+    async def get(self, id_hash: bytes) -> SessionRecord | None:
+        """The session whose id has this digest."""
+        row = await self._session.get(WebSession, id_hash)
+        return _record(row) if row else None
+
+    async def touch(self, id_hash: bytes, at: datetime, *, min_interval: timedelta) -> bool:
+        """Record activity at `at`, unless the recorded activity is more recent than
+        `min_interval`; keeps busy sessions from writing on every request."""
+        result = await self._session.execute(
+            update(WebSession)
+            .where(WebSession.id_hash == id_hash, WebSession.last_seen_at <= at - min_interval)
+            .values(last_seen_at=at)
+        )
+        return bool(result.rowcount)  # type: ignore[attr-defined]
+
+    async def claim_refresh(self, id_hash: bytes, now: datetime, until: datetime) -> bool:
+        """Take the right to refresh the provider tokens: true for one caller only while the
+        refresh is due, which moves the due time to `until` (a lease)."""
+        result = await self._session.execute(
+            update(WebSession)
+            .where(
+                WebSession.id_hash == id_hash,
+                WebSession.refresh_at.is_not(None),
+                WebSession.refresh_at <= now,
+            )
+            .values(refresh_at=until)
+        )
+        return bool(result.rowcount)  # type: ignore[attr-defined]
+
+    async def update_tokens(
+        self,
+        id_hash: bytes,
+        *,
+        tokens: bytes | None,
+        external_groups: Sequence[str],
+        refresh_at: datetime | None,
+    ) -> bool:
+        """Store refreshed tokens and groups; false when the session is gone."""
+        result = await self._session.execute(
+            update(WebSession)
+            .where(WebSession.id_hash == id_hash)
+            .values(tokens=tokens, external_groups=list(external_groups), refresh_at=refresh_at)
+        )
+        return bool(result.rowcount)  # type: ignore[attr-defined]
+
+    async def delete(self, id_hash: bytes) -> bool:
+        """End a session; false when it did not exist."""
+        result = await self._session.execute(
+            delete(WebSession).where(WebSession.id_hash == id_hash)
+        )
+        return bool(result.rowcount)  # type: ignore[attr-defined]
+
+    async def delete_for_principal(self, principal_id: str) -> int:
+        """End every session of a principal; the number ended."""
+        result = await self._session.execute(
+            delete(WebSession).where(WebSession.principal_id == principal_id)
+        )
+        return int(result.rowcount)  # type: ignore[attr-defined]
+
+    async def purge_expired(self, now: datetime, idle_before: datetime, *, limit: int) -> int:
+        """Delete up to `limit` sessions past their absolute expiry or idle since before
+        `idle_before`; the number deleted."""
+        stale = (
+            select(WebSession.id_hash)
+            .where(or_(WebSession.expires_at <= now, WebSession.last_seen_at <= idle_before))
+            .limit(limit)
+            .scalar_subquery()
+        )
+        result = await self._session.execute(
+            delete(WebSession).where(WebSession.id_hash.in_(stale))
+        )
+        return int(result.rowcount)  # type: ignore[attr-defined]

lorebase_store/auth/unit.py

@@ -0,0 +1,53 @@
+"""One transaction over the auth repositories."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from lorebase_store.auth.api_keys import ApiKeyRepository
+from lorebase_store.auth.audit import AuditRepository, AuditStream
+from lorebase_store.auth.grants import GrantRepository
+from lorebase_store.auth.principals import GroupRepository, PrincipalRepository
+from lorebase_store.auth.sessions import SessionRepository
+from lorebase_store.auth.usage import UsageRepository
+
+
+@dataclass(frozen=True)
+class AuthUnit:
+    """The auth repositories sharing one session and transaction."""
+
+    principals: PrincipalRepository
+    groups: GroupRepository
+    api_keys: ApiKeyRepository
+    grants: GrantRepository
+    audit: AuditRepository
+    read_audit: AuditRepository
+    """The read stream of the audit log (`audit_read_log`)."""
+    usage: UsageRepository
+    sessions: SessionRepository
+
+
+class SqlAuthStore:
+    """Opens transactions over the auth repositories. Commits on success, rolls back on error."""
+
+    def __init__(self, sessions: async_sessionmaker[AsyncSession]) -> None:
+        self._sessions = sessions
+
+    @asynccontextmanager
+    async def unit(self) -> AsyncIterator[AuthUnit]:
+        """A transaction: leave the block normally to commit."""
+        async with self._sessions() as session, session.begin():
+            yield AuthUnit(
+                principals=PrincipalRepository(session),
+                groups=GroupRepository(session),
+                api_keys=ApiKeyRepository(session),
+                grants=GrantRepository(session),
+                audit=AuditRepository(session),
+                read_audit=AuditRepository(session, AuditStream.READS),
+                usage=UsageRepository(session),
+                sessions=SessionRepository(session),
+            )

lorebase_store/auth/usage.py

@@ -0,0 +1,70 @@
+"""Daily usage of API keys, the counters behind per-key quotas (design 13.6).
+
+Each key has one row per UTC day. Writers add deltas, so several Keeper processes serving the
+same key sum up correctly; the row returned by an addition holds the totals of all of them.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import date, datetime
+
+from sqlalchemy import select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_store.models import ApiKeyUsage
+
+
+@dataclass(frozen=True)
+class UsageRecord:
+    """What a key used on one day."""
+
+    api_key_id: str
+    day: date
+    requests: int
+    context_tokens: int
+
+
+class UsageRepository:
+    """Daily usage counters of API keys."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def get(self, api_key_id: str, day: date) -> UsageRecord:
+        """The usage of a key on a day; zeros when nothing was recorded."""
+        row = (
+            await self._session.execute(
+                select(ApiKeyUsage.requests, ApiKeyUsage.context_tokens).where(
+                    ApiKeyUsage.api_key_id == api_key_id, ApiKeyUsage.day == day
+                )
+            )
+        ).one_or_none()
+        requests, tokens = (0, 0) if row is None else (int(row[0]), int(row[1]))
+        return UsageRecord(api_key_id, day, requests, tokens)
+
+    async def add(
+        self, api_key_id: str, day: date, *, requests: int, context_tokens: int, at: datetime
+    ) -> UsageRecord:
+        """Add to the counters of a key and day and return the new totals.
+
+        `IntegrityError` when the key does not exist (it was deleted meanwhile).
+        """
+        insertion = insert(ApiKeyUsage).values(
+            api_key_id=api_key_id,
+            day=day,
+            requests=requests,
+            context_tokens=context_tokens,
+            updated_at=at,
+        )
+        statement = insertion.on_conflict_do_update(
+            index_elements=[ApiKeyUsage.api_key_id, ApiKeyUsage.day],
+            set_={
+                "requests": ApiKeyUsage.requests + insertion.excluded.requests,
+                "context_tokens": ApiKeyUsage.context_tokens + insertion.excluded.context_tokens,
+                "updated_at": insertion.excluded.updated_at,
+            },
+        ).returning(ApiKeyUsage.requests, ApiKeyUsage.context_tokens)
+        total = (await self._session.execute(statement)).one()
+        return UsageRecord(api_key_id, day, int(total[0]), int(total[1]))

lorebase_store/blobs.py

@@ -0,0 +1,73 @@
+"""Blob store interface for document bytes and its default Postgres backend (design 9.2, AC-24).
+
+Blobs are content addressed: the reference names the backend and the SHA-256 of the bytes, so
+storing the same bytes twice keeps one copy. The Postgres backend writes `document_blobs` in the
+caller's transaction (row-level security applies); another backend (for example an
+S3-compatible store) implements the same protocol and ignores the session.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from typing import Protocol
+
+from sqlalchemy import Table, select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_store.models import metadata
+
+BLOBS: Table = metadata.tables["document_blobs"]
+POSTGRES_SCHEME = "postgres"
+TEXT_REF_KEY = "text_ref"
+"""Document metadata member holding the blob reference of the parsed text (UTF-8)."""
+
+
+class BlobNotFoundError(LookupError):
+    """No blob under that reference."""
+
+
+class BlobStore(Protocol):
+    """Content-addressed storage of document bytes."""
+
+    async def put(self, session: AsyncSession, domain_id: str, data: bytes) -> str:
+        """Store `data` and return its reference (idempotent per content)."""
+        ...
+
+    async def get(self, session: AsyncSession, domain_id: str, ref: str) -> bytes:
+        """The bytes stored under `ref`; raises `BlobNotFoundError`."""
+        ...
+
+
+def sha256(data: bytes) -> bytes:
+    return hashlib.sha256(data).digest()
+
+
+class PostgresBlobStore:
+    """`BlobStore` keeping the bytes in `document_blobs` (references `postgres:<hex>`)."""
+
+    async def put(self, session: AsyncSession, domain_id: str, data: bytes) -> str:
+        digest = sha256(data)
+        await session.execute(
+            insert(BLOBS)
+            .values(domain_id=domain_id, sha256=digest, content=data, size=len(data))
+            .on_conflict_do_nothing(index_elements=["domain_id", "sha256"])
+        )
+        return f"{POSTGRES_SCHEME}:{digest.hex()}"
+
+    async def get(self, session: AsyncSession, domain_id: str, ref: str) -> bytes:
+        scheme, _, digest = ref.partition(":")
+        if scheme != POSTGRES_SCHEME or not digest:
+            raise BlobNotFoundError(f"not a Postgres blob reference: {ref!r}")
+        try:
+            key = bytes.fromhex(digest)
+        except ValueError:
+            raise BlobNotFoundError(f"invalid blob reference: {ref!r}") from None
+        content = (
+            await session.execute(
+                select(BLOBS.c.content).where(BLOBS.c.domain_id == domain_id, BLOBS.c.sha256 == key)
+            )
+        ).scalar_one_or_none()
+        if content is None:
+            raise BlobNotFoundError(f"no blob {ref!r} in Domain {domain_id!r}")
+        return bytes(content)