gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/environment/schema/batch.py

@@ -0,0 +1,755 @@
+"""Rows in bulk: streamed from files, cut into chunks, checked and applied as a batch.
+
+The runtime side of world data ingestion at scale. :func:`iter_rows` streams the
+three rows-file shapes (JSON object of entity to rows, tagged JSONL, bare rows with
+an entity) without loading a file; :func:`chunk_rows` cuts a row stream into gzipped
+JSONL chunks on row boundaries; :func:`validate_chunks` checks every row's shape
+against the contract; :func:`build_batch` applies chunks onto a parent snapshot in
+dependency order and saves the new snapshot. Refusals everywhere are the one record
+of the contract (section 2)::
+
+    {"chunk": 0, "index": 17, "entity": "users", "primaryKey": {"id": "u17"},
+     "field": "email", "pointer": "/users/17/email", "detail": "..."}
+
+Standard library only, so the bridge and the api bundle can load it.
+"""
+
+from __future__ import annotations
+
+import gzip
+import hashlib
+import io
+import json
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import IO, Any, Iterable, Iterator
+
+from .compiler import Contract, quote
+from .store import SQLiteStore
+from .validation import ContractError, fail, validate
+
+CHUNK_BYTES = 64 * 1024 * 1024
+REPORT_CAP = 10_000
+_READ = 256 * 1024
+#: A first line longer than this is a one-line document, never a JSONL row.
+_LINE_CAP = 1024 * 1024
+
+
+def encode_row(value: Any) -> str:
+    """One JSONL line, the exact serialisation every client must match: compact
+    separators, non-ASCII kept as is, key order kept as given."""
+    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
+
+
+# ---------------------------------------------------------------- streaming rows
+
+
+class _Reader:
+    """A text stream with a growing buffer, for decoding JSON values one at a time.
+
+    The buffer only shrinks on ``compact()``, so a caller may note a position,
+    read ahead, and come back to it."""
+
+    def __init__(self, stream: IO[str]):
+        self._stream = stream
+        self.buffer = ""
+        self.pos = 0
+        self.eof = False
+
+    def fill(self) -> bool:
+        if self.eof:
+            return False
+        more = self._stream.read(_READ)
+        if not more:
+            self.eof = True
+            return False
+        self.buffer += more
+        return True
+
+    def compact(self) -> None:
+        self.buffer = self.buffer[self.pos :]
+        self.pos = 0
+
+    def peek(self) -> str | None:
+        """The next non-whitespace character, or None at the end."""
+        while True:
+            while self.pos < len(self.buffer) and self.buffer[self.pos] in " \t\r\n\ufeff":
+                self.pos += 1
+            if self.pos < len(self.buffer):
+                return self.buffer[self.pos]
+            if not self.fill():
+                return None
+
+    def take(self, char: str, where: str) -> None:
+        if self.peek() != char:
+            fail("rows", f"{where}: expected {char!r} in the JSON document", where)
+        self.pos += 1
+
+    def value(self, where: str) -> Any:
+        """Decode the next JSON value, reading more of the stream as needed."""
+        decoder = json.JSONDecoder()
+        self.peek()
+        while True:
+            try:
+                parsed, end = decoder.raw_decode(self.buffer, self.pos)
+            except ValueError:
+                if self.fill():
+                    continue
+                fail("rows", f"{where}: not valid JSON", where)
+            if end == len(self.buffer) and self.fill():
+                # The value may continue (a number, a literal) past the buffer's edge.
+                continue
+            self.pos = end
+            return parsed
+
+    def line(self, cap: int | None = None) -> str | None:
+        """The next line without its newline, or None at the end. With ``cap``, a line
+        longer than it is handed back unconsumed (the caller treats it as a document)."""
+        while True:
+            rest = self.buffer[self.pos :]
+            newline = rest.find("\n")
+            if newline >= 0:
+                self.pos += newline + 1
+                return rest[:newline]
+            if cap is not None and len(rest) > cap:
+                return rest
+            if not self.fill():
+                break
+        rest = self.buffer[self.pos :]
+        if not rest:
+            return None
+        self.pos = len(self.buffer)
+        return rest
+
+
+def _open_text(source: str | Path) -> IO[str]:
+    if str(source) == "-":
+        return io.TextIOWrapper(sys.stdin.buffer, encoding="utf-8")
+    path = Path(source)
+    if path.suffix == ".gz":
+        return io.TextIOWrapper(gzip.open(path, "rb"), encoding="utf-8")
+    return open(path, "r", encoding="utf-8")
+
+
+def _row(value: Any, where: str) -> dict:
+    if not isinstance(value, dict):
+        fail("rows", f"{where}: a row must be a JSON object", where)
+    return value
+
+
+def _stream_array(reader: _Reader, where: str) -> Iterator[dict]:
+    reader.take("[", where)
+    index = 0
+    while True:
+        if reader.peek() == "]":
+            reader.pos += 1
+            return
+        yield _row(reader.value(f"{where}[{index}]"), f"{where}[{index}]")
+        reader.compact()
+        index += 1
+        if reader.peek() == ",":
+            reader.pos += 1
+        elif reader.peek() != "]":
+            fail("rows", f"{where}: expected ',' or ']' after row {index - 1}", where)
+
+
+def _stream_document(
+    reader: _Reader, where: str, entity: str | None, seen: dict[str, int] | None = None
+) -> Iterator[tuple[str, dict]]:
+    """``{"<entity>": [row, ...], ...}`` streamed one row at a time. ``seen`` collects
+    every entity the document names, an empty array included (a replace names it)."""
+    reader.take("{", where)
+    while True:
+        if reader.peek() == "}":
+            reader.pos += 1
+            return
+        name = reader.value(where)
+        if not isinstance(name, str):
+            fail("rows", f"{where}: object keys must be entity names", where)
+        reader.take(":", where)
+        if reader.peek() != "[":
+            fail("rows", f'{where}: {name} must map to an array of rows (JSONL lines are {{"entity", "row"}})', where)
+        if entity is not None and name != entity:
+            fail("rows", f"{where} maps entities; drop --entity or make the file a bare array", where)
+        if seen is not None:
+            seen.setdefault(name, 0)
+        for row in _stream_array(reader, f"{where}:{name}"):
+            if seen is not None:
+                seen[name] += 1
+            yield name, row
+        if reader.peek() == ",":
+            reader.pos += 1
+        elif reader.peek() != "}":
+            fail("rows", f"{where}: expected ',' or '}}' after {name}", where)
+
+
+def _stream_jsonl(reader: _Reader, first: str, where: str, entity: str | None) -> Iterator[tuple[str, dict]]:
+    number = 0
+    line: str | None = first
+    while line is not None:
+        number += 1
+        if line.strip():
+            try:
+                item = json.loads(line)
+            except ValueError:
+                fail("rows", f"{where}:{number} is not JSON", where)
+            if entity is not None:
+                yield entity, _row(item, f"{where}:{number}")
+            elif not isinstance(item, dict) or set(item) != {"entity", "row"} or not isinstance(item["entity"], str):
+                fail(
+                    "rows",
+                    f'{where}:{number}: expected {{"entity": name, "row": {{…}}}} (or pass --entity for bare rows)',
+                    where,
+                )
+            else:
+                yield item["entity"], _row(item["row"], f"{where}:{number}")
+        reader.compact()
+        line = reader.line()
+
+
+def iter_file_rows(
+    stream: IO[str], where: str, *, entity: str | None = None, seen: dict[str, int] | None = None
+) -> Iterator[tuple[str, dict]]:
+    """Stream ``(entity, row)`` from one rows file of any of the three shapes. ``seen``,
+    when given, collects entity names to row counts, an entity with no rows included."""
+    reader = _Reader(stream)
+    head = reader.peek()
+    if head is None:
+        return
+    if head == "[":
+        if entity is None:
+            fail("rows", f"{where} is a bare array; say which entity it holds with --entity", where)
+        if seen is not None:
+            seen.setdefault(entity, 0)
+        for row in _stream_array(reader, where):
+            if seen is not None:
+                seen[entity] += 1
+            yield entity, row
+        return
+    if head != "{":
+        fail("rows", f"{where} is neither a JSON object of entity -> rows nor an array of rows", where)
+    start = reader.pos
+    first = reader.line(cap=_LINE_CAP)
+    parsed: Any = None
+    complete = False
+    if first is not None and len(first) <= _LINE_CAP:
+        try:
+            parsed = json.loads(first)
+            complete = True
+        except ValueError:
+            complete = False
+    if complete and isinstance(parsed, dict):
+        is_row = (entity is None and set(parsed) == {"entity", "row"}) or (
+            entity is not None and not (set(parsed) == {entity} and isinstance(parsed[entity], list))
+        )
+        if is_row:
+            for name, row in _stream_jsonl(reader, first, where, entity):
+                if seen is not None:
+                    seen[name] = seen.get(name, 0) + 1
+                yield name, row
+            return
+    # A document: rewind to its first character and stream it.
+    reader.pos = start
+    yield from _stream_document(reader, where, entity, seen)
+
+
+def iter_rows(
+    sources: Iterable[str | Path], *, entity: str | None = None, seen: dict[str, int] | None = None
+) -> Iterator[tuple[str, dict]]:
+    """Stream ``(entity, row)`` from rows files (``-`` is stdin, ``.gz`` is gunzipped)
+    in order, never holding a whole file."""
+    for source in sources:
+        where = str(source)
+        try:
+            stream = _open_text(source)
+        except OSError:
+            fail("rows", f"Cannot read {where}", where)
+        with stream:
+            try:
+                yield from iter_file_rows(stream, where, entity=entity, seen=seen)
+            except UnicodeError:
+                fail("rows", f"Cannot read {where} as UTF-8 text", where)
+
+
+def count_rows(sources: Iterable[str | Path], *, entity: str | None = None, cap: int | None = None) -> int:
+    """How many rows the files hold, stopping early past ``cap``."""
+    total = 0
+    for _ in iter_rows(sources, entity=entity):
+        total += 1
+        if cap is not None and total > cap:
+            break
+    return total
+
+
+# ---------------------------------------------------------------- chunks
+
+
+@dataclass
+class Chunk:
+    ordinal: int
+    path: Path
+    digest: str  # sha256 of the gzipped bytes: the chunk's address
+    bytes: int  # gzipped size
+    rows: int
+    entity: str | None
+    content_bytes: int  # the JSONL before gzip
+    content_sha256: str
+    first_row: int  # index of the chunk's first row in the input stream
+
+    def manifest_entry(self) -> dict:
+        return {
+            "ordinal": self.ordinal,
+            "digest": self.digest,
+            "bytes": self.bytes,
+            "rows": self.rows,
+            "entity": self.entity,
+            "chunkId": None,
+            "uploaded": False,
+        }
+
+
+class _ChunkWriter:
+    def __init__(self, path: Path, ordinal: int, entity: str | None, first_row: int):
+        self.path = path
+        self.ordinal = ordinal
+        self.entity = entity
+        self.first_row = first_row
+        self.rows = 0
+        self.content_bytes = 0
+        self._hash = hashlib.sha256()
+        self._raw = open(path, "wb")
+        self._gz = gzip.GzipFile(fileobj=self._raw, mode="wb", mtime=0)
+
+    def write(self, line: bytes) -> None:
+        self._gz.write(line)
+        self._hash.update(line)
+        self.rows += 1
+        self.content_bytes += len(line)
+
+    def close(self) -> Chunk:
+        self._gz.close()
+        self._raw.close()
+        digest = hashlib.sha256()
+        with open(self.path, "rb") as stream:
+            for block in iter(lambda: stream.read(_READ), b""):
+                digest.update(block)
+        return Chunk(
+            ordinal=self.ordinal,
+            path=self.path,
+            digest=digest.hexdigest(),
+            bytes=self.path.stat().st_size,
+            rows=self.rows,
+            entity=self.entity,
+            content_bytes=self.content_bytes,
+            content_sha256=self._hash.hexdigest(),
+            first_row=self.first_row,
+        )
+
+
+def chunk_rows(
+    rows: Iterable[tuple[str, dict]],
+    workdir: str | Path,
+    *,
+    chunk_bytes: int = CHUNK_BYTES,
+    chunk_rows: int | None = None,
+    entity: str | None = None,
+) -> Iterator[Chunk]:
+    """Cut a row stream into gzipped JSONL chunk files under ``workdir``.
+
+    Lines are tagged ``{"entity", "row"}`` unless ``entity`` names the one entity
+    every row belongs to, in which case they are bare rows. A chunk closes when the
+    next line would take it past ``chunk_bytes`` of JSONL (never mid-row; a single
+    row larger than the bound is a chunk of its own) or when it holds ``chunk_rows``
+    rows. Chunks are yielded as they close, so one chunk is on disk at a time."""
+    workdir = Path(workdir)
+    workdir.mkdir(parents=True, exist_ok=True)
+    if chunk_bytes < 1 or (chunk_rows is not None and chunk_rows < 1):
+        fail("chunk", "chunk bounds must be positive", "/chunk")
+    writer: _ChunkWriter | None = None
+    ordinal = 0
+    seen = 0
+    for name, row in rows:
+        if entity is not None and name != entity:
+            fail("rows", f"Row {seen} belongs to {name}; a chunk declared for {entity} holds only its rows", "/rows")
+        line = (encode_row(row) if entity is not None else encode_row({"entity": name, "row": row})) + "\n"
+        encoded = line.encode("utf-8")
+        if writer is not None and (
+            (chunk_rows is not None and writer.rows >= chunk_rows)
+            or (writer.content_bytes > 0 and writer.content_bytes + len(encoded) > chunk_bytes)
+        ):
+            yield writer.close()
+            writer = None
+        if writer is None:
+            writer = _ChunkWriter(workdir / f"chunk-{ordinal:05d}.jsonl.gz", ordinal, entity, seen)
+            ordinal += 1
+        writer.write(encoded)
+        seen += 1
+    if writer is not None:
+        yield writer.close()
+
+
+def iter_chunk_rows(path: str | Path, entity: str | None = None) -> Iterator[tuple[str, dict]]:
+    """Stream ``(entity, row)`` from one chunk file (gzip over JSONL; a plain file
+    is read as is), tagged lines or bare rows of the declared entity."""
+    path = Path(path)
+    with open(path, "rb") as probe:
+        compressed = probe.read(2) == b"\x1f\x8b"
+    opener = gzip.open if compressed else open
+    with opener(path, "rt", encoding="utf-8") as stream:  # type: ignore[call-overload]
+        reader = _Reader(stream)
+        first = reader.line()
+        if first is not None:
+            yield from _stream_jsonl(reader, first, str(path), entity)
+
+
+# ---------------------------------------------------------------- refusals
+
+
+def _field_of(error: ContractError) -> str | None:
+    path = error.detail.get("path") or ""
+    return path[1:] if path.startswith("/") and len(path) > 1 else None
+
+
+def refusal(
+    chunk: int | None,
+    index: int | None,
+    entity: str | None,
+    row: Any,
+    field_name: str | None,
+    detail: str,
+    contract: Contract | None = None,
+) -> dict:
+    """One refusal record. ``primaryKey`` is the row's key fields when it carries
+    them all, ``field`` the property refused (None for a row-level refusal)."""
+    key = None
+    if contract is not None and isinstance(row, dict) and entity in contract.entities:
+        fields = contract.entities[entity]["primary_key"]
+        if all(name in row for name in fields):
+            key = {name: row[name] for name in fields}
+    pointer = "/" + (entity or "")
+    if index is not None:
+        pointer += f"/{index}"
+    if field_name:
+        pointer += f"/{field_name}"
+    return {
+        "chunk": chunk,
+        "index": index,
+        "entity": entity,
+        "primaryKey": key,
+        "field": field_name,
+        "pointer": pointer,
+        "detail": detail,
+    }
+
+
+def shape_refusal(contract: Contract, chunk: int, index: int, entity: str, row: Any) -> dict | None:
+    """A row's shape against the contract: type, bounds, required, unknown property
+    and primary key present. No reference checks. None when the row conforms."""
+    if entity not in contract.entities:
+        return refusal(
+            chunk,
+            index,
+            entity,
+            row,
+            None,
+            f"Unknown entity {entity}; the world has {', '.join(sorted(contract.entities))}",
+        )
+    if not isinstance(row, dict):
+        return refusal(chunk, index, entity, row, None, "A row must be an object")
+    spec = contract.entities[entity]
+    for name in spec["primary_key"]:
+        if name not in row:
+            return refusal(chunk, index, entity, row, name, f"Row is missing primary key field {name!r}", contract)
+    try:
+        validate(spec["schema"], row, location=f"{spec['location']}/schema")
+    except ContractError as error:
+        return refusal(chunk, index, entity, row, _field_of(error), error.detail["message"], contract)
+    return None
+
+
+def write_report(records: Iterable[dict], destination: IO[str], *, cap: int = REPORT_CAP) -> int:
+    """Refusal records as JSONL, at most ``cap`` of them followed by
+    ``{"truncated": true, "total": N}`` when there were more. Returns the total."""
+    total = 0
+    for record in records:
+        total += 1
+        if total <= cap:
+            destination.write(encode_row(record) + "\n")
+    if total > cap:
+        destination.write(encode_row({"truncated": True, "total": total}) + "\n")
+    return total
+
+
+# ---------------------------------------------------------------- batches
+
+
+@dataclass
+class ChunkSpec:
+    ordinal: int
+    path: Path
+    entity: str | None = None
+
+
+@dataclass
+class _Staged:
+    """The valid rows of a batch after the in-batch merge, keyed the way the store keys."""
+
+    rows: dict[str, dict[str, dict]] = field(default_factory=dict)  # entity -> key -> row
+    origin: dict[str, dict[str, tuple[int, int]]] = field(default_factory=dict)  # entity -> key -> (chunk, index)
+    refusals: list[dict] = field(default_factory=list)
+    read: int = 0
+
+
+def dependency_order(contract: Contract) -> list[str]:
+    """Entities with the targets of relationships first; a cycle keeps name order."""
+    remaining = dict(sorted(contract.entities.items()))
+    ordered: list[str] = []
+    while remaining:
+        ready = [
+            name
+            for name, entity in remaining.items()
+            if all(rel["target"] == name or rel["target"] not in remaining for rel in entity["relationships"])
+        ]
+        if not ready:
+            ready = [next(iter(remaining))]
+        for name in ready:
+            ordered.append(name)
+            del remaining[name]
+    return ordered
+
+
+def _stage(
+    contract: Contract,
+    store: SQLiteStore,
+    chunks: Iterable[ChunkSpec],
+    merge: str,
+    staged: _Staged | None = None,
+) -> _Staged:
+    """Shape-check every row of the chunks and merge them by primary key."""
+    staged = staged or _Staged()
+    for chunk in chunks:
+        for index, (entity, row) in enumerate(iter_chunk_rows(chunk.path, chunk.entity)):
+            staged.read += 1
+            refused = shape_refusal(contract, chunk.ordinal, index, entity, row)
+            if refused is not None:
+                staged.refusals.append(refused)
+                continue
+            key = store._key(entity, row)
+            held = staged.rows.setdefault(entity, {})
+            origin = staged.origin.setdefault(entity, {})
+            if key in held and merge == "fill":
+                first = held[key]
+                for name, value in row.items():
+                    if name not in first:
+                        first[name] = value
+                continue
+            held[key] = row
+            origin[key] = (chunk.ordinal, index)
+    return staged
+
+
+def _reference_refusals(contract: Contract, store: SQLiteStore, staged: _Staged, mode: str) -> list[dict]:
+    """Every row whose relationship points at a key neither the parent snapshot
+    (minus replaced entities) nor the batch holds."""
+    refusals: list[dict] = []
+    known: dict[str, set[str]] = {}
+
+    def keys_of(target: str) -> set[str]:
+        if target not in known:
+            held = set(staged.rows.get(target, {}))
+            if not (mode == "replace" and target in staged.rows):
+                with store._trusted():
+                    held |= {row[0] for row in store._conn.execute(f"SELECT _key FROM {quote(target)}")}
+            known[target] = held
+        return known[target]
+
+    for entity, rows in staged.rows.items():
+        relationships = contract.entities[entity]["relationships"]
+        if not relationships:
+            continue
+        dropped = []
+        for key, row in rows.items():
+            for rel in relationships:
+                values = [row.get(name) for name in rel["fields"]]
+                if all(value is None for value in values):
+                    continue
+                chunk, index = staged.origin[entity][key]
+                if any(value is None for value in values):
+                    refusals.append(
+                        refusal(
+                            chunk,
+                            index,
+                            entity,
+                            row,
+                            rel["fields"][0],
+                            "Composite references must be wholly present or wholly absent/null",
+                            contract,
+                        )
+                    )
+                    dropped.append(key)
+                    break
+                target = rel["target"]
+                try:
+                    wanted = store._key(target, dict(zip(contract.entities[target]["primary_key"], values)))
+                except ContractError as error:
+                    refusals.append(
+                        refusal(chunk, index, entity, row, rel["fields"][0], error.detail["message"], contract)
+                    )
+                    dropped.append(key)
+                    break
+                if wanted not in keys_of(target):
+                    refusals.append(
+                        refusal(
+                            chunk,
+                            index,
+                            entity,
+                            row,
+                            rel["fields"][0],
+                            f"Relationship target does not exist in {target}",
+                            contract,
+                        )
+                    )
+                    dropped.append(key)
+                    break
+        for key in dropped:
+            del rows[key]
+    return refusals
+
+
+def _apply(
+    contract: Contract, store: SQLiteStore, staged: _Staged, mode: str, batch_id: str
+) -> tuple[int, dict | None]:
+    """Apply the staged rows as one batch; the count applied, or the store's refusal."""
+    from .data import seed_batch
+
+    ordered = {name: list(staged.rows[name].values()) for name in dependency_order(contract) if name in staged.rows}
+    batch = seed_batch(store, contract, ordered, mode, batch_id=batch_id)
+    origins = {id(row): staged.origin[entity][key] for entity, rows in staged.rows.items() for key, row in rows.items()}
+    try:
+        store.apply_batch(batch)
+    except ContractError as error:
+        detail = error.detail
+        mutation = None
+        index = detail.get("index")
+        if isinstance(index, int) and 0 <= index < len(batch["mutations"]):
+            mutation = batch["mutations"][index]
+        origin = origins.get(id(mutation.get("value"))) if mutation else None
+        entity = detail.get("entity") if isinstance(detail.get("entity"), str) else None
+        row = mutation.get("value") if mutation else None
+        chunk, position = origin if origin else (None, None)
+        return 0, refusal(chunk, position, entity, row, _field_of(error), detail["message"], contract)
+    applied = sum(1 for mutation in batch["mutations"] if mutation["op"] != "delete")
+    return applied, None
+
+
+def validate_chunks(contract: Contract, chunks: list[ChunkSpec]) -> dict:
+    """``validate-rows``: shape-only, every row of every chunk."""
+    rows = 0
+    refusals: list[dict] = []
+    entities: dict[str, int] = {}
+    for chunk in chunks:
+        for index, (entity, row) in enumerate(iter_chunk_rows(chunk.path, chunk.entity)):
+            rows += 1
+            refused = shape_refusal(contract, chunk.ordinal, index, entity, row)
+            if refused is not None:
+                refusals.append(refused)
+                continue
+            entities[entity] = entities.get(entity, 0) + 1
+    return {"rows": rows, "refusals": refusals, "entities": entities}
+
+
+def build_batch(
+    contract: Contract,
+    parent: Path,
+    chunks: list[ChunkSpec],
+    *,
+    mode: str = "append",
+    merge: str = "upsert",
+    atomic: bool = True,
+    batch_id: str,
+    parent_digest: str,
+    dry_run: bool = False,
+) -> dict:
+    """Apply ``chunks`` onto the snapshot at ``parent`` (a private copy: it is written
+    in place, so the caller passes the file the result should be) and describe the
+    result. Atomic: one transaction, any refusal means nothing applied and every
+    refusal found. Per chunk: one transaction per chunk, a refused chunk is left out
+    and the rest applied. ``dry_run`` finds the refusals and applies nothing."""
+    from .snapshot import batch_digest, file_sha256, write_sidecar, contract_hash
+
+    if mode not in ("append", "replace"):
+        fail("mode", "mode must be append or replace", "/mode")
+    if merge not in ("upsert", "fill"):
+        fail("merge", "merge must be upsert or fill", "/merge")
+    refusals: list[dict] = []
+    applied_rows = 0
+    applied_chunks: list[int] = []
+    applied_digests: list[str] = []
+    counts: dict[str, int] = {}
+    with SQLiteStore.open(parent, contract, audit=False) as store:
+        if atomic:
+            staged = _stage(contract, store, chunks, merge)
+            refusals.extend(staged.refusals)
+            refusals.extend(_reference_refusals(contract, store, staged, mode))
+            applied, refused = _apply(contract, store, staged, mode, batch_id)
+            if refused is not None:
+                refusals.append(refused)
+            if not refusals and not dry_run:
+                applied_rows = applied
+                applied_chunks = [chunk.ordinal for chunk in chunks]
+                applied_digests = [file_sha256(chunk.path) for chunk in chunks]
+        else:
+            for chunk in chunks:
+                staged = _stage(contract, store, [chunk], merge)
+                found = list(staged.refusals) + _reference_refusals(contract, store, staged, mode)
+                applied, refused = _apply(contract, store, staged, mode, f"{batch_id}/{chunk.ordinal}")
+                if refused is not None:
+                    found.append(refused)
+                refusals.extend(found)
+                if found and not dry_run:
+                    # The store rolled the chunk's transaction back; the rows it did
+                    # accept are gone with it, so the chunk is refused as a whole.
+                    continue
+                if not dry_run:
+                    applied_rows += applied
+                    applied_chunks.append(chunk.ordinal)
+                    applied_digests.append(file_sha256(chunk.path))
+        counts = store.counts()
+    digest = batch_digest(parent_digest, batch_id, applied_digests)
+    result: dict[str, Any] = {
+        "contractHash": contract_hash(contract),
+        "batchDigest": digest,
+        "entityCounts": counts,
+        "appliedRows": applied_rows,
+        "refusals": refusals[:REPORT_CAP],
+        "refusalCount": len(refusals),
+        "appliedChunks": applied_chunks,
+    }
+    if dry_run or (atomic and refusals):
+        parent.unlink(missing_ok=True)
+        return result
+    write_sidecar(parent, result["contractHash"], digest, counts)
+    result["snapshot"] = parent
+    result["fileSha256"] = file_sha256(parent)
+    return result
+
+
+__all__ = [
+    "CHUNK_BYTES",
+    "REPORT_CAP",
+    "Chunk",
+    "ChunkSpec",
+    "build_batch",
+    "chunk_rows",
+    "count_rows",
+    "dependency_order",
+    "encode_row",
+    "iter_chunk_rows",
+    "iter_file_rows",
+    "iter_rows",
+    "refusal",
+    "shape_refusal",
+    "validate_chunks",
+    "write_report",
+]