feat(preprocess): publish incremental Evidence corpus

harness/tht/corpus/store.py

@@ -0,0 +1,156 @@
+"""Durable immutable corpus generations and an atomic ACTIVE pointer."""
+
+from __future__ import annotations
+
+import json
+import fcntl
+import os
+import re
+import stat
+import uuid
+from pathlib import Path
+from contextlib import contextmanager
+
+from tht.corpus.models import CorpusManifest
+
+
+_GENERATION = re.compile(r"^gen:[0-9a-f]{32}$")
+
+
+class UnsafeCorpusPath(RuntimeError):
+    pass
+
+
+def _atomic_write(path: Path, payload: bytes) -> None:
+    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+    fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
+    try:
+        with os.fdopen(fd, "wb") as stream:
+            stream.write(payload)
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+        directory = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
+        try:
+            os.fsync(directory)
+        finally:
+            os.close(directory)
+    except BaseException:
+        temporary.unlink(missing_ok=True)
+        raise
+
+
+class CorpusStore:
+    def __init__(self, root: Path) -> None:
+        self.root = Path(root)
+        self.active_path = self.root / "ACTIVE"
+        self._replace = os.replace
+        self._ensure_root()
+
+    def _ensure_root(self) -> None:
+        if self.root.is_symlink():
+            raise UnsafeCorpusPath("corpus root must not be a symlink")
+        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
+        info = self.root.lstat()
+        if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid():
+            raise UnsafeCorpusPath("corpus root is unsafe")
+
+    @contextmanager
+    def writer_lock(self):
+        lock_path = self.root / ".writer.lock"
+        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600)
+        try:
+            info = os.fstat(fd)
+            if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1:
+                raise UnsafeCorpusPath("corpus writer lock is unsafe")
+            fcntl.flock(fd, fcntl.LOCK_EX)
+            yield
+        finally:
+            fcntl.flock(fd, fcntl.LOCK_UN)
+            os.close(fd)
+
+    def generation_path(self, generation: str) -> Path:
+        if not _GENERATION.fullmatch(generation):
+            raise UnsafeCorpusPath("invalid corpus generation")
+        path = self.root / generation.replace(":", "-")
+        if path.is_symlink():
+            raise UnsafeCorpusPath("generation must not be a symlink")
+        return path
+
+    def stage(
+        self,
+        manifest: CorpusManifest,
+        materialized: dict[str, str],
+        *,
+        generation: str | None = None,
+    ) -> str:
+        generation = generation or f"gen:{uuid.uuid4().hex}"
+        path = self.generation_path(generation)
+        try:
+            path.mkdir(mode=0o700)
+        except FileExistsError:
+            raise UnsafeCorpusPath("generation already exists") from None
+        documents = path / "documents"
+        documents.mkdir(mode=0o700)
+        files: dict[str, str] = {}
+        for document in manifest.documents:
+            relative = f"documents/{document.document_id.removeprefix('doc:')}.md"
+            _atomic_write(path / relative, materialized[document.document_id].encode("utf-8"))
+            files[document.document_id] = relative
+        payload = json.loads(manifest.model_dump_json())
+        metadata = payload["metadata"]
+        metadata["files"] = files
+        payload.update({"manifest_id": generation, "metadata": metadata})
+        staged = CorpusManifest.model_validate(payload)
+        _atomic_write(path / "manifest.json", (staged.model_dump_json(indent=2) + "\n").encode())
+        return generation
+
+    def publish(self, generation: str) -> str:
+        manifest = self.manifest(generation)
+        if manifest.manifest_id != generation:
+            raise UnsafeCorpusPath("manifest generation mismatch")
+        temporary = self.active_path.with_name(f".ACTIVE.{uuid.uuid4().hex}.tmp")
+        _atomic_write(temporary, (generation + "\n").encode())
+        self._replace(temporary, self.active_path)
+        directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
+        try:
+            os.fsync(directory)
+        finally:
+            os.close(directory)
+        return generation
+
+    def active_generation(self) -> str | None:
+        try:
+            if self.active_path.is_symlink():
+                raise UnsafeCorpusPath("ACTIVE must not be a symlink")
+            value = self.active_path.read_text(encoding="ascii").strip()
+        except FileNotFoundError:
+            return None
+        if not _GENERATION.fullmatch(value):
+            raise UnsafeCorpusPath("ACTIVE contains an invalid generation")
+        return value
+
+    def manifest(self, generation: str) -> CorpusManifest:
+        path = self.generation_path(generation)
+        manifest_path = path / "manifest.json"
+        if manifest_path.is_symlink():
+            raise UnsafeCorpusPath("manifest must not be a symlink")
+        return CorpusManifest.model_validate_json(manifest_path.read_text(encoding="utf-8"))
+
+    def active_manifest(self) -> CorpusManifest | None:
+        generation = self.active_generation()
+        return self.manifest(generation) if generation else None
+
+    def resolve_document(self, document_id: str, generation: str | None = None) -> Path | None:
+        generation = generation or self.active_generation()
+        if generation is None:
+            return None
+        manifest = self.manifest(generation)
+        relative = manifest.metadata.get("files", {}).get(document_id)
+        if not isinstance(relative, str):
+            return None
+        base = self.generation_path(generation).resolve()
+        path = (base / relative).resolve()
+        if not path.is_relative_to(base) or path.is_symlink():
+            raise UnsafeCorpusPath("materialized document escapes its generation")
+        return path