data_studio/identity.py

from __future__ import annotations

import hashlib
import json
from typing import Any

from data_studio.governance import RawRecordV1, RecordIdentityV1, SourceRegistryV1, SourceState


def canonical_json(payload: Any) -> str:
    return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)


def payload_hash(payload: Any) -> str:
    return hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest()


def stable_record_id(
    source: SourceRegistryV1,
    source_native_id: str,
) -> str:
    parts = (
        source.source_id,
        source.hf_revision,
        source.hf_config or "",
        source.hf_split,
        source_native_id,
    )
    return hashlib.sha256("\n".join(parts).encode("utf-8")).hexdigest()


def build_raw_record(
    source: SourceRegistryV1,
    source_native_id: str,
    raw_payload: dict[str, Any],
    normalized_payload: dict[str, Any],
) -> RawRecordV1:
    if source.state != SourceState.APPROVED:
        raise ValueError("Only approved source revisions can create governed raw records.")
    identity = RecordIdentityV1(
        record_id=stable_record_id(source, source_native_id),
        source_id=source.source_id,
        hf_revision=source.hf_revision,
        hf_config=source.hf_config,
        hf_split=source.hf_split,
        source_native_id=source_native_id,
    )
    return RawRecordV1(
        identity=identity,
        raw_payload=raw_payload,
        raw_payload_hash=payload_hash(raw_payload),
        normalized_payload=normalized_payload,
        normalized_payload_hash=payload_hash(normalized_payload),
    )