Source code for numeraire_dataset.zones.lock

"""``data.lock.json`` — the provenance ledger tying built tables to their recipe + raw vintages.

Each artifact (a pinned ``raw`` vintage or a built ``clean`` table) gets one lock entry. The entry
records its recipe hash (the transform) and content digest (the frame's schema, index, and values),
so a partial or corrupted build is detectable independently of the recipe.

The **contract with numeraire**: a result's ``data_vintage`` string is derived *directly* from a
clean artifact's lock entry — ``f"{name}@{recipe_hash-prefix}"`` — so provenance lives in exactly
one place and a result traces back to the committed step, explicit overrides, and raw vintages.
"""

from __future__ import annotations

import copy
import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from numeraire_dataset.zones.raw import PIT_STATUSES
from numeraire_dataset.zones.steps import (
    CONTENT_DIGEST_VERSION,
    RECIPE_HASH_VERSION,
    Built,
    _recipe_hash_from_canonical,
    content_digest,
)

LOCK_VERSION = 2
_VINTAGE_HASH_LEN = 12  # chars of the recipe-hash hex kept in the data_vintage string


def _hash_prefix(recipe_hash: str) -> str:
    """The hex-digest prefix used in a data_vintage string (drops the ``sha256:`` label)."""
    digest = recipe_hash.split(":", 1)[-1]
    return digest[:_VINTAGE_HASH_LEN]


[docs] def data_vintage_of(name: str, recipe_hash: str) -> str: """The numeraire ``data_vintage`` for a clean artifact: ``<name>@<recipe-hash-prefix>``.""" return f"{name}@{_hash_prefix(recipe_hash)}"
def _raw_identity_component(value: str, *, label: str) -> str: if not value or "@" in value or "#" in value: raise ValueError(f"raw {label} must be non-empty and exclude reserved '@'/'#' separators") return value
[docs] @dataclass class DataLock: """An in-memory ``data.lock.json`` with typed add/read helpers.""" artifacts: dict[str, dict[str, Any]] = field(default_factory=dict)
[docs] def add_raw( self, source: str, vintage: str, content_digest: str, *, query_hash: str = "", pit_status: str = "vintage", content_digest_version: int = CONTENT_DIGEST_VERSION, ) -> str: """Record a pinned raw vintage; returns its artifact key ``<source>@<vintage>``.""" source = _raw_identity_component(source, label="source") vintage = _raw_identity_component(vintage, label="vintage") key = f"{source}@{vintage}" entry: dict[str, Any] = { "kind": "raw", "source": source, "vintage": vintage, "content_digest": content_digest, "pit_status": pit_status, } if pit_status not in PIT_STATUSES: raise ValueError(f"pit_status must be one of {PIT_STATUSES}; got {pit_status!r}") if query_hash: entry["query_hash"] = query_hash if not isinstance(content_digest_version, int) or content_digest_version < 1: raise ValueError("content_digest_version must be a positive integer") entry["content_digest_version"] = content_digest_version self.artifacts[key] = entry return key
[docs] def add_clean( self, built: Built, *, inputs: Mapping[str, str], step_version: int | None = None, ) -> str: """Record an auditable clean recipe with role-bound artifact labels and hashes.""" if not isinstance(inputs, Mapping): raise TypeError( "clean inputs must be a role-to-artifact mapping; the legacy label-list form " "cannot represent recipe-hash contract version 2" ) if step_version is not None and step_version != built.step_version: raise ValueError( f"step version mismatch: built {built.step_version}, supplied {step_version}" ) if built.recipe_hash_version != RECIPE_HASH_VERSION: raise ValueError( f"unsupported built recipe hash version {built.recipe_hash_version}; " f"expected {RECIPE_HASH_VERSION}" ) expected_hash = _recipe_hash_from_canonical( built.step_name, built.step_version, built.canonical_params, built.input_hashes, ) if built.recipe_hash != expected_hash: raise ValueError("built recipe metadata does not match its recipe hash") observed_content_digest = content_digest(built.frame) if built.content_digest != observed_content_digest: raise ValueError( "built frame content does not match its recorded content digest; " "rerun the clean step instead of locking a mutated output" ) expected_roles = set(built.input_hashes) observed_roles = set(inputs) if observed_roles != expected_roles: raise ValueError( "clean input labels must match the built recipe roles; " f"expected {sorted(expected_roles)}, observed {sorted(observed_roles)}" ) bindings = { role: {"artifact": inputs[role], "hash": built.input_hashes[role]} for role in sorted(expected_roles) } self.artifacts[built.name] = { "kind": "clean", "step": built.step_name, "step_version": built.step_version, "recipe_hash": built.recipe_hash, "recipe_hash_version": built.recipe_hash_version, "params": copy.deepcopy(built.canonical_params), "content_digest": built.content_digest, "content_digest_version": CONTENT_DIGEST_VERSION, "inputs": bindings, "rows": len(built.frame), "data_vintage": data_vintage_of(built.name, built.recipe_hash), } return built.name
[docs] def data_vintage(self, name: str) -> str: """The ``data_vintage`` string of a recorded clean artifact (the numeraire-facing stamp).""" entry = self.artifacts.get(name) if entry is None: raise KeyError(f"no artifact {name!r} in the lock") if entry["kind"] != "clean": raise ValueError(f"artifact {name!r} is {entry['kind']!r}, not a clean table") return str(entry["data_vintage"])
[docs] def to_json(self) -> str: """Serialize to canonical (sorted-key) JSON.""" return json.dumps( {"version": LOCK_VERSION, "artifacts": self.artifacts}, sort_keys=True, indent=2, )
[docs] def write(self, path: str | Path) -> Path: """Write ``data.lock.json`` (or the given path) and return it.""" p = Path(path) if p.is_dir(): p = p / "data.lock.json" p.write_text(self.to_json(), encoding="utf-8") return p
[docs] @classmethod def read(cls, path: str | Path) -> DataLock: """Load a lock file (accepts a directory containing ``data.lock.json``).""" p = Path(path) if p.is_dir(): p = p / "data.lock.json" payload = json.loads(p.read_text(encoding="utf-8")) if payload.get("version") != LOCK_VERSION: raise ValueError(f"unsupported data.lock.json version {payload.get('version')!r}") return cls(artifacts=dict(payload.get("artifacts", {})))