lorebase-store 0.2.3__py3-none-any.whl

lorebase_store/repositories/revisions.py

@@ -0,0 +1,273 @@
+"""Revisions, their ordered parents and generation numbers (design 6.4)."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import Table, insert, select, text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+
+from lorebase_core.revision import (
+    Revision,
+    UnknownRevisionError,
+    next_generation,
+    revision_id,
+)
+from lorebase_store.models import metadata
+
+REVISIONS: Table = metadata.tables["revisions"]
+REVISION_PARENTS: Table = metadata.tables["revision_parents"]
+
+_FIRST_PARENT_CHAIN = text(
+    """
+    WITH RECURSIVE chain(id, depth) AS (
+        SELECT CAST(:head AS bytea), 0
+        UNION ALL
+        SELECT p.parent_id, chain.depth + 1
+        FROM chain
+        JOIN revision_parents p
+          ON p.domain_id = :domain AND p.revision_id = chain.id AND p.position = 0
+    )
+    SELECT id FROM chain ORDER BY depth
+    """
+)
+
+
+_CONTAINED = text(
+    """
+    WITH RECURSIVE walk(id) AS (
+        SELECT CAST(:head AS bytea)
+        UNION
+        SELECT p.parent_id
+        FROM walk
+        JOIN revision_parents p ON p.domain_id = :domain AND p.revision_id = walk.id
+        JOIN revisions r
+          ON r.domain_id = :domain AND r.id = p.parent_id AND r.generation >= :generation
+    )
+    SELECT id FROM walk WHERE id = ANY(CAST(:candidates AS bytea[]))
+    """
+)
+"""Candidates reachable from a head; the walk stops below the lowest candidate generation."""
+
+
+_AHEAD_BEHIND = text(
+    """
+    WITH RECURSIVE
+    left_side(id) AS (
+        SELECT CAST(:left AS bytea)
+        UNION
+        SELECT p.parent_id
+        FROM left_side
+        JOIN revision_parents p ON p.domain_id = :domain AND p.revision_id = left_side.id
+    ),
+    right_side(id) AS (
+        SELECT CAST(:right AS bytea)
+        UNION
+        SELECT p.parent_id
+        FROM right_side
+        JOIN revision_parents p ON p.domain_id = :domain AND p.revision_id = right_side.id
+    )
+    SELECT
+        (SELECT count(*) FROM left_side WHERE id NOT IN (SELECT id FROM right_side)),
+        (SELECT count(*) FROM right_side WHERE id NOT IN (SELECT id FROM left_side))
+    """
+)
+
+
+@dataclass(frozen=True, slots=True)
+class RevisionRow:
+    id: bytes
+    root: bytes
+    generation: int
+    parents: tuple[bytes, ...]
+    author: str
+    on_behalf_of: str | None
+    message: str
+    origin: dict[str, Any]
+    stats: dict[str, Any]
+    created_at: datetime
+
+
+class RevisionRepository:
+    """Revisions of one Domain."""
+
+    def __init__(self, session: AsyncSession, domain_id: str) -> None:
+        self._session = session
+        self._domain = domain_id
+
+    async def insert(self, revision: Revision) -> RevisionRow:
+        """Store a Revision whose parents are stored; returns it with id and generation."""
+        if revision.domain != self._domain:
+            raise ValueError(f"Revision of Domain {revision.domain!r} in {self._domain!r}")
+        generations = await self.generations(revision.parents)
+        missing = [p for p in revision.parents if p not in generations]
+        if missing:
+            raise UnknownRevisionError(missing[0].hex())
+        row = RevisionRow(
+            id=revision_id(revision),
+            root=revision.root,
+            generation=next_generation(generations.values()),
+            parents=tuple(revision.parents),
+            author=revision.author.principal,
+            on_behalf_of=revision.author.on_behalf_of,
+            message=revision.message,
+            origin=revision.origin.model_dump(),
+            stats=revision.stats.model_dump(),
+            created_at=revision.created_at,
+        )
+        await self._session.execute(
+            insert(REVISIONS).values(
+                domain_id=self._domain,
+                id=row.id,
+                root=row.root,
+                generation=row.generation,
+                author=row.author,
+                on_behalf_of=row.on_behalf_of,
+                message=row.message,
+                origin=row.origin,
+                stats=row.stats,
+                created_at=row.created_at,
+            )
+        )
+        if row.parents:
+            await self._session.execute(
+                insert(REVISION_PARENTS),
+                [
+                    {
+                        "domain_id": self._domain,
+                        "revision_id": row.id,
+                        "position": position,
+                        "parent_id": parent,
+                    }
+                    for position, parent in enumerate(row.parents)
+                ],
+            )
+        return row
+
+    async def get(self, revision: bytes) -> RevisionRow | None:
+        result = await self._session.execute(
+            select(REVISIONS).where(
+                REVISIONS.c.domain_id == self._domain, REVISIONS.c.id == revision
+            )
+        )
+        found = result.mappings().one_or_none()
+        if found is None:
+            return None
+        parents = await self.parents(revision)
+        values = {k: v for k, v in found.items() if k != "domain_id"}
+        return RevisionRow(parents=parents, **values)
+
+    async def require(self, revision: bytes) -> RevisionRow:
+        found = await self.get(revision)
+        if found is None:
+            raise UnknownRevisionError(revision.hex())
+        return found
+
+    async def parents(self, revision: bytes) -> tuple[bytes, ...]:
+        result = await self._session.execute(
+            select(REVISION_PARENTS.c.parent_id)
+            .where(
+                REVISION_PARENTS.c.domain_id == self._domain,
+                REVISION_PARENTS.c.revision_id == revision,
+            )
+            .order_by(REVISION_PARENTS.c.position)
+        )
+        return tuple(result.scalars())
+
+    async def generations(self, revisions: Iterable[bytes]) -> dict[bytes, int]:
+        wanted = list(dict.fromkeys(revisions))
+        if not wanted:
+            return {}
+        result = await self._session.execute(
+            select(REVISIONS.c.id, REVISIONS.c.generation).where(
+                REVISIONS.c.domain_id == self._domain, REVISIONS.c.id.in_(wanted)
+            )
+        )
+        return {rev: int(generation) for rev, generation in result}
+
+    async def ahead_behind(self, left: bytes, right: bytes) -> tuple[int, int]:
+        """Revisions reachable from `left` but not from `right`, and the other way round."""
+        if left == right:
+            return 0, 0
+        result = await self._session.execute(
+            _AHEAD_BEHIND, {"left": left, "right": right, "domain": self._domain}
+        )
+        ahead, behind = result.one()
+        return int(ahead), int(behind)
+
+    async def contained(self, head: bytes, candidates: Iterable[bytes]) -> set[bytes]:
+        """The candidates reachable from `head` (`head` included): its ancestors or itself.
+        Only Revisions at least as recent (by generation) as the oldest candidate are
+        walked."""
+        generations = await self.generations(candidates)
+        if not generations:
+            return set()
+        result = await self._session.execute(
+            _CONTAINED,
+            {
+                "head": head,
+                "domain": self._domain,
+                "generation": min(generations.values()),
+                "candidates": list(generations),
+            },
+        )
+        return set(result.scalars())
+
+    async def first_parent_chain(self, head: bytes) -> list[bytes]:
+        """`head` and its first-parent ancestors, newest first."""
+        result = await self._session.execute(
+            _FIRST_PARENT_CHAIN, {"head": head, "domain": self._domain}
+        )
+        return list(result.scalars())
+
+
+class PgRevisionGraph:
+    """Synchronous core `RevisionGraph` over a synchronous `Session` (see `PgNodeStore`)."""
+
+    def __init__(self, session: Session, domain_id: str) -> None:
+        self._session = session
+        self._domain = domain_id
+        self._parents: dict[bytes, tuple[bytes, ...]] = {}
+        self._generations: dict[bytes, int] = {}
+        self._roots: dict[bytes, bytes] = {}
+
+    def parents(self, revision_id: bytes) -> Sequence[bytes]:
+        if revision_id not in self._parents:
+            self._load(revision_id)
+        return self._parents[revision_id]
+
+    def generation(self, revision_id: bytes) -> int:
+        if revision_id not in self._generations:
+            self._load(revision_id)
+        return self._generations[revision_id]
+
+    def root(self, revision_id: bytes) -> bytes:
+        """Tree root of a Revision (the `revision_root` callable of a merge input)."""
+        if revision_id not in self._roots:
+            self._load(revision_id)
+        return self._roots[revision_id]
+
+    def _load(self, revision: bytes) -> None:
+        found = self._session.execute(
+            select(REVISIONS.c.generation, REVISIONS.c.root).where(
+                REVISIONS.c.domain_id == self._domain, REVISIONS.c.id == revision
+            )
+        ).one_or_none()
+        if found is None:
+            raise UnknownRevisionError(revision.hex())
+        generation, root = found
+        self._roots[revision] = root
+        parents = self._session.execute(
+            select(REVISION_PARENTS.c.parent_id)
+            .where(
+                REVISION_PARENTS.c.domain_id == self._domain,
+                REVISION_PARENTS.c.revision_id == revision,
+            )
+            .order_by(REVISION_PARENTS.c.position)
+        ).scalars()
+        self._generations[revision] = int(generation)
+        self._parents[revision] = tuple(parents)

lorebase_store/repositories/scope.py

@@ -0,0 +1,33 @@
+"""Row-level security scope of the current transaction (store schema notes, RLS)."""
+
+from __future__ import annotations
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_store.auth.scope import checked_domain_id
+
+DOMAIN_IDS_SETTING = "lorebase.domain_ids"
+"""Transaction setting listing, comma-separated, the Domains the transaction may touch."""
+
+_ADD_DOMAIN = text(
+    """
+    SELECT set_config(
+        :name,
+        CASE
+            WHEN coalesce(current_setting(:name, true), '') = '' THEN :domain
+            WHEN :domain = ANY(string_to_array(current_setting(:name, true), ','))
+                THEN current_setting(:name, true)
+            ELSE current_setting(:name, true) || ',' || :domain
+        END,
+        true
+    )
+    """
+)
+
+
+async def scope_to_domain(session: AsyncSession, domain_id: str) -> None:
+    """Add `domain_id` to the Domains visible in the current transaction."""
+    await session.execute(
+        _ADD_DOMAIN, {"name": DOMAIN_IDS_SETTING, "domain": checked_domain_id(domain_id)}
+    )

lorebase_store/repositories/source_runs.py

@@ -0,0 +1,123 @@
+"""Synchronization runs of structured Sources (design 7.4, 9.1)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import Table, and_, insert, or_, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_store.models import metadata
+from lorebase_store.models.enums import RunStatus
+
+SOURCE_RUNS: Table = metadata.tables["source_runs"]
+
+
+@dataclass(frozen=True, slots=True)
+class SourceRunRow:
+    id: str
+    source_id: str
+    status: RunStatus
+    started_at: datetime | None
+    finished_at: datetime | None
+    stats: dict[str, Any]
+    revision_id: bytes | None = None
+    proposal_id: str | None = None
+    error: dict[str, Any] | None = None
+
+
+def _row(found: Any) -> SourceRunRow:
+    return SourceRunRow(**{k: v for k, v in found.items() if k != "domain_id"})
+
+
+class SourceRunRepository:
+    """Source runs of one Domain, in the caller's transaction (after `scope_to_domain`)."""
+
+    def __init__(self, session: AsyncSession, domain_id: str) -> None:
+        self._session = session
+        self._domain = domain_id
+
+    async def insert(self, row: SourceRunRow) -> None:
+        await self._session.execute(
+            insert(SOURCE_RUNS).values(
+                domain_id=self._domain,
+                id=row.id,
+                source_id=row.source_id,
+                status=row.status,
+                started_at=row.started_at,
+                finished_at=row.finished_at,
+                stats=row.stats,
+                revision_id=row.revision_id,
+                proposal_id=row.proposal_id,
+                error=row.error,
+            )
+        )
+
+    async def get(self, run_id: str) -> SourceRunRow | None:
+        found = (
+            (
+                await self._session.execute(
+                    select(SOURCE_RUNS).where(
+                        SOURCE_RUNS.c.domain_id == self._domain, SOURCE_RUNS.c.id == run_id
+                    )
+                )
+            )
+            .mappings()
+            .one_or_none()
+        )
+        return None if found is None else _row(found)
+
+    async def finish(
+        self,
+        run_id: str,
+        *,
+        status: RunStatus,
+        finished_at: datetime,
+        stats: dict[str, Any],
+        revision_id: bytes | None = None,
+        proposal_id: str | None = None,
+        error: dict[str, Any] | None = None,
+    ) -> None:
+        await self._session.execute(
+            update(SOURCE_RUNS)
+            .where(SOURCE_RUNS.c.domain_id == self._domain, SOURCE_RUNS.c.id == run_id)
+            .values(
+                status=status,
+                finished_at=finished_at,
+                stats=stats,
+                revision_id=revision_id,
+                proposal_id=proposal_id,
+                error=error,
+            )
+        )
+
+    async def latest(
+        self, source_id: str, *, status: RunStatus | None = None
+    ) -> SourceRunRow | None:
+        """The most recently started run of a Source (with `status`, when given)."""
+        query = select(SOURCE_RUNS).where(
+            SOURCE_RUNS.c.domain_id == self._domain, SOURCE_RUNS.c.source_id == source_id
+        )
+        if status is not None:
+            query = query.where(SOURCE_RUNS.c.status == status)
+        query = query.order_by(SOURCE_RUNS.c.started_at.desc(), SOURCE_RUNS.c.id.desc()).limit(1)
+        found = (await self._session.execute(query)).mappings().one_or_none()
+        return None if found is None else _row(found)
+
+    async def page(
+        self, source_id: str, *, limit: int, after: tuple[datetime, str] | None = None
+    ) -> list[SourceRunRow]:
+        """Runs of a Source, newest first, after the keyset `(started_at, id)`."""
+        table = SOURCE_RUNS
+        query = select(table).where(
+            table.c.domain_id == self._domain, table.c.source_id == source_id
+        )
+        if after is not None:
+            at, last_id = after
+            query = query.where(
+                or_(table.c.started_at < at, and_(table.c.started_at == at, table.c.id < last_id))
+            )
+        query = query.order_by(table.c.started_at.desc(), table.c.id.desc()).limit(limit)
+        return [_row(row) for row in (await self._session.execute(query)).mappings()]