artifacts.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952
  1. """Digest-bound, bounded Parquet artifacts owned by the DataOps Runner."""
  2. from __future__ import annotations
  3. import hashlib
  4. import json
  5. import os
  6. import re
  7. import tempfile
  8. from collections.abc import Mapping
  9. from contextlib import contextmanager, suppress
  10. from datetime import UTC, datetime, timedelta
  11. from typing import Any
  12. import polars as pl
  13. import pyarrow.parquet as pq
  14. from sqlalchemy import text
  15. from app.core.common.identifiers import (
  16. ensure_governance_uid,
  17. new_governance_uid,
  18. )
  19. from app.core.data_rules.execution_contracts import canonical_schema_hash
  20. PARQUET_CONTENT_TYPE = "application/x-parquet"
  21. _DIGEST = re.compile(r"^[0-9a-f]{64}$")
  22. def _parquet_footer_bounds(path: str) -> tuple[int, int]:
  23. try:
  24. metadata = pq.ParquetFile(path).metadata
  25. except Exception as exc:
  26. raise ValueError("artifact Parquet footer is invalid") from exc
  27. if metadata is None or metadata.num_row_groups < 1:
  28. raise ValueError("artifact Parquet footer is incomplete")
  29. uncompressed = sum(
  30. metadata.row_group(index).total_byte_size
  31. for index in range(metadata.num_row_groups)
  32. )
  33. if metadata.num_rows < 0 or uncompressed < 0:
  34. raise ValueError("artifact Parquet footer is invalid")
  35. return metadata.num_rows, uncompressed
  36. def _now_utc(clock) -> datetime:
  37. value = clock()
  38. if not isinstance(value, datetime):
  39. raise ValueError("artifact clock must return a datetime")
  40. if value.tzinfo is None:
  41. value = value.replace(tzinfo=UTC)
  42. return value.astimezone(UTC)
  43. def _timestamp(value: datetime) -> str:
  44. return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
  45. def _parse_timestamp(value: Any) -> datetime:
  46. if not isinstance(value, str) or not value.endswith("Z"):
  47. raise ValueError("artifact expiry metadata is invalid")
  48. try:
  49. parsed = datetime.fromisoformat(value[:-1] + "+00:00")
  50. except ValueError as exc:
  51. raise ValueError("artifact expiry metadata is invalid") from exc
  52. return parsed.astimezone(UTC)
  53. def _uid(value: Any, label: str) -> str:
  54. try:
  55. return ensure_governance_uid({"uid": str(value)})
  56. except ValueError as exc:
  57. raise ValueError(f"{label} must be a valid UUIDv7") from exc
  58. def _inferred_schema_fields(
  59. frame: pl.DataFrame | pl.LazyFrame,
  60. ) -> list[dict[str, Any]]:
  61. schema = (
  62. frame.collect_schema()
  63. if isinstance(frame, pl.LazyFrame)
  64. else frame.schema
  65. )
  66. fields = []
  67. for name, dtype in schema.items():
  68. field: dict[str, Any] = {
  69. "name": name,
  70. "nullable": True,
  71. }
  72. if dtype == pl.Boolean:
  73. field["type"] = "boolean"
  74. elif dtype == pl.Date:
  75. field["type"] = "date"
  76. elif dtype == pl.String:
  77. field["type"] = "string"
  78. elif dtype.is_integer():
  79. field["type"] = "integer"
  80. elif dtype == pl.Float32:
  81. field["type"] = "float"
  82. elif dtype == pl.Float64:
  83. field["type"] = "double"
  84. elif dtype.is_decimal():
  85. field.update(
  86. {
  87. "type": "decimal",
  88. "precision": dtype.precision,
  89. "scale": dtype.scale,
  90. }
  91. )
  92. elif isinstance(dtype, pl.Datetime):
  93. field["type"] = (
  94. "timestamptz" if dtype.time_zone else "timestamp"
  95. )
  96. if dtype.time_zone:
  97. field["timezone"] = dtype.time_zone
  98. else:
  99. raise ValueError(f"unsupported artifact dtype for {name}")
  100. fields.append(field)
  101. return sorted(fields, key=lambda item: item["name"])
  102. def _normalized_schema_fields(value: Any) -> list[dict[str, Any]]:
  103. canonical_schema_hash(value)
  104. return sorted(
  105. [dict(field) for field in value],
  106. key=lambda item: item["name"],
  107. )
  108. def _schema_contract(value: Any) -> tuple[list[dict[str, Any]], str]:
  109. fields = _normalized_schema_fields(value)
  110. return fields, canonical_schema_hash(fields)
  111. def _validate_frame_schema(
  112. frame: pl.DataFrame | pl.LazyFrame,
  113. fields: list[dict[str, Any]],
  114. ) -> None:
  115. schema = (
  116. frame.collect_schema()
  117. if isinstance(frame, pl.LazyFrame)
  118. else frame.schema
  119. )
  120. expected_names = {field["name"] for field in fields}
  121. if set(schema.names()) != expected_names:
  122. raise ValueError("artifact schema fields do not match")
  123. for field in fields:
  124. dtype = schema[field["name"]]
  125. field_type = field["type"]
  126. matches = (
  127. (field_type == "boolean" and dtype == pl.Boolean)
  128. or (field_type == "date" and dtype == pl.Date)
  129. or (field_type == "string" and dtype == pl.String)
  130. or (field_type == "integer" and dtype.is_integer())
  131. or (field_type == "float" and dtype == pl.Float32)
  132. or (field_type == "double" and dtype == pl.Float64)
  133. or (
  134. field_type == "decimal"
  135. and dtype.is_decimal()
  136. and dtype.precision == field.get("precision")
  137. and dtype.scale == field.get("scale")
  138. )
  139. or (
  140. field_type == "timestamp"
  141. and isinstance(dtype, pl.Datetime)
  142. and dtype.time_zone is None
  143. )
  144. or (
  145. field_type == "timestamptz"
  146. and isinstance(dtype, pl.Datetime)
  147. and dtype.time_zone == field.get("timezone")
  148. )
  149. )
  150. if not matches:
  151. raise ValueError(
  152. f"artifact schema type for {field['name']} does not match"
  153. )
  154. if isinstance(frame, pl.DataFrame):
  155. for field in fields:
  156. if not field["nullable"] and frame[field["name"]].null_count():
  157. raise ValueError(
  158. f"artifact nullable contract for {field['name']} does not match"
  159. )
  160. def _metadata(value: Any) -> dict[str, str]:
  161. if not isinstance(value, Mapping):
  162. raise ValueError("artifact content metadata is missing")
  163. normalized = {}
  164. for key, item in value.items():
  165. name = str(key).lower()
  166. if name.startswith("x-amz-meta-"):
  167. name = name[len("x-amz-meta-") :]
  168. if name in {
  169. "sha256",
  170. "row-count",
  171. "schema-sha256",
  172. "expires-at",
  173. "artifact-bytes",
  174. }:
  175. normalized[name] = str(item)
  176. required = {
  177. "sha256",
  178. "row-count",
  179. "schema-sha256",
  180. "expires-at",
  181. "artifact-bytes",
  182. }
  183. if set(normalized) != required:
  184. raise ValueError("artifact content metadata is incomplete")
  185. if sum(len(key) + len(item) for key, item in normalized.items()) > 2_048:
  186. raise ValueError("artifact content metadata exceeds the safe limit")
  187. return normalized
  188. class ArtifactStore:
  189. """Read and write only server-owned, bounded Parquet artifacts."""
  190. def __init__(
  191. self,
  192. client,
  193. *,
  194. bucket: str,
  195. max_artifact_bytes: int,
  196. max_rows: int,
  197. memory_limit_bytes: int,
  198. max_ttl_seconds: int = 86400,
  199. clock=None,
  200. ):
  201. if not re.fullmatch(r"[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]", bucket):
  202. raise ValueError("artifact bucket name is invalid")
  203. self.client = client
  204. self.bucket = bucket
  205. self.max_artifact_bytes = int(max_artifact_bytes)
  206. self.max_rows = int(max_rows)
  207. self.memory_limit_bytes = int(memory_limit_bytes)
  208. self.max_ttl_seconds = int(max_ttl_seconds)
  209. self.clock = clock or (lambda: datetime.now(UTC))
  210. if (
  211. self.max_artifact_bytes < 1024
  212. or self.max_rows < 1
  213. or self.memory_limit_bytes < self.max_artifact_bytes
  214. or self.max_ttl_seconds < 1
  215. ):
  216. raise ValueError("artifact resource limits are invalid")
  217. if not self.client.bucket_exists(self.bucket):
  218. raise ValueError("artifact bucket does not exist")
  219. def _limits(self, value: Any = None) -> dict[str, int]:
  220. configured = {
  221. "max_rows": self.max_rows,
  222. "max_artifact_bytes": self.max_artifact_bytes,
  223. "memory_limit_bytes": self.memory_limit_bytes,
  224. }
  225. if value is None:
  226. return configured
  227. if not isinstance(value, dict) or set(value) != set(configured):
  228. raise ValueError("artifact limits must have a closed shape")
  229. result = {}
  230. for key, ceiling in configured.items():
  231. item = value[key]
  232. if (
  233. isinstance(item, bool)
  234. or not isinstance(item, int)
  235. or item < 1
  236. or item > ceiling
  237. ):
  238. raise ValueError(
  239. f"artifact {key} exceeds the configured ceiling"
  240. )
  241. result[key] = item
  242. return result
  243. def _parse_ref(self, ref: Any) -> str:
  244. prefix = f"minio://{self.bucket}/"
  245. if not isinstance(ref, str) or not ref.startswith(prefix):
  246. raise ValueError("artifact reference is not owned by this store")
  247. key = ref[len(prefix) :]
  248. match = re.fullmatch(
  249. r"rules/([0-9a-f-]{36})/([0-9a-f-]{36})\.parquet",
  250. key,
  251. )
  252. if match is None:
  253. raise ValueError("artifact reference is invalid")
  254. _uid(match.group(1), "artifact correlation id")
  255. _uid(match.group(2), "artifact id")
  256. return key
  257. def _validated_stat(
  258. self,
  259. key: str,
  260. *,
  261. expected_digest: str | None = None,
  262. limits: dict[str, int] | None = None,
  263. ) -> tuple[Any, dict[str, str]]:
  264. effective = self._limits(limits)
  265. stat = self.client.stat_object(self.bucket, key)
  266. size = int(getattr(stat, "size", -1))
  267. if size < 1 or size > effective["max_artifact_bytes"]:
  268. raise ValueError("artifact size exceeds the configured limit")
  269. if size > effective["memory_limit_bytes"]:
  270. raise ValueError("artifact download exceeds the memory limit")
  271. if str(getattr(stat, "content_type", "")).lower() != PARQUET_CONTENT_TYPE:
  272. raise ValueError("artifact content type is invalid")
  273. metadata = _metadata(getattr(stat, "metadata", None))
  274. digest = metadata["sha256"]
  275. if _DIGEST.fullmatch(digest) is None:
  276. raise ValueError("artifact digest metadata is invalid")
  277. if expected_digest is not None and digest != expected_digest:
  278. raise ValueError("artifact digest does not match")
  279. try:
  280. row_count = int(metadata["row-count"])
  281. metadata_size = int(metadata["artifact-bytes"])
  282. except (TypeError, ValueError) as exc:
  283. raise ValueError("artifact count metadata is invalid") from exc
  284. if row_count < 0 or row_count > effective["max_rows"]:
  285. raise ValueError("artifact row count exceeds the configured limit")
  286. if metadata_size != size:
  287. raise ValueError("artifact size metadata does not match")
  288. if _DIGEST.fullmatch(metadata["schema-sha256"]) is None:
  289. raise ValueError("artifact schema metadata is invalid")
  290. if _parse_timestamp(metadata["expires-at"]) <= _now_utc(self.clock):
  291. raise ValueError("artifact has expired")
  292. return stat, metadata
  293. def write(
  294. self,
  295. frame: pl.LazyFrame | pl.DataFrame,
  296. correlation_id: str,
  297. ttl_seconds: int,
  298. *,
  299. schema_fields: list[dict[str, Any]] | None = None,
  300. limits: dict[str, int] | None = None,
  301. ) -> dict[str, Any]:
  302. effective = self._limits(limits)
  303. correlation = _uid(correlation_id, "correlation_id")
  304. if (
  305. isinstance(ttl_seconds, bool)
  306. or not isinstance(ttl_seconds, int)
  307. or ttl_seconds < 1
  308. or ttl_seconds > self.max_ttl_seconds
  309. ):
  310. raise ValueError("artifact TTL is outside the configured limit")
  311. if isinstance(frame, pl.DataFrame):
  312. lazy = frame.lazy()
  313. elif isinstance(frame, pl.LazyFrame):
  314. lazy = frame
  315. else:
  316. raise ValueError("artifact frame must be a Polars frame")
  317. collected = lazy.head(effective["max_rows"] + 1).collect(
  318. engine="streaming"
  319. )
  320. if collected.height > effective["max_rows"]:
  321. raise ValueError("artifact row count exceeds the configured limit")
  322. if collected.estimated_size() > effective["memory_limit_bytes"]:
  323. raise ValueError("artifact frame exceeds the configured memory limit")
  324. fields, schema_digest = _schema_contract(
  325. schema_fields or _inferred_schema_fields(collected)
  326. )
  327. _validate_frame_schema(collected, fields)
  328. expires_at = _timestamp(
  329. _now_utc(self.clock) + timedelta(seconds=ttl_seconds)
  330. )
  331. artifact_id = new_governance_uid()
  332. key = f"rules/{correlation}/{artifact_id}.parquet"
  333. path = None
  334. uploaded = False
  335. try:
  336. with tempfile.NamedTemporaryFile(
  337. prefix="dataops-rule-artifact-",
  338. suffix=".parquet",
  339. delete=False,
  340. ) as handle:
  341. path = handle.name
  342. collected.write_parquet(path)
  343. size = os.path.getsize(path)
  344. if size < 1 or size > effective["max_artifact_bytes"]:
  345. raise ValueError("artifact size exceeds the configured limit")
  346. if size > effective["memory_limit_bytes"]:
  347. raise ValueError("serialized artifact exceeds the memory limit")
  348. if (
  349. size + collected.estimated_size()
  350. > effective["memory_limit_bytes"]
  351. ):
  352. raise ValueError(
  353. "artifact serialization exceeds the memory limit"
  354. )
  355. digest = hashlib.sha256()
  356. with open(path, "rb") as handle:
  357. while chunk := handle.read(1024 * 1024):
  358. digest.update(chunk)
  359. digest_hex = digest.hexdigest()
  360. with open(path, "rb") as handle:
  361. self.client.put_object(
  362. self.bucket,
  363. key,
  364. handle,
  365. size,
  366. content_type=PARQUET_CONTENT_TYPE,
  367. metadata={
  368. "sha256": digest_hex,
  369. "row-count": str(collected.height),
  370. "schema-sha256": schema_digest,
  371. "expires-at": expires_at,
  372. "artifact-bytes": str(size),
  373. },
  374. )
  375. uploaded = True
  376. self._validated_stat(
  377. key,
  378. expected_digest=digest_hex,
  379. limits=effective,
  380. )
  381. artifact_ref = f"minio://{self.bucket}/{key}"
  382. with self.stage(
  383. artifact_ref,
  384. digest_hex,
  385. expected_schema_fields=fields,
  386. limits=effective,
  387. ):
  388. pass
  389. except Exception:
  390. if uploaded:
  391. with suppress(Exception):
  392. self.client.remove_object(self.bucket, key)
  393. raise
  394. finally:
  395. if path is not None:
  396. with suppress(FileNotFoundError):
  397. os.unlink(path)
  398. return {
  399. "artifact_ref": artifact_ref,
  400. "digest": digest_hex,
  401. "row_count": collected.height,
  402. "schema_hash": schema_digest,
  403. "schema_fields": fields,
  404. "expires_at": expires_at,
  405. }
  406. def write_path(
  407. self,
  408. path: str,
  409. correlation_id: str,
  410. ttl_seconds: int,
  411. *,
  412. schema_fields: list[dict[str, Any]],
  413. limits: dict[str, int] | None = None,
  414. ) -> dict[str, Any]:
  415. """Upload a worker-produced Parquet file without collecting it."""
  416. effective = self._limits(limits)
  417. correlation = _uid(correlation_id, "correlation_id")
  418. if (
  419. isinstance(ttl_seconds, bool)
  420. or not isinstance(ttl_seconds, int)
  421. or ttl_seconds < 1
  422. or ttl_seconds > self.max_ttl_seconds
  423. ):
  424. raise ValueError("artifact TTL is outside the configured limit")
  425. if not isinstance(path, str) or not os.path.isfile(path):
  426. raise ValueError("worker artifact path is invalid")
  427. size = os.path.getsize(path)
  428. if size < 1 or size > effective["max_artifact_bytes"]:
  429. raise ValueError("artifact size exceeds the configured limit")
  430. rows, uncompressed = _parquet_footer_bounds(path)
  431. if rows > effective["max_rows"]:
  432. raise ValueError("artifact row count exceeds the configured limit")
  433. if uncompressed > effective["memory_limit_bytes"]:
  434. raise ValueError(
  435. "Parquet footer uncompressed size exceeds the memory limit"
  436. )
  437. fields, schema_digest = _schema_contract(schema_fields)
  438. _validate_frame_schema(pl.scan_parquet(path), fields)
  439. digest = hashlib.sha256()
  440. with open(path, "rb") as handle:
  441. while chunk := handle.read(1024 * 1024):
  442. digest.update(chunk)
  443. digest_hex = digest.hexdigest()
  444. expires_at = _timestamp(
  445. _now_utc(self.clock) + timedelta(seconds=ttl_seconds)
  446. )
  447. artifact_id = new_governance_uid()
  448. key = f"rules/{correlation}/{artifact_id}.parquet"
  449. uploaded = False
  450. try:
  451. with open(path, "rb") as handle:
  452. self.client.put_object(
  453. self.bucket,
  454. key,
  455. handle,
  456. size,
  457. content_type=PARQUET_CONTENT_TYPE,
  458. metadata={
  459. "sha256": digest_hex,
  460. "row-count": str(rows),
  461. "schema-sha256": schema_digest,
  462. "expires-at": expires_at,
  463. "artifact-bytes": str(size),
  464. },
  465. )
  466. uploaded = True
  467. artifact_ref = f"minio://{self.bucket}/{key}"
  468. with self.stage(
  469. artifact_ref,
  470. digest_hex,
  471. expected_schema_fields=fields,
  472. limits=effective,
  473. ):
  474. pass
  475. except Exception:
  476. if uploaded:
  477. with suppress(Exception):
  478. self.client.remove_object(self.bucket, key)
  479. raise
  480. return {
  481. "artifact_ref": artifact_ref,
  482. "digest": digest_hex,
  483. "row_count": rows,
  484. "schema_hash": schema_digest,
  485. "schema_fields": fields,
  486. "expires_at": expires_at,
  487. }
  488. def describe(self, ref: str) -> dict[str, Any]:
  489. """Return validated object metadata without exposing MinIO credentials."""
  490. key = self._parse_ref(ref)
  491. _stat, metadata = self._validated_stat(key)
  492. return {
  493. "artifact_ref": ref,
  494. "digest": metadata["sha256"],
  495. "row_count": int(metadata["row-count"]),
  496. "schema_hash": metadata["schema-sha256"],
  497. "expires_at": metadata["expires-at"],
  498. }
  499. def read(
  500. self,
  501. ref: str,
  502. expected_digest: str,
  503. *,
  504. expected_schema_fields: list[dict[str, Any]] | None = None,
  505. limits: dict[str, int] | None = None,
  506. ) -> pl.LazyFrame:
  507. with self.stage(
  508. ref,
  509. expected_digest,
  510. expected_schema_fields=expected_schema_fields,
  511. limits=limits,
  512. ) as path:
  513. effective = self._limits(limits)
  514. try:
  515. frame = pl.read_parquet(
  516. path,
  517. n_rows=effective["max_rows"] + 1,
  518. memory_map=False,
  519. )
  520. except Exception as exc:
  521. raise ValueError("artifact is not valid Parquet") from exc
  522. if frame.estimated_size() > effective["memory_limit_bytes"]:
  523. raise ValueError("artifact frame exceeds the configured memory limit")
  524. _validate_frame_schema(
  525. frame, _normalized_schema_fields(expected_schema_fields)
  526. )
  527. return frame.lazy()
  528. @contextmanager
  529. def stage(
  530. self,
  531. ref: str,
  532. expected_digest: str,
  533. *,
  534. expected_schema_fields: list[dict[str, Any]] | None = None,
  535. limits: dict[str, int] | None = None,
  536. ):
  537. effective = self._limits(limits)
  538. if _DIGEST.fullmatch(str(expected_digest or "")) is None:
  539. raise ValueError("expected artifact digest is invalid")
  540. key = self._parse_ref(ref)
  541. _stat, metadata = self._validated_stat(
  542. key,
  543. expected_digest=expected_digest,
  544. limits=effective,
  545. )
  546. if expected_schema_fields is None:
  547. raise ValueError("expected artifact schema fields are required")
  548. fields = _normalized_schema_fields(expected_schema_fields)
  549. if canonical_schema_hash(fields) != metadata["schema-sha256"]:
  550. raise ValueError("artifact schema contract is not expected")
  551. response = self.client.get_object(self.bucket, key)
  552. digest = hashlib.sha256()
  553. path = None
  554. size = 0
  555. try:
  556. try:
  557. with tempfile.NamedTemporaryFile(
  558. prefix="dataops-rule-stage-",
  559. suffix=".parquet",
  560. delete=False,
  561. ) as handle:
  562. path = handle.name
  563. while chunk := response.read(1024 * 1024):
  564. size += len(chunk)
  565. if size > effective["max_artifact_bytes"]:
  566. raise ValueError(
  567. "artifact size exceeds the configured limit"
  568. )
  569. digest.update(chunk)
  570. handle.write(chunk)
  571. finally:
  572. response.close()
  573. release = getattr(response, "release_conn", None)
  574. if callable(release):
  575. release()
  576. if digest.hexdigest() != expected_digest:
  577. raise ValueError("artifact digest does not match content")
  578. footer_rows, uncompressed = _parquet_footer_bounds(path)
  579. if footer_rows != int(metadata["row-count"]):
  580. raise ValueError("artifact row count does not match metadata")
  581. if footer_rows > effective["max_rows"]:
  582. raise ValueError(
  583. "artifact row count exceeds the configured limit"
  584. )
  585. if uncompressed > effective["memory_limit_bytes"]:
  586. raise ValueError(
  587. "Parquet footer uncompressed size exceeds the memory limit"
  588. )
  589. lazy = pl.scan_parquet(path)
  590. _validate_frame_schema(lazy, fields)
  591. yield path
  592. finally:
  593. if path is not None:
  594. with suppress(FileNotFoundError):
  595. os.unlink(path)
  596. def cleanup_expired(self, correlation_id: str) -> int:
  597. correlation = _uid(correlation_id, "correlation_id")
  598. prefix = f"rules/{correlation}/"
  599. removed = 0
  600. for item in self.client.list_objects(
  601. self.bucket,
  602. prefix=prefix,
  603. recursive=True,
  604. ):
  605. key = str(getattr(item, "object_name", ""))
  606. if not key.startswith(prefix):
  607. continue
  608. try:
  609. self._parse_ref(f"minio://{self.bucket}/{key}")
  610. stat = self.client.stat_object(self.bucket, key)
  611. metadata = _metadata(getattr(stat, "metadata", None))
  612. expired = _parse_timestamp(
  613. metadata["expires-at"]
  614. ) <= _now_utc(self.clock)
  615. except ValueError:
  616. continue
  617. if expired:
  618. self.client.remove_object(self.bucket, key)
  619. removed += 1
  620. return removed
  621. def delete(self, ref: str) -> None:
  622. """Delete one exact store-owned artifact after validating its key."""
  623. key = self._parse_ref(ref)
  624. self.client.remove_object(self.bucket, key)
  625. class PostgresArtifactResolver:
  626. """Resolve a canonical artifact binding without accepting caller paths."""
  627. def __init__(self, engine, artifact_store: ArtifactStore):
  628. self.engine = engine
  629. self.artifact_store = artifact_store
  630. def resolve(
  631. self,
  632. *,
  633. binding_id: str,
  634. correlation_id: str,
  635. kind: str,
  636. ) -> dict[str, Any]:
  637. binding = _uid(binding_id, "artifact binding id")
  638. correlation = _uid(correlation_id, "artifact correlation id")
  639. if kind not in {"input", "lookup", "output"}:
  640. raise ValueError("artifact kind is invalid")
  641. statement = text(
  642. """
  643. SELECT
  644. a.artifact_ref,
  645. a.artifact_digest,
  646. a.row_count,
  647. a.schema_hash,
  648. a.schema_fields,
  649. a.expires_at,
  650. b.binding_hash
  651. FROM public.rule_run_artifacts a
  652. JOIN public.dataflow_dataset_bindings b
  653. ON b.id = a.binding_id
  654. WHERE a.binding_id = CAST(:binding_id AS uuid)
  655. AND a.correlation_id = CAST(:correlation_id AS uuid)
  656. AND a.artifact_kind = :artifact_kind
  657. AND a.expires_at > CURRENT_TIMESTAMP
  658. AND b.object_kind = 'parquet_artifact'
  659. AND b.access_mode IN ('read', 'read_write')
  660. ORDER BY a.created_at DESC
  661. LIMIT 1
  662. """
  663. )
  664. with self.engine.connect() as connection:
  665. row = connection.execute(
  666. statement,
  667. {
  668. "binding_id": binding,
  669. "correlation_id": correlation,
  670. "artifact_kind": kind,
  671. },
  672. ).mappings().one_or_none()
  673. if row is None:
  674. raise ValueError("canonical artifact binding was not found")
  675. artifact_ref = str(row["artifact_ref"])
  676. key = self.artifact_store._parse_ref(artifact_ref)
  677. if not key.startswith(f"rules/{correlation}/"):
  678. raise ValueError(
  679. "catalog artifact does not match the execution correlation"
  680. )
  681. described = self.artifact_store.describe(artifact_ref)
  682. row_fields = row["schema_fields"]
  683. if isinstance(row_fields, str):
  684. row_fields = json.loads(row_fields)
  685. if (
  686. described["digest"] != str(row["artifact_digest"])
  687. or described["row_count"] != int(row["row_count"])
  688. or described["schema_hash"] != str(row["schema_hash"])
  689. ):
  690. raise ValueError("catalog artifact metadata does not match storage")
  691. return {
  692. **described,
  693. "schema_fields": _normalized_schema_fields(row_fields),
  694. "binding_hash": str(row["binding_hash"]),
  695. }
  696. def attest_binding(
  697. self,
  698. *,
  699. binding_id: str,
  700. binding_hash: str,
  701. access_mode: str,
  702. ) -> dict[str, str]:
  703. binding = _uid(binding_id, "artifact binding id")
  704. if _DIGEST.fullmatch(str(binding_hash or "")) is None:
  705. raise ValueError("artifact binding hash is invalid")
  706. allowed = {
  707. "read": {"read", "read_write"},
  708. "write": {"write", "read_write"},
  709. }.get(access_mode)
  710. if allowed is None:
  711. raise ValueError("artifact access mode is invalid")
  712. with self.engine.connect() as connection:
  713. row = connection.execute(
  714. text(
  715. """
  716. SELECT binding_hash, access_mode, object_kind
  717. FROM public.dataflow_dataset_bindings
  718. WHERE id = CAST(:binding_id AS uuid)
  719. """
  720. ),
  721. {"binding_id": binding},
  722. ).mappings().one_or_none()
  723. if (
  724. row is None
  725. or row["object_kind"] != "parquet_artifact"
  726. or row["access_mode"] not in allowed
  727. or str(row["binding_hash"]) != binding_hash
  728. ):
  729. raise ValueError("canonical artifact binding no longer matches")
  730. return {"binding_hash": str(row["binding_hash"])}
  731. def register(
  732. self,
  733. *,
  734. binding_id: str,
  735. correlation_id: str,
  736. artifact: dict[str, Any],
  737. kind: str,
  738. binding_hash: str,
  739. ) -> dict[str, Any]:
  740. binding = _uid(binding_id, "artifact binding id")
  741. correlation = _uid(correlation_id, "artifact correlation id")
  742. if kind not in {"input", "lookup", "output"}:
  743. raise ValueError("artifact kind is invalid")
  744. self.attest_binding(
  745. binding_id=binding,
  746. binding_hash=binding_hash,
  747. access_mode="write" if kind == "output" else "read",
  748. )
  749. if not isinstance(artifact, dict):
  750. raise ValueError("artifact metadata is invalid")
  751. artifact_ref = artifact.get("artifact_ref")
  752. key = self.artifact_store._parse_ref(artifact_ref)
  753. if not key.startswith(f"rules/{correlation}/"):
  754. raise ValueError(
  755. "artifact does not match the execution correlation"
  756. )
  757. described = self.artifact_store.describe(artifact_ref)
  758. for key in (
  759. "artifact_ref",
  760. "digest",
  761. "row_count",
  762. "schema_hash",
  763. "expires_at",
  764. ):
  765. if described[key] != artifact.get(key):
  766. raise ValueError("artifact metadata does not match storage")
  767. fields = _normalized_schema_fields(artifact.get("schema_fields"))
  768. if canonical_schema_hash(fields) != described["schema_hash"]:
  769. raise ValueError("artifact schema contract does not match storage")
  770. inserted = None
  771. with self.engine.begin() as connection:
  772. inserted = connection.execute(
  773. text(
  774. """
  775. INSERT INTO public.rule_run_artifacts (
  776. id, correlation_id, binding_id, artifact_ref,
  777. artifact_digest, row_count, schema_hash, schema_fields,
  778. artifact_kind, expires_at
  779. ) VALUES (
  780. CAST(:id AS uuid), CAST(:correlation_id AS uuid),
  781. CAST(:binding_id AS uuid), :artifact_ref,
  782. :artifact_digest, :row_count, :schema_hash,
  783. CAST(:schema_fields AS jsonb), :artifact_kind,
  784. CAST(:expires_at AS timestamptz)
  785. )
  786. ON CONFLICT (
  787. correlation_id, binding_id, artifact_kind
  788. ) DO NOTHING
  789. RETURNING artifact_ref, artifact_digest, row_count,
  790. schema_hash, schema_fields, expires_at
  791. """
  792. ),
  793. {
  794. "id": new_governance_uid(),
  795. "correlation_id": correlation,
  796. "binding_id": binding,
  797. "artifact_ref": described["artifact_ref"],
  798. "artifact_digest": described["digest"],
  799. "row_count": described["row_count"],
  800. "schema_hash": described["schema_hash"],
  801. "schema_fields": json.dumps(
  802. fields,
  803. sort_keys=True,
  804. separators=(",", ":"),
  805. ),
  806. "artifact_kind": kind,
  807. "expires_at": described["expires_at"],
  808. },
  809. ).mappings().one_or_none()
  810. if inserted is None:
  811. inserted = connection.execute(
  812. text(
  813. """
  814. SELECT artifact_ref, artifact_digest, row_count,
  815. schema_hash, schema_fields, expires_at
  816. FROM public.rule_run_artifacts
  817. WHERE correlation_id = CAST(:correlation_id AS uuid)
  818. AND binding_id = CAST(:binding_id AS uuid)
  819. AND artifact_kind = :artifact_kind
  820. FOR UPDATE
  821. """
  822. ),
  823. {
  824. "correlation_id": correlation,
  825. "binding_id": binding,
  826. "artifact_kind": kind,
  827. },
  828. ).mappings().one_or_none()
  829. if inserted is None:
  830. self.artifact_store.delete(artifact_ref)
  831. raise ValueError("immutable artifact catalog handoff was not found")
  832. existing = dict(inserted)
  833. existing_fields = existing["schema_fields"]
  834. if isinstance(existing_fields, str):
  835. existing_fields = json.loads(existing_fields)
  836. registered = {
  837. "artifact_ref": str(existing["artifact_ref"]),
  838. "digest": str(existing["artifact_digest"]),
  839. "row_count": int(existing["row_count"]),
  840. "schema_hash": str(existing["schema_hash"]),
  841. "schema_fields": _normalized_schema_fields(existing_fields),
  842. "expires_at": (
  843. _timestamp(existing["expires_at"])
  844. if isinstance(existing["expires_at"], datetime)
  845. else str(existing["expires_at"])
  846. ),
  847. }
  848. if registered["digest"] != described["digest"]:
  849. self.artifact_store.delete(artifact_ref)
  850. raise ValueError(
  851. "immutable artifact catalog digest conflicts with retry"
  852. )
  853. if registered["artifact_ref"] != artifact_ref:
  854. self.artifact_store.delete(artifact_ref)
  855. stored = self.artifact_store.describe(
  856. registered["artifact_ref"]
  857. )
  858. for key in (
  859. "digest",
  860. "row_count",
  861. "schema_hash",
  862. "expires_at",
  863. ):
  864. if stored[key] != registered[key]:
  865. raise ValueError(
  866. "immutable artifact catalog does not match storage"
  867. )
  868. return registered
  869. def cleanup_expired(self, *, limit: int = 100) -> int:
  870. if (
  871. isinstance(limit, bool)
  872. or not isinstance(limit, int)
  873. or limit < 1
  874. or limit > 1_000
  875. ):
  876. raise ValueError("artifact cleanup limit is invalid")
  877. removed = 0
  878. with self.engine.begin() as connection:
  879. rows = connection.execute(
  880. text(
  881. """
  882. SELECT id::text AS id, artifact_ref
  883. FROM public.rule_run_artifacts
  884. WHERE expires_at <= CURRENT_TIMESTAMP
  885. ORDER BY expires_at, id
  886. LIMIT :limit
  887. FOR UPDATE SKIP LOCKED
  888. """
  889. ),
  890. {"limit": limit},
  891. ).mappings().all()
  892. for row in rows:
  893. self.artifact_store.delete(str(row["artifact_ref"]))
  894. connection.execute(
  895. text(
  896. """
  897. DELETE FROM public.rule_run_artifacts
  898. WHERE id = CAST(:id AS uuid)
  899. """
  900. ),
  901. {"id": str(row["id"])},
  902. )
  903. removed += 1
  904. return removed