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 logging
import re
import threading
import time
from asyncio import CancelledError
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
from app.runner.rule_evidence import validate_public_rule_result

CONFIG_KEYS = {
    "component_binding_id",
    "rule_version_id",
    "execution_plan_hash",
    "provenance",
}
IDEMPOTENCY_STRATEGIES = {
    "partition_replace",
    "upsert",
    "deduplication_key",
}
LOGGER = logging.getLogger(__name__)


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],
        evidence_writer=None,
    ):
        self.repository = repository
        self.adapters = dict(adapters)
        self.evidence_writer = evidence_writer

    @staticmethod
    def _redacted_sample(value: Any) -> list[dict[str, Any]]:
        if not isinstance(value, list):
            return []
        sample = []
        for row in value[:100]:
            if not isinstance(row, dict):
                continue
            sample.append(
                {
                    str(key): (
                        None if item is None else "[REDACTED]"
                    )
                    for key, item in sorted(row.items())
                }
            )
        return sample

    def _start_evidence(
        self,
        *,
        component_binding_id: str,
        rule_version_id: str,
        plan_hash: str,
        execution_context: Mapping[str, Any],
    ):
        if self.evidence_writer is None:
            return None, None
        correlation_id = execution_context.get("correlation_id")
        dataflow_uid = execution_context.get("dataflow_uid")
        deployment_id = execution_context.get("deployment_id")
        environment = execution_context.get("environment")
        workflow_version = execution_context.get("workflow_version")
        node_id = execution_context.get("node_id")
        task_jti = execution_context.get("task_jti")
        if (
            not correlation_id
            or not dataflow_uid
            or not deployment_id
            or environment
            not in {"development", "test", "production"}
            or isinstance(workflow_version, bool)
            or not isinstance(workflow_version, int)
            or workflow_version < 1
            or not isinstance(node_id, str)
            or not node_id
        ):
            raise NodeExecutionError(
                "trusted rule evidence context is incomplete"
            )
        try:
            rule_run_id = self.evidence_writer.start(
                component_binding_id=component_binding_id,
                rule_version_id=rule_version_id,
                plan_hash=plan_hash,
                correlation_id=correlation_id,
                dataflow_uid=dataflow_uid,
                deployment_id=deployment_id,
                environment=environment,
                workflow_version=workflow_version,
                node_id=node_id,
                lease_owner=task_jti,
            )
            replay = self.evidence_writer.replay(rule_run_id)
        except Exception as exc:
            LOGGER.exception("governed rule evidence start failed")
            raise NodeExecutionError(
                "rule execution evidence is unavailable"
            ) from exc
        return rule_run_id, replay

    def _finish_evidence(self, rule_run_id, result):
        if rule_run_id is None:
            return
        try:
            self.evidence_writer.finish(rule_run_id, result)
        except Exception as exc:
            raise NodeExecutionError(
                "rule execution evidence commit outcome is unknown",
                commit_outcome="unknown",
            ) from exc

    def replay_task(
        self,
        *,
        node,
        correlation_id,
        deployment_id,
        task_jti,
    ):
        if self.evidence_writer is None:
            return None
        config = node.get("config")
        if not isinstance(config, dict):
            return None
        try:
            replay = self.evidence_writer.reconcile_expired_lease(
                lease_owner=task_jti,
                deployment_id=deployment_id,
                correlation_id=correlation_id,
                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 ""
                ),
            )
        except Exception as exc:
            raise NodeExecutionError(
                "rule execution evidence is unavailable"
            ) from exc
        return replay

    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_parameters = parameters
        staging_input = None
        if backend == "sql_pushdown" and parameters not in ({}, None):
            if (
                not isinstance(parameters, dict)
                or set(parameters) != {"input_artifact"}
                or not isinstance(parameters["input_artifact"], str)
            ):
                raise NodeExecutionError(
                    "SQL rule handoff must contain one opaque receipt"
                )
            staging_input = parameters["input_artifact"]
            adapter_parameters = {}
        rule_run_id, replay = self._start_evidence(
            component_binding_id=component_binding_id,
            rule_version_id=rule_version_id,
            plan_hash=plan_hash,
            execution_context=execution_context,
        )
        if isinstance(replay, dict):
            status = replay.get("status")
            if status == "success":
                return {
                    key: value
                    for key, value in replay.items()
                    if key != "status"
                }
            if status in {"failed", "unknown", "cancelled"}:
                raise NodeExecutionError(
                    "rule execution was already finalized",
                    commit_outcome=str(
                        replay.get("commit_outcome")
                        or "not_applicable"
                    ),
                )
            raise NodeExecutionError("rule execution is already in progress")
        if self.evidence_writer is not None:
            try:
                self.evidence_writer.heartbeat(
                    rule_run_id,
                    execution_context.get("task_jti"),
                )
            except Exception as exc:
                raise NodeExecutionError(
                    "rule execution lease is unavailable"
                ) from exc
        if staging_input is not None:
            if self.evidence_writer is None:
                raise NodeExecutionError(
                    "SQL staging receipt resolver is unavailable"
                )
            try:
                self.evidence_writer.resolve_sql_staging(
                    staging_input,
                    deployment_id=execution_context.get(
                        "deployment_id"
                    ),
                    correlation_id=execution_context.get(
                        "correlation_id"
                    ),
                    input_binding_id=plan["input_binding_id"],
                )
            except Exception as exc:
                self._finish_evidence(
                    rule_run_id,
                    {
                        "status": "failed",
                        "commit_outcome": "not_committed",
                        "timings": {"duration_ms": 0},
                    },
                )
                raise NodeExecutionError(
                    "SQL staging receipt is not executable"
                ) from exc
        adapter_context = {}
        if backend == "polars_batch":
            adapter_context["correlation_id"] = execution_context.get(
                "correlation_id"
            )
        started = time.monotonic()
        heartbeat_stop = threading.Event()
        heartbeat_failures = []
        heartbeat_thread = None
        if self.evidence_writer is not None:
            interval = float(
                getattr(
                    self.evidence_writer,
                    "heartbeat_interval_seconds",
                    10.0,
                )
            )

            def renew_lease():
                while not heartbeat_stop.wait(interval):
                    try:
                        self.evidence_writer.heartbeat(
                            rule_run_id,
                            execution_context.get("task_jti"),
                        )
                    except Exception as exc:
                        heartbeat_failures.append(exc)
                        heartbeat_stop.set()

            heartbeat_thread = threading.Thread(
                target=renew_lease,
                name=f"rule-heartbeat-{rule_run_id}",
                daemon=True,
            )
            heartbeat_thread.start()
        try:
            result = adapter.execute(
                plan=record["plan"],
                node=node,
                parameters=adapter_parameters,
                write_authorized=write_authorized,
                **adapter_context,
            )
        except CancelledError:
            self._finish_evidence(
                rule_run_id,
                {
                    "status": "cancelled",
                    "commit_outcome": "not_applicable",
                    "timings": {
                        "duration_ms": int(
                            (time.monotonic() - started) * 1000
                        )
                    },
                },
            )
            raise
        except NodeExecutionError as exc:
            self._finish_evidence(
                rule_run_id,
                {
                    "status": (
                        "unknown"
                        if exc.commit_outcome == "unknown"
                        else "failed"
                    ),
                    "commit_outcome": exc.commit_outcome,
                    "timings": {
                        "duration_ms": int(
                            (time.monotonic() - started) * 1000
                        )
                    },
                },
            )
            raise
        except Exception as exc:
            self._finish_evidence(
                rule_run_id,
                {
                    "status": "failed",
                    "commit_outcome": "not_committed",
                    "timings": {
                        "duration_ms": int(
                            (time.monotonic() - started) * 1000
                        )
                    },
                },
            )
            raise NodeExecutionError(
                "rule plan execution failed",
                commit_outcome="not_committed",
            ) from exc
        finally:
            heartbeat_stop.set()
            if heartbeat_thread is not None:
                heartbeat_thread.join(timeout=5)
        if heartbeat_failures:
            result_outcome = (
                result.get("commit_outcome")
                if isinstance(result, dict)
                else None
            )
            evidence_outcome = (
                "committed"
                if result_outcome == "committed"
                else "unknown"
            )
            self._finish_evidence(
                rule_run_id,
                {
                    "status": "unknown",
                    "commit_outcome": evidence_outcome,
                    "timings": {
                        "duration_ms": int(
                            (time.monotonic() - started) * 1000
                        )
                    },
                },
            )
            raise NodeExecutionError(
                "rule execution lease outcome is unknown",
                commit_outcome="unknown",
            )
        if not isinstance(result, dict):
            error = NodeExecutionError(
                "rule plan result must be an object"
            )
            self._finish_evidence(
                rule_run_id,
                {
                    "status": "failed",
                    "commit_outcome": error.commit_outcome,
                    "timings": {
                        "duration_ms": int(
                            (time.monotonic() - started) * 1000
                        )
                    },
                },
            )
            raise error
        adapter_commit_outcome = str(
            result.get("commit_outcome", "not_applicable")
        )
        staging_output = None
        if backend == "sql_pushdown" and self.evidence_writer is not None:
            try:
                staging_output = self.evidence_writer.stage_sql_output(
                    rule_run_id,
                    output_binding_id=plan["output_binding_id"],
                )
            except Exception as exc:
                self._finish_evidence(
                    rule_run_id,
                    {
                        "status": "unknown",
                        "commit_outcome": (
                            "committed"
                            if adapter_commit_outcome == "committed"
                            else "unknown"
                        ),
                        "timings": {
                            "duration_ms": int(
                                (time.monotonic() - started) * 1000
                            )
                        },
                    },
                )
                raise NodeExecutionError(
                    "SQL staging receipt outcome is unknown",
                    commit_outcome="unknown",
                ) from exc
        try:
            sample = self._redacted_sample(
                result.pop("_violation_sample", None)
            )
            public_result = {
                **result,
                "component_binding_id": component_binding_id,
                "rule_version_id": rule_version_id,
                "execution_plan_hash": plan_hash,
            }
            for metric in (
                "rows_in",
                "rows_out",
                "rows_rejected",
                "rows_quarantined",
            ):
                value = public_result.get(metric, 0)
                if isinstance(value, bool):
                    raise ValueError("boolean metric")
                public_result[metric] = max(0, int(value))
            artifact_ref = public_result.get("artifact_ref")
            if isinstance(artifact_ref, str):
                public_result["output_artifact"] = artifact_ref
            elif (
                backend == "sql_pushdown"
                and staging_output is not None
            ):
                public_result["output_artifact"] = staging_output
            public_result = validate_public_rule_result(
                public_result,
                backend=backend,
            )
        except (TypeError, ValueError) as exc:
            post_commit_unknown = adapter_commit_outcome in {
                "committed",
                "unknown",
            }
            self._finish_evidence(
                rule_run_id,
                {
                    "status": (
                        "unknown"
                        if post_commit_unknown
                        else "failed"
                    ),
                    "commit_outcome": (
                        adapter_commit_outcome
                        if post_commit_unknown
                        else "not_committed"
                    ),
                    "timings": {
                        "duration_ms": int(
                            (time.monotonic() - started) * 1000
                        )
                    },
                },
            )
            raise NodeExecutionError(
                "rule plan result is invalid",
                commit_outcome=(
                    "unknown"
                    if post_commit_unknown
                    else "not_committed"
                ),
            ) from exc
        evidence_result = {
            "status": "success",
            "rows_in": max(0, int(public_result.get("rows_in", 0))),
            "rows_out": max(0, int(public_result.get("rows_out", 0))),
            "rows_rejected": max(
                0, int(public_result.get("rows_rejected", 0))
            ),
            "rows_quarantined": max(
                0, int(public_result.get("rows_quarantined", 0))
            ),
            "commit_outcome": str(
                public_result.get("commit_outcome", "not_applicable")
            ),
            "timings": {
                "duration_ms": int(
                    (time.monotonic() - started) * 1000
                )
            },
            "public_result": public_result,
        }
        if sample:
            evidence_result.update(
                {
                    "violation_sample": sample,
                    "sample_count": len(sample),
                    "redaction_policy": (
                        "rule-violation-default-v1"
                    ),
                }
            )
        self._finish_evidence(rule_run_id, evidence_result)
        return public_result