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 typing import Any, Mapping

from sqlalchemy import text

from app.core.common.identifiers import ensure_governance_uid
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.plan,
                p.plan_hash,
                p.status AS plan_status,
                r.status AS rule_status
            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
            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 SqlRulePlanAdapter:
    """Execute a compiled, parameterized SQL plan using existing governed nodes."""

    def __init__(self, *, query_executor, write_executor):
        self.query_executor = query_executor
        self.write_executor = write_executor

    def execute(
        self,
        *,
        plan,
        node,
        parameters,
        write_authorized,
    ):
        if not isinstance(plan, dict):
            raise NodeExecutionError("published SQL rule plan is invalid")
        unknown = set(plan) - {
            "statement",
            "parameters",
            "data_source_uid",
        }
        if unknown:
            raise NodeExecutionError(
                "published SQL rule plan contains unsupported fields"
            )
        compiled_node = {
            "id": node.get("id"),
            "data_source_uid": _uid(
                plan.get("data_source_uid"), "plan data_source_uid"
            ),
            "purpose": node.get("purpose"),
            "config": {
                "statement": plan.get("statement"),
                "parameters": plan.get("parameters", {}),
            },
        }
        if node.get("type") == "rule.apply":
            compiled_node["idempotency"] = node.get("idempotency")
            return self.write_executor.execute(
                compiled_node,
                parameters,
                write_authorized=write_authorized,
            )
        return self.query_executor.execute(compiled_node, parameters)


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,
        **_kwargs,
    ):
        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")
        backend = record.get("backend")
        adapter = self.adapters.get(backend)
        if adapter is None or not callable(getattr(adapter, "execute", None)):
            raise NodeExecutionError("rule plan backend is not registered")
        result = adapter.execute(
            plan=record["plan"],
            node=node,
            parameters=parameters,
            write_authorized=write_authorized,
        )
        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,
        }