rule_evidence.py 61 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551
  1. """Transactional, server-attested evidence for governed rule execution."""
  2. from __future__ import annotations
  3. import hashlib
  4. import json
  5. import os
  6. import re
  7. import tempfile
  8. import uuid
  9. from contextlib import suppress
  10. from typing import Any
  11. import polars as pl
  12. from sqlalchemy import text
  13. from app.core.common.identifiers import (
  14. ensure_governance_uid,
  15. new_governance_uid,
  16. )
  17. _DIGEST = re.compile(r"^[0-9a-f]{64}$")
  18. _FINAL_STATUSES = {"success", "failed", "unknown", "cancelled"}
  19. _COMMIT_OUTCOMES = {
  20. "not_applicable",
  21. "not_committed",
  22. "committed",
  23. "unknown",
  24. }
  25. _FINISH_KEYS = {
  26. "status",
  27. "rows_in",
  28. "rows_out",
  29. "rows_rejected",
  30. "rows_quarantined",
  31. "commit_outcome",
  32. "timings",
  33. "public_result",
  34. "violation_sample",
  35. "sample_count",
  36. "redaction_policy",
  37. }
  38. _PUBLIC_RESULT_KEYS = {
  39. "affected_rows",
  40. "artifact_ref",
  41. "commit_outcome",
  42. "component_binding_id",
  43. "digest",
  44. "execution_plan_hash",
  45. "expires_at",
  46. "output_artifact",
  47. "row_count",
  48. "rows_aggregated",
  49. "rows_deduplicated",
  50. "rows_filtered",
  51. "rows_in",
  52. "rows_join_dropped",
  53. "rows_out",
  54. "rows_quarantined",
  55. "rows_rejected",
  56. "rule_version_id",
  57. "schema_hash",
  58. "violation_count",
  59. "violations",
  60. }
  61. _ATTESTATION_RESULT_KEYS = {
  62. "component_binding_id",
  63. "execution_plan_hash",
  64. "rule_version_id",
  65. }
  66. _BACKEND_PUBLIC_RESULT_KEYS = {
  67. "sql_pushdown": _ATTESTATION_RESULT_KEYS
  68. | {
  69. "commit_outcome",
  70. "output_artifact",
  71. "rows_in",
  72. "rows_out",
  73. "rows_quarantined",
  74. "rows_rejected",
  75. },
  76. "polars_batch": _ATTESTATION_RESULT_KEYS
  77. | {
  78. "artifact_ref",
  79. "commit_outcome",
  80. "digest",
  81. "expires_at",
  82. "output_artifact",
  83. "row_count",
  84. "rows_aggregated",
  85. "rows_deduplicated",
  86. "rows_filtered",
  87. "rows_in",
  88. "rows_join_dropped",
  89. "rows_out",
  90. "rows_quarantined",
  91. "rows_rejected",
  92. "schema_hash",
  93. "violation_count",
  94. "violations",
  95. },
  96. "quality_check": _ATTESTATION_RESULT_KEYS
  97. | {
  98. "commit_outcome",
  99. "rows_in",
  100. "rows_out",
  101. "rows_quarantined",
  102. "rows_rejected",
  103. "violation_count",
  104. "violations",
  105. },
  106. }
  107. def _uid(value: Any, label: str) -> str:
  108. try:
  109. return ensure_governance_uid({"uid": str(value)})
  110. except ValueError as exc:
  111. raise ValueError(f"{label} is invalid") from exc
  112. def _uuid(value: Any, label: str) -> str:
  113. try:
  114. return str(uuid.UUID(str(value)))
  115. except (TypeError, ValueError, AttributeError) as exc:
  116. raise ValueError(f"{label} is invalid") from exc
  117. def _canonical_digest(value: Any) -> str:
  118. encoded = json.dumps(
  119. value,
  120. sort_keys=True,
  121. separators=(",", ":"),
  122. ensure_ascii=False,
  123. ).encode("utf-8")
  124. return hashlib.sha256(encoded).hexdigest()
  125. def _bounded_count(value: Any, label: str) -> int:
  126. if isinstance(value, bool):
  127. raise ValueError(f"{label} is invalid")
  128. try:
  129. normalized = int(value or 0)
  130. except (TypeError, ValueError) as exc:
  131. raise ValueError(f"{label} is invalid") from exc
  132. if normalized < 0 or normalized > 10_000_000_000:
  133. raise ValueError(f"{label} is outside the evidence limit")
  134. return normalized
  135. def _sample_fields(sample: list[dict[str, Any]]) -> list[dict[str, Any]]:
  136. names = sorted({str(key) for row in sample for key in row})
  137. if not names:
  138. raise ValueError("violation sample has no fields")
  139. return [
  140. {"name": name, "type": "string", "nullable": True}
  141. for name in names
  142. ]
  143. def validate_public_rule_result(
  144. value: Any,
  145. *,
  146. backend: str | None = None,
  147. ) -> dict[str, Any]:
  148. allowed_keys = (
  149. _BACKEND_PUBLIC_RESULT_KEYS.get(backend)
  150. if backend is not None
  151. else _PUBLIC_RESULT_KEYS
  152. )
  153. if (
  154. not isinstance(value, dict)
  155. or allowed_keys is None
  156. or set(value) - allowed_keys
  157. or any(str(key).startswith("_") for key in value)
  158. or len(
  159. json.dumps(
  160. value,
  161. sort_keys=True,
  162. separators=(",", ":"),
  163. ensure_ascii=False,
  164. ).encode("utf-8")
  165. )
  166. > 32_768
  167. ):
  168. raise ValueError("public rule result is not evidence safe")
  169. for key, item in value.items():
  170. if key == "violations":
  171. if not isinstance(item, list) or len(item) > 100:
  172. raise ValueError(
  173. "public rule result is not evidence safe"
  174. )
  175. for summary in item:
  176. if (
  177. not isinstance(summary, dict)
  178. or set(summary) != {"step_id", "count"}
  179. or not isinstance(summary["step_id"], str)
  180. or isinstance(summary["count"], bool)
  181. or not isinstance(summary["count"], int)
  182. or summary["count"] < 0
  183. ):
  184. raise ValueError(
  185. "public rule result is not evidence safe"
  186. )
  187. elif isinstance(item, (dict, list, tuple, set)):
  188. raise ValueError("public rule result is not evidence safe")
  189. return dict(value)
  190. class PostgresRuleEvidenceWriter:
  191. """Persist one immutable run and at most one expiring violation sample."""
  192. def __init__(
  193. self,
  194. engine,
  195. artifact_store,
  196. *,
  197. sample_ttl_seconds: int = 3600,
  198. lease_seconds: int = 300,
  199. cleanup_claim_seconds: int = 300,
  200. ):
  201. self.engine = engine
  202. self.artifact_store = artifact_store
  203. self.sample_ttl_seconds = int(sample_ttl_seconds)
  204. self.lease_seconds = int(lease_seconds)
  205. self.cleanup_claim_seconds = int(cleanup_claim_seconds)
  206. if (
  207. self.sample_ttl_seconds < 1
  208. or self.sample_ttl_seconds
  209. > self.artifact_store.max_ttl_seconds
  210. ):
  211. raise ValueError("violation sample TTL is invalid")
  212. if self.lease_seconds < 30 or self.lease_seconds > 900:
  213. raise ValueError("rule execution lease is invalid")
  214. if (
  215. self.cleanup_claim_seconds < 30
  216. or self.cleanup_claim_seconds > 3600
  217. ):
  218. raise ValueError("rule cleanup lease is invalid")
  219. self.heartbeat_interval_seconds = max(
  220. 5.0,
  221. min(60.0, self.lease_seconds / 3),
  222. )
  223. def start(
  224. self,
  225. *,
  226. component_binding_id: str,
  227. rule_version_id: str,
  228. plan_hash: str,
  229. correlation_id: str,
  230. dataflow_uid: str,
  231. deployment_id: str,
  232. environment: str,
  233. workflow_version: int,
  234. node_id: str,
  235. lease_owner: str,
  236. ) -> str:
  237. component = _uid(component_binding_id, "component binding id")
  238. rule = _uid(rule_version_id, "rule version id")
  239. correlation = _uid(correlation_id, "correlation id")
  240. dataflow = _uid(dataflow_uid, "dataflow id")
  241. deployment = _uid(deployment_id, "deployment id")
  242. owner = _uuid(lease_owner, "lease owner")
  243. if environment not in {"development", "test", "production"}:
  244. raise ValueError("deployment environment is invalid")
  245. if _DIGEST.fullmatch(str(plan_hash or "")) is None:
  246. raise ValueError("plan hash is invalid")
  247. if (
  248. isinstance(workflow_version, bool)
  249. or not isinstance(workflow_version, int)
  250. or workflow_version < 1
  251. ):
  252. raise ValueError("workflow version is invalid")
  253. if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_-]{0,99}", str(node_id)):
  254. raise ValueError("node id is invalid")
  255. evidence_key = _canonical_digest(
  256. {
  257. "component_binding_id": component,
  258. "correlation_id": correlation,
  259. "dataflow_uid": dataflow,
  260. "deployment_id": deployment,
  261. "environment": environment,
  262. "node_id": node_id,
  263. "plan_hash": plan_hash,
  264. "rule_version_id": rule,
  265. "workflow_version": workflow_version,
  266. }
  267. )
  268. with self.engine.begin() as connection:
  269. canonical = connection.execute(
  270. text(
  271. """
  272. SELECT d.id::text AS deployment_id
  273. FROM public.dataflow_component_bindings b
  274. JOIN public.dataflow_versions v
  275. ON v.id = b.dataflow_version_id
  276. JOIN public.rule_execution_plans p
  277. ON p.component_binding_id = b.id
  278. JOIN public.dataflow_deployments d
  279. ON d.dataflow_version_id = v.id
  280. WHERE b.id = CAST(:component_binding_id AS uuid)
  281. AND b.rule_version_id =
  282. CAST(:rule_version_id AS uuid)
  283. AND p.plan_hash = :plan_hash
  284. AND p.status = 'published'
  285. AND b.component_id = :node_id
  286. AND v.dataflow_uid = CAST(:dataflow_uid AS uuid)
  287. AND v.version_no = :workflow_version
  288. AND d.id = CAST(:deployment_id AS uuid)
  289. AND d.environment = :environment
  290. AND d.status IN ('canary','active')
  291. """
  292. ),
  293. {
  294. "component_binding_id": component,
  295. "rule_version_id": rule,
  296. "plan_hash": plan_hash,
  297. "dataflow_uid": dataflow,
  298. "deployment_id": deployment,
  299. "environment": environment,
  300. "workflow_version": workflow_version,
  301. "node_id": node_id,
  302. },
  303. ).mappings().one_or_none()
  304. if canonical is None:
  305. raise ValueError("canonical rule deployment does not match")
  306. existing = connection.execute(
  307. text(
  308. """
  309. SELECT id::text, status, lease_owner::text,
  310. lease_expires_at
  311. FROM public.rule_runs
  312. WHERE evidence_key = :evidence_key
  313. FOR UPDATE
  314. """
  315. ),
  316. {"evidence_key": evidence_key},
  317. ).mappings().one_or_none()
  318. if existing is not None:
  319. if (
  320. existing["status"] == "running"
  321. and existing["lease_expires_at"] is not None
  322. ):
  323. expired = connection.execute(
  324. text(
  325. """
  326. SELECT :lease_expires_at <= CURRENT_TIMESTAMP
  327. """
  328. ),
  329. {
  330. "lease_expires_at": existing[
  331. "lease_expires_at"
  332. ]
  333. },
  334. ).scalar_one()
  335. if expired:
  336. connection.execute(
  337. text(
  338. """
  339. UPDATE public.rule_runs
  340. SET status = 'unknown',
  341. commit_outcome = 'unknown',
  342. failure_code = 'execution_lease_expired',
  343. finished_at = CURRENT_TIMESTAMP,
  344. updated_at = CURRENT_TIMESTAMP
  345. WHERE id = CAST(:id AS uuid)
  346. AND status = 'running'
  347. """
  348. ),
  349. {"id": existing["id"]},
  350. )
  351. return str(existing["id"])
  352. rule_run_id = new_governance_uid()
  353. selected = connection.execute(
  354. text(
  355. """
  356. INSERT INTO public.rule_runs (
  357. id, deployment_id, component_binding_id,
  358. rule_version_id, plan_hash, status,
  359. correlation_id, evidence_key, started_at
  360. , attempt_no, lease_owner, lease_expires_at,
  361. heartbeat_at
  362. ) VALUES (
  363. CAST(:id AS uuid), CAST(:deployment_id AS uuid),
  364. CAST(:component_binding_id AS uuid),
  365. CAST(:rule_version_id AS uuid), :plan_hash,
  366. 'running', CAST(:correlation_id AS uuid),
  367. :evidence_key, CURRENT_TIMESTAMP, 1,
  368. CAST(:lease_owner AS uuid),
  369. CURRENT_TIMESTAMP
  370. + make_interval(secs => :lease_seconds),
  371. CURRENT_TIMESTAMP
  372. )
  373. ON CONFLICT (evidence_key) DO NOTHING
  374. RETURNING id::text
  375. """
  376. ),
  377. {
  378. "id": rule_run_id,
  379. "deployment_id": canonical["deployment_id"],
  380. "component_binding_id": component,
  381. "rule_version_id": rule,
  382. "plan_hash": plan_hash,
  383. "correlation_id": correlation,
  384. "evidence_key": evidence_key,
  385. "lease_owner": owner,
  386. "lease_seconds": self.lease_seconds,
  387. },
  388. ).scalar_one_or_none()
  389. if selected is None:
  390. selected = connection.execute(
  391. text(
  392. """
  393. SELECT id::text
  394. FROM public.rule_runs
  395. WHERE evidence_key = :evidence_key
  396. """
  397. ),
  398. {"evidence_key": evidence_key},
  399. ).scalar_one()
  400. return str(selected)
  401. def heartbeat(self, rule_run_id: str, lease_owner: str) -> None:
  402. run_id = _uid(rule_run_id, "rule run id")
  403. owner = _uuid(lease_owner, "lease owner")
  404. with self.engine.begin() as connection:
  405. updated = connection.execute(
  406. text(
  407. """
  408. UPDATE public.rule_runs
  409. SET heartbeat_at = CURRENT_TIMESTAMP,
  410. lease_expires_at = CURRENT_TIMESTAMP
  411. + make_interval(secs => :lease_seconds),
  412. updated_at = CURRENT_TIMESTAMP
  413. WHERE id = CAST(:id AS uuid)
  414. AND status = 'running'
  415. AND lease_owner = CAST(:lease_owner AS uuid)
  416. """
  417. ),
  418. {
  419. "id": run_id,
  420. "lease_owner": owner,
  421. "lease_seconds": self.lease_seconds,
  422. },
  423. )
  424. if updated.rowcount != 1:
  425. raise ValueError("rule execution lease is not owned")
  426. def stage_sql_output(
  427. self,
  428. rule_run_id: str,
  429. *,
  430. output_binding_id: str,
  431. ttl_seconds: int | None = None,
  432. ) -> str:
  433. run_id = _uid(rule_run_id, "rule run id")
  434. binding_id = _uid(output_binding_id, "output binding id")
  435. if ttl_seconds is None:
  436. ttl_seconds = min(3600, self.sample_ttl_seconds)
  437. if (
  438. isinstance(ttl_seconds, bool)
  439. or not isinstance(ttl_seconds, int)
  440. or ttl_seconds < 1
  441. or ttl_seconds > self.sample_ttl_seconds
  442. ):
  443. raise ValueError("SQL staging TTL is invalid")
  444. receipt_id = new_governance_uid()
  445. with self.engine.begin() as connection:
  446. row = connection.execute(
  447. text(
  448. """
  449. SELECT r.deployment_id::text, r.correlation_id::text,
  450. b.binding_hash, b.object_ref, b.object_kind,
  451. b.access_mode
  452. FROM public.rule_runs r
  453. JOIN public.dataflow_dataset_bindings b
  454. ON b.id = CAST(:binding_id AS uuid)
  455. AND b.dataflow_deployment_id = r.deployment_id
  456. WHERE r.id = CAST(:run_id AS uuid)
  457. AND r.status = 'running'
  458. FOR SHARE OF r, b
  459. """
  460. ),
  461. {"run_id": run_id, "binding_id": binding_id},
  462. ).mappings().one_or_none()
  463. if (
  464. row is None
  465. or row["object_kind"] not in {"table", "view"}
  466. or row["access_mode"] not in {"write", "read_write"}
  467. ):
  468. raise ValueError("SQL staging output binding is invalid")
  469. relation_digest = _canonical_digest(
  470. {
  471. "binding_hash": str(row["binding_hash"]),
  472. "object_kind": str(row["object_kind"]),
  473. "object_ref": str(row["object_ref"]),
  474. }
  475. )
  476. selected = connection.execute(
  477. text(
  478. """
  479. INSERT INTO public.rule_sql_staging_receipts (
  480. id, producer_rule_run_id, deployment_id,
  481. correlation_id, output_binding_id,
  482. output_binding_hash, relation_ref,
  483. relation_digest, commit_outcome, status,
  484. expires_at
  485. ) VALUES (
  486. CAST(:id AS uuid), CAST(:run_id AS uuid),
  487. CAST(:deployment_id AS uuid),
  488. CAST(:correlation_id AS uuid),
  489. CAST(:binding_id AS uuid), :binding_hash,
  490. :relation_ref, :relation_digest, 'committed',
  491. 'pending', CURRENT_TIMESTAMP
  492. + make_interval(secs => :ttl_seconds)
  493. )
  494. ON CONFLICT (
  495. producer_rule_run_id, output_binding_id
  496. ) DO NOTHING
  497. RETURNING id::text
  498. """
  499. ),
  500. {
  501. "id": receipt_id,
  502. "run_id": run_id,
  503. "deployment_id": row["deployment_id"],
  504. "correlation_id": row["correlation_id"],
  505. "binding_id": binding_id,
  506. "binding_hash": str(row["binding_hash"]),
  507. "relation_ref": str(row["object_ref"]),
  508. "relation_digest": relation_digest,
  509. "ttl_seconds": ttl_seconds,
  510. },
  511. ).scalar_one_or_none()
  512. if selected is None:
  513. selected = connection.execute(
  514. text(
  515. """
  516. SELECT id::text
  517. FROM public.rule_sql_staging_receipts
  518. WHERE producer_rule_run_id = CAST(:run_id AS uuid)
  519. AND output_binding_id =
  520. CAST(:binding_id AS uuid)
  521. """
  522. ),
  523. {"run_id": run_id, "binding_id": binding_id},
  524. ).scalar_one()
  525. return f"dataops-staging://{selected}"
  526. def resolve_sql_staging(
  527. self,
  528. receipt_ref: str,
  529. *,
  530. deployment_id: str,
  531. correlation_id: str,
  532. input_binding_id: str,
  533. ) -> dict[str, str]:
  534. match = re.fullmatch(
  535. r"dataops-staging://([0-9a-f-]{36})",
  536. str(receipt_ref or ""),
  537. )
  538. if match is None:
  539. raise ValueError("SQL staging receipt is invalid")
  540. receipt_id = _uid(match.group(1), "SQL staging receipt id")
  541. deployment = _uid(deployment_id, "deployment id")
  542. correlation = _uid(correlation_id, "correlation id")
  543. binding_id = _uid(input_binding_id, "input binding id")
  544. with self.engine.connect() as connection:
  545. row = connection.execute(
  546. text(
  547. """
  548. SELECT s.relation_ref, s.relation_digest,
  549. s.output_binding_hash, b.binding_hash,
  550. b.object_ref, b.object_kind, b.access_mode
  551. FROM public.rule_sql_staging_receipts s
  552. JOIN public.rule_runs r
  553. ON r.id = s.producer_rule_run_id
  554. JOIN public.dataflow_dataset_bindings b
  555. ON b.id = s.output_binding_id
  556. WHERE s.id = CAST(:id AS uuid)
  557. AND s.deployment_id =
  558. CAST(:deployment_id AS uuid)
  559. AND s.correlation_id =
  560. CAST(:correlation_id AS uuid)
  561. AND s.output_binding_id =
  562. CAST(:input_binding_id AS uuid)
  563. AND s.status = 'ready'
  564. AND s.expires_at > CURRENT_TIMESTAMP
  565. AND s.commit_outcome = 'committed'
  566. AND r.status = 'success'
  567. AND r.commit_outcome = 'committed'
  568. AND b.binding_hash = s.output_binding_hash
  569. AND b.access_mode IN ('read','read_write')
  570. """
  571. ),
  572. {
  573. "id": receipt_id,
  574. "deployment_id": deployment,
  575. "correlation_id": correlation,
  576. "input_binding_id": binding_id,
  577. },
  578. ).mappings().one_or_none()
  579. if row is None:
  580. raise ValueError("SQL staging receipt is not executable")
  581. expected_digest = _canonical_digest(
  582. {
  583. "binding_hash": str(row["binding_hash"]),
  584. "object_kind": str(row["object_kind"]),
  585. "object_ref": str(row["object_ref"]),
  586. }
  587. )
  588. if (
  589. expected_digest != str(row["relation_digest"])
  590. or str(row["relation_ref"]) != str(row["object_ref"])
  591. ):
  592. raise ValueError("SQL staging receipt attestation does not match")
  593. return {
  594. "relation_ref": str(row["relation_ref"]),
  595. "relation_digest": str(row["relation_digest"]),
  596. }
  597. def replay(self, rule_run_id: str) -> dict[str, Any] | None:
  598. run_id = _uid(rule_run_id, "rule run id")
  599. with self.engine.connect() as connection:
  600. row = connection.execute(
  601. text(
  602. """
  603. SELECT status, commit_outcome, public_result
  604. FROM public.rule_runs
  605. WHERE id = CAST(:id AS uuid)
  606. """
  607. ),
  608. {"id": run_id},
  609. ).mappings().one_or_none()
  610. if row is None:
  611. raise ValueError("rule run was not found")
  612. if row["status"] in {"queued", "running"}:
  613. return None
  614. result = row["public_result"]
  615. if isinstance(result, str):
  616. result = json.loads(result)
  617. return {
  618. **(dict(result) if isinstance(result, dict) else {}),
  619. "status": str(row["status"]),
  620. "commit_outcome": str(row["commit_outcome"]),
  621. }
  622. def replay_by_lease_owner(
  623. self,
  624. *,
  625. lease_owner: str,
  626. deployment_id: str,
  627. correlation_id: str,
  628. component_binding_id: str,
  629. rule_version_id: str,
  630. plan_hash: str,
  631. ) -> dict[str, Any] | None:
  632. owner = _uuid(lease_owner, "lease owner")
  633. deployment = _uid(deployment_id, "deployment id")
  634. correlation = _uid(correlation_id, "correlation id")
  635. component = _uid(
  636. component_binding_id,
  637. "component binding id",
  638. )
  639. rule = _uid(rule_version_id, "rule version id")
  640. if _DIGEST.fullmatch(str(plan_hash or "")) is None:
  641. raise ValueError("plan hash is invalid")
  642. with self.engine.connect() as connection:
  643. row = connection.execute(
  644. text(
  645. """
  646. SELECT id::text
  647. FROM public.rule_runs
  648. WHERE lease_owner = CAST(:lease_owner AS uuid)
  649. AND deployment_id =
  650. CAST(:deployment_id AS uuid)
  651. AND correlation_id =
  652. CAST(:correlation_id AS uuid)
  653. AND component_binding_id =
  654. CAST(:component_binding_id AS uuid)
  655. AND rule_version_id =
  656. CAST(:rule_version_id AS uuid)
  657. AND plan_hash = :plan_hash
  658. """
  659. ),
  660. {
  661. "lease_owner": owner,
  662. "deployment_id": deployment,
  663. "correlation_id": correlation,
  664. "component_binding_id": component,
  665. "rule_version_id": rule,
  666. "plan_hash": plan_hash,
  667. },
  668. ).scalar_one_or_none()
  669. if row is None:
  670. return None
  671. return self.replay(str(row))
  672. def reconcile_expired_lease(
  673. self,
  674. *,
  675. lease_owner: str,
  676. deployment_id: str,
  677. correlation_id: str,
  678. component_binding_id: str,
  679. rule_version_id: str,
  680. plan_hash: str,
  681. ) -> dict[str, Any]:
  682. owner = _uuid(lease_owner, "lease owner")
  683. deployment = _uid(deployment_id, "deployment id")
  684. correlation = _uid(correlation_id, "correlation id")
  685. component = _uid(
  686. component_binding_id,
  687. "component binding id",
  688. )
  689. rule = _uid(rule_version_id, "rule version id")
  690. if _DIGEST.fullmatch(str(plan_hash or "")) is None:
  691. raise ValueError("plan hash is invalid")
  692. parameters = {
  693. "lease_owner": owner,
  694. "deployment_id": deployment,
  695. "correlation_id": correlation,
  696. "component_binding_id": component,
  697. "rule_version_id": rule,
  698. "plan_hash": plan_hash,
  699. }
  700. with self.engine.begin() as connection:
  701. row = connection.execute(
  702. text(
  703. """
  704. SELECT id::text, status, lease_expires_at,
  705. evidence_digest
  706. FROM public.rule_runs
  707. WHERE lease_owner = CAST(:lease_owner AS uuid)
  708. AND deployment_id =
  709. CAST(:deployment_id AS uuid)
  710. AND correlation_id =
  711. CAST(:correlation_id AS uuid)
  712. AND component_binding_id =
  713. CAST(:component_binding_id AS uuid)
  714. AND rule_version_id =
  715. CAST(:rule_version_id AS uuid)
  716. AND plan_hash = :plan_hash
  717. FOR UPDATE
  718. """
  719. ),
  720. parameters,
  721. ).mappings().one_or_none()
  722. if row is None:
  723. return {"state": "missing"}
  724. if (
  725. row["status"] == "running"
  726. and row["lease_expires_at"] is not None
  727. ):
  728. expired = connection.execute(
  729. text(
  730. "SELECT :lease_expires_at <= CURRENT_TIMESTAMP"
  731. ),
  732. {"lease_expires_at": row["lease_expires_at"]},
  733. ).scalar_one()
  734. if expired:
  735. connection.execute(
  736. text(
  737. """
  738. UPDATE public.rule_runs
  739. SET status = 'unknown',
  740. commit_outcome = 'unknown',
  741. failure_code =
  742. 'execution_lease_expired',
  743. lease_expires_at = NULL,
  744. finished_at = CURRENT_TIMESTAMP,
  745. updated_at = CURRENT_TIMESTAMP
  746. WHERE id = CAST(:id AS uuid)
  747. AND status = 'running'
  748. """
  749. ),
  750. {"id": row["id"]},
  751. )
  752. connection.execute(
  753. text(
  754. """
  755. UPDATE public.rule_sql_staging_receipts
  756. SET status = 'failed',
  757. commit_outcome = 'unknown',
  758. updated_at = CURRENT_TIMESTAMP
  759. WHERE producer_rule_run_id =
  760. CAST(:id AS uuid)
  761. AND status = 'pending'
  762. """
  763. ),
  764. {"id": row["id"]},
  765. )
  766. row = {**row, "status": "unknown"}
  767. if row["status"] == "running":
  768. return {"state": "running"}
  769. run_id = str(row["id"])
  770. replay = self.replay(run_id)
  771. result = {
  772. key: value
  773. for key, value in replay.items()
  774. if key != "status"
  775. }
  776. return {
  777. "state": "terminal",
  778. "status": str(replay["status"]),
  779. "commit_outcome": str(replay["commit_outcome"]),
  780. "result": result,
  781. "result_digest": _canonical_digest(result),
  782. "evidence_digest": str(row["evidence_digest"] or ""),
  783. }
  784. @staticmethod
  785. def _validate_finish(result: Any) -> dict[str, Any]:
  786. if not isinstance(result, dict) or set(result) - _FINISH_KEYS:
  787. raise ValueError("rule evidence result has unsupported fields")
  788. status = result.get("status")
  789. commit_outcome = result.get("commit_outcome", "not_applicable")
  790. if status not in _FINAL_STATUSES:
  791. raise ValueError("rule evidence status is invalid")
  792. if commit_outcome not in _COMMIT_OUTCOMES:
  793. raise ValueError("rule evidence commit outcome is invalid")
  794. timings = result.get("timings", {})
  795. if (
  796. not isinstance(timings, dict)
  797. or set(timings) != {"duration_ms"}
  798. or isinstance(timings.get("duration_ms"), bool)
  799. or not isinstance(timings.get("duration_ms"), int)
  800. or timings["duration_ms"] < 0
  801. or timings["duration_ms"] > 86_400_000
  802. ):
  803. raise ValueError("rule evidence timings are invalid")
  804. public_result = result.get("public_result")
  805. if public_result is not None:
  806. validate_public_rule_result(public_result)
  807. sample = result.get("violation_sample")
  808. if sample is not None:
  809. if (
  810. status != "success"
  811. or not isinstance(sample, list)
  812. or not 1 <= len(sample) <= 100
  813. or result.get("sample_count") != len(sample)
  814. or result.get("redaction_policy")
  815. != "rule-violation-default-v1"
  816. ):
  817. raise ValueError("violation sample is invalid")
  818. for row in sample:
  819. if not isinstance(row, dict):
  820. raise ValueError("violation sample row is invalid")
  821. for value in row.values():
  822. if value not in {None, "[REDACTED]"}:
  823. raise ValueError(
  824. "violation sample contains unredacted values"
  825. )
  826. return {
  827. **result,
  828. "rows_in": _bounded_count(result.get("rows_in"), "rows_in"),
  829. "rows_out": _bounded_count(result.get("rows_out"), "rows_out"),
  830. "rows_rejected": _bounded_count(
  831. result.get("rows_rejected"), "rows_rejected"
  832. ),
  833. "rows_quarantined": _bounded_count(
  834. result.get("rows_quarantined"), "rows_quarantined"
  835. ),
  836. "commit_outcome": commit_outcome,
  837. }
  838. def _prepare_sample(
  839. self,
  840. run_id: str,
  841. correlation_id: str,
  842. sample: list[dict[str, Any]],
  843. redaction_policy: str,
  844. ) -> tuple[str, dict[str, Any], str]:
  845. fields = _sample_fields(sample)
  846. frame = pl.DataFrame(
  847. {
  848. field["name"]: [
  849. row.get(field["name"]) for row in sample
  850. ]
  851. for field in fields
  852. },
  853. schema={field["name"]: pl.String for field in fields},
  854. )
  855. with tempfile.NamedTemporaryFile(
  856. prefix="dataops-rule-violation-",
  857. suffix=".parquet",
  858. delete=False,
  859. ) as handle:
  860. path = handle.name
  861. sample_id = None
  862. prepared = None
  863. uploaded = False
  864. try:
  865. frame.write_parquet(path)
  866. prepared = self.artifact_store.prepare_path(
  867. path,
  868. correlation_id,
  869. self.sample_ttl_seconds,
  870. schema_fields=fields,
  871. limits={
  872. "max_rows": min(100, self.artifact_store.max_rows),
  873. "max_artifact_bytes": min(
  874. 4 * 1024 * 1024,
  875. self.artifact_store.max_artifact_bytes,
  876. ),
  877. "memory_limit_bytes": min(
  878. 16 * 1024 * 1024,
  879. self.artifact_store.memory_limit_bytes,
  880. ),
  881. },
  882. )
  883. sample_id = new_governance_uid()
  884. with self.engine.begin() as connection:
  885. row = connection.execute(
  886. text(
  887. """
  888. SELECT id::text, artifact_ref, artifact_digest,
  889. schema_hash, sample_count,
  890. redaction_policy, expires_at, handoff_status
  891. FROM public.rule_violation_samples
  892. WHERE rule_run_id = CAST(:rule_run_id AS uuid)
  893. FOR UPDATE
  894. """
  895. ),
  896. {"rule_run_id": run_id},
  897. ).mappings().one_or_none()
  898. if row is None:
  899. connection.execute(
  900. text(
  901. """
  902. INSERT INTO public.rule_violation_samples (
  903. id, rule_run_id, artifact_ref,
  904. artifact_digest, schema_hash, sample_count,
  905. redaction_policy, expires_at,
  906. handoff_status, schema_fields
  907. ) VALUES (
  908. CAST(:id AS uuid),
  909. CAST(:rule_run_id AS uuid), :artifact_ref,
  910. :artifact_digest, :schema_hash,
  911. :sample_count, :redaction_policy,
  912. CAST(:expires_at AS timestamptz), 'pending',
  913. CAST(:schema_fields AS jsonb)
  914. )
  915. """
  916. ),
  917. {
  918. "id": sample_id,
  919. "rule_run_id": run_id,
  920. "artifact_ref": prepared["artifact_ref"],
  921. "artifact_digest": prepared["digest"],
  922. "schema_hash": prepared["schema_hash"],
  923. "sample_count": len(sample),
  924. "redaction_policy": redaction_policy,
  925. "expires_at": prepared["expires_at"],
  926. "schema_fields": json.dumps(fields),
  927. },
  928. )
  929. else:
  930. if (
  931. str(row["artifact_digest"]) != prepared["digest"]
  932. or int(row["sample_count"]) != len(sample)
  933. or str(row["redaction_policy"])
  934. != redaction_policy
  935. ):
  936. raise ValueError(
  937. "violation sample evidence is immutable"
  938. )
  939. sample_id = str(row["id"])
  940. prepared.update(
  941. {
  942. "artifact_ref": str(row["artifact_ref"]),
  943. "digest": str(row["artifact_digest"]),
  944. "schema_hash": str(row["schema_hash"]),
  945. "expires_at": str(row["expires_at"]),
  946. }
  947. )
  948. if row["handoff_status"] == "ready":
  949. stored = self.artifact_store.describe(
  950. prepared["artifact_ref"]
  951. )
  952. if (
  953. stored["digest"] != prepared["digest"]
  954. or stored["schema_hash"]
  955. != prepared["schema_hash"]
  956. or stored["row_count"] != len(sample)
  957. ):
  958. raise ValueError(
  959. "ready violation sample does not match storage"
  960. )
  961. return sample_id, prepared, path
  962. if row["handoff_status"] == "failed":
  963. raise ValueError(
  964. "violation sample handoff already failed"
  965. )
  966. self.artifact_store.upload_path(
  967. path,
  968. prepared,
  969. limits={
  970. "max_rows": min(100, self.artifact_store.max_rows),
  971. "max_artifact_bytes": min(
  972. 4 * 1024 * 1024,
  973. self.artifact_store.max_artifact_bytes,
  974. ),
  975. "memory_limit_bytes": min(
  976. 16 * 1024 * 1024,
  977. self.artifact_store.memory_limit_bytes,
  978. ),
  979. },
  980. )
  981. uploaded = True
  982. with self.engine.begin() as connection:
  983. updated = connection.execute(
  984. text(
  985. """
  986. UPDATE public.rule_violation_samples
  987. SET handoff_status = 'ready',
  988. updated_at = CURRENT_TIMESTAMP
  989. WHERE id = CAST(:id AS uuid)
  990. AND handoff_status = 'pending'
  991. AND artifact_digest = :artifact_digest
  992. """
  993. ),
  994. {
  995. "id": sample_id,
  996. "artifact_digest": prepared["digest"],
  997. },
  998. )
  999. if updated.rowcount != 1:
  1000. state = connection.execute(
  1001. text(
  1002. """
  1003. SELECT handoff_status
  1004. FROM public.rule_violation_samples
  1005. WHERE id = CAST(:id AS uuid)
  1006. """
  1007. ),
  1008. {"id": sample_id},
  1009. ).scalar_one_or_none()
  1010. if state != "ready":
  1011. raise RuntimeError(
  1012. "violation sample finalize outcome is unknown"
  1013. )
  1014. return sample_id, prepared, path
  1015. except Exception:
  1016. if sample_id is not None and prepared is not None:
  1017. try:
  1018. with self.engine.connect() as connection:
  1019. committed = connection.execute(
  1020. text(
  1021. """
  1022. SELECT handoff_status, artifact_digest
  1023. FROM public.rule_violation_samples
  1024. WHERE id = CAST(:id AS uuid)
  1025. """
  1026. ),
  1027. {"id": sample_id},
  1028. ).mappings().one_or_none()
  1029. if (
  1030. committed is not None
  1031. and committed["handoff_status"] == "ready"
  1032. and str(committed["artifact_digest"])
  1033. == prepared["digest"]
  1034. ):
  1035. return sample_id, prepared, path
  1036. except Exception:
  1037. pass
  1038. with self.engine.begin() as connection:
  1039. connection.execute(
  1040. text(
  1041. """
  1042. UPDATE public.rule_violation_samples
  1043. SET handoff_status = :handoff_status,
  1044. failure_code = :failure_code,
  1045. updated_at = CURRENT_TIMESTAMP
  1046. WHERE rule_run_id = CAST(:rule_run_id AS uuid)
  1047. AND handoff_status = 'pending'
  1048. """
  1049. ),
  1050. {
  1051. "rule_run_id": run_id,
  1052. "handoff_status": (
  1053. "unknown" if uploaded else "failed"
  1054. ),
  1055. "failure_code": (
  1056. "sample_finalize_unknown"
  1057. if uploaded
  1058. else "sample_handoff_failed"
  1059. ),
  1060. },
  1061. )
  1062. raise
  1063. finally:
  1064. with suppress(FileNotFoundError):
  1065. os.unlink(path)
  1066. def finish(self, rule_run_id: str, result: Any) -> None:
  1067. run_id = _uid(rule_run_id, "rule run id")
  1068. normalized = self._validate_finish(result)
  1069. evidence_digest = _canonical_digest(normalized)
  1070. sample_path = None
  1071. try:
  1072. with self.engine.connect() as connection:
  1073. current = connection.execute(
  1074. text(
  1075. """
  1076. SELECT status, correlation_id::text,
  1077. evidence_digest
  1078. FROM public.rule_runs
  1079. WHERE id = CAST(:id AS uuid)
  1080. """
  1081. ),
  1082. {"id": run_id},
  1083. ).mappings().one_or_none()
  1084. if current is None:
  1085. raise ValueError("rule run was not found")
  1086. if current["status"] not in {"queued", "running"}:
  1087. if (
  1088. current["status"] == normalized["status"]
  1089. and str(current["evidence_digest"] or "")
  1090. == evidence_digest
  1091. ):
  1092. return
  1093. raise ValueError("rule run evidence is immutable")
  1094. if normalized.get("violation_sample"):
  1095. _sample_id, _prepared, sample_path = self._prepare_sample(
  1096. run_id,
  1097. str(current["correlation_id"]),
  1098. normalized["violation_sample"],
  1099. normalized["redaction_policy"],
  1100. )
  1101. try:
  1102. with self.engine.begin() as connection:
  1103. updated = connection.execute(
  1104. text(
  1105. """
  1106. UPDATE public.rule_runs
  1107. SET rows_in = :rows_in,
  1108. rows_out = :rows_out,
  1109. rows_rejected = :rows_rejected,
  1110. rows_quarantined = :rows_quarantined,
  1111. status = :status,
  1112. timings = CAST(:timings AS jsonb),
  1113. commit_outcome = :commit_outcome,
  1114. public_result = CAST(:public_result AS jsonb),
  1115. failure_code = :failure_code,
  1116. evidence_digest = :evidence_digest,
  1117. lease_expires_at = NULL,
  1118. finished_at = CURRENT_TIMESTAMP,
  1119. updated_at = CURRENT_TIMESTAMP
  1120. WHERE id = CAST(:id AS uuid)
  1121. AND status IN ('queued','running')
  1122. """
  1123. ),
  1124. {
  1125. "id": run_id,
  1126. "rows_in": normalized["rows_in"],
  1127. "rows_out": normalized["rows_out"],
  1128. "rows_rejected": normalized[
  1129. "rows_rejected"
  1130. ],
  1131. "rows_quarantined": normalized[
  1132. "rows_quarantined"
  1133. ],
  1134. "status": normalized["status"],
  1135. "timings": json.dumps(
  1136. normalized["timings"]
  1137. ),
  1138. "commit_outcome": normalized[
  1139. "commit_outcome"
  1140. ],
  1141. "public_result": (
  1142. json.dumps(
  1143. normalized.get("public_result")
  1144. )
  1145. if normalized.get("public_result")
  1146. is not None
  1147. else None
  1148. ),
  1149. "failure_code": (
  1150. None
  1151. if normalized["status"] == "success"
  1152. else (
  1153. f"execution_{normalized['status']}"
  1154. )
  1155. ),
  1156. "evidence_digest": evidence_digest,
  1157. },
  1158. )
  1159. if updated.rowcount != 1:
  1160. state = connection.execute(
  1161. text(
  1162. """
  1163. SELECT status, commit_outcome,
  1164. evidence_digest
  1165. FROM public.rule_runs
  1166. WHERE id = CAST(:id AS uuid)
  1167. """
  1168. ),
  1169. {"id": run_id},
  1170. ).mappings().one_or_none()
  1171. if (
  1172. state is None
  1173. or state["status"] != normalized["status"]
  1174. or state["commit_outcome"]
  1175. != normalized["commit_outcome"]
  1176. or str(state["evidence_digest"] or "")
  1177. != evidence_digest
  1178. ):
  1179. raise RuntimeError(
  1180. "rule run finalize outcome is unknown"
  1181. )
  1182. receipt_status = (
  1183. "ready"
  1184. if normalized["status"] == "success"
  1185. and normalized["commit_outcome"] == "committed"
  1186. else "failed"
  1187. )
  1188. connection.execute(
  1189. text(
  1190. """
  1191. UPDATE public.rule_sql_staging_receipts
  1192. SET status = :status,
  1193. ready_at = CASE
  1194. WHEN :status = 'ready'
  1195. THEN CURRENT_TIMESTAMP
  1196. ELSE ready_at
  1197. END,
  1198. commit_outcome = CASE
  1199. WHEN :status = 'ready'
  1200. THEN 'committed'
  1201. ELSE 'unknown'
  1202. END,
  1203. updated_at = CURRENT_TIMESTAMP
  1204. WHERE producer_rule_run_id =
  1205. CAST(:run_id AS uuid)
  1206. AND status = 'pending'
  1207. """
  1208. ),
  1209. {
  1210. "run_id": run_id,
  1211. "status": receipt_status,
  1212. },
  1213. )
  1214. except Exception as exc:
  1215. try:
  1216. with self.engine.connect() as connection:
  1217. terminal = connection.execute(
  1218. text(
  1219. """
  1220. SELECT status, commit_outcome,
  1221. evidence_digest
  1222. FROM public.rule_runs
  1223. WHERE id = CAST(:id AS uuid)
  1224. """
  1225. ),
  1226. {"id": run_id},
  1227. ).mappings().one_or_none()
  1228. except Exception as recheck_exc:
  1229. raise RuntimeError(
  1230. "rule run finalize outcome is unknown"
  1231. ) from recheck_exc
  1232. if (
  1233. terminal is None
  1234. or terminal["status"] != normalized["status"]
  1235. or terminal["commit_outcome"]
  1236. != normalized["commit_outcome"]
  1237. or str(terminal["evidence_digest"] or "")
  1238. != evidence_digest
  1239. ):
  1240. raise RuntimeError(
  1241. "rule run finalize outcome is unknown"
  1242. ) from exc
  1243. finally:
  1244. if sample_path is not None:
  1245. with suppress(FileNotFoundError):
  1246. os.unlink(sample_path)
  1247. def reconcile_samples(self, *, limit: int = 100) -> dict[str, int]:
  1248. if isinstance(limit, bool) or not isinstance(limit, int):
  1249. raise ValueError("cleanup limit is invalid")
  1250. if limit < 1 or limit > 1000:
  1251. raise ValueError("cleanup limit is invalid")
  1252. claim = new_governance_uid()
  1253. with self.engine.begin() as connection:
  1254. candidates = connection.execute(
  1255. text(
  1256. """
  1257. SELECT id::text, artifact_ref, artifact_digest,
  1258. schema_hash, schema_fields, sample_count,
  1259. handoff_status,
  1260. expires_at <= CURRENT_TIMESTAMP AS expired
  1261. FROM public.rule_violation_samples
  1262. WHERE (
  1263. cleanup_claim IS NULL
  1264. OR cleanup_claim_expires_at <= CURRENT_TIMESTAMP
  1265. )
  1266. AND (
  1267. handoff_status IN ('pending','unknown')
  1268. OR (
  1269. expires_at <= CURRENT_TIMESTAMP
  1270. AND handoff_status IN (
  1271. 'legacy','ready','failed'
  1272. )
  1273. )
  1274. )
  1275. ORDER BY expires_at, id
  1276. FOR UPDATE SKIP LOCKED
  1277. LIMIT :limit
  1278. """
  1279. ),
  1280. {"limit": limit},
  1281. ).mappings().all()
  1282. rows = []
  1283. for row in candidates:
  1284. updated = connection.execute(
  1285. text(
  1286. """
  1287. UPDATE public.rule_violation_samples
  1288. SET cleanup_claim = CAST(:claim AS uuid),
  1289. cleanup_claim_expires_at =
  1290. CURRENT_TIMESTAMP
  1291. + make_interval(
  1292. secs => :claim_seconds
  1293. ),
  1294. updated_at = CURRENT_TIMESTAMP
  1295. WHERE id = CAST(:id AS uuid)
  1296. AND (
  1297. cleanup_claim IS NULL
  1298. OR cleanup_claim_expires_at
  1299. <= CURRENT_TIMESTAMP
  1300. )
  1301. """
  1302. ),
  1303. {
  1304. "id": str(row["id"]),
  1305. "claim": claim,
  1306. "claim_seconds": self.cleanup_claim_seconds,
  1307. },
  1308. )
  1309. if int(updated.rowcount or 0) == 1:
  1310. rows.append(row)
  1311. metrics = {
  1312. "claimed": len(rows),
  1313. "ready": 0,
  1314. "failed": 0,
  1315. "expired_deleted": 0,
  1316. }
  1317. for row in rows:
  1318. row_id = str(row["id"])
  1319. if bool(row["expired"]):
  1320. try:
  1321. self.artifact_store.delete(str(row["artifact_ref"]))
  1322. except Exception:
  1323. with self.engine.begin() as connection:
  1324. connection.execute(
  1325. text(
  1326. """
  1327. UPDATE public.rule_violation_samples
  1328. SET cleanup_claim = NULL,
  1329. cleanup_claim_expires_at = NULL,
  1330. updated_at = CURRENT_TIMESTAMP
  1331. WHERE id = CAST(:id AS uuid)
  1332. AND cleanup_claim =
  1333. CAST(:claim AS uuid)
  1334. """
  1335. ),
  1336. {"id": row_id, "claim": claim},
  1337. )
  1338. continue
  1339. with self.engine.begin() as connection:
  1340. deleted = connection.execute(
  1341. text(
  1342. """
  1343. DELETE FROM public.rule_violation_samples
  1344. WHERE id = CAST(:id AS uuid)
  1345. AND cleanup_claim = CAST(:claim AS uuid)
  1346. """
  1347. ),
  1348. {"id": row_id, "claim": claim},
  1349. )
  1350. metrics["expired_deleted"] += int(
  1351. deleted.rowcount or 0
  1352. )
  1353. continue
  1354. next_status = row["handoff_status"]
  1355. failure_code = None
  1356. try:
  1357. fields = row["schema_fields"]
  1358. if isinstance(fields, str):
  1359. fields = json.loads(fields)
  1360. if not isinstance(fields, list):
  1361. raise ValueError(
  1362. "violation sample schema is unavailable"
  1363. )
  1364. described = self.artifact_store.describe_optional(
  1365. str(row["artifact_ref"])
  1366. )
  1367. if described is None:
  1368. raise ValueError(
  1369. "violation sample object is missing"
  1370. )
  1371. if (
  1372. described["digest"]
  1373. != str(row["artifact_digest"])
  1374. or described["schema_hash"] != row["schema_hash"]
  1375. or described["row_count"]
  1376. != int(row["sample_count"])
  1377. ):
  1378. raise ValueError(
  1379. "violation sample attestation does not match"
  1380. )
  1381. with self.artifact_store.stage(
  1382. str(row["artifact_ref"]),
  1383. str(row["artifact_digest"]),
  1384. expected_schema_fields=fields,
  1385. limits={
  1386. "max_rows": min(
  1387. 100,
  1388. self.artifact_store.max_rows,
  1389. ),
  1390. "max_artifact_bytes": min(
  1391. 4 * 1024 * 1024,
  1392. self.artifact_store.max_artifact_bytes,
  1393. ),
  1394. "memory_limit_bytes": min(
  1395. 16 * 1024 * 1024,
  1396. self.artifact_store.memory_limit_bytes,
  1397. ),
  1398. },
  1399. ):
  1400. pass
  1401. next_status = "ready"
  1402. metrics["ready"] += 1
  1403. except ValueError:
  1404. next_status = "failed"
  1405. failure_code = "sample_reconcile_invalid"
  1406. metrics["failed"] += 1
  1407. except Exception:
  1408. next_status = row["handoff_status"]
  1409. with self.engine.begin() as connection:
  1410. connection.execute(
  1411. text(
  1412. """
  1413. UPDATE public.rule_violation_samples
  1414. SET handoff_status = :handoff_status,
  1415. failure_code = :failure_code,
  1416. cleanup_claim = NULL,
  1417. cleanup_claim_expires_at = NULL,
  1418. updated_at = CURRENT_TIMESTAMP
  1419. WHERE id = CAST(:id AS uuid)
  1420. AND cleanup_claim = CAST(:claim AS uuid)
  1421. """
  1422. ),
  1423. {
  1424. "id": row_id,
  1425. "claim": claim,
  1426. "handoff_status": next_status,
  1427. "failure_code": failure_code,
  1428. },
  1429. )
  1430. return metrics
  1431. def cleanup_sql_staging(self, *, limit: int = 100) -> int:
  1432. if (
  1433. isinstance(limit, bool)
  1434. or not isinstance(limit, int)
  1435. or limit < 1
  1436. or limit > 1000
  1437. ):
  1438. raise ValueError("cleanup limit is invalid")
  1439. claim = new_governance_uid()
  1440. with self.engine.begin() as connection:
  1441. candidates = connection.execute(
  1442. text(
  1443. """
  1444. SELECT id::text
  1445. FROM public.rule_sql_staging_receipts
  1446. WHERE expires_at <= CURRENT_TIMESTAMP
  1447. AND status IN ('pending','ready','failed')
  1448. AND (
  1449. cleanup_claim IS NULL
  1450. OR cleanup_claim_expires_at
  1451. <= CURRENT_TIMESTAMP
  1452. )
  1453. ORDER BY expires_at, id
  1454. FOR UPDATE SKIP LOCKED
  1455. LIMIT :limit
  1456. """
  1457. ),
  1458. {"limit": limit},
  1459. ).mappings().all()
  1460. rows = []
  1461. for row in candidates:
  1462. updated = connection.execute(
  1463. text(
  1464. """
  1465. UPDATE public.rule_sql_staging_receipts
  1466. SET cleanup_claim = CAST(:claim AS uuid),
  1467. cleanup_claim_expires_at =
  1468. CURRENT_TIMESTAMP
  1469. + make_interval(
  1470. secs => :claim_seconds
  1471. ),
  1472. updated_at = CURRENT_TIMESTAMP
  1473. WHERE id = CAST(:id AS uuid)
  1474. AND (
  1475. cleanup_claim IS NULL
  1476. OR cleanup_claim_expires_at
  1477. <= CURRENT_TIMESTAMP
  1478. )
  1479. """
  1480. ),
  1481. {
  1482. "id": str(row["id"]),
  1483. "claim": claim,
  1484. "claim_seconds": self.cleanup_claim_seconds,
  1485. },
  1486. )
  1487. if int(updated.rowcount or 0) == 1:
  1488. rows.append(row)
  1489. finalized = 0
  1490. for row in rows:
  1491. with self.engine.begin() as connection:
  1492. updated = connection.execute(
  1493. text(
  1494. """
  1495. UPDATE public.rule_sql_staging_receipts
  1496. SET status = 'expired',
  1497. cleanup_claim = NULL,
  1498. cleanup_claim_expires_at = NULL,
  1499. updated_at = CURRENT_TIMESTAMP
  1500. WHERE id = CAST(:id AS uuid)
  1501. AND cleanup_claim = CAST(:claim AS uuid)
  1502. AND expires_at <= CURRENT_TIMESTAMP
  1503. """
  1504. ),
  1505. {"id": str(row["id"]), "claim": claim},
  1506. )
  1507. finalized += int(updated.rowcount or 0)
  1508. return finalized
  1509. def cleanup_expired(self, *, limit: int = 100) -> int:
  1510. samples = self.reconcile_samples(limit=limit)
  1511. receipts = self.cleanup_sql_staging(limit=limit)
  1512. return samples["expired_deleted"] + receipts
  1513. __all__ = [
  1514. "PostgresRuleEvidenceWriter",
  1515. "validate_public_rule_result",
  1516. ]