lorebase-store 0.2.3__py3-none-any.whl

lorebase_store/migrations.py

@@ -0,0 +1,133 @@
+"""Run Alembic migrations programmatically and expose the expected head revisions.
+
+The Alembic environment lives in `packages/store/migrations` (shipped in the wheel as
+`lorebase_store/_migrations`). Settings not passed explicitly are loaded with `load_settings()`
+by the environment; revisions read `config.attributes["store_settings"]`.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+from alembic import command
+from alembic.config import Config
+from alembic.script import ScriptDirectory
+from alembic.util import CommandError
+
+from lorebase_store.settings import DatabaseSettings, StoreSettings
+
+DATABASE_ATTRIBUTE = "database_settings"
+STORE_ATTRIBUTE = "store_settings"
+URL_ATTRIBUTE = "database_url"
+
+DYNAMIC_INDEX_PREFIXES = ("embeddings_hnsw_", "mat_records_attr_", "ix_mat_records_bm25")
+"""Indexes created from configuration or schema at runtime, the optional BM25 index
+(`lorebase_store.bm25`) included; schema comparison skips them."""
+
+PROFILE_INDEXES = ("ix_mat_records_values", "ix_mat_lexicon_branch_word")
+"""Indexes of the models that a storage profile may leave out (`lorebase_store.storage`);
+schema comparison skips them whether they exist or not."""
+
+PARTITION_NAME = re.compile(r"^(?P<parent>[a-z_]+)_p(?P<remainder>\d+)$")
+"""Hash partitions are named `<table>_p<remainder>`."""
+
+_PACKAGED_SCRIPTS = Path(__file__).resolve().parent / "_migrations"
+_SOURCE_SCRIPTS = Path(__file__).resolve().parents[2] / "migrations"
+
+
+def script_location() -> Path:
+    """Directory of the Alembic environment: the packaged copy, else the source tree."""
+    return _PACKAGED_SCRIPTS if _PACKAGED_SCRIPTS.is_dir() else _SOURCE_SCRIPTS
+
+
+def include_object(
+    obj: Any, name: str | None, type_: str, reflected: bool, compare_to: Any
+) -> bool:
+    """Schema comparison filter: skip hash partitions, runtime-managed indexes and the
+    indexes a storage profile manages."""
+    if type_ == "index" and name in PROFILE_INDEXES:
+        return False
+    if not reflected or compare_to is not None or name is None:
+        return True
+    if type_ == "table":
+        return PARTITION_NAME.match(name) is None
+    if type_ == "index":
+        return not name.startswith(DYNAMIC_INDEX_PREFIXES)
+    return True
+
+
+def alembic_config(
+    database: DatabaseSettings | None = None,
+    store: StoreSettings | None = None,
+    *,
+    url: str | None = None,
+) -> Config:
+    """Alembic configuration; `url` overrides `database.dsn`."""
+    config = Config()
+    config.set_main_option("script_location", str(script_location()))
+    for key, value in (
+        (DATABASE_ATTRIBUTE, database),
+        (STORE_ATTRIBUTE, store),
+        (URL_ATTRIBUTE, url),
+    ):
+        if value is not None:
+            config.attributes[key] = value
+    return config
+
+
+def head_revisions() -> tuple[str, ...]:
+    """Alembic head revisions of this code, for readiness checks."""
+    return tuple(ScriptDirectory.from_config(alembic_config()).get_heads())
+
+
+class RevisionMismatchError(CommandError):
+    """The database is at a revision this code does not know: a newer version migrated it."""
+
+
+def unknown_revisions(script: ScriptDirectory, current: Iterable[str]) -> tuple[str, ...]:
+    """Revisions of `current` (the database's) that the migration scripts of this code lack."""
+    known = {revision.revision for revision in script.walk_revisions()}
+    return tuple(sorted({revision for revision in current if revision not in known}))
+
+
+def check_known_revisions(script: ScriptDirectory, current: Iterable[str]) -> None:
+    """Refuse to migrate a database that a newer version already migrated.
+
+    Alembic alone stops with "Can't locate revision"; this names both sides and the way out.
+    """
+    unknown = unknown_revisions(script, current)
+    if not unknown:
+        return
+    heads = ", ".join(sorted(script.get_heads())) or "none"
+    raise RevisionMismatchError(
+        f"the database is at revision {', '.join(unknown)}, which this code does not know "
+        f"(its newest revision is {heads}): a newer Lorebase migrated this database. Run the "
+        "version that did (with docker compose: `docker compose build` then `up`, or "
+        "`make up`, which rebuild the images from the checkout), or restore a backup taken "
+        "before that upgrade. Older code cannot downgrade it."
+    )
+
+
+def upgrade(
+    revision: str = "head",
+    *,
+    database: DatabaseSettings | None = None,
+    store: StoreSettings | None = None,
+    url: str | None = None,
+) -> None:
+    """Upgrade the database to `revision`."""
+    command.upgrade(alembic_config(database, store, url=url), revision)
+
+
+def downgrade(
+    revision: str = "base",
+    *,
+    database: DatabaseSettings | None = None,
+    store: StoreSettings | None = None,
+    url: str | None = None,
+) -> None:
+    """Downgrade the database to `revision`."""
+    command.downgrade(alembic_config(database, store, url=url), revision)

lorebase_store/models/__init__.py

@@ -0,0 +1,136 @@
+"""SQLAlchemy models of every Lorebase table (design 7). `metadata` is the migration target."""
+
+from __future__ import annotations
+
+from lorebase_store.models.base import Base
+from lorebase_store.models.governance import (
+    ApiKey,
+    ApiKeyUsage,
+    AuditEntry,
+    AuditReadEntry,
+    Grant,
+    Group,
+    GroupMember,
+    OperationDeferredEvent,
+    OperationInvocation,
+    Principal,
+    RateLimitBucket,
+    WebSession,
+)
+from lorebase_store.models.ingestion import (
+    Document,
+    DocumentBlob,
+    DocumentChunk,
+    ExtractionCacheEntry,
+    ExtractionJob,
+    InferenceEndpointMode,
+    InstallationSecret,
+    SchemaSuggestion,
+    SourceRun,
+    VaultEntry,
+    VaultLineage,
+    VaultTombstone,
+)
+from lorebase_store.models.materialized import (
+    BranchMaterialization,
+    Embedding,
+    EmbeddingCursor,
+    MatHistory,
+    MatLexicon,
+    MatLineageMap,
+    MatLink,
+    MatRecord,
+    MatSchema,
+    StorageTask,
+)
+from lorebase_store.models.proposals import (
+    Conflict,
+    ConflictResolution,
+    MergeAttempt,
+    Proposal,
+    ProposalItem,
+    ProposalReview,
+)
+from lorebase_store.models.queues import Job, OutboxConsumer, OutboxEvent, OutboxPending
+from lorebase_store.models.tracing import (
+    AgentTraceSession,
+    AgentTraceSessionDomain,
+    AgentTraceStep,
+)
+from lorebase_store.models.versioning import (
+    ChangesetEntry,
+    Domain,
+    Provenance,
+    Ref,
+    RefLogEntry,
+    Revision,
+    RevisionParent,
+    StoredObject,
+    TreeNode,
+)
+from lorebase_store.models.webhooks import Webhook, WebhookDelivery
+
+metadata = Base.metadata
+
+__all__ = [
+    "AgentTraceSession",
+    "AgentTraceSessionDomain",
+    "AgentTraceStep",
+    "ApiKey",
+    "ApiKeyUsage",
+    "AuditEntry",
+    "AuditReadEntry",
+    "Base",
+    "BranchMaterialization",
+    "ChangesetEntry",
+    "Conflict",
+    "ConflictResolution",
+    "Document",
+    "DocumentBlob",
+    "DocumentChunk",
+    "Domain",
+    "Embedding",
+    "EmbeddingCursor",
+    "ExtractionCacheEntry",
+    "ExtractionJob",
+    "Grant",
+    "Group",
+    "GroupMember",
+    "InferenceEndpointMode",
+    "InstallationSecret",
+    "Job",
+    "MatHistory",
+    "MatLexicon",
+    "MatLineageMap",
+    "MatLink",
+    "MatRecord",
+    "MatSchema",
+    "MergeAttempt",
+    "OperationDeferredEvent",
+    "OperationInvocation",
+    "OutboxConsumer",
+    "OutboxEvent",
+    "OutboxPending",
+    "Principal",
+    "Proposal",
+    "ProposalItem",
+    "ProposalReview",
+    "Provenance",
+    "RateLimitBucket",
+    "Ref",
+    "RefLogEntry",
+    "Revision",
+    "RevisionParent",
+    "SchemaSuggestion",
+    "SourceRun",
+    "StorageTask",
+    "StoredObject",
+    "TreeNode",
+    "VaultEntry",
+    "VaultLineage",
+    "VaultTombstone",
+    "WebSession",
+    "Webhook",
+    "WebhookDelivery",
+    "metadata",
+]

lorebase_store/models/base.py

@@ -0,0 +1,64 @@
+"""Declarative base, constraint naming convention and shared column helpers."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from enum import StrEnum
+from typing import Annotated, Any
+
+from sqlalchemy import (
+    CheckConstraint,
+    DateTime,
+    Enum,
+    ForeignKeyConstraint,
+    LargeBinary,
+    MetaData,
+    Text,
+)
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import DeclarativeBase, mapped_column
+
+NAMING_CONVENTION = {
+    "ix": "ix_%(table_name)s_%(column_0_N_name)s",
+    "uq": "uq_%(table_name)s_%(column_0_N_name)s",
+    "ck": "ck_%(table_name)s_%(constraint_name)s",
+    "fk": "fk_%(table_name)s_%(column_0_N_name)s_%(referred_table_name)s",
+    "pk": "pk_%(table_name)s",
+}
+
+HASH_PARTITIONED = {"postgresql_partition_by": "HASH (domain_id)"}
+"""Table option for tables hash-partitioned on `domain_id` (partitions come from migrations)."""
+
+
+class Base(DeclarativeBase):
+    """Base of every Lorebase table."""
+
+    metadata = MetaData(naming_convention=NAMING_CONVENTION)
+    type_annotation_map = {  # noqa: RUF012
+        dict[str, Any]: JSONB,
+        bytes: LargeBinary,
+        str: Text,
+        datetime: DateTime(timezone=True),
+    }
+
+
+H32 = bytes
+"""A 32-byte SHA-256 digest stored as `bytea`."""
+
+DomainKey = Annotated[str, mapped_column(Text, primary_key=True)]
+"""`domain_id` as the leading primary key column."""
+
+
+def h32_check(column: str) -> CheckConstraint:
+    """Check that `column` holds exactly 32 bytes."""
+    return CheckConstraint(f"octet_length({column}) = 32", name=f"{column}_h32")
+
+
+def domain_fk(*, ondelete: str = "CASCADE") -> ForeignKeyConstraint:
+    """`domain_id` references `domains.id`; rows go with their Domain."""
+    return ForeignKeyConstraint(["domain_id"], ["domains.id"], ondelete=ondelete)
+
+
+def pg_enum(enum: type[StrEnum], name: str) -> Enum:
+    """Native Postgres enum type `name` storing the member values of `enum`."""
+    return Enum(enum, name=name, values_callable=lambda members: [m.value for m in members])

lorebase_store/models/enums.py

@@ -0,0 +1,132 @@
+"""Closed value sets stored as native Postgres enum types."""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+from lorebase_store.models.base import pg_enum
+
+
+class RefKind(StrEnum):
+    """`branch` moves; `tag` is an immutable pointer (design 6.5); `snapshot` holds the
+    materialized rows of a past Revision for reads at it (`materialize.snapshot`)."""
+
+    BRANCH = "branch"
+    TAG = "tag"
+    SNAPSHOT = "snapshot"
+
+
+class Materialization(StrEnum):
+    """Branch materialization mode (design 6.11)."""
+
+    FULL = "full"
+    OVERLAY = "overlay"
+    NONE = "none"
+    BUILDING = "building"
+    """A full branch whose rows a background job is still writing from its tree; reads of
+    its head and commits wait for it (`materialize.background`)."""
+
+
+class ProposalOrigin(StrEnum):
+    """What opened a Proposal."""
+
+    MANUAL = "manual"
+    SOURCE_SYNC = "source_sync"
+    EXTRACTION = "extraction"
+    OPERATION = "operation"
+
+
+class ProposalStatus(StrEnum):
+    """Proposal lifecycle."""
+
+    OPEN = "open"
+    MERGED = "merged"
+    CLOSED = "closed"
+
+
+class ReviewVerdict(StrEnum):
+    """Outcome of one Proposal review."""
+
+    APPROVE = "approve"
+    REQUEST_CHANGES = "request_changes"
+    COMMENT = "comment"
+
+
+class ItemDecision(StrEnum):
+    """Decision on one extraction Proposal item."""
+
+    PENDING = "pending"
+    ACCEPTED = "accepted"
+    REJECTED = "rejected"
+
+
+class MergeAttemptStatus(StrEnum):
+    """Merge attempt lifecycle (design 6.8 step 5, 6.10)."""
+
+    CONFLICTED = "conflicted"
+    STALE = "stale"
+    RESOLVED = "resolved"
+    MERGED = "merged"
+    ABANDONED = "abandoned"
+
+
+class PrincipalType(StrEnum):
+    """Kind of authenticated actor (design 13.2)."""
+
+    USER = "user"
+    SERVICE = "service"
+    API_KEY = "api_key"
+
+
+class GrantScope(StrEnum):
+    """What a grant applies to (design 10.2)."""
+
+    DOMAIN = "domain"
+    KIND = "kind"
+    OPERATION = "operation"
+    BRANCH = "branch"
+
+
+class InvocationStatus(StrEnum):
+    """Operation invocation status (design 10.1)."""
+
+    PENDING = "pending"
+    COMMITTED = "committed"
+    PROPOSED = "proposed"
+    REJECTED = "rejected"
+    DRY_RUN = "dry_run"
+    FAILED = "failed"
+
+
+class JobStatus(StrEnum):
+    """Background job lifecycle; `queued` jobs are claimed with SKIP LOCKED."""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class RunStatus(StrEnum):
+    """Lifecycle of Source runs and extraction jobs."""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+REF_KIND = pg_enum(RefKind, "ref_kind")
+MATERIALIZATION = pg_enum(Materialization, "materialization")
+PROPOSAL_ORIGIN = pg_enum(ProposalOrigin, "proposal_origin")
+PROPOSAL_STATUS = pg_enum(ProposalStatus, "proposal_status")
+REVIEW_VERDICT = pg_enum(ReviewVerdict, "review_verdict")
+ITEM_DECISION = pg_enum(ItemDecision, "item_decision")
+MERGE_ATTEMPT_STATUS = pg_enum(MergeAttemptStatus, "merge_attempt_status")
+PRINCIPAL_TYPE = pg_enum(PrincipalType, "principal_type")
+GRANT_SCOPE = pg_enum(GrantScope, "grant_scope")
+INVOCATION_STATUS = pg_enum(InvocationStatus, "invocation_status")
+JOB_STATUS = pg_enum(JobStatus, "job_status")
+RUN_STATUS = pg_enum(RunStatus, "run_status")