"""Durable immutable corpus generations and an atomic ACTIVE pointer.""" from __future__ import annotations import json import fcntl import os import re import stat import shutil 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._fsync_directory = self._sync_root 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") if self.active_generation() == generation: return generation previous = self.active_generation() temporary = self.active_path.with_name(f".ACTIVE.{uuid.uuid4().hex}.tmp") replaced = False try: _atomic_write(temporary, (generation + "\n").encode()) self._replace(temporary, self.active_path) replaced = True self._fsync_directory() except BaseException: temporary.unlink(missing_ok=True) if replaced: if previous is None: self.active_path.unlink(missing_ok=True) else: rollback = self.active_path.with_name(f".ACTIVE.rollback.{uuid.uuid4().hex}.tmp") _atomic_write(rollback, (previous + "\n").encode()) self._replace(rollback, self.active_path) self._sync_root() raise return generation def _sync_root(self) -> None: directory = os.open(self.root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) try: os.fsync(directory) finally: os.close(directory) 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 discard(self, generation: str) -> None: path = self.generation_path(generation) if path.exists(): if path.is_symlink() or not stat.S_ISDIR(path.lstat().st_mode): raise UnsafeCorpusPath("generation cleanup target is unsafe") shutil.rmtree(path) 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