DataOps-platform/app/runner/rules.py

"""Load immutable rule plans and dispatch them through allowlisted adapters."""

from __future__ import annotations

import hashlib
import json
import re
from collections.abc import Mapping
from typing import Any

from sqlalchemy import text

from app.core.common.identifiers import ensure_governance_uid
from app.core.data_rules.compilers.polars import (
    COMPILER_VERSION as POLARS_COMPILER_VERSION,
)
from app.core.data_rules.compilers.polars import (
    validate_bound_polars_plan,
)
from app.core.data_rules.compilers.sql import (
    COMPILER_VERSION as SQL_COMPILER_VERSION,
)
from app.core.data_rules.compilers.sql import (
    validate_bound_sql_plan,
)
from app.runner.nodes import NodeExecutionError

CONFIG_KEYS = {
    "component_binding_id",
    "rule_version_id",
    "execution_plan_hash",
    "provenance",
}
IDEMPOTENCY_STRATEGIES = {
    "partition_replace",
    "upsert",
    "deduplication_key",
}


def _canonical_hash(value: Any) -> str:
    try:
        canonical = json.dumps(
            value,
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
        )
    except (TypeError, ValueError) as exc:
        raise NodeExecutionError(
            "published rule plan is not JSON serializable"
        ) from exc
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _uid(value: Any, label: str) -> str:
    try:
        return ensure_governance_uid({"uid": str(value)})
    except ValueError as exc:
        raise NodeExecutionError(f"{label} is invalid") from exc


class PostgresRulePlanRepository:
    """Read one plan through all immutable binding and version constraints."""

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

    def load(
        self,
        *,
        component_binding_id: str,
        rule_version_id: str,
        plan_hash: str,
    ) -> dict[str, Any] | None:
        statement = text(
            """
            SELECT
                p.component_binding_id::text AS component_binding_id,
                b.rule_version_id::text AS rule_version_id,
                p.backend,
                p.compiler_version,
                p.plan,
                p.plan_hash,
                p.schema_hashes,
                p.status AS plan_status,
                r.status AS rule_status,
                r.spec_hash AS canonical_rule_spec_hash,
                ins.id::text AS canonical_input_schema_snapshot_id,
                ins.schema_hash AS canonical_input_schema_hash,
                outs.id::text AS canonical_output_schema_snapshot_id,
                outs.schema_hash AS canonical_output_schema_hash,
                b.component_kind,
                b.idempotency AS binding_idempotency
            FROM public.rule_execution_plans p
            JOIN public.dataflow_component_bindings b
              ON b.id = p.component_binding_id
            JOIN public.data_rule_versions r
              ON r.id = b.rule_version_id
            LEFT JOIN public.dataflow_dataset_bindings ib
              ON ib.id = CAST(p.plan->>'input_binding_id' AS uuid)
            LEFT JOIN public.data_schema_snapshots ins
              ON ins.id = ib.schema_snapshot_id
            LEFT JOIN public.dataflow_dataset_bindings ob
              ON ob.id = CAST(p.plan->>'output_binding_id' AS uuid)
            LEFT JOIN public.data_schema_snapshots outs
              ON outs.id = ob.schema_snapshot_id
            WHERE p.component_binding_id = CAST(:component_binding_id AS uuid)
              AND b.rule_version_id = CAST(:rule_version_id AS uuid)
              AND p.plan_hash = :plan_hash
            """
        )
        with self.engine.connect() as connection:
            row = connection.execute(
                statement,
                {
                    "component_binding_id": component_binding_id,
                    "rule_version_id": rule_version_id,
                    "plan_hash": plan_hash,
                },
            ).mappings().one_or_none()
        return dict(row) if row is not None else None


class RulePlanExecutor:
    """Fail closed unless the exact published plan is still executable."""

    def __init__(self, repository, *, adapters: Mapping[str, Any]):
        self.repository = repository
        self.adapters = dict(adapters)

    def execute(
        self,
        node,
        parameters,
        *,
        write_authorized=False,
        **execution_context,
    ):
        if node.get("type") not in {"rule.apply", "quality.check"}:
            raise NodeExecutionError("unsupported governed rule node")
        config = node.get("config")
        if not isinstance(config, dict) or set(config) - CONFIG_KEYS:
            raise NodeExecutionError(
                "governed rule node contains inline or unsupported fields"
            )
        component_binding_id = _uid(
            config.get("component_binding_id"),
            "component_binding_id",
        )
        rule_version_id = _uid(
            config.get("rule_version_id"),
            "rule_version_id",
        )
        plan_hash = str(config.get("execution_plan_hash") or "")
        if not re.fullmatch(r"[0-9a-f]{64}", plan_hash):
            raise NodeExecutionError("execution plan hash is invalid")
        if node.get("type") == "rule.apply":
            idempotency = node.get("idempotency")
            if (
                node.get("purpose") != "write"
                or not write_authorized
                or not isinstance(idempotency, dict)
                or idempotency.get("strategy")
                not in IDEMPOTENCY_STRATEGIES
                or not str(idempotency.get("key") or "").strip()
            ):
                raise NodeExecutionError(
                    "governed write authorization and idempotency are required"
                )
        elif node.get("purpose") != "read":
            raise NodeExecutionError("quality check must be read only")

        record = self.repository.load(
            component_binding_id=component_binding_id,
            rule_version_id=rule_version_id,
            plan_hash=plan_hash,
        )
        if not isinstance(record, dict):
            raise NodeExecutionError("published rule plan was not found")
        if (
            record.get("component_binding_id") != component_binding_id
            or record.get("rule_version_id") != rule_version_id
            or record.get("plan_hash") != plan_hash
            or record.get("plan_status") != "published"
            or record.get("rule_status") != "published"
            or _canonical_hash(record.get("plan")) != plan_hash
        ):
            raise NodeExecutionError("published rule plan is not executable")
        if node.get("type") == "rule.apply" and (
            record.get("component_kind") != "rule.apply"
            or record.get("binding_idempotency") != node.get("idempotency")
        ):
            raise NodeExecutionError(
                "governed rule idempotency does not match its binding"
            )
        backend = record.get("backend")
        if backend == "sql_pushdown":
            try:
                plan = validate_bound_sql_plan(record.get("plan"))
            except ValueError as exc:
                raise NodeExecutionError(
                    "published rule plan is not executable"
                ) from exc
            expected_schema_hashes = {
                "rule_spec_hash": plan["rule_spec_hash"],
                "input_schema_snapshot_id": plan[
                    "input_schema_snapshot_id"
                ],
                "input_schema_hash": plan["input_schema_hash"],
                "output_schema_snapshot_id": plan[
                    "output_schema_snapshot_id"
                ],
                "output_schema_hash": plan["output_schema_hash"],
            }
            if (
                record.get("compiler_version") != SQL_COMPILER_VERSION
                or plan["compiler_version"] != SQL_COMPILER_VERSION
                or record.get("schema_hashes") != expected_schema_hashes
                or record.get("canonical_rule_spec_hash")
                != plan["rule_spec_hash"]
                or record.get("canonical_input_schema_snapshot_id")
                != plan["input_schema_snapshot_id"]
                or record.get("canonical_input_schema_hash")
                != plan["input_schema_hash"]
                or record.get("canonical_output_schema_snapshot_id")
                != plan["output_schema_snapshot_id"]
                or record.get("canonical_output_schema_hash")
                != plan["output_schema_hash"]
            ):
                raise NodeExecutionError(
                    "published rule plan canonical attestation does not match"
                )
        elif backend == "polars_batch":
            try:
                plan = validate_bound_polars_plan(record.get("plan"))
            except ValueError as exc:
                raise NodeExecutionError(
                    "published rule plan is not executable"
                ) from exc
            expected_schema_hashes = {
                "rule_spec_hash": plan["rule_spec_hash"],
                "input_schema_snapshot_id": plan[
                    "input_schema_snapshot_id"
                ],
                "input_schema_hash": plan["input_schema_hash"],
                "output_schema_snapshot_id": plan[
                    "output_schema_snapshot_id"
                ],
                "output_schema_hash": plan["output_schema_hash"],
            }
            if (
                record.get("compiler_version") != POLARS_COMPILER_VERSION
                or plan["compiler_version"] != POLARS_COMPILER_VERSION
                or record.get("schema_hashes") != expected_schema_hashes
                or record.get("canonical_rule_spec_hash")
                != plan["rule_spec_hash"]
                or record.get("canonical_input_schema_snapshot_id")
                != plan["input_schema_snapshot_id"]
                or record.get("canonical_input_schema_hash")
                != plan["input_schema_hash"]
                or record.get("canonical_output_schema_snapshot_id")
                != plan["output_schema_snapshot_id"]
                or record.get("canonical_output_schema_hash")
                != plan["output_schema_hash"]
            ):
                raise NodeExecutionError(
                    "published rule plan canonical attestation does not match"
                )
        adapter = self.adapters.get(backend)
        if adapter is None or not callable(getattr(adapter, "execute", None)):
            raise NodeExecutionError("rule plan backend is not registered")
        adapter_context = {}
        if backend == "polars_batch":
            adapter_context["correlation_id"] = execution_context.get(
                "correlation_id"
            )
        result = adapter.execute(
            plan=record["plan"],
            node=node,
            parameters=parameters,
            write_authorized=write_authorized,
            **adapter_context,
        )
        if not isinstance(result, dict):
            raise NodeExecutionError("rule plan result must be an object")
        return {
            **result,
            "component_binding_id": component_binding_id,
            "rule_version_id": rule_version_id,
            "execution_plan_hash": plan_hash,
        }