273 lines
11 KiB
Python
273 lines
11 KiB
Python
"""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
|
|
import hashlib
|
|
import threading
|
|
from datetime import UTC, datetime
|
|
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._lock_state = threading.local()
|
|
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):
|
|
depth = getattr(self._lock_state, "depth", 0)
|
|
if depth:
|
|
self._lock_state.depth = depth + 1
|
|
try:
|
|
yield
|
|
finally:
|
|
self._lock_state.depth -= 1
|
|
return
|
|
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)
|
|
self._lock_state.depth = 1
|
|
yield
|
|
finally:
|
|
self._lock_state.depth = 0
|
|
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()
|
|
published_marker = self.generation_path(generation) / "PUBLISHED"
|
|
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()
|
|
_atomic_write(
|
|
published_marker,
|
|
(datetime.now(UTC).isoformat().replace("+00:00", "Z") + "\n").encode("ascii"),
|
|
)
|
|
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 list_generations(self) -> list[str]:
|
|
values = []
|
|
for entry in self.root.iterdir():
|
|
match = re.fullmatch(r"gen-([0-9a-f]{32})", entry.name)
|
|
if match and not entry.is_symlink() and stat.S_ISDIR(entry.lstat().st_mode):
|
|
values.append(f"gen:{match.group(1)}")
|
|
return sorted(values, key=lambda value: self.generation_path(value).stat().st_mtime_ns)
|
|
|
|
def published_generations(self) -> list[str]:
|
|
active = self.active_generation()
|
|
published = []
|
|
for generation in self.list_generations():
|
|
path = self.generation_path(generation)
|
|
marker = path / "PUBLISHED"
|
|
if generation != active and not marker.is_file():
|
|
continue
|
|
try:
|
|
manifest = self.manifest(generation)
|
|
if manifest.manifest_id != generation:
|
|
continue
|
|
timestamp = marker.read_text(encoding="ascii").strip() if marker.is_file() else ""
|
|
key = (timestamp or manifest.created_at.isoformat(), generation)
|
|
published.append((key, generation))
|
|
except (OSError, ValueError):
|
|
continue
|
|
return [generation for _, generation in sorted(published)]
|
|
|
|
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
|
|
parts = Path(relative).parts
|
|
if Path(relative).is_absolute() or parts[:1] != ("documents",) or len(parts) != 2:
|
|
raise UnsafeCorpusPath("materialized document path is unsafe")
|
|
return self.generation_path(generation) / relative
|
|
|
|
def read_document(self, document_id: str, generation: str | None = None) -> str | None:
|
|
generation = generation or self.active_generation()
|
|
if generation is None:
|
|
return None
|
|
manifest = self.manifest(generation)
|
|
path = self.resolve_document(document_id, generation)
|
|
document = next((item for item in manifest.documents if item.document_id == document_id), None)
|
|
if path is None or document is None:
|
|
return None
|
|
generation_fd = os.open(self.generation_path(generation), os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
|
documents_fd = fd = None
|
|
try:
|
|
documents_fd = os.open("documents", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=generation_fd)
|
|
fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC, dir_fd=documents_fd)
|
|
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("materialized document is unsafe")
|
|
payload = os.read(fd, info.st_size + 1)
|
|
if len(payload) != info.st_size or "sha256:" + hashlib.sha256(payload).hexdigest() != document.content_hash:
|
|
raise UnsafeCorpusPath("materialized document hash mismatch")
|
|
return payload.decode("utf-8")
|
|
except (OSError, UnicodeError) as error:
|
|
raise UnsafeCorpusPath("materialized document read failed") from error
|
|
finally:
|
|
if fd is not None:
|
|
os.close(fd)
|
|
if documents_fd is not None:
|
|
os.close(documents_fd)
|
|
os.close(generation_fd)
|
|
|
|
def materialize_document(
|
|
self, document_id: str, destination: Path, generation: str | None = None,
|
|
) -> Path | None:
|
|
content = self.read_document(document_id, generation)
|
|
if content is None:
|
|
return None
|
|
destination = Path(destination)
|
|
destination.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
_atomic_write(destination, content.encode("utf-8"))
|
|
destination.chmod(0o400)
|
|
return destination
|