artifacts.py 62 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694
  1. """Digest-bound, bounded Parquet artifacts owned by the DataOps Runner."""
  2. from __future__ import annotations
  3. import hashlib
  4. import json
  5. import math
  6. import os
  7. import re
  8. import tempfile
  9. from collections.abc import Mapping
  10. from contextlib import contextmanager, suppress
  11. from datetime import UTC, datetime, timedelta
  12. from typing import Any
  13. import polars as pl
  14. import pyarrow.parquet as pq
  15. from minio.error import S3Error
  16. from sqlalchemy import text
  17. from app.core.common.identifiers import (
  18. ensure_governance_uid,
  19. new_governance_uid,
  20. )
  21. from app.core.data_rules.execution_contracts import canonical_schema_hash
  22. PARQUET_CONTENT_TYPE = "application/x-parquet"
  23. _DIGEST = re.compile(r"^[0-9a-f]{64}$")
  24. def _confirmed_object_missing(exc: BaseException) -> bool:
  25. return isinstance(exc, KeyError) or (
  26. isinstance(exc, S3Error)
  27. and exc.code in {"NoSuchKey", "NoSuchObject", "NotFound"}
  28. )
  29. class ArtifactCommitUnknown(RuntimeError):
  30. """A catalog transaction may have committed but cannot be confirmed."""
  31. class ArtifactHandoffPending(RuntimeError):
  32. """Another publisher owns the durable pending handoff."""
  33. def _parquet_footer_bounds(path: str) -> tuple[int, int]:
  34. try:
  35. metadata = pq.ParquetFile(path).metadata
  36. except Exception as exc:
  37. raise ValueError("artifact Parquet footer is invalid") from exc
  38. if metadata is None or metadata.num_row_groups < 1:
  39. raise ValueError("artifact Parquet footer is incomplete")
  40. uncompressed = sum(
  41. metadata.row_group(index).total_byte_size
  42. for index in range(metadata.num_row_groups)
  43. )
  44. if metadata.num_rows < 0 or uncompressed < 0:
  45. raise ValueError("artifact Parquet footer is invalid")
  46. return metadata.num_rows, uncompressed
  47. def _now_utc(clock) -> datetime:
  48. value = clock()
  49. if not isinstance(value, datetime):
  50. raise ValueError("artifact clock must return a datetime")
  51. if value.tzinfo is None:
  52. value = value.replace(tzinfo=UTC)
  53. return value.astimezone(UTC)
  54. def _timestamp(value: datetime) -> str:
  55. return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
  56. def _parse_timestamp(value: Any) -> datetime:
  57. if not isinstance(value, str) or not value.endswith("Z"):
  58. raise ValueError("artifact expiry metadata is invalid")
  59. try:
  60. parsed = datetime.fromisoformat(value[:-1] + "+00:00")
  61. except ValueError as exc:
  62. raise ValueError("artifact expiry metadata is invalid") from exc
  63. return parsed.astimezone(UTC)
  64. def _uid(value: Any, label: str) -> str:
  65. try:
  66. return ensure_governance_uid({"uid": str(value)})
  67. except ValueError as exc:
  68. raise ValueError(f"{label} must be a valid UUIDv7") from exc
  69. def _inferred_schema_fields(
  70. frame: pl.DataFrame | pl.LazyFrame,
  71. ) -> list[dict[str, Any]]:
  72. schema = (
  73. frame.collect_schema()
  74. if isinstance(frame, pl.LazyFrame)
  75. else frame.schema
  76. )
  77. fields = []
  78. for name, dtype in schema.items():
  79. field: dict[str, Any] = {
  80. "name": name,
  81. "nullable": True,
  82. }
  83. if dtype == pl.Boolean:
  84. field["type"] = "boolean"
  85. elif dtype == pl.Date:
  86. field["type"] = "date"
  87. elif dtype == pl.String:
  88. field["type"] = "string"
  89. elif dtype.is_integer():
  90. field["type"] = "integer"
  91. elif dtype == pl.Float32:
  92. field["type"] = "float"
  93. elif dtype == pl.Float64:
  94. field["type"] = "double"
  95. elif dtype.is_decimal():
  96. field.update(
  97. {
  98. "type": "decimal",
  99. "precision": dtype.precision,
  100. "scale": dtype.scale,
  101. }
  102. )
  103. elif isinstance(dtype, pl.Datetime):
  104. field["type"] = (
  105. "timestamptz" if dtype.time_zone else "timestamp"
  106. )
  107. if dtype.time_zone:
  108. field["timezone"] = dtype.time_zone
  109. else:
  110. raise ValueError(f"unsupported artifact dtype for {name}")
  111. fields.append(field)
  112. return sorted(fields, key=lambda item: item["name"])
  113. def _normalized_schema_fields(value: Any) -> list[dict[str, Any]]:
  114. canonical_schema_hash(value)
  115. return sorted(
  116. [dict(field) for field in value],
  117. key=lambda item: item["name"],
  118. )
  119. def _schema_contract(value: Any) -> tuple[list[dict[str, Any]], str]:
  120. fields = _normalized_schema_fields(value)
  121. return fields, canonical_schema_hash(fields)
  122. def _validate_frame_schema(
  123. frame: pl.DataFrame | pl.LazyFrame,
  124. fields: list[dict[str, Any]],
  125. ) -> None:
  126. schema = (
  127. frame.collect_schema()
  128. if isinstance(frame, pl.LazyFrame)
  129. else frame.schema
  130. )
  131. expected_names = {field["name"] for field in fields}
  132. if set(schema.names()) != expected_names:
  133. raise ValueError("artifact schema fields do not match")
  134. for field in fields:
  135. dtype = schema[field["name"]]
  136. field_type = field["type"]
  137. matches = (
  138. (field_type == "boolean" and dtype == pl.Boolean)
  139. or (field_type == "date" and dtype == pl.Date)
  140. or (field_type == "string" and dtype == pl.String)
  141. or (field_type == "integer" and dtype.is_integer())
  142. or (field_type == "float" and dtype == pl.Float32)
  143. or (field_type == "double" and dtype == pl.Float64)
  144. or (
  145. field_type == "decimal"
  146. and dtype.is_decimal()
  147. and dtype.precision == field.get("precision")
  148. and dtype.scale == field.get("scale")
  149. )
  150. or (
  151. field_type == "timestamp"
  152. and isinstance(dtype, pl.Datetime)
  153. and dtype.time_zone is None
  154. )
  155. or (
  156. field_type == "timestamptz"
  157. and isinstance(dtype, pl.Datetime)
  158. and dtype.time_zone == field.get("timezone")
  159. )
  160. )
  161. if not matches:
  162. raise ValueError(
  163. f"artifact schema type for {field['name']} does not match"
  164. )
  165. if isinstance(frame, pl.DataFrame):
  166. for field in fields:
  167. if not field["nullable"] and frame[field["name"]].null_count():
  168. raise ValueError(
  169. f"artifact nullable contract for {field['name']} does not match"
  170. )
  171. def _metadata(value: Any) -> dict[str, str]:
  172. if not isinstance(value, Mapping):
  173. raise ValueError("artifact content metadata is missing")
  174. normalized = {}
  175. for key, item in value.items():
  176. name = str(key).lower()
  177. if name.startswith("x-amz-meta-"):
  178. name = name[len("x-amz-meta-") :]
  179. if name in {
  180. "sha256",
  181. "row-count",
  182. "schema-sha256",
  183. "expires-at",
  184. "artifact-bytes",
  185. }:
  186. normalized[name] = str(item)
  187. required = {
  188. "sha256",
  189. "row-count",
  190. "schema-sha256",
  191. "expires-at",
  192. "artifact-bytes",
  193. }
  194. if set(normalized) != required:
  195. raise ValueError("artifact content metadata is incomplete")
  196. if sum(len(key) + len(item) for key, item in normalized.items()) > 2_048:
  197. raise ValueError("artifact content metadata exceeds the safe limit")
  198. return normalized
  199. class ArtifactStore:
  200. """Read and write only server-owned, bounded Parquet artifacts."""
  201. def __init__(
  202. self,
  203. client,
  204. *,
  205. bucket: str,
  206. max_artifact_bytes: int,
  207. max_rows: int,
  208. memory_limit_bytes: int,
  209. max_ttl_seconds: int = 86400,
  210. clock=None,
  211. ):
  212. if not re.fullmatch(r"[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]", bucket):
  213. raise ValueError("artifact bucket name is invalid")
  214. self.client = client
  215. self.bucket = bucket
  216. self.max_artifact_bytes = int(max_artifact_bytes)
  217. self.max_rows = int(max_rows)
  218. self.memory_limit_bytes = int(memory_limit_bytes)
  219. self.max_ttl_seconds = int(max_ttl_seconds)
  220. self.clock = clock or (lambda: datetime.now(UTC))
  221. if (
  222. self.max_artifact_bytes < 1024
  223. or self.max_rows < 1
  224. or self.memory_limit_bytes < self.max_artifact_bytes
  225. or self.max_ttl_seconds < 1
  226. ):
  227. raise ValueError("artifact resource limits are invalid")
  228. if not self.client.bucket_exists(self.bucket):
  229. raise ValueError("artifact bucket does not exist")
  230. def _limits(self, value: Any = None) -> dict[str, int]:
  231. configured = {
  232. "max_rows": self.max_rows,
  233. "max_artifact_bytes": self.max_artifact_bytes,
  234. "memory_limit_bytes": self.memory_limit_bytes,
  235. }
  236. if value is None:
  237. return configured
  238. if not isinstance(value, dict) or set(value) != set(configured):
  239. raise ValueError("artifact limits must have a closed shape")
  240. result = {}
  241. for key, ceiling in configured.items():
  242. item = value[key]
  243. if (
  244. isinstance(item, bool)
  245. or not isinstance(item, int)
  246. or item < 1
  247. or item > ceiling
  248. ):
  249. raise ValueError(
  250. f"artifact {key} exceeds the configured ceiling"
  251. )
  252. result[key] = item
  253. return result
  254. def _parse_ref(self, ref: Any) -> str:
  255. prefix = f"minio://{self.bucket}/"
  256. if not isinstance(ref, str) or not ref.startswith(prefix):
  257. raise ValueError("artifact reference is not owned by this store")
  258. key = ref[len(prefix) :]
  259. match = re.fullmatch(
  260. r"rules/([0-9a-f-]{36})/([0-9a-f-]{36})\.parquet",
  261. key,
  262. )
  263. if match is None:
  264. raise ValueError("artifact reference is invalid")
  265. _uid(match.group(1), "artifact correlation id")
  266. _uid(match.group(2), "artifact id")
  267. return key
  268. def _validated_stat(
  269. self,
  270. key: str,
  271. *,
  272. expected_digest: str | None = None,
  273. limits: dict[str, int] | None = None,
  274. ) -> tuple[Any, dict[str, str]]:
  275. effective = self._limits(limits)
  276. stat = self.client.stat_object(self.bucket, key)
  277. size = int(getattr(stat, "size", -1))
  278. if size < 1 or size > effective["max_artifact_bytes"]:
  279. raise ValueError("artifact size exceeds the configured limit")
  280. if size > effective["memory_limit_bytes"]:
  281. raise ValueError("artifact download exceeds the memory limit")
  282. if str(getattr(stat, "content_type", "")).lower() != PARQUET_CONTENT_TYPE:
  283. raise ValueError("artifact content type is invalid")
  284. metadata = _metadata(getattr(stat, "metadata", None))
  285. digest = metadata["sha256"]
  286. if _DIGEST.fullmatch(digest) is None:
  287. raise ValueError("artifact digest metadata is invalid")
  288. if expected_digest is not None and digest != expected_digest:
  289. raise ValueError("artifact digest does not match")
  290. try:
  291. row_count = int(metadata["row-count"])
  292. metadata_size = int(metadata["artifact-bytes"])
  293. except (TypeError, ValueError) as exc:
  294. raise ValueError("artifact count metadata is invalid") from exc
  295. if row_count < 0 or row_count > effective["max_rows"]:
  296. raise ValueError("artifact row count exceeds the configured limit")
  297. if metadata_size != size:
  298. raise ValueError("artifact size metadata does not match")
  299. if _DIGEST.fullmatch(metadata["schema-sha256"]) is None:
  300. raise ValueError("artifact schema metadata is invalid")
  301. if _parse_timestamp(metadata["expires-at"]) <= _now_utc(self.clock):
  302. raise ValueError("artifact has expired")
  303. return stat, metadata
  304. def write(
  305. self,
  306. frame: pl.LazyFrame | pl.DataFrame,
  307. correlation_id: str,
  308. ttl_seconds: int,
  309. *,
  310. schema_fields: list[dict[str, Any]] | None = None,
  311. limits: dict[str, int] | None = None,
  312. ) -> dict[str, Any]:
  313. effective = self._limits(limits)
  314. correlation = _uid(correlation_id, "correlation_id")
  315. if (
  316. isinstance(ttl_seconds, bool)
  317. or not isinstance(ttl_seconds, int)
  318. or ttl_seconds < 1
  319. or ttl_seconds > self.max_ttl_seconds
  320. ):
  321. raise ValueError("artifact TTL is outside the configured limit")
  322. if isinstance(frame, pl.DataFrame):
  323. lazy = frame.lazy()
  324. elif isinstance(frame, pl.LazyFrame):
  325. lazy = frame
  326. else:
  327. raise ValueError("artifact frame must be a Polars frame")
  328. collected = lazy.head(effective["max_rows"] + 1).collect(
  329. engine="streaming"
  330. )
  331. if collected.height > effective["max_rows"]:
  332. raise ValueError("artifact row count exceeds the configured limit")
  333. if collected.estimated_size() > effective["memory_limit_bytes"]:
  334. raise ValueError("artifact frame exceeds the configured memory limit")
  335. fields, schema_digest = _schema_contract(
  336. schema_fields or _inferred_schema_fields(collected)
  337. )
  338. _validate_frame_schema(collected, fields)
  339. expires_at = _timestamp(
  340. _now_utc(self.clock) + timedelta(seconds=ttl_seconds)
  341. )
  342. artifact_id = new_governance_uid()
  343. key = f"rules/{correlation}/{artifact_id}.parquet"
  344. path = None
  345. uploaded = False
  346. try:
  347. with tempfile.NamedTemporaryFile(
  348. prefix="dataops-rule-artifact-",
  349. suffix=".parquet",
  350. delete=False,
  351. ) as handle:
  352. path = handle.name
  353. collected.write_parquet(path)
  354. size = os.path.getsize(path)
  355. if size < 1 or size > effective["max_artifact_bytes"]:
  356. raise ValueError("artifact size exceeds the configured limit")
  357. if size > effective["memory_limit_bytes"]:
  358. raise ValueError("serialized artifact exceeds the memory limit")
  359. if (
  360. size + collected.estimated_size()
  361. > effective["memory_limit_bytes"]
  362. ):
  363. raise ValueError(
  364. "artifact serialization exceeds the memory limit"
  365. )
  366. digest = hashlib.sha256()
  367. with open(path, "rb") as handle:
  368. while chunk := handle.read(1024 * 1024):
  369. digest.update(chunk)
  370. digest_hex = digest.hexdigest()
  371. with open(path, "rb") as handle:
  372. self.client.put_object(
  373. self.bucket,
  374. key,
  375. handle,
  376. size,
  377. content_type=PARQUET_CONTENT_TYPE,
  378. metadata={
  379. "sha256": digest_hex,
  380. "row-count": str(collected.height),
  381. "schema-sha256": schema_digest,
  382. "expires-at": expires_at,
  383. "artifact-bytes": str(size),
  384. },
  385. )
  386. uploaded = True
  387. self._validated_stat(
  388. key,
  389. expected_digest=digest_hex,
  390. limits=effective,
  391. )
  392. artifact_ref = f"minio://{self.bucket}/{key}"
  393. with self.stage(
  394. artifact_ref,
  395. digest_hex,
  396. expected_schema_fields=fields,
  397. limits=effective,
  398. ):
  399. pass
  400. except Exception:
  401. if uploaded:
  402. with suppress(Exception):
  403. self.client.remove_object(self.bucket, key)
  404. raise
  405. finally:
  406. if path is not None:
  407. with suppress(FileNotFoundError):
  408. os.unlink(path)
  409. return {
  410. "artifact_ref": artifact_ref,
  411. "digest": digest_hex,
  412. "row_count": collected.height,
  413. "schema_hash": schema_digest,
  414. "schema_fields": fields,
  415. "expires_at": expires_at,
  416. }
  417. def prepare_path(
  418. self,
  419. path: str,
  420. correlation_id: str,
  421. ttl_seconds: int,
  422. *,
  423. schema_fields: list[dict[str, Any]],
  424. limits: dict[str, int] | None = None,
  425. ) -> dict[str, Any]:
  426. """Validate a local Parquet file and reserve its server-owned key."""
  427. effective = self._limits(limits)
  428. correlation = _uid(correlation_id, "correlation_id")
  429. if (
  430. isinstance(ttl_seconds, bool)
  431. or not isinstance(ttl_seconds, int)
  432. or ttl_seconds < 1
  433. or ttl_seconds > self.max_ttl_seconds
  434. ):
  435. raise ValueError("artifact TTL is outside the configured limit")
  436. if not isinstance(path, str) or not os.path.isfile(path):
  437. raise ValueError("worker artifact path is invalid")
  438. size = os.path.getsize(path)
  439. if size < 1 or size > effective["max_artifact_bytes"]:
  440. raise ValueError("artifact size exceeds the configured limit")
  441. rows, uncompressed = _parquet_footer_bounds(path)
  442. if rows > effective["max_rows"]:
  443. raise ValueError("artifact row count exceeds the configured limit")
  444. if uncompressed > effective["memory_limit_bytes"]:
  445. raise ValueError(
  446. "Parquet footer uncompressed size exceeds the memory limit"
  447. )
  448. fields, schema_digest = _schema_contract(schema_fields)
  449. _validate_frame_schema(pl.scan_parquet(path), fields)
  450. digest = hashlib.sha256()
  451. with open(path, "rb") as handle:
  452. while chunk := handle.read(1024 * 1024):
  453. digest.update(chunk)
  454. digest_hex = digest.hexdigest()
  455. expires_at = _timestamp(
  456. _now_utc(self.clock) + timedelta(seconds=ttl_seconds)
  457. )
  458. artifact_id = new_governance_uid()
  459. key = f"rules/{correlation}/{artifact_id}.parquet"
  460. return {
  461. "artifact_ref": f"minio://{self.bucket}/{key}",
  462. "digest": digest_hex,
  463. "row_count": rows,
  464. "schema_hash": schema_digest,
  465. "schema_fields": fields,
  466. "expires_at": expires_at,
  467. }
  468. def upload_path(
  469. self,
  470. path: str,
  471. artifact: dict[str, Any],
  472. *,
  473. limits: dict[str, int] | None = None,
  474. ) -> dict[str, Any]:
  475. """Upload to an already reserved exact key and verify the object."""
  476. if not isinstance(artifact, dict):
  477. raise ValueError("prepared artifact metadata is invalid")
  478. required = {
  479. "artifact_ref",
  480. "digest",
  481. "row_count",
  482. "schema_hash",
  483. "schema_fields",
  484. "expires_at",
  485. }
  486. if set(artifact) != required:
  487. raise ValueError("prepared artifact metadata has a closed shape")
  488. key = self._parse_ref(artifact["artifact_ref"])
  489. correlation = key.split("/", 2)[1]
  490. ttl_seconds = math.ceil(
  491. (
  492. _parse_timestamp(artifact["expires_at"])
  493. - _now_utc(self.clock)
  494. ).total_seconds()
  495. )
  496. if ttl_seconds < 1 or ttl_seconds > self.max_ttl_seconds:
  497. raise ValueError("prepared artifact TTL is invalid")
  498. expected = self.prepare_path(
  499. path,
  500. correlation,
  501. ttl_seconds,
  502. schema_fields=artifact["schema_fields"],
  503. limits=limits,
  504. )
  505. for name in (
  506. "digest",
  507. "row_count",
  508. "schema_hash",
  509. "schema_fields",
  510. ):
  511. if expected[name] != artifact[name]:
  512. raise ValueError("prepared artifact no longer matches its path")
  513. effective = self._limits(limits)
  514. size = os.path.getsize(path)
  515. uploaded = False
  516. try:
  517. with open(path, "rb") as handle:
  518. self.client.put_object(
  519. self.bucket,
  520. key,
  521. handle,
  522. size,
  523. content_type=PARQUET_CONTENT_TYPE,
  524. metadata={
  525. "sha256": artifact["digest"],
  526. "row-count": str(artifact["row_count"]),
  527. "schema-sha256": artifact["schema_hash"],
  528. "expires-at": artifact["expires_at"],
  529. "artifact-bytes": str(size),
  530. },
  531. )
  532. uploaded = True
  533. with self.stage(
  534. artifact["artifact_ref"],
  535. artifact["digest"],
  536. expected_schema_fields=artifact["schema_fields"],
  537. limits=effective,
  538. ):
  539. pass
  540. except Exception:
  541. if uploaded:
  542. with suppress(Exception):
  543. self.client.remove_object(self.bucket, key)
  544. raise
  545. return dict(artifact)
  546. def write_path(
  547. self,
  548. path: str,
  549. correlation_id: str,
  550. ttl_seconds: int,
  551. *,
  552. schema_fields: list[dict[str, Any]],
  553. limits: dict[str, int] | None = None,
  554. ) -> dict[str, Any]:
  555. """Prepare and upload a non-cataloged compatibility artifact."""
  556. artifact = self.prepare_path(
  557. path,
  558. correlation_id,
  559. ttl_seconds,
  560. schema_fields=schema_fields,
  561. limits=limits,
  562. )
  563. return self.upload_path(path, artifact, limits=limits)
  564. def describe(self, ref: str) -> dict[str, Any]:
  565. """Return validated object metadata without exposing MinIO credentials."""
  566. key = self._parse_ref(ref)
  567. _stat, metadata = self._validated_stat(key)
  568. return {
  569. "artifact_ref": ref,
  570. "digest": metadata["sha256"],
  571. "row_count": int(metadata["row-count"]),
  572. "schema_hash": metadata["schema-sha256"],
  573. "expires_at": metadata["expires-at"],
  574. }
  575. def describe_optional(self, ref: str) -> dict[str, Any] | None:
  576. """Return None only for a confirmed missing object."""
  577. try:
  578. return self.describe(ref)
  579. except Exception as exc:
  580. if _confirmed_object_missing(exc):
  581. return None
  582. raise
  583. def read(
  584. self,
  585. ref: str,
  586. expected_digest: str,
  587. *,
  588. expected_schema_fields: list[dict[str, Any]] | None = None,
  589. limits: dict[str, int] | None = None,
  590. ) -> pl.LazyFrame:
  591. with self.stage(
  592. ref,
  593. expected_digest,
  594. expected_schema_fields=expected_schema_fields,
  595. limits=limits,
  596. ) as path:
  597. effective = self._limits(limits)
  598. try:
  599. frame = pl.read_parquet(
  600. path,
  601. n_rows=effective["max_rows"] + 1,
  602. memory_map=False,
  603. )
  604. except Exception as exc:
  605. raise ValueError("artifact is not valid Parquet") from exc
  606. if frame.estimated_size() > effective["memory_limit_bytes"]:
  607. raise ValueError("artifact frame exceeds the configured memory limit")
  608. _validate_frame_schema(
  609. frame, _normalized_schema_fields(expected_schema_fields)
  610. )
  611. return frame.lazy()
  612. @contextmanager
  613. def stage(
  614. self,
  615. ref: str,
  616. expected_digest: str,
  617. *,
  618. expected_schema_fields: list[dict[str, Any]] | None = None,
  619. limits: dict[str, int] | None = None,
  620. ):
  621. effective = self._limits(limits)
  622. if _DIGEST.fullmatch(str(expected_digest or "")) is None:
  623. raise ValueError("expected artifact digest is invalid")
  624. key = self._parse_ref(ref)
  625. _stat, metadata = self._validated_stat(
  626. key,
  627. expected_digest=expected_digest,
  628. limits=effective,
  629. )
  630. if expected_schema_fields is None:
  631. raise ValueError("expected artifact schema fields are required")
  632. fields = _normalized_schema_fields(expected_schema_fields)
  633. if canonical_schema_hash(fields) != metadata["schema-sha256"]:
  634. raise ValueError("artifact schema contract is not expected")
  635. response = self.client.get_object(self.bucket, key)
  636. digest = hashlib.sha256()
  637. path = None
  638. size = 0
  639. try:
  640. try:
  641. with tempfile.NamedTemporaryFile(
  642. prefix="dataops-rule-stage-",
  643. suffix=".parquet",
  644. delete=False,
  645. ) as handle:
  646. path = handle.name
  647. while chunk := response.read(1024 * 1024):
  648. size += len(chunk)
  649. if size > effective["max_artifact_bytes"]:
  650. raise ValueError(
  651. "artifact size exceeds the configured limit"
  652. )
  653. digest.update(chunk)
  654. handle.write(chunk)
  655. finally:
  656. response.close()
  657. release = getattr(response, "release_conn", None)
  658. if callable(release):
  659. release()
  660. if digest.hexdigest() != expected_digest:
  661. raise ValueError("artifact digest does not match content")
  662. footer_rows, uncompressed = _parquet_footer_bounds(path)
  663. if footer_rows != int(metadata["row-count"]):
  664. raise ValueError("artifact row count does not match metadata")
  665. if footer_rows > effective["max_rows"]:
  666. raise ValueError(
  667. "artifact row count exceeds the configured limit"
  668. )
  669. if uncompressed > effective["memory_limit_bytes"]:
  670. raise ValueError(
  671. "Parquet footer uncompressed size exceeds the memory limit"
  672. )
  673. lazy = pl.scan_parquet(path)
  674. _validate_frame_schema(lazy, fields)
  675. yield path
  676. finally:
  677. if path is not None:
  678. with suppress(FileNotFoundError):
  679. os.unlink(path)
  680. def cleanup_expired(self, correlation_id: str) -> int:
  681. correlation = _uid(correlation_id, "correlation_id")
  682. prefix = f"rules/{correlation}/"
  683. removed = 0
  684. for item in self.client.list_objects(
  685. self.bucket,
  686. prefix=prefix,
  687. recursive=True,
  688. ):
  689. key = str(getattr(item, "object_name", ""))
  690. if not key.startswith(prefix):
  691. continue
  692. try:
  693. self._parse_ref(f"minio://{self.bucket}/{key}")
  694. stat = self.client.stat_object(self.bucket, key)
  695. metadata = _metadata(getattr(stat, "metadata", None))
  696. expired = _parse_timestamp(
  697. metadata["expires-at"]
  698. ) <= _now_utc(self.clock)
  699. except ValueError:
  700. continue
  701. if expired:
  702. self.client.remove_object(self.bucket, key)
  703. removed += 1
  704. return removed
  705. def delete(self, ref: str) -> None:
  706. """Delete one exact store-owned artifact after validating its key."""
  707. key = self._parse_ref(ref)
  708. self.client.remove_object(self.bucket, key)
  709. class PostgresArtifactResolver:
  710. """Resolve a canonical artifact binding without accepting caller paths."""
  711. def __init__(self, engine, artifact_store: ArtifactStore):
  712. self.engine = engine
  713. self.artifact_store = artifact_store
  714. def resolve(
  715. self,
  716. *,
  717. binding_id: str,
  718. correlation_id: str,
  719. kind: str,
  720. ) -> dict[str, Any]:
  721. binding = _uid(binding_id, "artifact binding id")
  722. correlation = _uid(correlation_id, "artifact correlation id")
  723. if kind not in {"input", "lookup", "output"}:
  724. raise ValueError("artifact kind is invalid")
  725. statement = text(
  726. """
  727. SELECT
  728. a.artifact_ref,
  729. a.artifact_digest,
  730. a.row_count,
  731. a.schema_hash,
  732. a.schema_fields,
  733. a.expires_at,
  734. a.binding_hash AS catalog_binding_hash,
  735. b.binding_hash AS current_binding_hash
  736. FROM public.rule_run_artifacts a
  737. JOIN public.dataflow_dataset_bindings b
  738. ON b.id = a.binding_id
  739. WHERE a.binding_id = CAST(:binding_id AS uuid)
  740. AND a.correlation_id = CAST(:correlation_id AS uuid)
  741. AND a.artifact_kind = :artifact_kind
  742. AND a.handoff_status = 'ready'
  743. AND a.expires_at > CURRENT_TIMESTAMP
  744. AND a.binding_hash = b.binding_hash
  745. AND b.object_kind = 'parquet_artifact'
  746. AND b.access_mode IN ('read', 'read_write')
  747. ORDER BY a.created_at DESC
  748. LIMIT 1
  749. """
  750. )
  751. with self.engine.connect() as connection:
  752. row = connection.execute(
  753. statement,
  754. {
  755. "binding_id": binding,
  756. "correlation_id": correlation,
  757. "artifact_kind": kind,
  758. },
  759. ).mappings().one_or_none()
  760. if row is None:
  761. raise ValueError("canonical artifact binding was not found")
  762. artifact_ref = str(row["artifact_ref"])
  763. key = self.artifact_store._parse_ref(artifact_ref)
  764. if not key.startswith(f"rules/{correlation}/"):
  765. raise ValueError(
  766. "catalog artifact does not match the execution correlation"
  767. )
  768. described = self.artifact_store.describe(artifact_ref)
  769. row_fields = row["schema_fields"]
  770. if isinstance(row_fields, str):
  771. row_fields = json.loads(row_fields)
  772. if (
  773. described["digest"] != str(row["artifact_digest"])
  774. or described["row_count"] != int(row["row_count"])
  775. or described["schema_hash"] != str(row["schema_hash"])
  776. ):
  777. raise ValueError("catalog artifact metadata does not match storage")
  778. return {
  779. **described,
  780. "schema_fields": _normalized_schema_fields(row_fields),
  781. "binding_hash": str(row["catalog_binding_hash"]),
  782. }
  783. def resolve_handoff(
  784. self,
  785. *,
  786. artifact_ref: str,
  787. correlation_id: str,
  788. ) -> dict[str, Any]:
  789. """Resolve only a ready upstream output from the same execution."""
  790. correlation = _uid(correlation_id, "artifact correlation id")
  791. key = self.artifact_store._parse_ref(artifact_ref)
  792. if not key.startswith(f"rules/{correlation}/"):
  793. raise ValueError(
  794. "artifact handoff does not match the execution correlation"
  795. )
  796. with self.engine.connect() as connection:
  797. row = connection.execute(
  798. text(
  799. """
  800. SELECT a.artifact_digest, a.row_count, a.schema_hash,
  801. a.schema_fields, a.expires_at
  802. FROM public.rule_run_artifacts a
  803. JOIN public.dataflow_dataset_bindings b
  804. ON b.id = a.binding_id
  805. WHERE a.artifact_ref = :artifact_ref
  806. AND a.correlation_id =
  807. CAST(:correlation_id AS uuid)
  808. AND a.artifact_kind = 'output'
  809. AND a.handoff_status = 'ready'
  810. AND a.expires_at > CURRENT_TIMESTAMP
  811. AND a.binding_hash = b.binding_hash
  812. """
  813. ),
  814. {
  815. "artifact_ref": artifact_ref,
  816. "correlation_id": correlation,
  817. },
  818. ).mappings().one_or_none()
  819. if row is None:
  820. raise ValueError("ready upstream artifact handoff was not found")
  821. fields = row["schema_fields"]
  822. if isinstance(fields, str):
  823. fields = json.loads(fields)
  824. described = self.artifact_store.describe(artifact_ref)
  825. if (
  826. described["digest"] != str(row["artifact_digest"])
  827. or described["row_count"] != int(row["row_count"])
  828. or described["schema_hash"] != str(row["schema_hash"])
  829. ):
  830. raise ValueError(
  831. "upstream artifact handoff metadata does not match"
  832. )
  833. return {
  834. **described,
  835. "schema_fields": _normalized_schema_fields(fields),
  836. }
  837. def attest_binding(
  838. self,
  839. *,
  840. binding_id: str,
  841. binding_hash: str,
  842. access_mode: str,
  843. ) -> dict[str, str]:
  844. binding = _uid(binding_id, "artifact binding id")
  845. if _DIGEST.fullmatch(str(binding_hash or "")) is None:
  846. raise ValueError("artifact binding hash is invalid")
  847. allowed = {
  848. "read": {"read", "read_write"},
  849. "write": {"write", "read_write"},
  850. }.get(access_mode)
  851. if allowed is None:
  852. raise ValueError("artifact access mode is invalid")
  853. with self.engine.connect() as connection:
  854. row = connection.execute(
  855. text(
  856. """
  857. SELECT binding_hash, access_mode, object_kind
  858. FROM public.dataflow_dataset_bindings
  859. WHERE id = CAST(:binding_id AS uuid)
  860. """
  861. ),
  862. {"binding_id": binding},
  863. ).mappings().one_or_none()
  864. if (
  865. row is None
  866. or row["object_kind"] != "parquet_artifact"
  867. or row["access_mode"] not in allowed
  868. or str(row["binding_hash"]) != binding_hash
  869. ):
  870. raise ValueError("canonical artifact binding no longer matches")
  871. return {"binding_hash": str(row["binding_hash"])}
  872. @staticmethod
  873. def _catalog_artifact(row: Mapping[str, Any]) -> dict[str, Any]:
  874. fields = row["schema_fields"]
  875. if isinstance(fields, str):
  876. fields = json.loads(fields)
  877. expires_at = row["expires_at"]
  878. return {
  879. "artifact_ref": str(row["artifact_ref"]),
  880. "digest": str(row["artifact_digest"]),
  881. "row_count": int(row["row_count"]),
  882. "schema_hash": str(row["schema_hash"]),
  883. "schema_fields": _normalized_schema_fields(fields),
  884. "expires_at": (
  885. _timestamp(expires_at)
  886. if isinstance(expires_at, datetime)
  887. else str(expires_at)
  888. ),
  889. }
  890. @staticmethod
  891. def _attest_binding_locked(
  892. connection,
  893. *,
  894. binding_id: str,
  895. binding_hash: str,
  896. kind: str,
  897. ) -> None:
  898. row = connection.execute(
  899. text(
  900. """
  901. SELECT binding_hash, access_mode, object_kind
  902. FROM public.dataflow_dataset_bindings
  903. WHERE id = CAST(:binding_id AS uuid)
  904. FOR SHARE
  905. """
  906. ),
  907. {"binding_id": binding_id},
  908. ).mappings().one_or_none()
  909. allowed = (
  910. {"write", "read_write"}
  911. if kind == "output"
  912. else {"read", "read_write"}
  913. )
  914. if (
  915. row is None
  916. or row["object_kind"] != "parquet_artifact"
  917. or row["access_mode"] not in allowed
  918. or str(row["binding_hash"]) != binding_hash
  919. ):
  920. raise ValueError("canonical artifact binding no longer matches")
  921. def _lookup_handoff(
  922. self,
  923. *,
  924. correlation_id: str,
  925. binding_id: str,
  926. kind: str,
  927. ) -> dict[str, Any] | None:
  928. with self.engine.connect() as connection:
  929. row = connection.execute(
  930. text(
  931. """
  932. SELECT id::text AS id, correlation_id::text,
  933. binding_id::text, artifact_ref, artifact_digest,
  934. row_count, schema_hash, schema_fields,
  935. artifact_kind, binding_hash, expires_at,
  936. handoff_status
  937. FROM public.rule_run_artifacts
  938. WHERE correlation_id = CAST(:correlation_id AS uuid)
  939. AND binding_id = CAST(:binding_id AS uuid)
  940. AND artifact_kind = :artifact_kind
  941. """
  942. ),
  943. {
  944. "correlation_id": correlation_id,
  945. "binding_id": binding_id,
  946. "artifact_kind": kind,
  947. },
  948. ).mappings().one_or_none()
  949. return dict(row) if row is not None else None
  950. def reserve(
  951. self,
  952. *,
  953. binding_id: str,
  954. correlation_id: str,
  955. artifact: dict[str, Any],
  956. kind: str,
  957. binding_hash: str,
  958. ) -> dict[str, Any]:
  959. """Atomically attest the binding and reserve one pending handoff."""
  960. binding = _uid(binding_id, "artifact binding id")
  961. correlation = _uid(correlation_id, "artifact correlation id")
  962. if kind not in {"input", "lookup", "output"}:
  963. raise ValueError("artifact kind is invalid")
  964. if _DIGEST.fullmatch(str(binding_hash or "")) is None:
  965. raise ValueError("artifact binding hash is invalid")
  966. if not isinstance(artifact, dict):
  967. raise ValueError("artifact metadata is invalid")
  968. artifact_ref = artifact.get("artifact_ref")
  969. key = self.artifact_store._parse_ref(artifact_ref)
  970. if not key.startswith(f"rules/{correlation}/"):
  971. raise ValueError(
  972. "artifact does not match the execution correlation"
  973. )
  974. fields = _normalized_schema_fields(artifact.get("schema_fields"))
  975. if canonical_schema_hash(fields) != artifact.get("schema_hash"):
  976. raise ValueError("artifact schema contract does not match")
  977. if _DIGEST.fullmatch(str(artifact.get("digest") or "")) is None:
  978. raise ValueError("artifact digest is invalid")
  979. reservation_id = new_governance_uid()
  980. parameters = {
  981. "id": reservation_id,
  982. "correlation_id": correlation,
  983. "binding_id": binding,
  984. "artifact_ref": artifact_ref,
  985. "artifact_digest": artifact["digest"],
  986. "row_count": int(artifact["row_count"]),
  987. "schema_hash": artifact["schema_hash"],
  988. "schema_fields": json.dumps(
  989. fields,
  990. sort_keys=True,
  991. separators=(",", ":"),
  992. ),
  993. "artifact_kind": kind,
  994. "binding_hash": binding_hash,
  995. "expires_at": artifact["expires_at"],
  996. }
  997. selected = None
  998. inserted = False
  999. try:
  1000. with self.engine.begin() as connection:
  1001. self._attest_binding_locked(
  1002. connection,
  1003. binding_id=binding,
  1004. binding_hash=binding_hash,
  1005. kind=kind,
  1006. )
  1007. selected = connection.execute(
  1008. text(
  1009. """
  1010. INSERT INTO public.rule_run_artifacts (
  1011. id, correlation_id, binding_id, artifact_ref,
  1012. artifact_digest, row_count, schema_hash,
  1013. schema_fields, artifact_kind, binding_hash,
  1014. handoff_status, expires_at
  1015. ) VALUES (
  1016. CAST(:id AS uuid),
  1017. CAST(:correlation_id AS uuid),
  1018. CAST(:binding_id AS uuid), :artifact_ref,
  1019. :artifact_digest, :row_count, :schema_hash,
  1020. CAST(:schema_fields AS jsonb), :artifact_kind,
  1021. :binding_hash, 'pending',
  1022. CAST(:expires_at AS timestamptz)
  1023. )
  1024. ON CONFLICT (
  1025. correlation_id, binding_id, artifact_kind
  1026. ) DO NOTHING
  1027. RETURNING id::text AS id, correlation_id::text,
  1028. binding_id::text, artifact_ref,
  1029. artifact_digest, row_count, schema_hash,
  1030. schema_fields, artifact_kind, binding_hash,
  1031. expires_at, handoff_status
  1032. """
  1033. ),
  1034. parameters,
  1035. ).mappings().one_or_none()
  1036. inserted = selected is not None
  1037. if selected is None:
  1038. selected = connection.execute(
  1039. text(
  1040. """
  1041. SELECT id::text AS id, correlation_id::text,
  1042. binding_id::text, artifact_ref,
  1043. artifact_digest, row_count, schema_hash,
  1044. schema_fields, artifact_kind, binding_hash,
  1045. expires_at, handoff_status
  1046. FROM public.rule_run_artifacts
  1047. WHERE correlation_id =
  1048. CAST(:correlation_id AS uuid)
  1049. AND binding_id = CAST(:binding_id AS uuid)
  1050. AND artifact_kind = :artifact_kind
  1051. FOR UPDATE
  1052. """
  1053. ),
  1054. parameters,
  1055. ).mappings().one_or_none()
  1056. except ValueError:
  1057. raise
  1058. except Exception as exc:
  1059. try:
  1060. selected = self._lookup_handoff(
  1061. correlation_id=correlation,
  1062. binding_id=binding,
  1063. kind=kind,
  1064. )
  1065. except Exception as recheck_exc:
  1066. raise ArtifactCommitUnknown(
  1067. "artifact reservation commit outcome is unknown"
  1068. ) from recheck_exc
  1069. if (
  1070. selected is None
  1071. or str(selected["artifact_ref"]) != artifact_ref
  1072. or str(selected["artifact_digest"]) != artifact["digest"]
  1073. or str(selected["binding_hash"]) != binding_hash
  1074. ):
  1075. raise ArtifactCommitUnknown(
  1076. "artifact reservation commit outcome is unknown"
  1077. ) from exc
  1078. inserted = True
  1079. if selected is None:
  1080. raise ArtifactCommitUnknown(
  1081. "artifact reservation commit outcome is unknown"
  1082. )
  1083. row = dict(selected)
  1084. if str(row["binding_hash"]) != binding_hash:
  1085. raise ValueError("artifact reservation binding hash conflicts")
  1086. if str(row["artifact_digest"]) != artifact["digest"]:
  1087. raise ValueError(
  1088. "immutable artifact catalog digest conflicts with retry"
  1089. )
  1090. status = str(row["handoff_status"])
  1091. if not inserted:
  1092. if status == "ready":
  1093. return {
  1094. **self._catalog_artifact(row),
  1095. "correlation_id": correlation,
  1096. "reservation_id": str(row["id"]),
  1097. "handoff_status": "ready",
  1098. "upload_required": False,
  1099. }
  1100. if status == "pending":
  1101. raise ArtifactHandoffPending(
  1102. "artifact handoff is already pending"
  1103. )
  1104. raise ValueError("artifact handoff has failed")
  1105. return {
  1106. **self._catalog_artifact(row),
  1107. "correlation_id": correlation,
  1108. "reservation_id": str(row["id"]),
  1109. "handoff_status": status,
  1110. "upload_required": status == "pending",
  1111. }
  1112. def _abort_pending(self, reservation_id: str) -> None:
  1113. with self.engine.begin() as connection:
  1114. connection.execute(
  1115. text(
  1116. """
  1117. DELETE FROM public.rule_run_artifacts
  1118. WHERE id = CAST(:id AS uuid)
  1119. AND handoff_status = 'pending'
  1120. """
  1121. ),
  1122. {"id": reservation_id},
  1123. )
  1124. def finalize(
  1125. self,
  1126. *,
  1127. reservation: dict[str, Any],
  1128. binding_id: str,
  1129. binding_hash: str,
  1130. kind: str,
  1131. ) -> dict[str, Any]:
  1132. reservation_id = _uid(
  1133. reservation.get("reservation_id"), "artifact reservation id"
  1134. )
  1135. binding = _uid(binding_id, "artifact binding id")
  1136. selected = None
  1137. try:
  1138. with self.engine.begin() as connection:
  1139. self._attest_binding_locked(
  1140. connection,
  1141. binding_id=binding,
  1142. binding_hash=binding_hash,
  1143. kind=kind,
  1144. )
  1145. selected = connection.execute(
  1146. text(
  1147. """
  1148. UPDATE public.rule_run_artifacts
  1149. SET handoff_status = 'ready',
  1150. ready_at = CURRENT_TIMESTAMP,
  1151. updated_at = CURRENT_TIMESTAMP,
  1152. failure_code = NULL,
  1153. failed_at = NULL
  1154. WHERE id = CAST(:id AS uuid)
  1155. AND binding_id = CAST(:binding_id AS uuid)
  1156. AND binding_hash = :binding_hash
  1157. AND artifact_digest = :artifact_digest
  1158. AND handoff_status = 'pending'
  1159. RETURNING id::text AS id, correlation_id::text,
  1160. binding_id::text, artifact_ref,
  1161. artifact_digest, row_count, schema_hash,
  1162. schema_fields, artifact_kind, binding_hash,
  1163. expires_at, handoff_status
  1164. """
  1165. ),
  1166. {
  1167. "id": reservation_id,
  1168. "binding_id": binding,
  1169. "binding_hash": binding_hash,
  1170. "artifact_digest": reservation["digest"],
  1171. },
  1172. ).mappings().one_or_none()
  1173. if selected is None:
  1174. raise ValueError(
  1175. "pending artifact handoff no longer matches"
  1176. )
  1177. except ValueError:
  1178. raise
  1179. except Exception as exc:
  1180. try:
  1181. selected = self._lookup_handoff(
  1182. correlation_id=_uid(
  1183. reservation["correlation_id"],
  1184. "artifact correlation id",
  1185. ),
  1186. binding_id=binding,
  1187. kind=kind,
  1188. )
  1189. except Exception as recheck_exc:
  1190. raise ArtifactCommitUnknown(
  1191. "artifact finalize commit outcome is unknown"
  1192. ) from recheck_exc
  1193. if (
  1194. selected is None
  1195. or str(selected["id"]) != reservation_id
  1196. or str(selected["artifact_digest"])
  1197. != reservation["digest"]
  1198. or str(selected["binding_hash"]) != binding_hash
  1199. or str(selected["handoff_status"]) != "ready"
  1200. ):
  1201. raise ArtifactCommitUnknown(
  1202. "artifact finalize commit outcome is unknown"
  1203. ) from exc
  1204. return self._catalog_artifact(selected)
  1205. def publish_path(
  1206. self,
  1207. path: str,
  1208. *,
  1209. binding_id: str,
  1210. binding_hash: str,
  1211. correlation_id: str,
  1212. kind: str,
  1213. ttl_seconds: int,
  1214. schema_fields: list[dict[str, Any]],
  1215. limits: dict[str, int] | None = None,
  1216. ) -> dict[str, Any]:
  1217. """Reserve, upload, and finalize one durable artifact handoff."""
  1218. prepared = self.artifact_store.prepare_path(
  1219. path,
  1220. correlation_id,
  1221. ttl_seconds,
  1222. schema_fields=schema_fields,
  1223. limits=limits,
  1224. )
  1225. reservation = self.reserve(
  1226. binding_id=binding_id,
  1227. correlation_id=correlation_id,
  1228. artifact=prepared,
  1229. kind=kind,
  1230. binding_hash=binding_hash,
  1231. )
  1232. if not reservation["upload_required"]:
  1233. stored = self.artifact_store.describe(
  1234. reservation["artifact_ref"]
  1235. )
  1236. if stored["digest"] != reservation["digest"]:
  1237. raise ValueError(
  1238. "ready artifact catalog does not match storage"
  1239. )
  1240. return {
  1241. key: reservation[key]
  1242. for key in (
  1243. "artifact_ref",
  1244. "digest",
  1245. "row_count",
  1246. "schema_hash",
  1247. "schema_fields",
  1248. "expires_at",
  1249. )
  1250. }
  1251. reserved_artifact = {
  1252. key: reservation[key]
  1253. for key in (
  1254. "artifact_ref",
  1255. "digest",
  1256. "row_count",
  1257. "schema_hash",
  1258. "schema_fields",
  1259. "expires_at",
  1260. )
  1261. }
  1262. try:
  1263. self.artifact_store.upload_path(
  1264. path,
  1265. reserved_artifact,
  1266. limits=limits,
  1267. )
  1268. except Exception:
  1269. with suppress(Exception):
  1270. self.artifact_store.delete(
  1271. reserved_artifact["artifact_ref"]
  1272. )
  1273. with suppress(Exception):
  1274. self._abort_pending(reservation["reservation_id"])
  1275. raise
  1276. return self.finalize(
  1277. reservation=reservation,
  1278. binding_id=binding_id,
  1279. binding_hash=binding_hash,
  1280. kind=kind,
  1281. )
  1282. def _mark_failed(
  1283. self,
  1284. *,
  1285. row_id: str,
  1286. expected_status: str,
  1287. failure_code: str,
  1288. ) -> bool:
  1289. if expected_status not in {"pending", "ready"}:
  1290. raise ValueError("artifact expected handoff status is invalid")
  1291. with self.engine.begin() as connection:
  1292. updated = connection.execute(
  1293. text(
  1294. """
  1295. UPDATE public.rule_run_artifacts
  1296. SET handoff_status = 'failed',
  1297. failed_at = CURRENT_TIMESTAMP,
  1298. updated_at = CURRENT_TIMESTAMP,
  1299. failure_code = :failure_code
  1300. WHERE id = CAST(:id AS uuid)
  1301. AND handoff_status = :expected_status
  1302. """
  1303. ),
  1304. {
  1305. "id": row_id,
  1306. "expected_status": expected_status,
  1307. "failure_code": failure_code,
  1308. },
  1309. )
  1310. return int(updated.rowcount or 0) == 1
  1311. def _lookup_handoff_by_id(
  1312. self,
  1313. row_id: str,
  1314. ) -> dict[str, Any] | None:
  1315. with self.engine.connect() as connection:
  1316. row = connection.execute(
  1317. text(
  1318. """
  1319. SELECT id::text AS id, correlation_id::text,
  1320. binding_id::text, artifact_ref, artifact_digest,
  1321. row_count, schema_hash, schema_fields,
  1322. artifact_kind, binding_hash, expires_at,
  1323. handoff_status
  1324. FROM public.rule_run_artifacts
  1325. WHERE id = CAST(:id AS uuid)
  1326. """
  1327. ),
  1328. {"id": row_id},
  1329. ).mappings().one_or_none()
  1330. return dict(row) if row is not None else None
  1331. @staticmethod
  1332. def _matching_ready(
  1333. snapshot: Mapping[str, Any],
  1334. current: Mapping[str, Any] | None,
  1335. ) -> bool:
  1336. if current is None or str(current.get("handoff_status")) != "ready":
  1337. return False
  1338. return all(
  1339. str(current.get(key)) == str(snapshot.get(key))
  1340. for key in (
  1341. "id",
  1342. "correlation_id",
  1343. "binding_id",
  1344. "artifact_ref",
  1345. "artifact_digest",
  1346. "row_count",
  1347. "schema_hash",
  1348. "artifact_kind",
  1349. "binding_hash",
  1350. )
  1351. )
  1352. def _verify_reconcile_object(
  1353. self,
  1354. row: Mapping[str, Any],
  1355. ) -> dict[str, Any] | None:
  1356. artifact_ref = str(row["artifact_ref"])
  1357. self.artifact_store._parse_ref(artifact_ref)
  1358. stored = self.artifact_store.describe_optional(artifact_ref)
  1359. if stored is None:
  1360. return None
  1361. if any(
  1362. (
  1363. stored["digest"] != str(row["artifact_digest"]),
  1364. stored["row_count"] != int(row["row_count"]),
  1365. stored["schema_hash"] != str(row["schema_hash"]),
  1366. )
  1367. ):
  1368. raise ValueError("catalog artifact metadata does not match storage")
  1369. fields = row["schema_fields"]
  1370. if isinstance(fields, str):
  1371. fields = json.loads(fields)
  1372. try:
  1373. with self.artifact_store.stage(
  1374. artifact_ref,
  1375. str(row["artifact_digest"]),
  1376. expected_schema_fields=fields,
  1377. ):
  1378. pass
  1379. except Exception as exc:
  1380. if _confirmed_object_missing(exc):
  1381. return None
  1382. raise
  1383. return stored
  1384. def reconcile(
  1385. self,
  1386. *,
  1387. limit: int = 100,
  1388. grace_seconds: int = 300,
  1389. ) -> dict[str, int]:
  1390. """Repair bounded catalog/store drift after the grace period."""
  1391. if (
  1392. isinstance(limit, bool)
  1393. or not isinstance(limit, int)
  1394. or limit < 1
  1395. or limit > 1_000
  1396. ):
  1397. raise ValueError("artifact reconciliation limit is invalid")
  1398. if (
  1399. isinstance(grace_seconds, bool)
  1400. or not isinstance(grace_seconds, int)
  1401. or grace_seconds < 30
  1402. or grace_seconds > 86_400
  1403. ):
  1404. raise ValueError("artifact reconciliation grace is invalid")
  1405. result = {
  1406. "pending_finalized": 0,
  1407. "pending_deleted": 0,
  1408. "ready_failed": 0,
  1409. "orphans_deleted": 0,
  1410. }
  1411. with self.engine.connect() as connection:
  1412. rows = connection.execute(
  1413. text(
  1414. """
  1415. SELECT id::text AS id, correlation_id::text,
  1416. binding_id::text, artifact_ref, artifact_digest,
  1417. row_count, schema_hash, schema_fields,
  1418. artifact_kind, binding_hash, expires_at,
  1419. handoff_status
  1420. FROM public.rule_run_artifacts
  1421. WHERE handoff_status IN ('pending','ready')
  1422. AND updated_at <= CURRENT_TIMESTAMP
  1423. - make_interval(secs => :grace_seconds)
  1424. ORDER BY updated_at, id
  1425. LIMIT :limit
  1426. """
  1427. ),
  1428. {
  1429. "grace_seconds": grace_seconds,
  1430. "limit": limit,
  1431. },
  1432. ).mappings().all()
  1433. for raw_row in rows:
  1434. row = dict(raw_row)
  1435. row_id = str(row["id"])
  1436. status = str(row["handoff_status"])
  1437. try:
  1438. stored = self._verify_reconcile_object(row)
  1439. except ValueError:
  1440. if status == "pending":
  1441. failed = self._mark_failed(
  1442. row_id=row_id,
  1443. expected_status="pending",
  1444. failure_code="pending_object_invalid",
  1445. )
  1446. if failed:
  1447. with suppress(Exception):
  1448. self.artifact_store.delete(row["artifact_ref"])
  1449. else:
  1450. failed = self._mark_failed(
  1451. row_id=row_id,
  1452. expected_status="ready",
  1453. failure_code="ready_object_invalid",
  1454. )
  1455. if failed:
  1456. result["ready_failed"] += 1
  1457. continue
  1458. except Exception:
  1459. # Authentication, timeout, transport, and server failures are
  1460. # not evidence that a cataloged object is invalid.
  1461. continue
  1462. if status == "ready":
  1463. if stored is None:
  1464. failed = self._mark_failed(
  1465. row_id=row_id,
  1466. expected_status="ready",
  1467. failure_code="ready_object_missing",
  1468. )
  1469. if failed:
  1470. result["ready_failed"] += 1
  1471. continue
  1472. if stored is None:
  1473. with self.engine.begin() as connection:
  1474. deleted = connection.execute(
  1475. text(
  1476. """
  1477. DELETE FROM public.rule_run_artifacts
  1478. WHERE id = CAST(:id AS uuid)
  1479. AND handoff_status = 'pending'
  1480. """
  1481. ),
  1482. {"id": row_id},
  1483. )
  1484. if int(deleted.rowcount or 0) == 1:
  1485. result["pending_deleted"] += 1
  1486. continue
  1487. with self.engine.begin() as connection:
  1488. finalized = connection.execute(
  1489. text(
  1490. """
  1491. UPDATE public.rule_run_artifacts a
  1492. SET handoff_status = 'ready',
  1493. ready_at = CURRENT_TIMESTAMP,
  1494. updated_at = CURRENT_TIMESTAMP,
  1495. failure_code = NULL,
  1496. failed_at = NULL
  1497. WHERE a.id = CAST(:id AS uuid)
  1498. AND a.handoff_status = 'pending'
  1499. AND EXISTS (
  1500. SELECT 1
  1501. FROM public.dataflow_dataset_bindings b
  1502. WHERE b.id = a.binding_id
  1503. AND b.binding_hash = a.binding_hash
  1504. AND b.object_kind = 'parquet_artifact'
  1505. AND b.access_mode IN (
  1506. 'read','write','read_write'
  1507. )
  1508. )
  1509. RETURNING a.id
  1510. """
  1511. ),
  1512. {"id": row_id},
  1513. )
  1514. if int(finalized.rowcount or 0) == 1:
  1515. result["pending_finalized"] += 1
  1516. else:
  1517. current = self._lookup_handoff_by_id(row_id)
  1518. if self._matching_ready(row, current):
  1519. result["pending_finalized"] += 1
  1520. elif (
  1521. current is not None
  1522. and str(current.get("handoff_status")) == "pending"
  1523. ):
  1524. self._mark_failed(
  1525. row_id=row_id,
  1526. expected_status="pending",
  1527. failure_code="pending_binding_changed",
  1528. )
  1529. remaining = limit - len(rows)
  1530. if remaining <= 0:
  1531. return result
  1532. now = _now_utc(self.artifact_store.clock)
  1533. candidates = []
  1534. scanned = 0
  1535. for item in self.artifact_store.client.list_objects(
  1536. self.artifact_store.bucket,
  1537. prefix="rules/",
  1538. recursive=True,
  1539. ):
  1540. scanned += 1
  1541. if scanned > limit * 10 or len(candidates) >= remaining:
  1542. break
  1543. key = str(getattr(item, "object_name", ""))
  1544. ref = f"minio://{self.artifact_store.bucket}/{key}"
  1545. try:
  1546. self.artifact_store._parse_ref(ref)
  1547. except ValueError:
  1548. continue
  1549. modified = getattr(item, "last_modified", None)
  1550. if not isinstance(modified, datetime):
  1551. continue
  1552. if modified.tzinfo is None:
  1553. modified = modified.replace(tzinfo=UTC)
  1554. if modified.astimezone(UTC) > now - timedelta(
  1555. seconds=grace_seconds
  1556. ):
  1557. continue
  1558. candidates.append(ref)
  1559. if not candidates:
  1560. return result
  1561. with self.engine.connect() as connection:
  1562. referenced = {
  1563. str(row["artifact_ref"])
  1564. for row in connection.execute(
  1565. text(
  1566. """
  1567. SELECT artifact_ref
  1568. FROM public.rule_run_artifacts
  1569. WHERE artifact_ref =
  1570. ANY(CAST(:artifact_refs AS text[]))
  1571. """
  1572. ),
  1573. {"artifact_refs": candidates},
  1574. ).mappings().all()
  1575. }
  1576. for ref in candidates:
  1577. if ref in referenced:
  1578. continue
  1579. self.artifact_store.delete(ref)
  1580. result["orphans_deleted"] += 1
  1581. return result
  1582. def cleanup_expired(self, *, limit: int = 100) -> int:
  1583. if (
  1584. isinstance(limit, bool)
  1585. or not isinstance(limit, int)
  1586. or limit < 1
  1587. or limit > 1_000
  1588. ):
  1589. raise ValueError("artifact cleanup limit is invalid")
  1590. removed = 0
  1591. with self.engine.begin() as connection:
  1592. rows = connection.execute(
  1593. text(
  1594. """
  1595. SELECT id::text AS id, artifact_ref
  1596. FROM public.rule_run_artifacts
  1597. WHERE expires_at <= CURRENT_TIMESTAMP
  1598. ORDER BY expires_at, id
  1599. LIMIT :limit
  1600. FOR UPDATE SKIP LOCKED
  1601. """
  1602. ),
  1603. {"limit": limit},
  1604. ).mappings().all()
  1605. for row in rows:
  1606. self.artifact_store.delete(str(row["artifact_ref"]))
  1607. connection.execute(
  1608. text(
  1609. """
  1610. DELETE FROM public.rule_run_artifacts
  1611. WHERE id = CAST(:id AS uuid)
  1612. """
  1613. ),
  1614. {"id": str(row["id"])},
  1615. )
  1616. removed += 1
  1617. return removed