publication.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845
  1. """Trusted generation receipts and rule publication lifecycle gates."""
  2. from __future__ import annotations
  3. import base64
  4. import copy
  5. import hashlib
  6. import hmac
  7. import json
  8. import os
  9. import re
  10. import tempfile
  11. from contextlib import suppress
  12. from datetime import UTC, datetime
  13. from typing import Any
  14. from app.core.common.identifiers import ensure_governance_uid, new_governance_uid
  15. from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec
  16. _DIGEST = re.compile(r"^[0-9a-f]{64}$")
  17. _RECEIPT_KEYS = {
  18. "version",
  19. "generation_run_id",
  20. "actor_uid",
  21. "source_text_hash",
  22. "candidate_hash",
  23. "rule_spec_hash",
  24. "model_hash",
  25. "prompt_hash",
  26. "context_hash",
  27. "expires_at",
  28. }
  29. _COMPILED_KEYS = {
  30. "backend",
  31. "compiler_version",
  32. "plan",
  33. "plan_hash",
  34. "schema_hashes",
  35. "binding_hashes",
  36. "capabilities",
  37. }
  38. _TEST_KEYS = {
  39. "status",
  40. "test_kind",
  41. "run_id",
  42. "plan_hash",
  43. "schema_hashes",
  44. "binding_hashes",
  45. "counts",
  46. "attestation",
  47. }
  48. class RuleValidationRejected(ValueError):
  49. """A deterministic validation failure whose audit row must be committed."""
  50. def _canonical(value: Any) -> bytes:
  51. try:
  52. encoded = json.dumps(
  53. value,
  54. sort_keys=True,
  55. separators=(",", ":"),
  56. ensure_ascii=False,
  57. )
  58. except (TypeError, ValueError) as exc:
  59. raise ValueError("publication value must be JSON serializable") from exc
  60. return encoded.encode("utf-8")
  61. def _digest(value: Any, label: str) -> str:
  62. normalized = str(value or "")
  63. if not _DIGEST.fullmatch(normalized):
  64. raise ValueError(f"{label} must be a sha256 digest")
  65. return normalized
  66. def _uid(value: Any, label: str) -> str:
  67. try:
  68. return ensure_governance_uid({"uid": str(value)})
  69. except ValueError as exc:
  70. raise ValueError(f"{label} must be a valid UUIDv7") from exc
  71. def _hash_text(value: Any, label: str, maximum: int = 20_000) -> str:
  72. if not isinstance(value, str) or not value.strip():
  73. raise ValueError(f"{label} is required")
  74. normalized = value.strip()
  75. if len(normalized) > maximum:
  76. raise ValueError(f"{label} exceeds {maximum} characters")
  77. return hashlib.sha256(normalized.encode("utf-8")).hexdigest()
  78. def _b64encode(value: bytes) -> str:
  79. return base64.urlsafe_b64encode(value).decode("ascii").rstrip("=")
  80. def _b64decode(value: str) -> bytes:
  81. if (
  82. not isinstance(value, str)
  83. or not value
  84. or re.fullmatch(r"[A-Za-z0-9_-]+", value) is None
  85. ):
  86. raise ValueError("generation receipt encoding is invalid")
  87. try:
  88. decoded = base64.urlsafe_b64decode(
  89. value + ("=" * (-len(value) % 4))
  90. )
  91. except Exception as exc:
  92. raise ValueError("generation receipt encoding is invalid") from exc
  93. if _b64encode(decoded) != value:
  94. raise ValueError("generation receipt encoding is not canonical")
  95. return decoded
  96. def generation_receipt_claims(
  97. *,
  98. generation_run_id: str,
  99. actor_uid: str,
  100. source_text: str,
  101. candidate_hash: str,
  102. rule_spec: dict[str, Any],
  103. model_hash: str,
  104. prompt_hash: str,
  105. context_hash: str,
  106. expires_at: datetime,
  107. ) -> dict[str, Any]:
  108. """Build the closed, non-secret claim set signed by ``/interpret``."""
  109. if not isinstance(expires_at, datetime) or expires_at.tzinfo is None:
  110. raise ValueError("generation receipt expiry must be timezone aware")
  111. spec = validate_rule_spec(rule_spec)
  112. return {
  113. "version": "1",
  114. "generation_run_id": _uid(
  115. generation_run_id, "generation_run_id"
  116. ),
  117. "actor_uid": _uid(actor_uid, "actor_uid"),
  118. "source_text_hash": _hash_text(source_text, "source_text"),
  119. "candidate_hash": _digest(candidate_hash, "candidate_hash"),
  120. "rule_spec_hash": rule_spec_hash(spec),
  121. "model_hash": _digest(model_hash, "model_hash"),
  122. "prompt_hash": _digest(prompt_hash, "prompt_hash"),
  123. "context_hash": _digest(context_hash, "context_hash"),
  124. "expires_at": int(expires_at.timestamp()),
  125. }
  126. class GenerationReceiptSigner:
  127. """Issue and verify compact HMAC receipts without exposing the secret."""
  128. def __init__(self, secret: str | bytes):
  129. secret_bytes = (
  130. secret.encode("utf-8") if isinstance(secret, str) else bytes(secret)
  131. )
  132. if len(secret_bytes) < 16:
  133. raise ValueError("generation receipt secret is too short")
  134. self._secret = secret_bytes
  135. def issue(self, claims: dict[str, Any]) -> str:
  136. normalized = self._validate_claims(claims)
  137. payload = _b64encode(_canonical(normalized))
  138. signature = hmac.new(
  139. self._secret, payload.encode("ascii"), hashlib.sha256
  140. ).digest()
  141. return f"{payload}.{_b64encode(signature)}"
  142. def verify(
  143. self,
  144. receipt: str,
  145. *,
  146. actor_uid: str,
  147. source_text: str,
  148. rule_spec: dict[str, Any],
  149. now: datetime | None = None,
  150. ) -> dict[str, Any]:
  151. if not isinstance(receipt, str) or len(receipt) > 4096:
  152. raise ValueError("generation receipt is invalid")
  153. try:
  154. payload, encoded_signature = receipt.split(".", 1)
  155. except ValueError as exc:
  156. raise ValueError("generation receipt is invalid") from exc
  157. expected = hmac.new(
  158. self._secret, payload.encode("ascii"), hashlib.sha256
  159. ).digest()
  160. provided = _b64decode(encoded_signature)
  161. if not hmac.compare_digest(expected, provided):
  162. raise ValueError("generation receipt signature is invalid")
  163. try:
  164. decoded = json.loads(_b64decode(payload))
  165. except (UnicodeDecodeError, json.JSONDecodeError) as exc:
  166. raise ValueError("generation receipt payload is invalid") from exc
  167. claims = self._validate_claims(decoded)
  168. current = now or datetime.now(UTC)
  169. if current.tzinfo is None:
  170. raise ValueError("generation receipt clock must be timezone aware")
  171. if int(current.timestamp()) >= claims["expires_at"]:
  172. raise ValueError("generation receipt has expired")
  173. if claims["actor_uid"] != _uid(actor_uid, "actor_uid"):
  174. raise ValueError("generation receipt actor does not match")
  175. if claims["source_text_hash"] != _hash_text(
  176. source_text, "source_text"
  177. ):
  178. raise ValueError("generation receipt source does not match")
  179. if claims["rule_spec_hash"] != rule_spec_hash(
  180. validate_rule_spec(rule_spec)
  181. ):
  182. raise ValueError("generation receipt RuleSpec does not match")
  183. return claims
  184. @staticmethod
  185. def _validate_claims(value: Any) -> dict[str, Any]:
  186. if not isinstance(value, dict) or set(value) != _RECEIPT_KEYS:
  187. raise ValueError("generation receipt claims are invalid")
  188. if value.get("version") != "1":
  189. raise ValueError("generation receipt version is unsupported")
  190. expires_at = value.get("expires_at")
  191. if (
  192. isinstance(expires_at, bool)
  193. or not isinstance(expires_at, int)
  194. or expires_at <= 0
  195. ):
  196. raise ValueError("generation receipt expiry is invalid")
  197. return {
  198. "version": "1",
  199. "generation_run_id": _uid(
  200. value.get("generation_run_id"), "generation_run_id"
  201. ),
  202. "actor_uid": _uid(value.get("actor_uid"), "actor_uid"),
  203. "source_text_hash": _digest(
  204. value.get("source_text_hash"), "source_text_hash"
  205. ),
  206. "candidate_hash": _digest(
  207. value.get("candidate_hash"), "candidate_hash"
  208. ),
  209. "rule_spec_hash": _digest(
  210. value.get("rule_spec_hash"), "rule_spec_hash"
  211. ),
  212. "model_hash": _digest(value.get("model_hash"), "model_hash"),
  213. "prompt_hash": _digest(value.get("prompt_hash"), "prompt_hash"),
  214. "context_hash": _digest(
  215. value.get("context_hash"), "context_hash"
  216. ),
  217. "expires_at": expires_at,
  218. }
  219. def _normalize_compiled(value: Any) -> dict[str, Any]:
  220. if not isinstance(value, dict) or set(value) != _COMPILED_KEYS:
  221. raise ValueError("trusted compiler returned an invalid result")
  222. if value.get("backend") not in {"sql_pushdown", "polars_batch"}:
  223. raise ValueError("trusted compiler returned an unsupported backend")
  224. compiler_version = str(value.get("compiler_version") or "").strip()
  225. if not compiler_version or len(compiler_version) > 80:
  226. raise ValueError("trusted compiler version is invalid")
  227. if not isinstance(value.get("plan"), dict):
  228. raise ValueError("trusted compiler plan is invalid")
  229. for key in ("schema_hashes", "binding_hashes", "capabilities"):
  230. if not isinstance(value.get(key), dict):
  231. raise ValueError(f"trusted compiler {key} are invalid")
  232. return {
  233. **copy.deepcopy(value),
  234. "compiler_version": compiler_version,
  235. "plan_hash": _digest(value.get("plan_hash"), "plan_hash"),
  236. }
  237. def _normalize_test(value: Any, context: dict[str, Any]) -> dict[str, Any]:
  238. if not isinstance(value, dict) or set(value) != _TEST_KEYS:
  239. raise ValueError("trusted test runner returned an invalid result")
  240. if value.get("status") != "success":
  241. raise ValueError("trusted rule dry-run did not succeed")
  242. if value.get("test_kind") not in {"dry_run", "sample", "golden"}:
  243. raise ValueError("trusted rule test kind is invalid")
  244. _uid(value.get("run_id"), "test run_id")
  245. if (
  246. value.get("plan_hash") != context.get("plan_hash")
  247. or value.get("schema_hashes") != context.get("schema_hashes")
  248. or value.get("binding_hashes") != context.get("binding_hashes")
  249. ):
  250. raise ValueError("test evidence does not match the exact compiled plan")
  251. counts = value.get("counts")
  252. if not isinstance(counts, dict) or any(
  253. isinstance(item, bool) or not isinstance(item, int) or item < 0
  254. for item in counts.values()
  255. ):
  256. raise ValueError("trusted test counts are invalid")
  257. attestation = value.get("attestation")
  258. if not isinstance(attestation, dict):
  259. raise ValueError("trusted test attestation is invalid")
  260. return copy.deepcopy(value)
  261. class RulePublicationService:
  262. """Coordinate only server-generated compile and test evidence."""
  263. def __init__(
  264. self,
  265. repository,
  266. *,
  267. receipt_signer: GenerationReceiptSigner | None,
  268. compiler,
  269. test_runner,
  270. ):
  271. self.repository = repository
  272. self.receipt_signer = receipt_signer
  273. self.compiler = compiler
  274. self.test_runner = test_runner
  275. def create_draft(
  276. self,
  277. *,
  278. rule_spec: dict[str, Any],
  279. source_text: str,
  280. actor_uid: str,
  281. generation_receipt: str,
  282. category: str,
  283. source_language: str,
  284. generated_kind: str,
  285. ) -> dict[str, Any]:
  286. if self.receipt_signer is None:
  287. raise RuntimeError("generation receipt signer is not configured")
  288. spec = validate_rule_spec(rule_spec)
  289. actor = _uid(actor_uid, "actor_uid")
  290. claims = self.receipt_signer.verify(
  291. generation_receipt,
  292. actor_uid=actor,
  293. source_text=source_text,
  294. rule_spec=spec,
  295. )
  296. receipt_hash = hashlib.sha256(
  297. generation_receipt.encode("utf-8")
  298. ).hexdigest()
  299. return self.repository.create_draft_from_generation(
  300. rule_spec=spec,
  301. source_text=source_text,
  302. created_by=actor,
  303. category=category,
  304. source_language=source_language,
  305. generated_kind=generated_kind,
  306. receipt_claims=claims,
  307. receipt_hash=receipt_hash,
  308. )
  309. def validate(
  310. self,
  311. version_id: str,
  312. actor_uid: str,
  313. ) -> dict[str, Any]:
  314. version = _uid(version_id, "version_id")
  315. actor = _uid(actor_uid, "actor_uid")
  316. context = self.repository.load_logical_validation_context(
  317. version_id=version,
  318. )
  319. try:
  320. compiled = _normalize_compiled(self.compiler.compile(context))
  321. except ValueError as exc:
  322. self.repository.record_validation_failure(
  323. version_id=version,
  324. actor_uid=actor,
  325. error_code="deterministic_validation_error",
  326. error_hash=hashlib.sha256(str(exc).encode("utf-8")).hexdigest(),
  327. )
  328. raise RuleValidationRejected(str(exc)) from exc
  329. return self.repository.persist_logical_compile(
  330. version_id=version,
  331. actor_uid=actor,
  332. context=context,
  333. compiled=compiled,
  334. )
  335. def test(
  336. self,
  337. version_id: str,
  338. actor_uid: str,
  339. *,
  340. plan_id: str,
  341. ) -> dict[str, Any]:
  342. version = _uid(version_id, "version_id")
  343. actor = _uid(actor_uid, "actor_uid")
  344. plan = _uid(plan_id, "plan_id")
  345. context = self.repository.load_logical_test_context(
  346. version_id=version,
  347. plan_id=plan,
  348. )
  349. if "terminal_result" in context:
  350. return copy.deepcopy(context["terminal_result"])
  351. evidence = _normalize_test(self.test_runner.run(context), context)
  352. return self.repository.persist_logical_test(
  353. version_id=version,
  354. plan_id=plan,
  355. actor_uid=actor,
  356. context=context,
  357. evidence=evidence,
  358. )
  359. def publish(self, version_id: str, actor_uid: str) -> dict[str, Any]:
  360. return self.repository.publish_validated_version(
  361. version_id=_uid(version_id, "version_id"),
  362. actor_uid=_uid(actor_uid, "actor_uid"),
  363. )
  364. def evidence(self, version_id: str) -> dict[str, Any]:
  365. return self.repository.get_publication_evidence(
  366. version_id=_uid(version_id, "version_id")
  367. )
  368. def catalog(self, *, query: str = "", limit: int = 50):
  369. if not isinstance(query, str) or len(query) > 200:
  370. raise ValueError("catalog query is invalid")
  371. if (
  372. isinstance(limit, bool)
  373. or not isinstance(limit, int)
  374. or limit < 1
  375. or limit > 100
  376. ):
  377. raise ValueError("catalog limit is invalid")
  378. return self.repository.search_published_rules(
  379. query=query.strip(), limit=limit
  380. )
  381. class PhysicalPlanPublicationService:
  382. """Test and publish deployment-bound plans independently of RuleVersion."""
  383. def __init__(self, repository, *, test_runner):
  384. self.repository = repository
  385. self.test_runner = test_runner
  386. def validate(self, plan_id: str, actor_uid: str) -> dict[str, Any]:
  387. return self.repository.attest_physical_compile(
  388. plan_id=_uid(plan_id, "plan_id"),
  389. actor_uid=_uid(actor_uid, "actor_uid"),
  390. )
  391. def test(self, plan_id: str, actor_uid: str) -> dict[str, Any]:
  392. plan = _uid(plan_id, "plan_id")
  393. actor = _uid(actor_uid, "actor_uid")
  394. context = self.repository.load_physical_test_context(plan_id=plan)
  395. if "terminal_result" in context:
  396. return copy.deepcopy(context["terminal_result"])
  397. evidence = _normalize_test(self.test_runner.run(context), context)
  398. return self.repository.persist_physical_test(
  399. plan_id=plan,
  400. actor_uid=actor,
  401. context=context,
  402. evidence=evidence,
  403. )
  404. def publish(self, plan_id: str, actor_uid: str) -> dict[str, Any]:
  405. return self.repository.publish_tested_physical_plan(
  406. plan_id=_uid(plan_id, "plan_id"),
  407. actor_uid=_uid(actor_uid, "actor_uid"),
  408. )
  409. class CanonicalBoundCompiler:
  410. """Compile a repository-loaded physical context through the registry."""
  411. def __init__(self, compiler_registry):
  412. self.compiler_registry = compiler_registry
  413. def compile(self, context: dict[str, Any]) -> dict[str, Any]:
  414. if not isinstance(context, dict):
  415. raise ValueError("canonical publication context is invalid")
  416. try:
  417. rule_version = context["rule_version"]
  418. input_schema = context["input_schema"]
  419. output_schema = context["output_schema"]
  420. input_binding = context["input_binding"]
  421. output_binding = context["output_binding"]
  422. backend = context["backend"]
  423. except KeyError as exc:
  424. raise ValueError(
  425. "canonical publication context is incomplete"
  426. ) from exc
  427. compiler = self.compiler_registry.select(
  428. rule_version.get("rule_spec"),
  429. input_binding,
  430. output_binding,
  431. )
  432. compiled = compiler.compile(
  433. rule_version=rule_version,
  434. input_schema=input_schema,
  435. output_schema=output_schema,
  436. input_binding=input_binding,
  437. output_binding=output_binding,
  438. backend=backend,
  439. )
  440. plan = compiled.get("plan")
  441. if not isinstance(plan, dict):
  442. raise ValueError("trusted compiler returned an invalid plan")
  443. return {
  444. "backend": compiled.get("backend"),
  445. "compiler_version": compiled.get("compiler_version"),
  446. "plan": plan,
  447. "plan_hash": compiled.get("plan_hash"),
  448. "schema_hashes": {
  449. "rule_spec_hash": rule_version.get("spec_hash"),
  450. "input_schema_snapshot_id": input_schema.get("id"),
  451. "input_schema_hash": input_schema.get("schema_hash"),
  452. "output_schema_snapshot_id": output_schema.get("id"),
  453. "output_schema_hash": output_schema.get("schema_hash"),
  454. },
  455. "binding_hashes": {
  456. "input": input_binding.get("binding_hash"),
  457. "output": output_binding.get("binding_hash"),
  458. },
  459. "capabilities": copy.deepcopy(backend),
  460. }
  461. class LogicalRuleCompiler:
  462. """Compile a draft RuleSpec against its immutable validation profile."""
  463. VERSION = "dataops-logical-rule-1.0"
  464. def compile(self, context: dict[str, Any]) -> dict[str, Any]:
  465. from app.core.data_rules.compilers.polars import PolarsRuleCompiler
  466. from app.core.data_rules.contracts import read_rule_spec
  467. if not isinstance(context, dict):
  468. raise ValueError("logical validation context is invalid")
  469. version = context.get("version")
  470. profile = context.get("profile")
  471. if not isinstance(version, dict) or version.get("status") != "draft":
  472. raise ValueError("only draft rule versions may be validated")
  473. if not isinstance(profile, dict):
  474. raise ValueError("trusted validation profile is unavailable")
  475. spec = read_rule_spec(version.get("rule_spec"))
  476. if rule_spec_hash(spec) != version.get("spec_hash"):
  477. raise ValueError("canonical RuleSpec hash does not match")
  478. fields = profile.get("fields")
  479. if not isinstance(fields, list) or not fields:
  480. raise ValueError("trusted validation schema fields are unavailable")
  481. schema_hash = _digest(profile.get("schema_hash"), "schema_hash")
  482. snapshot = {
  483. "id": _uid(
  484. profile.get("schema_snapshot_id"), "schema_snapshot_id"
  485. ),
  486. "schema_ref": str(profile.get("schema_ref") or ""),
  487. "schema_hash": schema_hash,
  488. "fields": copy.deepcopy(fields),
  489. "source_revision": str(profile.get("source_revision") or ""),
  490. }
  491. version_id = _uid(version.get("id"), "version_id")
  492. profile_id = _uid(profile.get("id"), "validation_profile_id")
  493. binding_base = {
  494. "data_source_uid": version_id,
  495. "object_kind": "parquet_artifact",
  496. "schema_snapshot_id": snapshot["id"],
  497. "access_mode": "read_write",
  498. "dialect": "polars",
  499. "write_mode": "append",
  500. }
  501. compiled = PolarsRuleCompiler().compile(
  502. rule_version={**version, "status": "published"},
  503. input_schema=snapshot,
  504. output_schema=snapshot,
  505. input_binding={
  506. "id": profile_id,
  507. **binding_base,
  508. "object_ref": "validation-input",
  509. },
  510. output_binding={
  511. "id": version_id,
  512. **binding_base,
  513. "object_ref": "validation-output",
  514. },
  515. backend={
  516. "max_rows": 100_000,
  517. "max_artifact_bytes": 32 * 1024 * 1024,
  518. "memory_limit_bytes": 256 * 1024 * 1024,
  519. "masking_policies": {
  520. "customer_mobile_last4": "preserve_last_4",
  521. "redact": "redact",
  522. },
  523. "lookup_bindings": {},
  524. },
  525. )
  526. plan = compiled["plan"]
  527. return {
  528. "backend": "polars_batch",
  529. "compiler_version": compiled["compiler_version"],
  530. "plan": plan,
  531. "plan_hash": compiled["plan_hash"],
  532. "schema_hashes": {"input": schema_hash},
  533. "binding_hashes": {},
  534. "capabilities": {
  535. "supported_execution_backends": ["polars"],
  536. "closed_rulespec": "2.0",
  537. "resource_limits": plan["resource_limits"],
  538. },
  539. }
  540. class ServerOwnedLogicalDryRunRunner:
  541. """Execute a logical Polars plan on a bounded server-owned artifact."""
  542. def __init__(self, artifact_store):
  543. self.artifact_store = artifact_store
  544. def run(self, context: dict[str, Any]) -> dict[str, Any]:
  545. from app.core.data_rules.compilers.polars import (
  546. validate_bound_polars_plan,
  547. )
  548. from app.runner.polars_worker import execute_isolated_polars_plan
  549. sample_ref = context.get("sample_artifact_ref")
  550. if not isinstance(sample_ref, str) or not sample_ref:
  551. raise ValueError("server-owned sample artifact is required")
  552. sample_digest = _digest(
  553. context.get("sample_artifact_digest"),
  554. "sample_artifact_digest",
  555. )
  556. plan = validate_bound_polars_plan(context.get("plan"))
  557. described = self.artifact_store.describe(sample_ref)
  558. if (
  559. described["digest"] != sample_digest
  560. or described["schema_hash"] != plan["input_schema_hash"]
  561. ):
  562. raise ValueError("sample artifact schema has drifted")
  563. output_path = None
  564. try:
  565. with self.artifact_store.stage(
  566. sample_ref,
  567. sample_digest,
  568. expected_schema_fields=plan["input_fields"],
  569. limits=plan["resource_limits"],
  570. ) as input_path:
  571. with tempfile.NamedTemporaryFile(
  572. prefix="dataops-logical-dry-run-",
  573. suffix=".parquet",
  574. delete=False,
  575. ) as handle:
  576. output_path = handle.name
  577. result = execute_isolated_polars_plan(
  578. {
  579. "plan": plan,
  580. "input_path": input_path,
  581. "lookup_paths": {},
  582. "output_path": output_path,
  583. "masking_policies": {
  584. "customer_mobile_last4": "preserve_last_4",
  585. "redact": "redact",
  586. },
  587. },
  588. memory_limit_bytes=plan["resource_limits"][
  589. "memory_limit_bytes"
  590. ],
  591. )
  592. prepared = self.artifact_store.prepare_path(
  593. output_path,
  594. new_governance_uid(),
  595. 300,
  596. schema_fields=plan["output_fields"],
  597. limits=plan["resource_limits"],
  598. )
  599. finally:
  600. if output_path is not None:
  601. with suppress(FileNotFoundError):
  602. os.unlink(output_path)
  603. count_keys = (
  604. "rows_in",
  605. "rows_out",
  606. "rows_rejected",
  607. "rows_quarantined",
  608. "rows_filtered",
  609. "rows_deduplicated",
  610. "rows_join_dropped",
  611. "rows_aggregated",
  612. "violation_count",
  613. )
  614. return {
  615. "status": "success",
  616. "test_kind": "dry_run",
  617. "run_id": new_governance_uid(),
  618. "plan_hash": context["plan_hash"],
  619. "schema_hashes": context["schema_hashes"],
  620. "binding_hashes": context["binding_hashes"],
  621. "counts": {key: int(result.get(key, 0)) for key in count_keys},
  622. "attestation": {
  623. "input_digest": described["digest"],
  624. "output_digest": prepared["digest"],
  625. "output_schema_hash": prepared["schema_hash"],
  626. "violation_digest": hashlib.sha256(
  627. _canonical(result.get("violations", []))
  628. ).hexdigest(),
  629. },
  630. }
  631. class ServerOwnedPhysicalPreflightRunner:
  632. """Execute an exact physical Polars plan on its latest trusted input."""
  633. def __init__(self, artifact_store, datasource_manager=None):
  634. self.artifact_store = artifact_store
  635. self.datasource_manager = datasource_manager
  636. def run(self, context: dict[str, Any]) -> dict[str, Any]:
  637. from app.core.data_rules.compilers.polars import (
  638. validate_bound_polars_plan,
  639. )
  640. from app.runner.polars_worker import execute_isolated_polars_plan
  641. if context.get("backend") == "sql_pushdown":
  642. return self._run_sql(context)
  643. if context.get("backend") != "polars_batch":
  644. raise ValueError("physical preflight backend is unsupported")
  645. sample = context.get("sample_artifact")
  646. if (
  647. not isinstance(sample, dict)
  648. or set(sample) != {
  649. "artifact_ref",
  650. "digest",
  651. "schema_fields",
  652. }
  653. ):
  654. raise ValueError("server-owned physical sample is required")
  655. plan = validate_bound_polars_plan(context.get("plan"))
  656. output_path = None
  657. try:
  658. with self.artifact_store.stage(
  659. sample["artifact_ref"],
  660. sample["digest"],
  661. expected_schema_fields=sample["schema_fields"],
  662. limits=plan["resource_limits"],
  663. ) as input_path:
  664. with tempfile.NamedTemporaryFile(
  665. prefix="dataops-physical-preflight-",
  666. suffix=".parquet",
  667. delete=False,
  668. ) as handle:
  669. output_path = handle.name
  670. result = execute_isolated_polars_plan(
  671. {
  672. "plan": plan,
  673. "input_path": input_path,
  674. "lookup_paths": {},
  675. "output_path": output_path,
  676. "masking_policies": {
  677. "customer_mobile_last4": "preserve_last_4",
  678. "redact": "redact",
  679. },
  680. },
  681. memory_limit_bytes=plan["resource_limits"][
  682. "memory_limit_bytes"
  683. ],
  684. )
  685. prepared = self.artifact_store.prepare_path(
  686. output_path,
  687. new_governance_uid(),
  688. 300,
  689. schema_fields=plan["output_fields"],
  690. limits=plan["resource_limits"],
  691. )
  692. finally:
  693. if output_path is not None:
  694. with suppress(FileNotFoundError):
  695. os.unlink(output_path)
  696. count_keys = (
  697. "rows_in",
  698. "rows_out",
  699. "rows_rejected",
  700. "rows_quarantined",
  701. "rows_filtered",
  702. "rows_deduplicated",
  703. "rows_join_dropped",
  704. "rows_aggregated",
  705. "violation_count",
  706. )
  707. return {
  708. "status": "success",
  709. "test_kind": "dry_run",
  710. "run_id": new_governance_uid(),
  711. "plan_hash": context["plan_hash"],
  712. "schema_hashes": context["schema_hashes"],
  713. "binding_hashes": context["binding_hashes"],
  714. "counts": {key: int(result.get(key, 0)) for key in count_keys},
  715. "attestation": {
  716. "input_digest": sample["digest"],
  717. "output_digest": prepared["digest"],
  718. "output_schema_hash": prepared["schema_hash"],
  719. "violation_digest": hashlib.sha256(
  720. _canonical(result.get("violations", []))
  721. ).hexdigest(),
  722. },
  723. }
  724. def _run_sql(self, context: dict[str, Any]) -> dict[str, Any]:
  725. from sqlalchemy import text
  726. from app.core.data_rules.compilers.sql import validate_bound_sql_plan
  727. if self.datasource_manager is None:
  728. raise ValueError(
  729. "server-owned SQL preflight executor is not configured"
  730. )
  731. plan = validate_bound_sql_plan(context.get("plan"))
  732. statement = plan["statements"][0]
  733. explain_sql = f"EXPLAIN {statement['sql']}"
  734. with self.datasource_manager.connect(
  735. plan["data_source_uid"], purpose="dataflow_read"
  736. ) as connection:
  737. rows = connection.execute(
  738. text(explain_sql), statement["parameters"]
  739. ).fetchall()
  740. if not rows:
  741. raise ValueError("physical SQL preflight returned no explain plan")
  742. explain_digest = hashlib.sha256(
  743. _canonical([list(row) for row in rows])
  744. ).hexdigest()
  745. return {
  746. "status": "success",
  747. "test_kind": "dry_run",
  748. "run_id": new_governance_uid(),
  749. "plan_hash": context["plan_hash"],
  750. "schema_hashes": context["schema_hashes"],
  751. "binding_hashes": context["binding_hashes"],
  752. "counts": {
  753. "rows_in": 0,
  754. "rows_out": 0,
  755. "rows_rejected": 0,
  756. },
  757. "attestation": {
  758. "dialect": plan["dialect"],
  759. "explain_digest": explain_digest,
  760. "statement_digest": hashlib.sha256(
  761. statement["sql"].encode("utf-8")
  762. ).hexdigest(),
  763. },
  764. }
  765. __all__ = [
  766. "CanonicalBoundCompiler",
  767. "GenerationReceiptSigner",
  768. "LogicalRuleCompiler",
  769. "PhysicalPlanPublicationService",
  770. "RulePublicationService",
  771. "RuleValidationRejected",
  772. "ServerOwnedLogicalDryRunRunner",
  773. "ServerOwnedPhysicalPreflightRunner",
  774. "generation_receipt_claims",
  775. ]