"""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 ( bound_sql_plan_relations, 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, ib.binding_hash AS canonical_input_binding_hash, ib.data_source_uid::text AS canonical_input_data_source_uid, ib.object_kind AS canonical_input_object_kind, ib.object_ref AS canonical_input_object_ref, ib.dialect AS canonical_input_dialect, ib.access_mode AS canonical_input_access_mode, ib.write_mode AS canonical_input_write_mode, outs.id::text AS canonical_output_schema_snapshot_id, outs.schema_hash AS canonical_output_schema_hash, ob.binding_hash AS canonical_output_binding_hash, ob.data_source_uid::text AS canonical_output_data_source_uid, ob.object_kind AS canonical_output_object_kind, ob.object_ref AS canonical_output_object_ref, ob.dialect AS canonical_output_dialect, ob.access_mode AS canonical_output_access_mode, ob.write_mode AS canonical_output_write_mode, b.component_kind, b.idempotency AS binding_idempotency, EXISTS ( SELECT 1 FROM public.rule_publication_audits pa WHERE pa.rule_execution_plan_id = p.id AND pa.rule_version_id = r.id AND pa.action = 'published' AND pa.to_status = 'published' AND pa.evidence_hash = p.plan_hash ) AS publication_audit_trusted, EXISTS ( SELECT 1 FROM public.rule_logical_plans lp JOIN public.rule_logical_compile_evidence lce ON lce.logical_plan_id = lp.id JOIN public.rule_logical_test_evidence lte ON lte.logical_plan_id = lp.id WHERE lp.rule_version_id = r.id AND lp.status = 'published' AND lce.status = 'success' AND lce.compiler_version = lp.compiler_version AND lce.plan_hash = lp.plan_hash AND lce.schema_hashes = lp.schema_hashes AND lce.capabilities = lp.capabilities AND lte.status = 'success' AND lte.plan_hash = lp.plan_hash AND lte.schema_hashes = lp.schema_hashes ) AS logical_evidence_trusted, EXISTS ( SELECT 1 FROM public.rule_compile_evidence pce JOIN public.rule_test_evidence pte ON pte.rule_execution_plan_id = pce.rule_execution_plan_id WHERE pce.rule_execution_plan_id = p.id AND pce.status = 'success' AND pce.legacy_untrusted = FALSE AND pce.compiler_version = p.compiler_version AND pce.plan_hash = p.plan_hash AND pce.schema_hashes = p.schema_hashes AND pce.binding_hashes = jsonb_build_object( 'input', ib.binding_hash, 'output', ob.binding_hash ) AND pce.capabilities = CASE WHEN p.backend = 'polars_batch' THEN jsonb_build_object( 'resource_limits', p.plan->'resource_limits' ) ELSE p.plan->'capabilities' END AND pte.status = 'success' AND pte.legacy_untrusted = FALSE AND pte.plan_hash = p.plan_hash AND pte.schema_hashes = p.schema_hashes AND pte.binding_hashes = jsonb_build_object( 'input', ib.binding_hash, 'output', ob.binding_hash ) ) AS physical_evidence_trusted 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 record.get("publication_audit_trusted") is not True or record.get("logical_evidence_trusted") is not True or record.get("physical_evidence_trusted") is not True or _canonical_hash(record.get("plan")) != plan_hash ): raise NodeExecutionError("published rule plan is not executable") if node.get("type") == "rule.apply": if ( record.get("component_kind") != "rule.apply" or record.get("binding_idempotency") != node.get("idempotency") ): raise NodeExecutionError( "governed rule idempotency does not match its binding" ) elif ( record.get("component_kind") != "quality.check" or record.get("binding_idempotency") is not None ): raise NodeExecutionError( "governed quality binding must be read only" ) backend = record.get("backend") if backend == "sql_pushdown": try: plan = validate_bound_sql_plan(record.get("plan")) relations = bound_sql_plan_relations(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"] or record.get("canonical_input_data_source_uid") != plan["data_source_uid"] or record.get("canonical_output_data_source_uid") != plan["data_source_uid"] or record.get("canonical_input_object_ref") != relations["input_object_ref"] or record.get("canonical_output_object_ref") != relations["output_object_ref"] or record.get("canonical_input_access_mode") != "read" or record.get("canonical_input_write_mode") is not None or record.get("canonical_output_access_mode") not in {"write", "read_write"} or record.get("canonical_output_write_mode") != "append" or ( "postgresql" if record.get("canonical_input_dialect") == "postgres" else record.get("canonical_input_dialect") ) != plan["dialect"] or ( "postgresql" if record.get("canonical_output_dialect") == "postgres" else record.get("canonical_output_dialect") ) != plan["dialect"] ): 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"] or record.get("canonical_input_object_kind") != "parquet_artifact" or record.get("canonical_output_object_kind") != "parquet_artifact" or record.get("canonical_input_binding_hash") != plan["input_binding_hash"] or record.get("canonical_output_binding_hash") != plan["output_binding_hash"] ): raise NodeExecutionError( "published rule plan canonical attestation does not match" ) adapter_backend = ( "quality_check" if ( node.get("type") == "quality.check" and backend in {"sql_pushdown", "quality_check"} ) else backend ) adapter = self.adapters.get(adapter_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 ( node.get("type") == "rule.apply" and 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 evidence_backend = ( "quality_check" if node.get("type") == "quality.check" else adapter_backend ) public_result = validate_public_rule_result( public_result, backend=evidence_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