atif-sql 0.2.0__py3-none-any.whl

atif_corpus/infrastructure/atomic.py

@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Atomic file writes: tmp-in-same-dir + fsync + rename, never a partial artifact.
+
+atif-duck reads the corpus with no coordination — no locks, no journal — so
+the ONLY correctness lever the writer has is that every artifact path either
+holds a complete previous version or a complete new version. ``os.replace``
+of a sibling tmp file gives exactly that on POSIX (same filesystem, atomic
+rename); a failure mid-serialization leaves the destination untouched and
+the tmp file removed.
+
+Durability, not just visibility: write ORDER is not persistence order across
+a kernel crash or power loss — a journalling filesystem can replay the
+renames while the data blocks are still dirty, publishing a ``meta.json``
+(which atif-duck treats as "this session is complete") next to a
+zero-length ``trajectory.json``. Every writer here therefore fsyncs the tmp
+file BEFORE the rename, and the directory holding the renamed entries after,
+so persistence order matches write order. The corpus is regenerable, so this
+is cheap insurance rather than a hot-path cost.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Iterator
+    from pathlib import Path
+
+#: Compact separators for trajectory.json, fixed by CONTRACT.md.
+COMPACT_SEPARATORS: tuple[str, str] = (",", ":")
+
+#: zstd level for the compressed artifacts. The library default: on the
+#: copied corpora it stores the JSON artifacts at about a tenth of their size
+#: at a few hundred MB/s on one core, and the level is part of what makes a
+#: pool worker and ``--workers 1`` write the same bytes, so it is pinned here
+#: rather than left to the library.
+ZSTD_LEVEL: int = 3
+
+
+def _tmp_sibling(path: Path) -> Path:
+    """A tmp path in the SAME directory (rename across filesystems isn't atomic)."""
+    return path.with_name(f".{path.name}.tmp-{os.getpid()}")
+
+
+def fsync_dir(path: Path) -> None:
+    """Flush ``path``'s directory entries so renames into it survive power loss.
+
+    A rename is metadata: without this the entry can be lost even though the
+    file it names was fsynced. Best-effort — a filesystem that refuses
+    ``O_RDONLY`` fsync on a directory must not fail a materialization pass
+    over a durability upgrade.
+    """
+    try:
+        fd = os.open(path, os.O_RDONLY)
+    except OSError as error:
+        logger.debug("fsync_dir: cannot open {} ({}); skipping", path, error)
+        return
+    try:
+        os.fsync(fd)
+    except OSError as error:
+        logger.debug("fsync_dir: fsync of {} failed ({}); skipping", path, error)
+    finally:
+        os.close(fd)
+
+
+def write_json_atomic(
+    path: Path,
+    obj: Any,
+    *,
+    compact: bool = False,
+    indent: int | None = None,
+) -> None:
+    """Serialize ``obj`` as JSON to ``path`` atomically and durably.
+
+    Streams ``json.dump`` into the tmp file — an object that fails to
+    serialize partway (the mid-write failure the tests pin) leaves a partial
+    TMP file, which is unlinked, and never a partial destination. The tmp
+    file is fsynced before the rename so its bytes are on stable storage
+    before the name that publishes them, and the parent directory after so
+    the rename itself survives power loss.
+    """
+    tmp = _tmp_sibling(path)
+    try:
+        with tmp.open("w", encoding="utf-8") as handle:
+            if compact:
+                json.dump(obj, handle, separators=COMPACT_SEPARATORS)
+            else:
+                json.dump(obj, handle, indent=indent)
+            handle.write("\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        tmp.replace(path)
+    except BaseException:
+        logger.debug("atomic write failed; removing tmp {}", tmp)
+        tmp.unlink(missing_ok=True)
+        raise
+    fsync_dir(path.parent)
+
+
+def _json_chunks(obj: Any, *, compact: bool) -> Iterator[str]:
+    """``obj`` serialized in pieces, exactly as :func:`write_json_atomic` serializes it."""
+    encoder = json.JSONEncoder(separators=COMPACT_SEPARATORS if compact else None)
+    yield from encoder.iterencode(obj)
+    yield "\n"
+
+
+def _zstd_compressor() -> Any:
+    import zstandard
+
+    return zstandard.ZstdCompressor(level=ZSTD_LEVEL, write_content_size=True, threads=0)
+
+
+def _write_zstd_chunks(path: Path, chunks: Callable[[], Iterable[bytes]], size: int) -> None:
+    """Compress ``chunks()`` (``size`` bytes in all) into ``path`` atomically and durably.
+
+    The frame records ``size`` as its content size, so a reader learns the
+    document's size from the first bytes of the file. Same tmp, fsync, rename,
+    fsync-dir discipline as :func:`write_json_atomic`.
+    """
+    tmp = _tmp_sibling(path)
+    try:
+        with tmp.open("wb") as handle:
+            with _zstd_compressor().stream_writer(handle, size=size, closefd=False) as writer:
+                for chunk in chunks():
+                    writer.write(chunk)
+            handle.flush()
+            os.fsync(handle.fileno())
+        tmp.replace(path)
+    except BaseException:
+        logger.debug("atomic write failed; removing tmp {}", tmp)
+        tmp.unlink(missing_ok=True)
+        raise
+    fsync_dir(path.parent)
+
+
+def write_json_zstd_atomic(path: Path, obj: Any, *, compact: bool = False) -> None:
+    """Write what :func:`write_json_atomic` would write, zstd-compressed, to ``path``.
+
+    Serializes twice: once to count the bytes (the frame header carries the
+    content size, which readers size their JSON parser from) and once into
+    the compressor, so no serialized copy of a large document is ever held in
+    memory. Decompressed, the file is byte for byte the plain artifact.
+    """
+
+    def chunks() -> Iterator[bytes]:
+        for piece in _json_chunks(obj, compact=compact):
+            yield piece.encode("utf-8")
+
+    size = sum(len(chunk) for chunk in chunks())
+    _write_zstd_chunks(path, chunks, size)
+
+
+def write_text_zstd_atomic(path: Path, text: str) -> None:
+    """Write ``text`` zstd-compressed to ``path`` atomically and durably (see :func:`write_text_atomic`)."""
+    data = text.encode("utf-8")
+    _write_zstd_chunks(path, lambda: (data,), len(data))
+
+
+def replace_dir_atomic(tmp_dir: Path, dst_dir: Path) -> None:
+    """Swap a fully-written ``tmp_dir`` into place at ``dst_dir``.
+
+    The directory rename is the atomicity boundary for an artifact SET: a
+    reader either sees the complete previous generation at ``dst_dir`` or
+    the complete new one, never a mix. ``tmp_dir`` must be on the same
+    filesystem AND must live OUTSIDE any directory readers glob (DuckDB's
+    ``read_json`` glob matches dot-dirs, verified on 1.5) — atif-corpus
+    stages under ``<corpus_root>/.staging/``, which no reader glob touches.
+
+    ``os.replace`` cannot overwrite a non-empty directory, so an existing
+    ``dst_dir`` is first renamed aside NEXT TO ``tmp_dir`` (outside the
+    globbed tree) and removed after the swap; an EXCEPTION between the two
+    renames triggers the restore below.
+
+    What this does NOT survive: a SIGKILL or power loss in the window
+    between the aside-rename and the swap leaves NO generation at
+    ``dst_dir`` — no in-process handler can run there, and POSIX offers no
+    atomic directory exchange (``renameat2(RENAME_EXCHANGE)`` is
+    Linux-specific and unexposed). Recovery is owned by the next pass, not
+    by this function:
+    :func:`atif_corpus.application.materialize.materialize` sweeps
+    ``.staging/`` at pass start and force-replans any session the watermark
+    claims is current but whose dir is absent. A reader that catches the
+    window sees a missing session, never a torn one.
+    """
+    old = tmp_dir.with_name(f"{tmp_dir.name}.old-{os.getpid()}")
+    shutil.rmtree(old, ignore_errors=True)
+    fsync_dir(tmp_dir)
+    if dst_dir.is_dir():
+        dst_dir.rename(old)
+    try:
+        tmp_dir.replace(dst_dir)
+    except BaseException:
+        logger.debug("dir swap failed; restoring previous generation at {}", dst_dir)
+        if old.is_dir() and not dst_dir.exists():
+            old.rename(dst_dir)
+        raise
+    fsync_dir(dst_dir.parent)
+    shutil.rmtree(old, ignore_errors=True)
+
+
+#: Stored blobs are read-only, like the columnar artifacts: nothing rewrites
+#: a content-addressed file in place.
+_BLOB_MODE: int = 0o444
+
+
+def write_blob_atomic(path: Path, data: bytes) -> bool:
+    """Store a content-addressed blob at ``path`` unless it is already there.
+
+    Returns ``True`` when this call wrote the file, ``False`` when a file of
+    the same size already held the name (the name IS the content hash, so a
+    same-size file under it is the same bytes, written by an earlier session
+    or a concurrent worker). A file of another size under the name is a torn
+    or foreign file and is replaced.
+
+    Same discipline as :func:`write_json_atomic`: a pid-suffixed tmp sibling,
+    fsynced before the rename that publishes it, the directory fsynced after.
+    Two workers storing the same blob at once each rename a complete file of
+    identical bytes into place, so neither can expose a partial one.
+    """
+    try:
+        if path.stat().st_size == len(data):
+            return False
+    except FileNotFoundError:
+        pass
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = _tmp_sibling(path)
+    try:
+        with tmp.open("wb") as handle:
+            handle.write(data)
+            handle.flush()
+            os.fsync(handle.fileno())
+        tmp.chmod(_BLOB_MODE)
+        tmp.replace(path)
+    except BaseException:
+        logger.debug("blob write failed; removing tmp {}", tmp)
+        tmp.unlink(missing_ok=True)
+        raise
+    fsync_dir(path.parent)
+    return True
+
+
+def write_text_atomic(path: Path, text: str) -> None:
+    """Write ``text`` to ``path`` atomically and durably.
+
+    Same discipline as :func:`write_json_atomic`: tmp sibling, fsync the data
+    before the rename that publishes it, fsync the parent directory after so
+    the rename survives power loss.
+    """
+    tmp = _tmp_sibling(path)
+    try:
+        with tmp.open("w", encoding="utf-8") as handle:
+            handle.write(text)
+            handle.flush()
+            os.fsync(handle.fileno())
+        tmp.replace(path)
+    except BaseException:
+        logger.debug("atomic write failed; removing tmp {}", tmp)
+        tmp.unlink(missing_ok=True)
+        raise
+    fsync_dir(path.parent)

atif_corpus/infrastructure/compress_artifacts.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Compress an existing session's plain JSON artifacts in place (``atif-sql corpus slim``).
+
+A session written before materialize stored its JSON artifacts compressed
+holds ``trajectory.json``, ``edges.jsonl`` and ``session_events.jsonl`` as
+plain files. :func:`compress_in_place` turns one of them into ``<name>.zst``
+without a moment where the session lacks a readable copy:
+
+1. stream the plain file through zstd into a pid-suffixed tmp sibling, the
+   frame recording the plain size as its content size (what the materialize
+   writer records too), then fsync it;
+2. read the tmp file back, decompressing, and compare its sha256 with the
+   plain file's; a difference unlinks the tmp file and raises, leaving the
+   plain file as it was;
+3. give the tmp file the plain file's mtime, then rename it to
+   ``<name>.zst`` and fsync the directory;
+4. unlink the plain file and fsync the directory again.
+
+Between 3 and 4 the session holds both spellings with the same content, and
+every reader prefers the compressed one. The mtime is carried over because
+analyze's checkpoints bound each session by its trajectory's mtime: a
+compressed copy with a new mtime would make every session look rewritten and
+re-run every pipeline on it.
+
+The same compressor settings as the materialize writer
+(:data:`atif_corpus.infrastructure.atomic.ZSTD_LEVEL`, single-threaded), so a
+slimmed session and a freshly converted one store the same bytes for the same
+document.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from loguru import logger
+
+from atif_corpus.domain.layout import COMPRESSED_ARTIFACT_FILENAMES, COMPRESSED_SUFFIX
+from atif_corpus.infrastructure.atomic import ZSTD_LEVEL, fsync_dir
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from pathlib import Path
+
+#: Read and write chunk size.
+_CHUNK_BYTES: int = 1024 * 1024
+
+
+class CompressionMismatchError(RuntimeError):
+    """The compressed copy did not decompress to the plain file's bytes; nothing was replaced."""
+
+
+@dataclass(frozen=True, slots=True)
+class Compressed:
+    """One artifact compressed (or, on a dry run, measured)."""
+
+    #: The plain file (removed unless this was a dry run).
+    plain: Path
+    #: Its size.
+    plain_bytes: int
+    #: The compressed file's size (on a dry run, what it would be).
+    compressed_bytes: int
+
+
+def plain_artifacts(session_dir: Path) -> list[Path]:
+    """The session's artifacts still stored plain, in :data:`COMPRESSED_ARTIFACT_FILENAMES` order."""
+    return [
+        session_dir / name
+        for name in COMPRESSED_ARTIFACT_FILENAMES
+        if (session_dir / name).is_file()
+    ]
+
+
+def _compressed_chunks(plain: Path, size: int) -> Iterator[bytes]:
+    """``plain``'s bytes as one zstd frame, in pieces (a dry run counts them, a real run writes them)."""
+    import zstandard
+
+    compressor = zstandard.ZstdCompressor(level=ZSTD_LEVEL, write_content_size=True, threads=0)
+    stream = compressor.compressobj(size=size)
+    with plain.open("rb") as source:
+        while chunk := source.read(_CHUNK_BYTES):
+            if out := stream.compress(chunk):
+                yield out
+    yield stream.flush()
+
+
+def _sha256_plain(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _sha256_decompressed(path: Path) -> str:
+    import zstandard
+
+    digest = hashlib.sha256()
+    with path.open("rb") as handle, zstandard.ZstdDecompressor().stream_reader(handle) as reader:
+        while chunk := reader.read(_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _check_round_trip(compressed: Path, plain: Path) -> None:
+    if _sha256_decompressed(compressed) != _sha256_plain(plain):
+        msg = f"{compressed} does not decompress to {plain}; left {plain} as it was"
+        raise CompressionMismatchError(msg)
+
+
+def measure_compressed(plain: Path) -> Compressed:
+    """What :func:`compress_in_place` would store for ``plain``, without writing anything."""
+    size = plain.stat().st_size
+    compressed = sum(len(chunk) for chunk in _compressed_chunks(plain, size))
+    return Compressed(plain=plain, plain_bytes=size, compressed_bytes=compressed)
+
+
+def compress_in_place(plain: Path) -> Compressed:
+    """Replace ``plain`` with ``<plain>.zst`` holding the same bytes (see the module docstring).
+
+    Raises
+    ------
+    CompressionMismatchError
+        The compressed copy read back differently; ``plain`` is untouched.
+    OSError
+        The file or its directory went away (a materialize pass swapped the
+        session's directory meanwhile) or could not be written.
+    """
+    stat = plain.stat()
+    target = plain.with_name(f"{plain.name}{COMPRESSED_SUFFIX}")
+    tmp = plain.with_name(f".{target.name}.tmp-{os.getpid()}")
+    try:
+        with tmp.open("wb") as handle:
+            for chunk in _compressed_chunks(plain, stat.st_size):
+                handle.write(chunk)
+            handle.flush()
+            os.fsync(handle.fileno())
+        _check_round_trip(tmp, plain)
+        os.utime(tmp, ns=(stat.st_atime_ns, stat.st_mtime_ns))
+        compressed_bytes = tmp.stat().st_size
+        tmp.replace(target)
+    except BaseException:
+        tmp.unlink(missing_ok=True)
+        raise
+    fsync_dir(plain.parent)
+    plain.unlink()
+    fsync_dir(plain.parent)
+    logger.debug("slim: {} -> {} ({} -> {} bytes)", plain, target, stat.st_size, compressed_bytes)
+    return Compressed(plain=plain, plain_bytes=stat.st_size, compressed_bytes=compressed_bytes)
+
+
+__all__ = [
+    "Compressed",
+    "CompressionMismatchError",
+    "compress_in_place",
+    "measure_compressed",
+    "plain_artifacts",
+]

atif_corpus/infrastructure/fake_converter.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Test adapter for :class:`atif_corpus.domain.ports.ConverterPort`.
+
+Lives in ``infrastructure`` (not ``tests/``) on purpose: atif-cli's tests
+and future integration harnesses need the same fake, and a fake that ships
+with the package is type-checked against the port on every ``ty`` run
+instead of drifting in a conftest.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import time
+from typing import TYPE_CHECKING, Any
+
+import zstandard
+
+from atif_corpus.domain.ports import ArchivedSource, ConversionOutput, EmptySourceError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+def _archive_sources(session_jsonl: Path, archive_dir: Path) -> tuple[ArchivedSource, ...]:
+    """Write a zstd copy of the main file and every file under its side dir.
+
+    The fake reads the files itself; the real adapter takes the bytes from the
+    converter's verifying read instead (see ``atif_converter.infrastructure
+    .source_archive``). The on-disk shape is the same, which is all the use
+    case sees.
+    """
+    base = session_jsonl.parent
+    side_dir = base / session_jsonl.stem
+    paths = [session_jsonl]
+    if side_dir.is_dir():
+        paths.extend(p for p in sorted(side_dir.rglob("*")) if p.is_file() and not p.is_symlink())
+    archived: list[ArchivedSource] = []
+    for path in paths:
+        data = path.read_bytes()
+        relative = path.relative_to(base).as_posix()
+        target = archive_dir / f"{relative}.zst"
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_bytes(zstandard.ZstdCompressor(level=3).compress(data))
+        archived.append(
+            ArchivedSource(
+                relative_path=relative, size=len(data), sha256=hashlib.sha256(data).hexdigest()
+            )
+        )
+    return tuple(archived)
+
+
+#: This process's simulated clock, in seconds. Only :class:`FakeConverter`
+#: with ``simulated_seconds`` moves it, and only inside ``convert``.
+_simulated_now: float = 0.0
+
+
+def simulated_clock() -> float:
+    """A clock that advances only while a :class:`FakeConverter` converts.
+
+    Hand it to ``materialize(convert_clock=...)`` with
+    ``FakeConverter(simulated_seconds=s)`` and every conversion measures
+    exactly ``s``, in whichever process (inline or pool worker) runs it,
+    however long the pass really takes. Module-level, so it pickles into a
+    spawned worker by reference; each process keeps its own count.
+    """
+    return _simulated_now
+
+
+class FakeConverter:
+    """A deterministic in-memory converter with scriptable failures.
+
+    ``convert`` fabricates a minimal-but-contract-shaped output keyed by the
+    session stem, records every call in ``converted``, and raises for any
+    session listed in ``fail_sessions`` — which is how the "a failing
+    session is recorded and skipped" behavior gets exercised.
+
+    Two knobs exist for the process-pool tests. ``record_pid`` stamps the
+    converting process's pid into the trajectory as ``worker_pid``, so a test
+    can read back from disk WHICH process converted each session — a pool
+    that quietly ran inline shows the parent's pid on every one.
+    ``delay_seconds`` holds each conversion open long enough for the pool's
+    other workers to start and take a session of their own; without it one
+    fast worker can drain a small fixture before the second exists. Both are
+    off by default, and ``converted`` is only meaningful on the inline path:
+    a pool worker's copy of this object never comes back.
+
+    ``simulated_seconds`` advances :func:`simulated_clock` by that much on
+    every conversion, so a pass timed with that clock reports a known
+    duration per session.
+
+    ``empty_sessions`` raises :class:`~atif_corpus.domain.ports.EmptySourceError`
+    for the named sessions, the port's "nothing to convert" verdict.
+    ``output_tag`` stamps ``converter_tag`` into every trajectory so a test can
+    read back which converter build wrote a session. Handed an
+    ``archive_dir``, the fake writes a real zstd archive of the session's files.
+    """
+
+    def __init__(
+        self,
+        *,
+        fail_sessions: frozenset[str] = frozenset(),
+        empty_sessions: frozenset[str] = frozenset(),
+        record_pid: bool = False,
+        delay_seconds: float = 0.0,
+        simulated_seconds: float = 0.0,
+        output_tag: str | None = None,
+    ) -> None:
+        #: Session ids whose conversion should raise.
+        self.fail_sessions = fail_sessions
+        #: Session ids whose conversion raises :class:`EmptySourceError`.
+        self.empty_sessions = empty_sessions
+        #: Stamped into the trajectory as ``converter_tag``, so a test can see
+        #: which converter wrote the artifacts on disk (a version bump's re-run).
+        self.output_tag = output_tag
+        #: Stamp ``os.getpid()`` into the trajectory as ``worker_pid``.
+        self.record_pid = record_pid
+        #: Seconds each conversion sleeps before returning.
+        self.delay_seconds = delay_seconds
+        #: Seconds each conversion advances :func:`simulated_clock` by.
+        self.simulated_seconds = simulated_seconds
+        #: Every session path convert() was asked about, in call order.
+        self.converted: list[Path] = []
+
+    def convert(self, session_jsonl: Path, *, archive_dir: Path | None = None) -> ConversionOutput:
+        """Fabricate contract-shaped artifacts for ``session_jsonl``."""
+        session_id = session_jsonl.stem
+        self.converted.append(session_jsonl)
+        if session_id in self.fail_sessions:
+            msg = f"scripted failure for {session_id}"
+            raise RuntimeError(msg)
+        if session_id in self.empty_sessions:
+            msg = f"no convertible events in {session_jsonl}"
+            raise EmptySourceError(msg)
+        if self.delay_seconds:
+            time.sleep(self.delay_seconds)
+        if self.simulated_seconds:
+            global _simulated_now  # noqa: PLW0603 — the clock is per process by design
+            _simulated_now += self.simulated_seconds
+        trajectory: dict[str, Any] = {
+            "schema_version": "ATIF-v1.7",
+            "session_id": session_id,
+            "steps": [],
+        }
+        if self.record_pid:
+            trajectory["worker_pid"] = os.getpid()
+        if self.output_tag is not None:
+            trajectory["converter_tag"] = self.output_tag
+        loss_report: dict[str, Any] = {
+            "records_converted": 0,
+            "records_dropped": 0,
+            "gaps_observed": [],
+        }
+        edge: dict[str, Any] = {
+            "uuid": f"{session_id}-u1",
+            "parent_uuid": None,
+            "message_id": None,
+            "type": "user",
+            "ts": "2026-01-01T00:00:00Z",
+            "is_sidechain": False,
+            "is_compact_summary": False,
+            "source_file": str(session_jsonl),
+            "tool_use_ids": [],
+        }
+        return ConversionOutput(
+            trajectory_dict=trajectory,
+            loss_report_dict=loss_report,
+            edges_lines=[json.dumps(edge, separators=(",", ":"))],
+            source_archive=(
+                () if archive_dir is None else _archive_sources(session_jsonl, archive_dir)
+            ),
+        )