rules.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882
  1. """Load immutable rule plans and dispatch them through allowlisted adapters."""
  2. from __future__ import annotations
  3. import hashlib
  4. import json
  5. import logging
  6. import re
  7. import threading
  8. import time
  9. from asyncio import CancelledError
  10. from collections.abc import Mapping
  11. from typing import Any
  12. from sqlalchemy import text
  13. from app.core.common.identifiers import ensure_governance_uid
  14. from app.core.data_rules.compilers.polars import (
  15. COMPILER_VERSION as POLARS_COMPILER_VERSION,
  16. )
  17. from app.core.data_rules.compilers.polars import (
  18. validate_bound_polars_plan,
  19. )
  20. from app.core.data_rules.compilers.sql import (
  21. COMPILER_VERSION as SQL_COMPILER_VERSION,
  22. )
  23. from app.core.data_rules.compilers.sql import (
  24. bound_sql_plan_relations,
  25. validate_bound_sql_plan,
  26. )
  27. from app.runner.nodes import NodeExecutionError
  28. from app.runner.rule_evidence import validate_public_rule_result
  29. CONFIG_KEYS = {
  30. "component_binding_id",
  31. "rule_version_id",
  32. "execution_plan_hash",
  33. "provenance",
  34. }
  35. IDEMPOTENCY_STRATEGIES = {
  36. "partition_replace",
  37. "upsert",
  38. "deduplication_key",
  39. }
  40. LOGGER = logging.getLogger(__name__)
  41. def _canonical_hash(value: Any) -> str:
  42. try:
  43. canonical = json.dumps(
  44. value,
  45. sort_keys=True,
  46. separators=(",", ":"),
  47. ensure_ascii=False,
  48. )
  49. except (TypeError, ValueError) as exc:
  50. raise NodeExecutionError(
  51. "published rule plan is not JSON serializable"
  52. ) from exc
  53. return hashlib.sha256(canonical.encode("utf-8")).hexdigest()
  54. def _uid(value: Any, label: str) -> str:
  55. try:
  56. return ensure_governance_uid({"uid": str(value)})
  57. except ValueError as exc:
  58. raise NodeExecutionError(f"{label} is invalid") from exc
  59. class PostgresRulePlanRepository:
  60. """Read one plan through all immutable binding and version constraints."""
  61. def __init__(self, engine):
  62. self.engine = engine
  63. def load(
  64. self,
  65. *,
  66. component_binding_id: str,
  67. rule_version_id: str,
  68. plan_hash: str,
  69. ) -> dict[str, Any] | None:
  70. statement = text(
  71. """
  72. SELECT
  73. p.component_binding_id::text AS component_binding_id,
  74. b.rule_version_id::text AS rule_version_id,
  75. p.backend,
  76. p.compiler_version,
  77. p.plan,
  78. p.plan_hash,
  79. p.schema_hashes,
  80. p.status AS plan_status,
  81. r.status AS rule_status,
  82. r.spec_hash AS canonical_rule_spec_hash,
  83. ins.id::text AS canonical_input_schema_snapshot_id,
  84. ins.schema_hash AS canonical_input_schema_hash,
  85. ib.binding_hash AS canonical_input_binding_hash,
  86. ib.data_source_uid::text AS canonical_input_data_source_uid,
  87. ib.object_kind AS canonical_input_object_kind,
  88. ib.object_ref AS canonical_input_object_ref,
  89. ib.dialect AS canonical_input_dialect,
  90. ib.access_mode AS canonical_input_access_mode,
  91. ib.write_mode AS canonical_input_write_mode,
  92. outs.id::text AS canonical_output_schema_snapshot_id,
  93. outs.schema_hash AS canonical_output_schema_hash,
  94. ob.binding_hash AS canonical_output_binding_hash,
  95. ob.data_source_uid::text AS canonical_output_data_source_uid,
  96. ob.object_kind AS canonical_output_object_kind,
  97. ob.object_ref AS canonical_output_object_ref,
  98. ob.dialect AS canonical_output_dialect,
  99. ob.access_mode AS canonical_output_access_mode,
  100. ob.write_mode AS canonical_output_write_mode,
  101. b.component_kind,
  102. b.idempotency AS binding_idempotency,
  103. EXISTS (
  104. SELECT 1 FROM public.rule_publication_audits pa
  105. WHERE pa.rule_execution_plan_id = p.id
  106. AND pa.rule_version_id = r.id
  107. AND pa.action = 'published'
  108. AND pa.to_status = 'published'
  109. AND pa.evidence_hash = p.plan_hash
  110. ) AS publication_audit_trusted,
  111. EXISTS (
  112. SELECT 1
  113. FROM public.rule_logical_plans lp
  114. JOIN public.rule_logical_compile_evidence lce
  115. ON lce.logical_plan_id = lp.id
  116. JOIN public.rule_logical_test_evidence lte
  117. ON lte.logical_plan_id = lp.id
  118. WHERE lp.rule_version_id = r.id
  119. AND lp.status = 'published'
  120. AND lce.status = 'success'
  121. AND lce.compiler_version = lp.compiler_version
  122. AND lce.plan_hash = lp.plan_hash
  123. AND lce.schema_hashes = lp.schema_hashes
  124. AND lce.capabilities = lp.capabilities
  125. AND lte.status = 'success'
  126. AND lte.plan_hash = lp.plan_hash
  127. AND lte.schema_hashes = lp.schema_hashes
  128. ) AS logical_evidence_trusted,
  129. EXISTS (
  130. SELECT 1
  131. FROM public.rule_compile_evidence pce
  132. JOIN public.rule_test_evidence pte
  133. ON pte.rule_execution_plan_id =
  134. pce.rule_execution_plan_id
  135. WHERE pce.rule_execution_plan_id = p.id
  136. AND pce.status = 'success'
  137. AND pce.legacy_untrusted = FALSE
  138. AND pce.compiler_version = p.compiler_version
  139. AND pce.plan_hash = p.plan_hash
  140. AND pce.schema_hashes = p.schema_hashes
  141. AND pce.binding_hashes = jsonb_build_object(
  142. 'input', ib.binding_hash,
  143. 'output', ob.binding_hash
  144. )
  145. AND pce.capabilities = CASE
  146. WHEN p.backend = 'polars_batch' THEN
  147. jsonb_build_object(
  148. 'resource_limits',
  149. p.plan->'resource_limits'
  150. )
  151. ELSE p.plan->'capabilities'
  152. END
  153. AND pte.status = 'success'
  154. AND pte.legacy_untrusted = FALSE
  155. AND pte.plan_hash = p.plan_hash
  156. AND pte.schema_hashes = p.schema_hashes
  157. AND pte.binding_hashes = jsonb_build_object(
  158. 'input', ib.binding_hash,
  159. 'output', ob.binding_hash
  160. )
  161. ) AS physical_evidence_trusted
  162. FROM public.rule_execution_plans p
  163. JOIN public.dataflow_component_bindings b
  164. ON b.id = p.component_binding_id
  165. JOIN public.data_rule_versions r
  166. ON r.id = b.rule_version_id
  167. LEFT JOIN public.dataflow_dataset_bindings ib
  168. ON ib.id = CAST(p.plan->>'input_binding_id' AS uuid)
  169. LEFT JOIN public.data_schema_snapshots ins
  170. ON ins.id = ib.schema_snapshot_id
  171. LEFT JOIN public.dataflow_dataset_bindings ob
  172. ON ob.id = CAST(p.plan->>'output_binding_id' AS uuid)
  173. LEFT JOIN public.data_schema_snapshots outs
  174. ON outs.id = ob.schema_snapshot_id
  175. WHERE p.component_binding_id = CAST(:component_binding_id AS uuid)
  176. AND b.rule_version_id = CAST(:rule_version_id AS uuid)
  177. AND p.plan_hash = :plan_hash
  178. """
  179. )
  180. with self.engine.connect() as connection:
  181. row = connection.execute(
  182. statement,
  183. {
  184. "component_binding_id": component_binding_id,
  185. "rule_version_id": rule_version_id,
  186. "plan_hash": plan_hash,
  187. },
  188. ).mappings().one_or_none()
  189. return dict(row) if row is not None else None
  190. class RulePlanExecutor:
  191. """Fail closed unless the exact published plan is still executable."""
  192. def __init__(
  193. self,
  194. repository,
  195. *,
  196. adapters: Mapping[str, Any],
  197. evidence_writer=None,
  198. ):
  199. self.repository = repository
  200. self.adapters = dict(adapters)
  201. self.evidence_writer = evidence_writer
  202. @staticmethod
  203. def _redacted_sample(value: Any) -> list[dict[str, Any]]:
  204. if not isinstance(value, list):
  205. return []
  206. sample = []
  207. for row in value[:100]:
  208. if not isinstance(row, dict):
  209. continue
  210. sample.append(
  211. {
  212. str(key): (
  213. None if item is None else "[REDACTED]"
  214. )
  215. for key, item in sorted(row.items())
  216. }
  217. )
  218. return sample
  219. def _start_evidence(
  220. self,
  221. *,
  222. component_binding_id: str,
  223. rule_version_id: str,
  224. plan_hash: str,
  225. execution_context: Mapping[str, Any],
  226. ):
  227. if self.evidence_writer is None:
  228. return None, None
  229. correlation_id = execution_context.get("correlation_id")
  230. dataflow_uid = execution_context.get("dataflow_uid")
  231. deployment_id = execution_context.get("deployment_id")
  232. environment = execution_context.get("environment")
  233. workflow_version = execution_context.get("workflow_version")
  234. node_id = execution_context.get("node_id")
  235. task_jti = execution_context.get("task_jti")
  236. if (
  237. not correlation_id
  238. or not dataflow_uid
  239. or not deployment_id
  240. or environment
  241. not in {"development", "test", "production"}
  242. or isinstance(workflow_version, bool)
  243. or not isinstance(workflow_version, int)
  244. or workflow_version < 1
  245. or not isinstance(node_id, str)
  246. or not node_id
  247. ):
  248. raise NodeExecutionError(
  249. "trusted rule evidence context is incomplete"
  250. )
  251. try:
  252. rule_run_id = self.evidence_writer.start(
  253. component_binding_id=component_binding_id,
  254. rule_version_id=rule_version_id,
  255. plan_hash=plan_hash,
  256. correlation_id=correlation_id,
  257. dataflow_uid=dataflow_uid,
  258. deployment_id=deployment_id,
  259. environment=environment,
  260. workflow_version=workflow_version,
  261. node_id=node_id,
  262. lease_owner=task_jti,
  263. )
  264. replay = self.evidence_writer.replay(rule_run_id)
  265. except Exception as exc:
  266. LOGGER.exception("governed rule evidence start failed")
  267. raise NodeExecutionError(
  268. "rule execution evidence is unavailable"
  269. ) from exc
  270. return rule_run_id, replay
  271. def _finish_evidence(self, rule_run_id, result):
  272. if rule_run_id is None:
  273. return
  274. try:
  275. self.evidence_writer.finish(rule_run_id, result)
  276. except Exception as exc:
  277. raise NodeExecutionError(
  278. "rule execution evidence commit outcome is unknown",
  279. commit_outcome="unknown",
  280. ) from exc
  281. def replay_task(
  282. self,
  283. *,
  284. node,
  285. correlation_id,
  286. deployment_id,
  287. task_jti,
  288. ):
  289. if self.evidence_writer is None:
  290. return None
  291. config = node.get("config")
  292. if not isinstance(config, dict):
  293. return None
  294. try:
  295. replay = self.evidence_writer.reconcile_expired_lease(
  296. lease_owner=task_jti,
  297. deployment_id=deployment_id,
  298. correlation_id=correlation_id,
  299. component_binding_id=_uid(
  300. config.get("component_binding_id"),
  301. "component_binding_id",
  302. ),
  303. rule_version_id=_uid(
  304. config.get("rule_version_id"),
  305. "rule_version_id",
  306. ),
  307. plan_hash=str(
  308. config.get("execution_plan_hash") or ""
  309. ),
  310. )
  311. except Exception as exc:
  312. raise NodeExecutionError(
  313. "rule execution evidence is unavailable"
  314. ) from exc
  315. return replay
  316. def execute(
  317. self,
  318. node,
  319. parameters,
  320. *,
  321. write_authorized=False,
  322. **execution_context,
  323. ):
  324. if node.get("type") not in {"rule.apply", "quality.check"}:
  325. raise NodeExecutionError("unsupported governed rule node")
  326. config = node.get("config")
  327. if not isinstance(config, dict) or set(config) - CONFIG_KEYS:
  328. raise NodeExecutionError(
  329. "governed rule node contains inline or unsupported fields"
  330. )
  331. component_binding_id = _uid(
  332. config.get("component_binding_id"),
  333. "component_binding_id",
  334. )
  335. rule_version_id = _uid(
  336. config.get("rule_version_id"),
  337. "rule_version_id",
  338. )
  339. plan_hash = str(config.get("execution_plan_hash") or "")
  340. if not re.fullmatch(r"[0-9a-f]{64}", plan_hash):
  341. raise NodeExecutionError("execution plan hash is invalid")
  342. if node.get("type") == "rule.apply":
  343. idempotency = node.get("idempotency")
  344. if (
  345. node.get("purpose") != "write"
  346. or not write_authorized
  347. or not isinstance(idempotency, dict)
  348. or idempotency.get("strategy")
  349. not in IDEMPOTENCY_STRATEGIES
  350. or not str(idempotency.get("key") or "").strip()
  351. ):
  352. raise NodeExecutionError(
  353. "governed write authorization and idempotency are required"
  354. )
  355. elif node.get("purpose") != "read":
  356. raise NodeExecutionError("quality check must be read only")
  357. record = self.repository.load(
  358. component_binding_id=component_binding_id,
  359. rule_version_id=rule_version_id,
  360. plan_hash=plan_hash,
  361. )
  362. if not isinstance(record, dict):
  363. raise NodeExecutionError("published rule plan was not found")
  364. if (
  365. record.get("component_binding_id") != component_binding_id
  366. or record.get("rule_version_id") != rule_version_id
  367. or record.get("plan_hash") != plan_hash
  368. or record.get("plan_status") != "published"
  369. or record.get("rule_status") != "published"
  370. or record.get("publication_audit_trusted") is not True
  371. or record.get("logical_evidence_trusted") is not True
  372. or record.get("physical_evidence_trusted") is not True
  373. or _canonical_hash(record.get("plan")) != plan_hash
  374. ):
  375. raise NodeExecutionError("published rule plan is not executable")
  376. if node.get("type") == "rule.apply":
  377. if (
  378. record.get("component_kind") != "rule.apply"
  379. or record.get("binding_idempotency")
  380. != node.get("idempotency")
  381. ):
  382. raise NodeExecutionError(
  383. "governed rule idempotency does not match its binding"
  384. )
  385. elif (
  386. record.get("component_kind") != "quality.check"
  387. or record.get("binding_idempotency") is not None
  388. ):
  389. raise NodeExecutionError(
  390. "governed quality binding must be read only"
  391. )
  392. backend = record.get("backend")
  393. if backend == "sql_pushdown":
  394. try:
  395. plan = validate_bound_sql_plan(record.get("plan"))
  396. relations = bound_sql_plan_relations(plan)
  397. except ValueError as exc:
  398. raise NodeExecutionError(
  399. "published rule plan is not executable"
  400. ) from exc
  401. expected_schema_hashes = {
  402. "rule_spec_hash": plan["rule_spec_hash"],
  403. "input_schema_snapshot_id": plan[
  404. "input_schema_snapshot_id"
  405. ],
  406. "input_schema_hash": plan["input_schema_hash"],
  407. "output_schema_snapshot_id": plan[
  408. "output_schema_snapshot_id"
  409. ],
  410. "output_schema_hash": plan["output_schema_hash"],
  411. }
  412. if (
  413. record.get("compiler_version") != SQL_COMPILER_VERSION
  414. or plan["compiler_version"] != SQL_COMPILER_VERSION
  415. or record.get("schema_hashes") != expected_schema_hashes
  416. or record.get("canonical_rule_spec_hash")
  417. != plan["rule_spec_hash"]
  418. or record.get("canonical_input_schema_snapshot_id")
  419. != plan["input_schema_snapshot_id"]
  420. or record.get("canonical_input_schema_hash")
  421. != plan["input_schema_hash"]
  422. or record.get("canonical_output_schema_snapshot_id")
  423. != plan["output_schema_snapshot_id"]
  424. or record.get("canonical_output_schema_hash")
  425. != plan["output_schema_hash"]
  426. or record.get("canonical_input_data_source_uid")
  427. != plan["data_source_uid"]
  428. or record.get("canonical_output_data_source_uid")
  429. != plan["data_source_uid"]
  430. or record.get("canonical_input_object_ref")
  431. != relations["input_object_ref"]
  432. or record.get("canonical_output_object_ref")
  433. != relations["output_object_ref"]
  434. or record.get("canonical_input_access_mode") != "read"
  435. or record.get("canonical_input_write_mode") is not None
  436. or record.get("canonical_output_access_mode")
  437. not in {"write", "read_write"}
  438. or record.get("canonical_output_write_mode") != "append"
  439. or (
  440. "postgresql"
  441. if record.get("canonical_input_dialect") == "postgres"
  442. else record.get("canonical_input_dialect")
  443. )
  444. != plan["dialect"]
  445. or (
  446. "postgresql"
  447. if record.get("canonical_output_dialect") == "postgres"
  448. else record.get("canonical_output_dialect")
  449. )
  450. != plan["dialect"]
  451. ):
  452. raise NodeExecutionError(
  453. "published rule plan canonical attestation does not match"
  454. )
  455. elif backend == "polars_batch":
  456. try:
  457. plan = validate_bound_polars_plan(record.get("plan"))
  458. except ValueError as exc:
  459. raise NodeExecutionError(
  460. "published rule plan is not executable"
  461. ) from exc
  462. expected_schema_hashes = {
  463. "rule_spec_hash": plan["rule_spec_hash"],
  464. "input_schema_snapshot_id": plan[
  465. "input_schema_snapshot_id"
  466. ],
  467. "input_schema_hash": plan["input_schema_hash"],
  468. "output_schema_snapshot_id": plan[
  469. "output_schema_snapshot_id"
  470. ],
  471. "output_schema_hash": plan["output_schema_hash"],
  472. }
  473. if (
  474. record.get("compiler_version") != POLARS_COMPILER_VERSION
  475. or plan["compiler_version"] != POLARS_COMPILER_VERSION
  476. or record.get("schema_hashes") != expected_schema_hashes
  477. or record.get("canonical_rule_spec_hash")
  478. != plan["rule_spec_hash"]
  479. or record.get("canonical_input_schema_snapshot_id")
  480. != plan["input_schema_snapshot_id"]
  481. or record.get("canonical_input_schema_hash")
  482. != plan["input_schema_hash"]
  483. or record.get("canonical_output_schema_snapshot_id")
  484. != plan["output_schema_snapshot_id"]
  485. or record.get("canonical_output_schema_hash")
  486. != plan["output_schema_hash"]
  487. or record.get("canonical_input_object_kind")
  488. != "parquet_artifact"
  489. or record.get("canonical_output_object_kind")
  490. != "parquet_artifact"
  491. or record.get("canonical_input_binding_hash")
  492. != plan["input_binding_hash"]
  493. or record.get("canonical_output_binding_hash")
  494. != plan["output_binding_hash"]
  495. ):
  496. raise NodeExecutionError(
  497. "published rule plan canonical attestation does not match"
  498. )
  499. adapter_backend = (
  500. "quality_check"
  501. if (
  502. node.get("type") == "quality.check"
  503. and backend in {"sql_pushdown", "quality_check"}
  504. )
  505. else backend
  506. )
  507. adapter = self.adapters.get(adapter_backend)
  508. if adapter is None or not callable(getattr(adapter, "execute", None)):
  509. raise NodeExecutionError("rule plan backend is not registered")
  510. adapter_parameters = parameters
  511. staging_input = None
  512. if backend == "sql_pushdown" and parameters not in ({}, None):
  513. if (
  514. not isinstance(parameters, dict)
  515. or set(parameters) != {"input_artifact"}
  516. or not isinstance(parameters["input_artifact"], str)
  517. ):
  518. raise NodeExecutionError(
  519. "SQL rule handoff must contain one opaque receipt"
  520. )
  521. staging_input = parameters["input_artifact"]
  522. adapter_parameters = {}
  523. rule_run_id, replay = self._start_evidence(
  524. component_binding_id=component_binding_id,
  525. rule_version_id=rule_version_id,
  526. plan_hash=plan_hash,
  527. execution_context=execution_context,
  528. )
  529. if isinstance(replay, dict):
  530. status = replay.get("status")
  531. if status == "success":
  532. return {
  533. key: value
  534. for key, value in replay.items()
  535. if key != "status"
  536. }
  537. if status in {"failed", "unknown", "cancelled"}:
  538. raise NodeExecutionError(
  539. "rule execution was already finalized",
  540. commit_outcome=str(
  541. replay.get("commit_outcome")
  542. or "not_applicable"
  543. ),
  544. )
  545. raise NodeExecutionError("rule execution is already in progress")
  546. if self.evidence_writer is not None:
  547. try:
  548. self.evidence_writer.heartbeat(
  549. rule_run_id,
  550. execution_context.get("task_jti"),
  551. )
  552. except Exception as exc:
  553. raise NodeExecutionError(
  554. "rule execution lease is unavailable"
  555. ) from exc
  556. if staging_input is not None:
  557. if self.evidence_writer is None:
  558. raise NodeExecutionError(
  559. "SQL staging receipt resolver is unavailable"
  560. )
  561. try:
  562. self.evidence_writer.resolve_sql_staging(
  563. staging_input,
  564. deployment_id=execution_context.get(
  565. "deployment_id"
  566. ),
  567. correlation_id=execution_context.get(
  568. "correlation_id"
  569. ),
  570. input_binding_id=plan["input_binding_id"],
  571. )
  572. except Exception as exc:
  573. self._finish_evidence(
  574. rule_run_id,
  575. {
  576. "status": "failed",
  577. "commit_outcome": "not_committed",
  578. "timings": {"duration_ms": 0},
  579. },
  580. )
  581. raise NodeExecutionError(
  582. "SQL staging receipt is not executable"
  583. ) from exc
  584. adapter_context = {}
  585. if backend == "polars_batch":
  586. adapter_context["correlation_id"] = execution_context.get(
  587. "correlation_id"
  588. )
  589. started = time.monotonic()
  590. heartbeat_stop = threading.Event()
  591. heartbeat_failures = []
  592. heartbeat_thread = None
  593. if self.evidence_writer is not None:
  594. interval = float(
  595. getattr(
  596. self.evidence_writer,
  597. "heartbeat_interval_seconds",
  598. 10.0,
  599. )
  600. )
  601. def renew_lease():
  602. while not heartbeat_stop.wait(interval):
  603. try:
  604. self.evidence_writer.heartbeat(
  605. rule_run_id,
  606. execution_context.get("task_jti"),
  607. )
  608. except Exception as exc:
  609. heartbeat_failures.append(exc)
  610. heartbeat_stop.set()
  611. heartbeat_thread = threading.Thread(
  612. target=renew_lease,
  613. name=f"rule-heartbeat-{rule_run_id}",
  614. daemon=True,
  615. )
  616. heartbeat_thread.start()
  617. try:
  618. result = adapter.execute(
  619. plan=record["plan"],
  620. node=node,
  621. parameters=adapter_parameters,
  622. write_authorized=write_authorized,
  623. **adapter_context,
  624. )
  625. except CancelledError:
  626. self._finish_evidence(
  627. rule_run_id,
  628. {
  629. "status": "cancelled",
  630. "commit_outcome": "not_applicable",
  631. "timings": {
  632. "duration_ms": int(
  633. (time.monotonic() - started) * 1000
  634. )
  635. },
  636. },
  637. )
  638. raise
  639. except NodeExecutionError as exc:
  640. self._finish_evidence(
  641. rule_run_id,
  642. {
  643. "status": (
  644. "unknown"
  645. if exc.commit_outcome == "unknown"
  646. else "failed"
  647. ),
  648. "commit_outcome": exc.commit_outcome,
  649. "timings": {
  650. "duration_ms": int(
  651. (time.monotonic() - started) * 1000
  652. )
  653. },
  654. },
  655. )
  656. raise
  657. except Exception as exc:
  658. self._finish_evidence(
  659. rule_run_id,
  660. {
  661. "status": "failed",
  662. "commit_outcome": "not_committed",
  663. "timings": {
  664. "duration_ms": int(
  665. (time.monotonic() - started) * 1000
  666. )
  667. },
  668. },
  669. )
  670. raise NodeExecutionError(
  671. "rule plan execution failed",
  672. commit_outcome="not_committed",
  673. ) from exc
  674. finally:
  675. heartbeat_stop.set()
  676. if heartbeat_thread is not None:
  677. heartbeat_thread.join(timeout=5)
  678. if heartbeat_failures:
  679. result_outcome = (
  680. result.get("commit_outcome")
  681. if isinstance(result, dict)
  682. else None
  683. )
  684. evidence_outcome = (
  685. "committed"
  686. if result_outcome == "committed"
  687. else "unknown"
  688. )
  689. self._finish_evidence(
  690. rule_run_id,
  691. {
  692. "status": "unknown",
  693. "commit_outcome": evidence_outcome,
  694. "timings": {
  695. "duration_ms": int(
  696. (time.monotonic() - started) * 1000
  697. )
  698. },
  699. },
  700. )
  701. raise NodeExecutionError(
  702. "rule execution lease outcome is unknown",
  703. commit_outcome="unknown",
  704. )
  705. if not isinstance(result, dict):
  706. error = NodeExecutionError(
  707. "rule plan result must be an object"
  708. )
  709. self._finish_evidence(
  710. rule_run_id,
  711. {
  712. "status": "failed",
  713. "commit_outcome": error.commit_outcome,
  714. "timings": {
  715. "duration_ms": int(
  716. (time.monotonic() - started) * 1000
  717. )
  718. },
  719. },
  720. )
  721. raise error
  722. adapter_commit_outcome = str(
  723. result.get("commit_outcome", "not_applicable")
  724. )
  725. staging_output = None
  726. if (
  727. node.get("type") == "rule.apply"
  728. and backend == "sql_pushdown"
  729. and self.evidence_writer is not None
  730. ):
  731. try:
  732. staging_output = self.evidence_writer.stage_sql_output(
  733. rule_run_id,
  734. output_binding_id=plan["output_binding_id"],
  735. )
  736. except Exception as exc:
  737. self._finish_evidence(
  738. rule_run_id,
  739. {
  740. "status": "unknown",
  741. "commit_outcome": (
  742. "committed"
  743. if adapter_commit_outcome == "committed"
  744. else "unknown"
  745. ),
  746. "timings": {
  747. "duration_ms": int(
  748. (time.monotonic() - started) * 1000
  749. )
  750. },
  751. },
  752. )
  753. raise NodeExecutionError(
  754. "SQL staging receipt outcome is unknown",
  755. commit_outcome="unknown",
  756. ) from exc
  757. try:
  758. sample = self._redacted_sample(
  759. result.pop("_violation_sample", None)
  760. )
  761. public_result = {
  762. **result,
  763. "component_binding_id": component_binding_id,
  764. "rule_version_id": rule_version_id,
  765. "execution_plan_hash": plan_hash,
  766. }
  767. for metric in (
  768. "rows_in",
  769. "rows_out",
  770. "rows_rejected",
  771. "rows_quarantined",
  772. ):
  773. value = public_result.get(metric, 0)
  774. if isinstance(value, bool):
  775. raise ValueError("boolean metric")
  776. public_result[metric] = max(0, int(value))
  777. artifact_ref = public_result.get("artifact_ref")
  778. if isinstance(artifact_ref, str):
  779. public_result["output_artifact"] = artifact_ref
  780. elif (
  781. backend == "sql_pushdown"
  782. and staging_output is not None
  783. ):
  784. public_result["output_artifact"] = staging_output
  785. evidence_backend = (
  786. "quality_check"
  787. if node.get("type") == "quality.check"
  788. else adapter_backend
  789. )
  790. public_result = validate_public_rule_result(
  791. public_result,
  792. backend=evidence_backend,
  793. )
  794. except (TypeError, ValueError) as exc:
  795. post_commit_unknown = adapter_commit_outcome in {
  796. "committed",
  797. "unknown",
  798. }
  799. self._finish_evidence(
  800. rule_run_id,
  801. {
  802. "status": (
  803. "unknown"
  804. if post_commit_unknown
  805. else "failed"
  806. ),
  807. "commit_outcome": (
  808. adapter_commit_outcome
  809. if post_commit_unknown
  810. else "not_committed"
  811. ),
  812. "timings": {
  813. "duration_ms": int(
  814. (time.monotonic() - started) * 1000
  815. )
  816. },
  817. },
  818. )
  819. raise NodeExecutionError(
  820. "rule plan result is invalid",
  821. commit_outcome=(
  822. "unknown"
  823. if post_commit_unknown
  824. else "not_committed"
  825. ),
  826. ) from exc
  827. evidence_result = {
  828. "status": "success",
  829. "rows_in": max(0, int(public_result.get("rows_in", 0))),
  830. "rows_out": max(0, int(public_result.get("rows_out", 0))),
  831. "rows_rejected": max(
  832. 0, int(public_result.get("rows_rejected", 0))
  833. ),
  834. "rows_quarantined": max(
  835. 0, int(public_result.get("rows_quarantined", 0))
  836. ),
  837. "commit_outcome": str(
  838. public_result.get("commit_outcome", "not_applicable")
  839. ),
  840. "timings": {
  841. "duration_ms": int(
  842. (time.monotonic() - started) * 1000
  843. )
  844. },
  845. "public_result": public_result,
  846. }
  847. if sample:
  848. evidence_result.update(
  849. {
  850. "violation_sample": sample,
  851. "sample_count": len(sample),
  852. "redaction_policy": (
  853. "rule-violation-default-v1"
  854. ),
  855. }
  856. )
  857. self._finish_evidence(rule_run_id, evidence_result)
  858. return public_result