atlas-iac/dockerfiles/hermes-hux-foundation/hux/store.py

"""Per-tenant durable storage on the tenant PVC.

Every path is derived from the caller's tenant slot and hashed subject, so
there is no shared cross-tenant store and no way to address another tenant's
data. Records are JSON documents with a ``revision`` for optimistic
concurrency; ledgers are append-only JSONL. Writes are atomic (temp file +
fsync + rename) and guarded by one lock per family per tenant.
"""

from __future__ import annotations

import json
import os
import re
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from collections.abc import Iterator

from hux.errors import Conflict, Invalid, NotFound, TooLarge
from hux.identity import Identity

ID_RE = re.compile(r"^[a-z]{2,6}_[A-Za-z0-9._-]{4,80}$")
LAYOUT_VERSION = 1
MAX_RECORD_BYTES = 256 * 1024
MAX_FAMILY_RECORDS = 20000


def now_iso() -> str:
    """RFC 3339 UTC timestamp with second precision and a Z suffix."""
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def new_id(prefix: str) -> str:
    """Opaque, prefix-typed identifier matching common.schema.json#/$defs/id."""
    return f"{prefix}_{os.urandom(8).hex()}{int(time.time() * 1000) % 100000:05d}"


def check_id(value: Any) -> str:
    """Return ``value`` when it is a well-formed record id, else raise Invalid."""
    if not isinstance(value, str) or not ID_RE.match(value):
        raise Invalid("malformed id")
    return value


def _atomic_write(path: Path, data: bytes) -> None:
    tmp = path.with_name(f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp")
    with open(tmp, "wb") as handle:
        handle.write(data)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(tmp, path)


class TenantStore:
    """Files for exactly one (tenant slot, subject) pair."""

    _locks: dict[str, threading.RLock] = {}
    _locks_guard = threading.Lock()

    def __init__(self, root: Path, identity: Identity) -> None:
        self.identity = identity
        self.root = Path(root) / "hux" / f"v{LAYOUT_VERSION}" / identity.tenant_slot / identity.subject
        self.root.mkdir(parents=True, exist_ok=True)

    # -- locking -----------------------------------------------------------
    def lock(self, family: str) -> threading.RLock:
        """One re-entrant lock per family per tenant directory."""
        key = f"{self.root}:{family}"
        with self._locks_guard:
            return self._locks.setdefault(key, threading.RLock())

    # -- documents ---------------------------------------------------------
    def _doc_path(self, family: str, record_id: str) -> Path:
        directory = self.root / family
        directory.mkdir(parents=True, exist_ok=True)
        return directory / f"{check_id(record_id)}.json"

    def get(self, family: str, record_id: str) -> dict[str, Any]:
        """Read one document or raise NotFound."""
        path = self._doc_path(family, record_id)
        if not path.exists():
            raise NotFound(f"{family} {record_id} not found")
        return json.loads(path.read_bytes())

    def exists(self, family: str, record_id: str) -> bool:
        """True when the document is present."""
        return self._doc_path(family, record_id).exists()

    def put(self, family: str, record: dict[str, Any], expected_revision: int | None = None) -> dict[str, Any]:
        """Create or replace a document, bumping ``revision``.

        ``expected_revision`` implements If-Match: it must equal the stored
        revision (or be None for a create of a new id) or Conflict is raised.
        """
        record_id = check_id(record.get("id"))
        with self.lock(family):
            path = self._doc_path(family, record_id)
            current = json.loads(path.read_bytes()) if path.exists() else None
            if current is None:
                if expected_revision not in (None, 0):
                    raise Conflict("record does not exist yet")
                if self.count(family) >= MAX_FAMILY_RECORDS:
                    raise TooLarge(f"{family} is full")
                revision = 1
            else:
                if expected_revision is not None and expected_revision != current.get("revision"):
                    raise Conflict(f"revision {expected_revision} does not match current revision {current.get('revision')}")
                revision = int(current.get("revision", 0)) + 1
            stored = {**record, "revision": revision}
            data = json.dumps(stored, sort_keys=True, separators=(",", ":")).encode()
            if len(data) > MAX_RECORD_BYTES:
                raise TooLarge("record exceeds size bound")
            _atomic_write(path, data)
            return stored

    def delete(self, family: str, record_id: str) -> None:
        """Remove a document; missing is not an error."""
        path = self._doc_path(family, record_id)
        with self.lock(family):
            if path.exists():
                path.unlink()

    def count(self, family: str) -> int:
        """Number of documents in a family."""
        directory = self.root / family
        return sum(1 for p in directory.glob("*.json")) if directory.exists() else 0

    def scan(self, family: str) -> Iterator[dict[str, Any]]:
        """Yield every document in a family, oldest file first."""
        directory = self.root / family
        if not directory.exists():
            return
        for path in sorted(directory.glob("*.json"), key=lambda p: (p.stat().st_mtime_ns, p.name)):
            yield json.loads(path.read_bytes())

    # -- ledgers -----------------------------------------------------------
    def _ledger_path(self, family: str, name: str) -> Path:
        directory = self.root / family
        directory.mkdir(parents=True, exist_ok=True)
        if not re.match(r"^[A-Za-z0-9._-]{1,120}$", name):
            raise Invalid("malformed ledger name")
        return directory / f"{name}.jsonl"

    def append(self, family: str, name: str, record: dict[str, Any]) -> None:
        """Append one JSON line with fsync."""
        line = json.dumps(record, sort_keys=True, separators=(",", ":")).encode() + b"\n"
        if len(line) > MAX_RECORD_BYTES:
            raise TooLarge("ledger record exceeds size bound")
        with self.lock(family):
            path = self._ledger_path(family, name)
            with open(path, "ab") as handle:
                handle.write(line)
                handle.flush()
                os.fsync(handle.fileno())

    def read(self, family: str, name: str) -> list[dict[str, Any]]:
        """Read a whole ledger; a torn trailing line from a crash is dropped."""
        path = self._ledger_path(family, name)
        if not path.exists():
            return []
        rows: list[dict[str, Any]] = []
        for line in path.read_bytes().split(b"\n"):
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                break
        return rows

    def ledgers(self, family: str) -> list[str]:
        """Names of the ledgers in a family."""
        directory = self.root / family
        return sorted(p.stem for p in directory.glob("*.jsonl")) if directory.exists() else []

    def rewrite(self, family: str, name: str, rows: list[dict[str, Any]]) -> None:
        """Replace a ledger atomically (used by retention purges)."""
        data = b"".join(json.dumps(r, sort_keys=True, separators=(",", ":")).encode() + b"\n" for r in rows)
        with self.lock(family):
            _atomic_write(self._ledger_path(family, name), data)

    # -- blobs -------------------------------------------------------------
    def put_blob(self, digest: str, data: bytes) -> Path:
        """Store content-addressed bytes; the caller has already verified the digest."""
        if not re.match(r"^[0-9a-f]{64}$", digest):
            raise Invalid("malformed digest")
        directory = self.root / "blobs" / digest[:2]
        directory.mkdir(parents=True, exist_ok=True)
        path = directory / digest
        if not path.exists():
            _atomic_write(path, data)
        return path

    def get_blob(self, digest: str) -> bytes:
        """Read content-addressed bytes or raise NotFound."""
        if not re.match(r"^[0-9a-f]{64}$", digest):
            raise Invalid("malformed digest")
        path = self.root / "blobs" / digest[:2] / digest
        if not path.exists():
            raise NotFound("blob not found")
        return path.read_bytes()

    # -- manifest ----------------------------------------------------------
    def manifest(self, contract_version: str) -> dict[str, Any]:
        """Read or create the ``hux.manifest.v1`` for this tenant directory."""
        path = self.root / "MANIFEST.json"
        with self.lock("manifest"):
            if path.exists():
                return json.loads(path.read_bytes())
            stamp = now_iso()
            record = {
                "schema": "hux.manifest.v1",
                "contract_version": contract_version,
                "data_layout_version": LAYOUT_VERSION,
                "min_reader_contract_version": "1.0.0",
                "created_at": stamp,
                "updated_at": stamp,
            }
            _atomic_write(path, json.dumps(record, sort_keys=True).encode())
            return record