DataOps-platform/app/core/data_rules/publication.py

"""Trusted generation receipts and rule publication lifecycle gates."""

from __future__ import annotations

import base64
import copy
import hashlib
import hmac
import json
import os
import re
import tempfile
from contextlib import ExitStack, suppress
from datetime import UTC, datetime
from typing import Any

from app.core.common.identifiers import ensure_governance_uid, new_governance_uid
from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec

_DIGEST = re.compile(r"^[0-9a-f]{64}$")
_RECEIPT_KEYS = {
    "version",
    "generation_run_id",
    "actor_uid",
    "source_text_hash",
    "candidate_hash",
    "rule_spec_hash",
    "model_hash",
    "prompt_hash",
    "context_hash",
    "issued_at",
    "expires_at",
}
_COMPILED_KEYS = {
    "backend",
    "compiler_version",
    "plan",
    "plan_hash",
    "schema_hashes",
    "binding_hashes",
    "capabilities",
}
_TEST_KEYS = {
    "status",
    "test_kind",
    "run_id",
    "plan_hash",
    "schema_hashes",
    "binding_hashes",
    "counts",
    "attestation",
}


class RuleValidationRejected(ValueError):
    """A deterministic validation failure whose audit row must be committed."""


def _canonical(value: Any) -> bytes:
    try:
        encoded = json.dumps(
            value,
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
        )
    except (TypeError, ValueError) as exc:
        raise ValueError("publication value must be JSON serializable") from exc
    return encoded.encode("utf-8")


def _digest(value: Any, label: str) -> str:
    normalized = str(value or "")
    if not _DIGEST.fullmatch(normalized):
        raise ValueError(f"{label} must be a sha256 digest")
    return normalized


def _uid(value: Any, label: str) -> str:
    try:
        return ensure_governance_uid({"uid": str(value)})
    except ValueError as exc:
        raise ValueError(f"{label} must be a valid UUIDv7") from exc


def _hash_text(value: Any, label: str, maximum: int = 20_000) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{label} is required")
    normalized = value.strip()
    if len(normalized) > maximum:
        raise ValueError(f"{label} exceeds {maximum} characters")
    return hashlib.sha256(normalized.encode("utf-8")).hexdigest()


def _b64encode(value: bytes) -> str:
    return base64.urlsafe_b64encode(value).decode("ascii").rstrip("=")


def _b64decode(value: str) -> bytes:
    if (
        not isinstance(value, str)
        or not value
        or re.fullmatch(r"[A-Za-z0-9_-]+", value) is None
    ):
        raise ValueError("generation receipt encoding is invalid")
    try:
        decoded = base64.urlsafe_b64decode(
            value + ("=" * (-len(value) % 4))
        )
    except Exception as exc:
        raise ValueError("generation receipt encoding is invalid") from exc
    if _b64encode(decoded) != value:
        raise ValueError("generation receipt encoding is not canonical")
    return decoded


def generation_receipt_claims(
    *,
    generation_run_id: str,
    actor_uid: str,
    source_text: str,
    candidate_hash: str,
    rule_spec: dict[str, Any],
    model_hash: str,
    prompt_hash: str,
    context_hash: str,
    issued_at: datetime | None = None,
    expires_at: datetime,
) -> dict[str, Any]:
    """Build the closed, non-secret claim set signed by ``/interpret``."""

    issued = issued_at or datetime.now(UTC)
    if not isinstance(issued, datetime) or issued.tzinfo is None:
        raise ValueError("generation receipt issue time must be timezone aware")
    if not isinstance(expires_at, datetime) or expires_at.tzinfo is None:
        raise ValueError("generation receipt expiry must be timezone aware")
    spec = validate_rule_spec(rule_spec)
    return {
        "version": "1",
        "generation_run_id": _uid(
            generation_run_id, "generation_run_id"
        ),
        "actor_uid": _uid(actor_uid, "actor_uid"),
        "source_text_hash": _hash_text(source_text, "source_text"),
        "candidate_hash": _digest(candidate_hash, "candidate_hash"),
        "rule_spec_hash": rule_spec_hash(spec),
        "model_hash": _digest(model_hash, "model_hash"),
        "prompt_hash": _digest(prompt_hash, "prompt_hash"),
        "context_hash": _digest(context_hash, "context_hash"),
        "issued_at": int(issued.timestamp()),
        "expires_at": int(expires_at.timestamp()),
    }


class GenerationReceiptSigner:
    """Issue and verify compact HMAC receipts without exposing the secret."""

    def __init__(
        self,
        secret: str | bytes,
        *,
        clock=None,
        max_ttl_seconds: int = 600,
        clock_skew_seconds: int = 30,
    ):
        secret_bytes = (
            secret.encode("utf-8") if isinstance(secret, str) else bytes(secret)
        )
        if len(secret_bytes) < 16:
            raise ValueError("generation receipt secret is too short")
        if (
            isinstance(max_ttl_seconds, bool)
            or not isinstance(max_ttl_seconds, int)
            or max_ttl_seconds < 1
            or max_ttl_seconds > 3600
        ):
            raise ValueError("generation receipt max TTL is invalid")
        if (
            isinstance(clock_skew_seconds, bool)
            or not isinstance(clock_skew_seconds, int)
            or clock_skew_seconds < 0
            or clock_skew_seconds > 300
        ):
            raise ValueError("generation receipt clock skew is invalid")
        self._secret = secret_bytes
        self._clock = clock or (lambda: datetime.now(UTC))
        self.max_ttl_seconds = max_ttl_seconds
        self.clock_skew_seconds = clock_skew_seconds

    def issue(self, claims: dict[str, Any]) -> str:
        normalized = self._validate_claims(claims)
        self._validate_time(normalized, self._clock())
        payload = _b64encode(_canonical(normalized))
        signature = hmac.new(
            self._secret, payload.encode("ascii"), hashlib.sha256
        ).digest()
        return f"{payload}.{_b64encode(signature)}"

    def verify(
        self,
        receipt: str,
        *,
        actor_uid: str,
        source_text: str,
        rule_spec: dict[str, Any],
        now: datetime | None = None,
        allow_expired: bool = False,
    ) -> dict[str, Any]:
        if not isinstance(receipt, str) or len(receipt) > 4096:
            raise ValueError("generation receipt is invalid")
        try:
            payload, encoded_signature = receipt.split(".", 1)
        except ValueError as exc:
            raise ValueError("generation receipt is invalid") from exc
        expected = hmac.new(
            self._secret, payload.encode("ascii"), hashlib.sha256
        ).digest()
        provided = _b64decode(encoded_signature)
        if not hmac.compare_digest(expected, provided):
            raise ValueError("generation receipt signature is invalid")
        try:
            decoded = json.loads(_b64decode(payload))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise ValueError("generation receipt payload is invalid") from exc
        claims = self._validate_claims(decoded)
        current = now or self._clock()
        if current.tzinfo is None:
            raise ValueError("generation receipt clock must be timezone aware")
        self._validate_time(
            claims,
            current,
            allow_expired=allow_expired,
        )
        if claims["actor_uid"] != _uid(actor_uid, "actor_uid"):
            raise ValueError("generation receipt actor does not match")
        if claims["source_text_hash"] != _hash_text(
            source_text, "source_text"
        ):
            raise ValueError("generation receipt source does not match")
        if claims["rule_spec_hash"] != rule_spec_hash(
            validate_rule_spec(rule_spec)
        ):
            raise ValueError("generation receipt RuleSpec does not match")
        return claims

    def _validate_time(
        self,
        claims: dict[str, Any],
        current: datetime,
        *,
        allow_expired: bool = False,
    ) -> None:
        if not isinstance(current, datetime) or current.tzinfo is None:
            raise ValueError("generation receipt clock must be timezone aware")
        now = int(current.timestamp())
        if now >= claims["expires_at"] and not allow_expired:
            raise ValueError("generation receipt has expired")
        if claims["issued_at"] > now + self.clock_skew_seconds:
            raise ValueError("generation receipt issue time is in the future")
        if claims["expires_at"] <= claims["issued_at"]:
            raise ValueError("generation receipt expiry is invalid")
        if (
            claims["expires_at"] - claims["issued_at"]
            > self.max_ttl_seconds
        ):
            raise ValueError("generation receipt TTL exceeds the maximum")

    @staticmethod
    def _validate_claims(value: Any) -> dict[str, Any]:
        if not isinstance(value, dict) or set(value) != _RECEIPT_KEYS:
            raise ValueError("generation receipt claims are invalid")
        if value.get("version") != "1":
            raise ValueError("generation receipt version is unsupported")
        temporal = {}
        for key in ("issued_at", "expires_at"):
            item = value.get(key)
            if (
                isinstance(item, bool)
                or not isinstance(item, int)
                or item <= 0
            ):
                raise ValueError(
                    f"generation receipt {key} is invalid"
                )
            temporal[key] = item
        return {
            "version": "1",
            "generation_run_id": _uid(
                value.get("generation_run_id"), "generation_run_id"
            ),
            "actor_uid": _uid(value.get("actor_uid"), "actor_uid"),
            "source_text_hash": _digest(
                value.get("source_text_hash"), "source_text_hash"
            ),
            "candidate_hash": _digest(
                value.get("candidate_hash"), "candidate_hash"
            ),
            "rule_spec_hash": _digest(
                value.get("rule_spec_hash"), "rule_spec_hash"
            ),
            "model_hash": _digest(value.get("model_hash"), "model_hash"),
            "prompt_hash": _digest(value.get("prompt_hash"), "prompt_hash"),
            "context_hash": _digest(
                value.get("context_hash"), "context_hash"
            ),
            **temporal,
        }


def _normalize_compiled(value: Any) -> dict[str, Any]:
    if not isinstance(value, dict) or set(value) != _COMPILED_KEYS:
        raise ValueError("trusted compiler returned an invalid result")
    if value.get("backend") not in {"sql_pushdown", "polars_batch"}:
        raise ValueError("trusted compiler returned an unsupported backend")
    compiler_version = str(value.get("compiler_version") or "").strip()
    if not compiler_version or len(compiler_version) > 80:
        raise ValueError("trusted compiler version is invalid")
    if not isinstance(value.get("plan"), dict):
        raise ValueError("trusted compiler plan is invalid")
    for key in ("schema_hashes", "binding_hashes", "capabilities"):
        if not isinstance(value.get(key), dict):
            raise ValueError(f"trusted compiler {key} are invalid")
    return {
        **copy.deepcopy(value),
        "compiler_version": compiler_version,
        "plan_hash": _digest(value.get("plan_hash"), "plan_hash"),
    }


def _normalize_test(value: Any, context: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(value, dict) or set(value) != _TEST_KEYS:
        raise ValueError("trusted test runner returned an invalid result")
    if value.get("status") != "success":
        raise ValueError("trusted rule dry-run did not succeed")
    if value.get("test_kind") not in {"dry_run", "sample", "golden"}:
        raise ValueError("trusted rule test kind is invalid")
    _uid(value.get("run_id"), "test run_id")
    if (
        value.get("plan_hash") != context.get("plan_hash")
        or value.get("schema_hashes") != context.get("schema_hashes")
        or value.get("binding_hashes") != context.get("binding_hashes")
    ):
        raise ValueError("test evidence does not match the exact compiled plan")
    counts = value.get("counts")
    if not isinstance(counts, dict) or any(
        isinstance(item, bool) or not isinstance(item, int) or item < 0
        for item in counts.values()
    ):
        raise ValueError("trusted test counts are invalid")
    attestation = value.get("attestation")
    if not isinstance(attestation, dict):
        raise ValueError("trusted test attestation is invalid")
    return copy.deepcopy(value)


class RulePublicationService:
    """Coordinate only server-generated compile and test evidence."""

    def __init__(
        self,
        repository,
        *,
        receipt_signer: GenerationReceiptSigner | None,
        compiler,
        test_runner,
    ):
        self.repository = repository
        self.receipt_signer = receipt_signer
        self.compiler = compiler
        self.test_runner = test_runner

    def create_draft(
        self,
        *,
        rule_spec: dict[str, Any],
        source_text: str,
        actor_uid: str,
        generation_receipt: str,
        category: str,
        source_language: str,
        generated_kind: str,
    ) -> dict[str, Any]:
        if self.receipt_signer is None:
            raise RuntimeError("generation receipt signer is not configured")
        spec = validate_rule_spec(rule_spec)
        actor = _uid(actor_uid, "actor_uid")
        claims = self.receipt_signer.verify(
            generation_receipt,
            actor_uid=actor,
            source_text=source_text,
            rule_spec=spec,
            allow_expired=True,
        )
        receipt_hash = hashlib.sha256(
            generation_receipt.encode("utf-8")
        ).hexdigest()
        return self.repository.create_draft_from_generation(
            rule_spec=spec,
            source_text=source_text,
            created_by=actor,
            category=category,
            source_language=source_language,
            generated_kind=generated_kind,
            receipt_claims=claims,
            receipt_hash=receipt_hash,
        )

    def validate(
        self,
        version_id: str,
        actor_uid: str,
    ) -> dict[str, Any]:
        version = _uid(version_id, "version_id")
        actor = _uid(actor_uid, "actor_uid")
        context = self.repository.load_logical_validation_context(
            version_id=version,
        )
        try:
            compiled = _normalize_compiled(self.compiler.compile(context))
        except ValueError as exc:
            self.repository.record_validation_failure(
                version_id=version,
                actor_uid=actor,
                error_code="deterministic_validation_error",
                error_hash=hashlib.sha256(str(exc).encode("utf-8")).hexdigest(),
            )
            raise RuleValidationRejected(str(exc)) from exc
        return self.repository.persist_logical_compile(
            version_id=version,
            actor_uid=actor,
            context=context,
            compiled=compiled,
        )

    def test(
        self,
        version_id: str,
        actor_uid: str,
        *,
        plan_id: str,
    ) -> dict[str, Any]:
        version = _uid(version_id, "version_id")
        actor = _uid(actor_uid, "actor_uid")
        plan = _uid(plan_id, "plan_id")
        context = self.repository.load_logical_test_context(
            version_id=version,
            plan_id=plan,
        )
        if "terminal_result" in context:
            return copy.deepcopy(context["terminal_result"])
        evidence = _normalize_test(self.test_runner.run(context), context)
        return self.repository.persist_logical_test(
            version_id=version,
            plan_id=plan,
            actor_uid=actor,
            context=context,
            evidence=evidence,
        )

    def publish(self, version_id: str, actor_uid: str) -> dict[str, Any]:
        return self.repository.publish_validated_version(
            version_id=_uid(version_id, "version_id"),
            actor_uid=_uid(actor_uid, "actor_uid"),
        )

    def evidence(self, version_id: str) -> dict[str, Any]:
        return self.repository.get_publication_evidence(
            version_id=_uid(version_id, "version_id")
        )

    def catalog(self, *, query: str = "", limit: int = 50):
        if not isinstance(query, str) or len(query) > 200:
            raise ValueError("catalog query is invalid")
        if (
            isinstance(limit, bool)
            or not isinstance(limit, int)
            or limit < 1
            or limit > 100
        ):
            raise ValueError("catalog limit is invalid")
        return self.repository.search_published_rules(
            query=query.strip(), limit=limit
        )


class PhysicalPlanPublicationService:
    """Test and publish deployment-bound plans independently of RuleVersion."""

    def __init__(self, repository, *, test_runner):
        self.repository = repository
        self.test_runner = test_runner

    def validate(self, plan_id: str, actor_uid: str) -> dict[str, Any]:
        return self.repository.attest_physical_compile(
            plan_id=_uid(plan_id, "plan_id"),
            actor_uid=_uid(actor_uid, "actor_uid"),
        )

    def test(self, plan_id: str, actor_uid: str) -> dict[str, Any]:
        plan = _uid(plan_id, "plan_id")
        actor = _uid(actor_uid, "actor_uid")
        context = self.repository.load_physical_test_context(plan_id=plan)
        if "terminal_result" in context:
            return copy.deepcopy(context["terminal_result"])
        evidence = _normalize_test(self.test_runner.run(context), context)
        return self.repository.persist_physical_test(
            plan_id=plan,
            actor_uid=actor,
            context=context,
            evidence=evidence,
        )

    def publish(self, plan_id: str, actor_uid: str) -> dict[str, Any]:
        return self.repository.publish_tested_physical_plan(
            plan_id=_uid(plan_id, "plan_id"),
            actor_uid=_uid(actor_uid, "actor_uid"),
        )


class CanonicalBoundCompiler:
    """Compile a repository-loaded physical context through the registry."""

    def __init__(self, compiler_registry):
        self.compiler_registry = compiler_registry

    def compile(self, context: dict[str, Any]) -> dict[str, Any]:
        if not isinstance(context, dict):
            raise ValueError("canonical publication context is invalid")
        try:
            rule_version = context["rule_version"]
            input_schema = context["input_schema"]
            output_schema = context["output_schema"]
            input_binding = context["input_binding"]
            output_binding = context["output_binding"]
            backend = context["backend"]
        except KeyError as exc:
            raise ValueError(
                "canonical publication context is incomplete"
            ) from exc
        compiler = self.compiler_registry.select(
            rule_version.get("rule_spec"),
            input_binding,
            output_binding,
        )
        compiled = compiler.compile(
            rule_version=rule_version,
            input_schema=input_schema,
            output_schema=output_schema,
            input_binding=input_binding,
            output_binding=output_binding,
            backend=backend,
        )
        plan = compiled.get("plan")
        if not isinstance(plan, dict):
            raise ValueError("trusted compiler returned an invalid plan")
        return {
            "backend": compiled.get("backend"),
            "compiler_version": compiled.get("compiler_version"),
            "plan": plan,
            "plan_hash": compiled.get("plan_hash"),
            "schema_hashes": {
                "rule_spec_hash": rule_version.get("spec_hash"),
                "input_schema_snapshot_id": input_schema.get("id"),
                "input_schema_hash": input_schema.get("schema_hash"),
                "output_schema_snapshot_id": output_schema.get("id"),
                "output_schema_hash": output_schema.get("schema_hash"),
            },
            "binding_hashes": {
                "input": input_binding.get("binding_hash"),
                "output": output_binding.get("binding_hash"),
            },
            "capabilities": copy.deepcopy(backend),
        }


class LogicalRuleCompiler:
    """Compile a draft RuleSpec against its immutable validation profile."""

    VERSION = "dataops-logical-rule-1.0"

    def compile(self, context: dict[str, Any]) -> dict[str, Any]:
        from app.core.data_rules.compilers.polars import PolarsRuleCompiler
        from app.core.data_rules.contracts import read_rule_spec

        if not isinstance(context, dict):
            raise ValueError("logical validation context is invalid")
        version = context.get("version")
        profile = context.get("profile")
        if not isinstance(version, dict) or version.get("status") != "draft":
            raise ValueError("only draft rule versions may be validated")
        if not isinstance(profile, dict):
            raise ValueError("trusted validation profile is unavailable")
        spec = read_rule_spec(version.get("rule_spec"))
        if rule_spec_hash(spec) != version.get("spec_hash"):
            raise ValueError("canonical RuleSpec hash does not match")
        snapshots = {}
        for prefix in ("input", "output"):
            fields = profile.get(f"{prefix}_fields")
            if not isinstance(fields, list) or not fields:
                raise ValueError(
                    f"trusted {prefix} validation fields are unavailable"
                )
            snapshots[prefix] = {
                "id": _uid(
                    profile.get(f"{prefix}_schema_snapshot_id"),
                    f"{prefix}_schema_snapshot_id",
                ),
                "schema_ref": str(
                    profile.get(f"{prefix}_schema_ref") or ""
                ),
                "schema_hash": _digest(
                    profile.get(f"{prefix}_schema_hash"),
                    f"{prefix}_schema_hash",
                ),
                "fields": copy.deepcopy(fields),
                "source_revision": str(
                    profile.get(f"{prefix}_source_revision") or ""
                ),
            }
        if (
            spec["input_schema_ref"] != snapshots["input"]["schema_ref"]
            or spec["output_schema_ref"]
            != snapshots["output"]["schema_ref"]
        ):
            raise ValueError(
                "RuleSpec schema references do not match validation profile"
            )
        version_id = _uid(version.get("id"), "version_id")
        profile_id = _uid(profile.get("id"), "validation_profile_id")
        binding_base = {
            "data_source_uid": version_id,
            "object_kind": "parquet_artifact",
            "dialect": "polars",
            "write_mode": "append",
        }
        compiled = PolarsRuleCompiler().compile(
            rule_version={**version, "status": "published"},
            input_schema=snapshots["input"],
            output_schema=snapshots["output"],
            input_binding={
                "id": profile_id,
                **binding_base,
                "schema_snapshot_id": snapshots["input"]["id"],
                "access_mode": "read",
                "object_ref": "validation-input",
            },
            output_binding={
                "id": version_id,
                **binding_base,
                "schema_snapshot_id": snapshots["output"]["id"],
                "access_mode": "write",
                "object_ref": "validation-output",
            },
            backend={
                "max_rows": 100_000,
                "max_artifact_bytes": 32 * 1024 * 1024,
                "memory_limit_bytes": 256 * 1024 * 1024,
                "masking_policies": {
                    "customer_mobile_last4": "preserve_last_4",
                    "redact": "redact",
                },
                "lookup_bindings": {},
            },
        )
        plan = compiled["plan"]
        return {
            "backend": "polars_batch",
            "compiler_version": compiled["compiler_version"],
            "plan": plan,
            "plan_hash": compiled["plan_hash"],
            "schema_hashes": {
                "input": snapshots["input"]["schema_hash"],
                "output": snapshots["output"]["schema_hash"],
            },
            "binding_hashes": {},
            "capabilities": {
                "supported_execution_backends": ["polars"],
                "closed_rulespec": "2.0",
                "resource_limits": plan["resource_limits"],
            },
        }


class ServerOwnedLogicalDryRunRunner:
    """Execute a logical Polars plan on a bounded server-owned artifact."""

    def __init__(self, artifact_store):
        self.artifact_store = artifact_store

    def run(self, context: dict[str, Any]) -> dict[str, Any]:
        from app.core.data_rules.compilers.polars import (
            validate_bound_polars_plan,
        )
        from app.runner.polars_worker import execute_isolated_polars_plan

        sample_ref = context.get("input_sample_artifact_ref")
        if not isinstance(sample_ref, str) or not sample_ref:
            raise ValueError("server-owned sample artifact is required")
        sample_digest = _digest(
            context.get("input_sample_artifact_digest"),
            "input_sample_artifact_digest",
        )
        plan = validate_bound_polars_plan(context.get("plan"))
        described = self.artifact_store.describe(sample_ref)
        if (
            described["digest"] != sample_digest
            or described["schema_hash"] != plan["input_schema_hash"]
        ):
            raise ValueError("sample artifact schema has drifted")
        golden_ref = context.get("golden_output_artifact_ref")
        golden_digest = context.get("golden_output_artifact_digest")
        if (golden_ref is None) != (golden_digest is None):
            raise ValueError("golden output artifact is incomplete")
        if golden_ref is not None:
            if not isinstance(golden_ref, str) or not golden_ref:
                raise ValueError("golden output artifact is invalid")
            golden_digest = _digest(
                golden_digest,
                "golden_output_artifact_digest",
            )
        output_path = None
        try:
            with ExitStack() as stack:
                input_path = stack.enter_context(
                    self.artifact_store.stage(
                        sample_ref,
                        sample_digest,
                        expected_schema_fields=plan["input_fields"],
                        limits=plan["resource_limits"],
                    )
                )
                golden_path = None
                if golden_ref is not None:
                    golden_path = stack.enter_context(
                        self.artifact_store.stage(
                            golden_ref,
                            golden_digest,
                            expected_schema_fields=plan["output_fields"],
                            limits=plan["resource_limits"],
                        )
                    )
                with tempfile.NamedTemporaryFile(
                    prefix="dataops-logical-dry-run-",
                    suffix=".parquet",
                    delete=False,
                ) as handle:
                    output_path = handle.name
                result = execute_isolated_polars_plan(
                    {
                        "plan": plan,
                        "input_path": input_path,
                        "lookup_paths": {},
                        "output_path": output_path,
                        "golden_path": golden_path,
                        "golden_digest": golden_digest,
                        "masking_policies": {
                            "customer_mobile_last4": "preserve_last_4",
                            "redact": "redact",
                        },
                    },
                    memory_limit_bytes=plan["resource_limits"][
                        "memory_limit_bytes"
                    ],
                )
            prepared = self.artifact_store.prepare_path(
                output_path,
                new_governance_uid(),
                300,
                schema_fields=plan["output_fields"],
                limits=plan["resource_limits"],
            )
            compared_output_digest = result.get("output_file_digest")
            if (
                compared_output_digest is not None
                and prepared["digest"] != compared_output_digest
            ):
                raise ValueError(
                    "golden-compared output artifact has drifted"
                )
        finally:
            if output_path is not None:
                with suppress(FileNotFoundError):
                    os.unlink(output_path)
        count_keys = (
            "rows_in",
            "rows_out",
            "rows_rejected",
            "rows_quarantined",
            "rows_filtered",
            "rows_deduplicated",
            "rows_join_dropped",
            "rows_aggregated",
            "violation_count",
        )
        return {
            "status": "success",
            "test_kind": "dry_run",
            "run_id": new_governance_uid(),
            "plan_hash": context["plan_hash"],
            "schema_hashes": context["schema_hashes"],
            "binding_hashes": context["binding_hashes"],
            "counts": {key: int(result.get(key, 0)) for key in count_keys},
            "attestation": {
                "input_digest": described["digest"],
                "output_digest": prepared["digest"],
                "output_schema_hash": prepared["schema_hash"],
                "violation_digest": hashlib.sha256(
                    _canonical(result.get("violations", []))
                ).hexdigest(),
                "golden_output_digest": result.get(
                    "golden_output_digest"
                ),
                "golden_rows_compared": result.get(
                    "golden_rows_compared"
                ),
            },
        }


class ServerOwnedPhysicalPreflightRunner:
    """Execute an exact physical Polars plan on its latest trusted input."""

    def __init__(self, artifact_store, datasource_manager=None):
        self.artifact_store = artifact_store
        self.datasource_manager = datasource_manager

    def run(self, context: dict[str, Any]) -> dict[str, Any]:
        from app.core.data_rules.compilers.polars import (
            validate_bound_polars_plan,
        )
        from app.runner.polars_worker import execute_isolated_polars_plan

        if context.get("backend") == "sql_pushdown":
            return self._run_sql(context)
        if context.get("backend") != "polars_batch":
            raise ValueError("physical preflight backend is unsupported")
        sample = context.get("sample_artifact")
        if (
            not isinstance(sample, dict)
            or set(sample) != {
                "artifact_ref",
                "digest",
                "schema_fields",
            }
        ):
            raise ValueError("server-owned physical sample is required")
        plan = validate_bound_polars_plan(context.get("plan"))
        output_path = None
        try:
            with self.artifact_store.stage(
                sample["artifact_ref"],
                sample["digest"],
                expected_schema_fields=sample["schema_fields"],
                limits=plan["resource_limits"],
            ) as input_path:
                with tempfile.NamedTemporaryFile(
                    prefix="dataops-physical-preflight-",
                    suffix=".parquet",
                    delete=False,
                ) as handle:
                    output_path = handle.name
                result = execute_isolated_polars_plan(
                    {
                        "plan": plan,
                        "input_path": input_path,
                        "lookup_paths": {},
                        "output_path": output_path,
                        "masking_policies": {
                            "customer_mobile_last4": "preserve_last_4",
                            "redact": "redact",
                        },
                    },
                    memory_limit_bytes=plan["resource_limits"][
                        "memory_limit_bytes"
                    ],
                )
            prepared = self.artifact_store.prepare_path(
                output_path,
                new_governance_uid(),
                300,
                schema_fields=plan["output_fields"],
                limits=plan["resource_limits"],
            )
        finally:
            if output_path is not None:
                with suppress(FileNotFoundError):
                    os.unlink(output_path)
        count_keys = (
            "rows_in",
            "rows_out",
            "rows_rejected",
            "rows_quarantined",
            "rows_filtered",
            "rows_deduplicated",
            "rows_join_dropped",
            "rows_aggregated",
            "violation_count",
        )
        return {
            "status": "success",
            "test_kind": "dry_run",
            "run_id": new_governance_uid(),
            "plan_hash": context["plan_hash"],
            "schema_hashes": context["schema_hashes"],
            "binding_hashes": context["binding_hashes"],
            "counts": {key: int(result.get(key, 0)) for key in count_keys},
            "attestation": {
                "input_digest": sample["digest"],
                "output_digest": prepared["digest"],
                "output_schema_hash": prepared["schema_hash"],
                "violation_digest": hashlib.sha256(
                    _canonical(result.get("violations", []))
                ).hexdigest(),
            },
        }

    def _run_sql(self, context: dict[str, Any]) -> dict[str, Any]:
        from sqlalchemy import text

        from app.core.data_rules.compilers.sql import validate_bound_sql_plan

        if self.datasource_manager is None:
            raise ValueError(
                "server-owned SQL preflight executor is not configured"
            )
        plan = validate_bound_sql_plan(context.get("plan"))
        statement = plan["statements"][0]
        explain_sql = f"EXPLAIN {statement['sql']}"
        with self.datasource_manager.connect(
            plan["data_source_uid"], purpose="dataflow_read"
        ) as connection:
            rows = connection.execute(
                text(explain_sql), statement["parameters"]
            ).fetchall()
        if not rows:
            raise ValueError("physical SQL preflight returned no explain plan")
        explain_digest = hashlib.sha256(
            _canonical([list(row) for row in rows])
        ).hexdigest()
        return {
            "status": "success",
            "test_kind": "dry_run",
            "run_id": new_governance_uid(),
            "plan_hash": context["plan_hash"],
            "schema_hashes": context["schema_hashes"],
            "binding_hashes": context["binding_hashes"],
            "counts": {
                "rows_in": 0,
                "rows_out": 0,
                "rows_rejected": 0,
            },
            "attestation": {
                "dialect": plan["dialect"],
                "explain_digest": explain_digest,
                "statement_digest": hashlib.sha256(
                    statement["sql"].encode("utf-8")
                ).hexdigest(),
            },
        }


__all__ = [
    "CanonicalBoundCompiler",
    "GenerationReceiptSigner",
    "LogicalRuleCompiler",
    "PhysicalPlanPublicationService",
    "RulePublicationService",
    "RuleValidationRejected",
    "ServerOwnedLogicalDryRunRunner",
    "ServerOwnedPhysicalPreflightRunner",
    "generation_receipt_claims",
]