mostlyright-data 0.9.0__py3-none-any.whl

mostlyright/data_harness/sources/catalog/fill_partitions.py

@@ -0,0 +1,459 @@
+"""Fixed identifier partitions and bounded external merge for Data.gov v4 passes."""
+
+from __future__ import annotations
+
+import sqlite3
+import tempfile
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from mostlyright.data_harness.canonical import canonical_sha256, sha256_bytes
+from mostlyright.data_harness.sources.catalog.fill_staging import CatalogFillRefused, FillStaging
+
+PARTITION_COUNT = 256
+RUN_MAX_ITEMS = 128
+MERGE_FAN_IN = 16
+
+
+@dataclass(frozen=True)
+class PassMapResult:
+    root_sha256: str
+    manifest_sha256: str
+    unique_records: int
+    conflicting_identifiers: int
+
+
+def write_partition_runs(staging: FillStaging, records: tuple[Any, ...]) -> list[dict[str, Any]]:
+    """Route exact identifier bytes by SHA-256 prefix into immutable bounded sorted runs."""
+
+    descriptors: list[dict[str, Any]] = []
+    for partition, entries, conflicting_identifiers in plan_partition_runs(records):
+        digest = _write_run(
+            staging,
+            partition=partition,
+            entries=entries,
+            conflicting_identifiers=conflicting_identifiers,
+        )
+        descriptors.append({"partition": partition, "run_sha256": digest})
+    return descriptors
+
+
+def plan_partition_runs(
+    records: Iterable[Any],
+) -> Iterable[tuple[str, list[tuple[str, str]], set[str]]]:
+    """Yield the exact bounded page runs used by both installation and validation."""
+
+    groups: dict[str, list[tuple[str, str]]] = {}
+    for record in records:
+        identifier = getattr(record, "record_id", None)
+        semantic_sha256 = getattr(record, "digest", None)
+        if not isinstance(identifier, str) or not _is_digest(semantic_sha256):
+            raise CatalogFillRefused(
+                "DATAGOV_PARTITION", "fill.records", "record identity/digest is invalid"
+            )
+        partition = sha256_bytes(identifier.encode("utf-8", errors="strict"))[:2]
+        groups.setdefault(partition, []).append((identifier, semantic_sha256))
+    for partition in sorted(groups):
+        entries, conflicting_identifiers = _deduplicate_sorted(groups[partition])
+        for offset in range(0, len(entries), RUN_MAX_ITEMS):
+            chunk = entries[offset : offset + RUN_MAX_ITEMS]
+            yield (
+                partition,
+                chunk,
+                conflicting_identifiers & {identifier for identifier, _semantic in chunk},
+            )
+
+
+def external_merge_pass(
+    staging: FillStaging,
+    run_descriptors: Iterable[dict[str, Any]],
+) -> PassMapResult:
+    """Merge a descriptor stream through bounded disk-backed fixed-size batches."""
+
+    try:
+        return _external_merge_pass(staging, run_descriptors)
+    except CatalogFillRefused:
+        raise
+    except (OSError, sqlite3.Error) as error:
+        raise CatalogFillRefused(
+            "DATAGOV_EXTERNAL_MERGE",
+            "fill.partition_runs",
+            "bounded external-merge scratch storage failed",
+        ) from error
+
+
+def _external_merge_pass(
+    staging: FillStaging,
+    run_descriptors: Iterable[dict[str, Any]],
+) -> PassMapResult:
+
+    with tempfile.TemporaryDirectory(
+        prefix="mostlyright-datagov-pass-", dir=staging.root.parent
+    ) as directory:
+        database_path = Path(directory) / "merge.sqlite3"
+        connection = sqlite3.connect(database_path)
+        try:
+            connection.execute("PRAGMA journal_mode=OFF")
+            connection.execute("PRAGMA synchronous=OFF")
+            connection.execute("PRAGMA temp_store=FILE")
+            connection.execute(
+                "CREATE TABLE merged ("
+                "partition TEXT NOT NULL, identifier TEXT NOT NULL, semantic TEXT NOT NULL, "
+                "conflicting INTEGER NOT NULL, PRIMARY KEY (partition, identifier)"
+                ") WITHOUT ROWID"
+            )
+            descriptor_count = 0
+            for descriptor in run_descriptors:
+                if (
+                    not isinstance(descriptor, dict)
+                    or set(descriptor) != {"partition", "run_sha256"}
+                    or not _valid_partition(descriptor["partition"])
+                    or not _is_digest(descriptor["run_sha256"])
+                ):
+                    raise CatalogFillRefused(
+                        "DATAGOV_PARTITION", "fill.partition_runs", "run descriptor is invalid"
+                    )
+                entries, inherited_conflicts = _validate_run(
+                    staging.read_shard("partition_run", descriptor["run_sha256"]),
+                    descriptor["partition"],
+                )
+                connection.executemany(
+                    "INSERT INTO merged(partition, identifier, semantic, conflicting) "
+                    "VALUES (?, ?, ?, ?) ON CONFLICT(partition, identifier) DO UPDATE SET "
+                    "conflicting = merged.conflicting OR excluded.conflicting "
+                    "OR merged.semantic <> excluded.semantic, "
+                    "semantic = min(merged.semantic, excluded.semantic)",
+                    (
+                        (
+                            descriptor["partition"],
+                            identifier,
+                            semantic,
+                            int(identifier in inherited_conflicts),
+                        )
+                        for identifier, semantic in entries
+                    ),
+                )
+                descriptor_count += 1
+                if descriptor_count % MERGE_FAN_IN == 0:
+                    connection.commit()
+            connection.commit()
+            partition_manifests, unique_records, conflicts = _write_spooled_partitions(
+                staging, connection
+            )
+        finally:
+            connection.close()
+    manifest = {
+        "partition_count": PARTITION_COUNT,
+        "partitions": partition_manifests,
+        "unique_records": unique_records,
+        "conflicting_identifiers": conflicts,
+    }
+    root = canonical_sha256(manifest)
+    manifest_sha256 = staging.write_shard("pass_map", {**manifest, "root_sha256": root})
+    return PassMapResult(
+        root_sha256=root,
+        manifest_sha256=manifest_sha256,
+        unique_records=unique_records,
+        conflicting_identifiers=conflicts,
+    )
+
+
+def _write_spooled_partitions(
+    staging: FillStaging, connection: sqlite3.Connection
+) -> tuple[list[dict[str, Any]], int, int]:
+    partition_manifests: list[dict[str, Any]] = []
+    unique_records = 0
+    conflicts = 0
+    partitions = connection.execute("SELECT DISTINCT partition FROM merged ORDER BY partition")
+    for (partition,) in partitions:
+        rows = connection.execute(
+            "SELECT identifier, semantic, conflicting FROM merged "
+            "WHERE partition = ? ORDER BY identifier",
+            (partition,),
+        )
+        run_sha256s: list[str] = []
+        partition_records = 0
+        partition_conflicts = 0
+        while chunk := rows.fetchmany(RUN_MAX_ITEMS):
+            entries = [(identifier, semantic) for identifier, semantic, _conflict in chunk]
+            conflicting_identifiers = {
+                identifier for identifier, _semantic, conflict in chunk if conflict
+            }
+            run_sha256s.append(
+                _write_run(
+                    staging,
+                    partition=partition,
+                    entries=entries,
+                    conflicting_identifiers=conflicting_identifiers,
+                )
+            )
+            partition_records += len(entries)
+            partition_conflicts += len(conflicting_identifiers)
+        unique_records += partition_records
+        conflicts += partition_conflicts
+        partition_manifests.append(
+            {
+                "partition": partition,
+                "run_sha256s": run_sha256s,
+                "record_count": partition_records,
+                "conflicting_identifiers": partition_conflicts,
+            }
+        )
+    return partition_manifests, unique_records, conflicts
+
+
+def dry_run_external_merge_pass(
+    staging: FillStaging,
+    run_descriptors: Iterable[dict[str, Any]],
+) -> PassMapResult:
+    """Recompute a pass in bounded scratch space without mutating authenticated staging."""
+
+    with tempfile.TemporaryDirectory(
+        prefix="mostlyright-datagov-merge-", dir=staging.root.parent
+    ) as directory:
+        scratch = FillStaging(Path(directory))
+        with scratch.locked():
+            adapter = _DryRunStaging(staging, scratch)
+            return external_merge_pass(adapter, run_descriptors)  # type: ignore[arg-type]
+
+
+class _DryRunStaging:
+    def __init__(self, source: FillStaging, scratch: FillStaging) -> None:
+        self._source = source
+        self._scratch = scratch
+        self._written: set[tuple[str, str]] = set()
+
+    @property
+    def root(self) -> Path:
+        return self._scratch.root
+
+    def write_shard(self, kind: str, payload: dict[str, Any]) -> str:
+        digest = self._scratch.write_shard(kind, payload)
+        self._written.add((kind, digest))
+        return digest
+
+    def read_shard(self, kind: str, digest: str) -> dict[str, Any]:
+        if (kind, digest) in self._written:
+            return self._scratch.read_shard(kind, digest)
+        return self._source.read_shard(kind, digest)
+
+
+def validate_partition_run(staging: FillStaging, descriptor: dict[str, Any]) -> tuple[int, int]:
+    """Authenticate one indexed run descriptor and its exact immutable shard."""
+
+    if (
+        not isinstance(descriptor, dict)
+        or set(descriptor) != {"partition", "run_sha256"}
+        or not _valid_partition(descriptor["partition"])
+        or not _is_digest(descriptor["run_sha256"])
+    ):
+        raise CatalogFillRefused(
+            "DATAGOV_PARTITION", "fill.partition_runs", "run descriptor is invalid"
+        )
+    entries, conflicts = _validate_run(
+        staging.read_shard("partition_run", descriptor["run_sha256"]),
+        descriptor["partition"],
+    )
+    return len(entries), len(conflicts)
+
+
+def read_partition_run(
+    staging: FillStaging, descriptor: dict[str, Any]
+) -> tuple[str, list[tuple[str, str]], set[str]]:
+    """Authenticate and return one bounded run for exact cross-stream reconciliation."""
+
+    if (
+        not isinstance(descriptor, dict)
+        or set(descriptor) != {"partition", "run_sha256"}
+        or not _valid_partition(descriptor["partition"])
+        or not _is_digest(descriptor["run_sha256"])
+    ):
+        raise CatalogFillRefused(
+            "DATAGOV_PARTITION", "fill.partition_runs", "run descriptor is invalid"
+        )
+    entries, conflicts = _validate_run(
+        staging.read_shard("partition_run", descriptor["run_sha256"]),
+        descriptor["partition"],
+    )
+    return descriptor["partition"], entries, set(conflicts)
+
+
+def validate_pass_map(
+    staging: FillStaging, manifest_sha256: str, expected_root_sha256: str
+) -> PassMapResult:
+    """Stream-authenticate the terminal pass manifest and every referenced final run."""
+
+    payload = staging.read_shard("pass_map", manifest_sha256)
+    if (
+        set(payload)
+        != {
+            "partition_count",
+            "partitions",
+            "unique_records",
+            "conflicting_identifiers",
+            "root_sha256",
+        }
+        or payload["partition_count"] != PARTITION_COUNT
+        or not isinstance(payload["partitions"], list)
+        or type(payload["unique_records"]) is not int
+        or type(payload["conflicting_identifiers"]) is not int
+        or payload["root_sha256"] != expected_root_sha256
+    ):
+        raise CatalogFillRefused("DATAGOV_PASS_MAP", "fill.pass_map", "pass manifest is invalid")
+    manifest = {key: value for key, value in payload.items() if key != "root_sha256"}
+    if canonical_sha256(manifest) != expected_root_sha256:
+        raise CatalogFillRefused("DATAGOV_PASS_MAP", "fill.pass_map", "pass manifest root differs")
+    prior_partition: str | None = None
+    records = 0
+    conflicts = 0
+    for partition in payload["partitions"]:
+        if (
+            not isinstance(partition, dict)
+            or set(partition)
+            != {
+                "partition",
+                "run_sha256s",
+                "record_count",
+                "conflicting_identifiers",
+            }
+            or not _valid_partition(partition["partition"])
+            or (prior_partition is not None and partition["partition"] <= prior_partition)
+            or not isinstance(partition["run_sha256s"], list)
+            or not partition["run_sha256s"]
+            or type(partition["record_count"]) is not int
+            or type(partition["conflicting_identifiers"]) is not int
+        ):
+            raise CatalogFillRefused(
+                "DATAGOV_PASS_MAP", "fill.pass_map", "partition manifest is invalid"
+            )
+        run_records = 0
+        run_conflicts = 0
+        for digest in partition["run_sha256s"]:
+            count, conflict_count = validate_partition_run(
+                staging,
+                {"partition": partition["partition"], "run_sha256": digest},
+            )
+            run_records += count
+            run_conflicts += conflict_count
+        if (
+            run_records != partition["record_count"]
+            or run_conflicts != partition["conflicting_identifiers"]
+        ):
+            raise CatalogFillRefused("DATAGOV_PASS_MAP", "fill.pass_map", "partition totals differ")
+        records += run_records
+        conflicts += run_conflicts
+        prior_partition = partition["partition"]
+    if records != payload["unique_records"] or conflicts != payload["conflicting_identifiers"]:
+        raise CatalogFillRefused("DATAGOV_PASS_MAP", "fill.pass_map", "pass totals differ")
+    return PassMapResult(
+        root_sha256=expected_root_sha256,
+        manifest_sha256=manifest_sha256,
+        unique_records=records,
+        conflicting_identifiers=conflicts,
+    )
+
+
+def _write_run(
+    staging: FillStaging,
+    *,
+    partition: str,
+    entries: list[tuple[str, str]],
+    conflicting_identifiers: set[str],
+) -> str:
+    if not _valid_partition(partition) or not 1 <= len(entries) <= RUN_MAX_ITEMS:
+        raise CatalogFillRefused(
+            "DATAGOV_PARTITION", "fill.partition_run", "bounded run is invalid"
+        )
+    return staging.write_shard(
+        "partition_run",
+        {
+            "partition": partition,
+            "entries": [
+                {"identifier": identifier, "semantic_sha256": semantic}
+                for identifier, semantic in entries
+            ],
+            "conflicting_identifiers": sorted(conflicting_identifiers),
+        },
+    )
+
+
+def _validate_run(
+    payload: dict[str, Any], expected_partition: str
+) -> tuple[list[tuple[str, str]], set[str]]:
+    if (
+        set(payload) != {"partition", "entries", "conflicting_identifiers"}
+        or payload["partition"] != expected_partition
+        or not isinstance(payload["entries"], list)
+        or not 1 <= len(payload["entries"]) <= RUN_MAX_ITEMS
+        or not isinstance(payload["conflicting_identifiers"], list)
+    ):
+        raise CatalogFillRefused(
+            "DATAGOV_PARTITION", "fill.partition_run", "run payload is invalid"
+        )
+    result: list[tuple[str, str]] = []
+    for item in payload["entries"]:
+        if (
+            not isinstance(item, dict)
+            or set(item) != {"identifier", "semantic_sha256"}
+            or not isinstance(item["identifier"], str)
+            or not _is_digest(item["semantic_sha256"])
+        ):
+            raise CatalogFillRefused(
+                "DATAGOV_PARTITION", "fill.partition_run", "run entry is invalid"
+            )
+        if (
+            sha256_bytes(item["identifier"].encode("utf-8", errors="strict"))[:2]
+            != expected_partition
+        ):
+            raise CatalogFillRefused(
+                "DATAGOV_PARTITION", "fill.partition_run", "entry is in the wrong partition"
+            )
+        result.append((item["identifier"], item["semantic_sha256"]))
+    if result != sorted(result) or len({identifier for identifier, _ in result}) != len(result):
+        raise CatalogFillRefused(
+            "DATAGOV_PARTITION", "fill.partition_run", "run is not unique and sorted"
+        )
+    conflicts = payload["conflicting_identifiers"]
+    identifiers = {identifier for identifier, _semantic in result}
+    if (
+        any(not isinstance(identifier, str) for identifier in conflicts)
+        or conflicts != sorted(set(conflicts))
+        or not set(conflicts) <= identifiers
+    ):
+        raise CatalogFillRefused(
+            "DATAGOV_PARTITION", "fill.partition_run", "conflict membership is invalid"
+        )
+    return result, set(conflicts)
+
+
+def _deduplicate_sorted(
+    entries: list[tuple[str, str]],
+) -> tuple[list[tuple[str, str]], set[str]]:
+    result: list[tuple[str, str]] = []
+    conflicts: set[str] = set()
+    for identifier, semantic in sorted(entries):
+        if result and result[-1][0] == identifier:
+            if result[-1][1] != semantic:
+                conflicts.add(identifier)
+                result[-1] = (identifier, min(result[-1][1], semantic))
+            continue
+        result.append((identifier, semantic))
+    return result, conflicts
+
+
+def _valid_partition(value: Any) -> bool:
+    return (
+        isinstance(value, str)
+        and len(value) == 2
+        and all(character in "0123456789abcdef" for character in value)
+    )
+
+
+def _is_digest(value: Any) -> bool:
+    return (
+        isinstance(value, str)
+        and len(value) == 64
+        and all(character in "0123456789abcdef" for character in value)
+    )