"""Digest-bound, bounded Parquet artifacts owned by the DataOps Runner.""" from __future__ import annotations import hashlib import json import os import re import tempfile from collections.abc import Mapping from contextlib import contextmanager, suppress from datetime import UTC, datetime, timedelta from typing import Any import polars as pl import pyarrow.parquet as pq from minio.error import S3Error from sqlalchemy import text from app.core.common.identifiers import ( ensure_governance_uid, new_governance_uid, ) from app.core.data_rules.execution_contracts import canonical_schema_hash PARQUET_CONTENT_TYPE = "application/x-parquet" _DIGEST = re.compile(r"^[0-9a-f]{64}$") class ArtifactCommitUnknown(RuntimeError): """A catalog transaction may have committed but cannot be confirmed.""" class ArtifactHandoffPending(RuntimeError): """Another publisher owns the durable pending handoff.""" def _parquet_footer_bounds(path: str) -> tuple[int, int]: try: metadata = pq.ParquetFile(path).metadata except Exception as exc: raise ValueError("artifact Parquet footer is invalid") from exc if metadata is None or metadata.num_row_groups < 1: raise ValueError("artifact Parquet footer is incomplete") uncompressed = sum( metadata.row_group(index).total_byte_size for index in range(metadata.num_row_groups) ) if metadata.num_rows < 0 or uncompressed < 0: raise ValueError("artifact Parquet footer is invalid") return metadata.num_rows, uncompressed def _now_utc(clock) -> datetime: value = clock() if not isinstance(value, datetime): raise ValueError("artifact clock must return a datetime") if value.tzinfo is None: value = value.replace(tzinfo=UTC) return value.astimezone(UTC) def _timestamp(value: datetime) -> str: return value.astimezone(UTC).isoformat().replace("+00:00", "Z") def _parse_timestamp(value: Any) -> datetime: if not isinstance(value, str) or not value.endswith("Z"): raise ValueError("artifact expiry metadata is invalid") try: parsed = datetime.fromisoformat(value[:-1] + "+00:00") except ValueError as exc: raise ValueError("artifact expiry metadata is invalid") from exc return parsed.astimezone(UTC) def _uid(value: Any, label: str) -> str: try: return ensure_governance_uid({"uid": str(value)}) except ValueError as exc: raise ValueError(f"{label} must be a valid UUIDv7") from exc def _inferred_schema_fields( frame: pl.DataFrame | pl.LazyFrame, ) -> list[dict[str, Any]]: schema = ( frame.collect_schema() if isinstance(frame, pl.LazyFrame) else frame.schema ) fields = [] for name, dtype in schema.items(): field: dict[str, Any] = { "name": name, "nullable": True, } if dtype == pl.Boolean: field["type"] = "boolean" elif dtype == pl.Date: field["type"] = "date" elif dtype == pl.String: field["type"] = "string" elif dtype.is_integer(): field["type"] = "integer" elif dtype == pl.Float32: field["type"] = "float" elif dtype == pl.Float64: field["type"] = "double" elif dtype.is_decimal(): field.update( { "type": "decimal", "precision": dtype.precision, "scale": dtype.scale, } ) elif isinstance(dtype, pl.Datetime): field["type"] = ( "timestamptz" if dtype.time_zone else "timestamp" ) if dtype.time_zone: field["timezone"] = dtype.time_zone else: raise ValueError(f"unsupported artifact dtype for {name}") fields.append(field) return sorted(fields, key=lambda item: item["name"]) def _normalized_schema_fields(value: Any) -> list[dict[str, Any]]: canonical_schema_hash(value) return sorted( [dict(field) for field in value], key=lambda item: item["name"], ) def _schema_contract(value: Any) -> tuple[list[dict[str, Any]], str]: fields = _normalized_schema_fields(value) return fields, canonical_schema_hash(fields) def _validate_frame_schema( frame: pl.DataFrame | pl.LazyFrame, fields: list[dict[str, Any]], ) -> None: schema = ( frame.collect_schema() if isinstance(frame, pl.LazyFrame) else frame.schema ) expected_names = {field["name"] for field in fields} if set(schema.names()) != expected_names: raise ValueError("artifact schema fields do not match") for field in fields: dtype = schema[field["name"]] field_type = field["type"] matches = ( (field_type == "boolean" and dtype == pl.Boolean) or (field_type == "date" and dtype == pl.Date) or (field_type == "string" and dtype == pl.String) or (field_type == "integer" and dtype.is_integer()) or (field_type == "float" and dtype == pl.Float32) or (field_type == "double" and dtype == pl.Float64) or ( field_type == "decimal" and dtype.is_decimal() and dtype.precision == field.get("precision") and dtype.scale == field.get("scale") ) or ( field_type == "timestamp" and isinstance(dtype, pl.Datetime) and dtype.time_zone is None ) or ( field_type == "timestamptz" and isinstance(dtype, pl.Datetime) and dtype.time_zone == field.get("timezone") ) ) if not matches: raise ValueError( f"artifact schema type for {field['name']} does not match" ) if isinstance(frame, pl.DataFrame): for field in fields: if not field["nullable"] and frame[field["name"]].null_count(): raise ValueError( f"artifact nullable contract for {field['name']} does not match" ) def _metadata(value: Any) -> dict[str, str]: if not isinstance(value, Mapping): raise ValueError("artifact content metadata is missing") normalized = {} for key, item in value.items(): name = str(key).lower() if name.startswith("x-amz-meta-"): name = name[len("x-amz-meta-") :] if name in { "sha256", "row-count", "schema-sha256", "expires-at", "artifact-bytes", }: normalized[name] = str(item) required = { "sha256", "row-count", "schema-sha256", "expires-at", "artifact-bytes", } if set(normalized) != required: raise ValueError("artifact content metadata is incomplete") if sum(len(key) + len(item) for key, item in normalized.items()) > 2_048: raise ValueError("artifact content metadata exceeds the safe limit") return normalized class ArtifactStore: """Read and write only server-owned, bounded Parquet artifacts.""" def __init__( self, client, *, bucket: str, max_artifact_bytes: int, max_rows: int, memory_limit_bytes: int, max_ttl_seconds: int = 86400, clock=None, ): if not re.fullmatch(r"[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]", bucket): raise ValueError("artifact bucket name is invalid") self.client = client self.bucket = bucket self.max_artifact_bytes = int(max_artifact_bytes) self.max_rows = int(max_rows) self.memory_limit_bytes = int(memory_limit_bytes) self.max_ttl_seconds = int(max_ttl_seconds) self.clock = clock or (lambda: datetime.now(UTC)) if ( self.max_artifact_bytes < 1024 or self.max_rows < 1 or self.memory_limit_bytes < self.max_artifact_bytes or self.max_ttl_seconds < 1 ): raise ValueError("artifact resource limits are invalid") if not self.client.bucket_exists(self.bucket): raise ValueError("artifact bucket does not exist") def _limits(self, value: Any = None) -> dict[str, int]: configured = { "max_rows": self.max_rows, "max_artifact_bytes": self.max_artifact_bytes, "memory_limit_bytes": self.memory_limit_bytes, } if value is None: return configured if not isinstance(value, dict) or set(value) != set(configured): raise ValueError("artifact limits must have a closed shape") result = {} for key, ceiling in configured.items(): item = value[key] if ( isinstance(item, bool) or not isinstance(item, int) or item < 1 or item > ceiling ): raise ValueError( f"artifact {key} exceeds the configured ceiling" ) result[key] = item return result def _parse_ref(self, ref: Any) -> str: prefix = f"minio://{self.bucket}/" if not isinstance(ref, str) or not ref.startswith(prefix): raise ValueError("artifact reference is not owned by this store") key = ref[len(prefix) :] match = re.fullmatch( r"rules/([0-9a-f-]{36})/([0-9a-f-]{36})\.parquet", key, ) if match is None: raise ValueError("artifact reference is invalid") _uid(match.group(1), "artifact correlation id") _uid(match.group(2), "artifact id") return key def _validated_stat( self, key: str, *, expected_digest: str | None = None, limits: dict[str, int] | None = None, ) -> tuple[Any, dict[str, str]]: effective = self._limits(limits) stat = self.client.stat_object(self.bucket, key) size = int(getattr(stat, "size", -1)) if size < 1 or size > effective["max_artifact_bytes"]: raise ValueError("artifact size exceeds the configured limit") if size > effective["memory_limit_bytes"]: raise ValueError("artifact download exceeds the memory limit") if str(getattr(stat, "content_type", "")).lower() != PARQUET_CONTENT_TYPE: raise ValueError("artifact content type is invalid") metadata = _metadata(getattr(stat, "metadata", None)) digest = metadata["sha256"] if _DIGEST.fullmatch(digest) is None: raise ValueError("artifact digest metadata is invalid") if expected_digest is not None and digest != expected_digest: raise ValueError("artifact digest does not match") try: row_count = int(metadata["row-count"]) metadata_size = int(metadata["artifact-bytes"]) except (TypeError, ValueError) as exc: raise ValueError("artifact count metadata is invalid") from exc if row_count < 0 or row_count > effective["max_rows"]: raise ValueError("artifact row count exceeds the configured limit") if metadata_size != size: raise ValueError("artifact size metadata does not match") if _DIGEST.fullmatch(metadata["schema-sha256"]) is None: raise ValueError("artifact schema metadata is invalid") if _parse_timestamp(metadata["expires-at"]) <= _now_utc(self.clock): raise ValueError("artifact has expired") return stat, metadata def write( self, frame: pl.LazyFrame | pl.DataFrame, correlation_id: str, ttl_seconds: int, *, schema_fields: list[dict[str, Any]] | None = None, limits: dict[str, int] | None = None, ) -> dict[str, Any]: effective = self._limits(limits) correlation = _uid(correlation_id, "correlation_id") if ( isinstance(ttl_seconds, bool) or not isinstance(ttl_seconds, int) or ttl_seconds < 1 or ttl_seconds > self.max_ttl_seconds ): raise ValueError("artifact TTL is outside the configured limit") if isinstance(frame, pl.DataFrame): lazy = frame.lazy() elif isinstance(frame, pl.LazyFrame): lazy = frame else: raise ValueError("artifact frame must be a Polars frame") collected = lazy.head(effective["max_rows"] + 1).collect( engine="streaming" ) if collected.height > effective["max_rows"]: raise ValueError("artifact row count exceeds the configured limit") if collected.estimated_size() > effective["memory_limit_bytes"]: raise ValueError("artifact frame exceeds the configured memory limit") fields, schema_digest = _schema_contract( schema_fields or _inferred_schema_fields(collected) ) _validate_frame_schema(collected, fields) expires_at = _timestamp( _now_utc(self.clock) + timedelta(seconds=ttl_seconds) ) artifact_id = new_governance_uid() key = f"rules/{correlation}/{artifact_id}.parquet" path = None uploaded = False try: with tempfile.NamedTemporaryFile( prefix="dataops-rule-artifact-", suffix=".parquet", delete=False, ) as handle: path = handle.name collected.write_parquet(path) size = os.path.getsize(path) if size < 1 or size > effective["max_artifact_bytes"]: raise ValueError("artifact size exceeds the configured limit") if size > effective["memory_limit_bytes"]: raise ValueError("serialized artifact exceeds the memory limit") if ( size + collected.estimated_size() > effective["memory_limit_bytes"] ): raise ValueError( "artifact serialization exceeds the memory limit" ) digest = hashlib.sha256() with open(path, "rb") as handle: while chunk := handle.read(1024 * 1024): digest.update(chunk) digest_hex = digest.hexdigest() with open(path, "rb") as handle: self.client.put_object( self.bucket, key, handle, size, content_type=PARQUET_CONTENT_TYPE, metadata={ "sha256": digest_hex, "row-count": str(collected.height), "schema-sha256": schema_digest, "expires-at": expires_at, "artifact-bytes": str(size), }, ) uploaded = True self._validated_stat( key, expected_digest=digest_hex, limits=effective, ) artifact_ref = f"minio://{self.bucket}/{key}" with self.stage( artifact_ref, digest_hex, expected_schema_fields=fields, limits=effective, ): pass except Exception: if uploaded: with suppress(Exception): self.client.remove_object(self.bucket, key) raise finally: if path is not None: with suppress(FileNotFoundError): os.unlink(path) return { "artifact_ref": artifact_ref, "digest": digest_hex, "row_count": collected.height, "schema_hash": schema_digest, "schema_fields": fields, "expires_at": expires_at, } def prepare_path( self, path: str, correlation_id: str, ttl_seconds: int, *, schema_fields: list[dict[str, Any]], limits: dict[str, int] | None = None, ) -> dict[str, Any]: """Validate a local Parquet file and reserve its server-owned key.""" effective = self._limits(limits) correlation = _uid(correlation_id, "correlation_id") if ( isinstance(ttl_seconds, bool) or not isinstance(ttl_seconds, int) or ttl_seconds < 1 or ttl_seconds > self.max_ttl_seconds ): raise ValueError("artifact TTL is outside the configured limit") if not isinstance(path, str) or not os.path.isfile(path): raise ValueError("worker artifact path is invalid") size = os.path.getsize(path) if size < 1 or size > effective["max_artifact_bytes"]: raise ValueError("artifact size exceeds the configured limit") rows, uncompressed = _parquet_footer_bounds(path) if rows > effective["max_rows"]: raise ValueError("artifact row count exceeds the configured limit") if uncompressed > effective["memory_limit_bytes"]: raise ValueError( "Parquet footer uncompressed size exceeds the memory limit" ) fields, schema_digest = _schema_contract(schema_fields) _validate_frame_schema(pl.scan_parquet(path), fields) digest = hashlib.sha256() with open(path, "rb") as handle: while chunk := handle.read(1024 * 1024): digest.update(chunk) digest_hex = digest.hexdigest() expires_at = _timestamp( _now_utc(self.clock) + timedelta(seconds=ttl_seconds) ) artifact_id = new_governance_uid() key = f"rules/{correlation}/{artifact_id}.parquet" return { "artifact_ref": f"minio://{self.bucket}/{key}", "digest": digest_hex, "row_count": rows, "schema_hash": schema_digest, "schema_fields": fields, "expires_at": expires_at, } def upload_path( self, path: str, artifact: dict[str, Any], *, limits: dict[str, int] | None = None, ) -> dict[str, Any]: """Upload to an already reserved exact key and verify the object.""" if not isinstance(artifact, dict): raise ValueError("prepared artifact metadata is invalid") required = { "artifact_ref", "digest", "row_count", "schema_hash", "schema_fields", "expires_at", } if set(artifact) != required: raise ValueError("prepared artifact metadata has a closed shape") key = self._parse_ref(artifact["artifact_ref"]) correlation = key.split("/", 2)[1] ttl_seconds = int( ( _parse_timestamp(artifact["expires_at"]) - _now_utc(self.clock) ).total_seconds() ) if ttl_seconds < 1 or ttl_seconds > self.max_ttl_seconds: raise ValueError("prepared artifact TTL is invalid") expected = self.prepare_path( path, correlation, ttl_seconds, schema_fields=artifact["schema_fields"], limits=limits, ) for name in ( "digest", "row_count", "schema_hash", "schema_fields", ): if expected[name] != artifact[name]: raise ValueError("prepared artifact no longer matches its path") effective = self._limits(limits) size = os.path.getsize(path) uploaded = False try: with open(path, "rb") as handle: self.client.put_object( self.bucket, key, handle, size, content_type=PARQUET_CONTENT_TYPE, metadata={ "sha256": artifact["digest"], "row-count": str(artifact["row_count"]), "schema-sha256": artifact["schema_hash"], "expires-at": artifact["expires_at"], "artifact-bytes": str(size), }, ) uploaded = True with self.stage( artifact["artifact_ref"], artifact["digest"], expected_schema_fields=artifact["schema_fields"], limits=effective, ): pass except Exception: if uploaded: with suppress(Exception): self.client.remove_object(self.bucket, key) raise return dict(artifact) def write_path( self, path: str, correlation_id: str, ttl_seconds: int, *, schema_fields: list[dict[str, Any]], limits: dict[str, int] | None = None, ) -> dict[str, Any]: """Prepare and upload a non-cataloged compatibility artifact.""" artifact = self.prepare_path( path, correlation_id, ttl_seconds, schema_fields=schema_fields, limits=limits, ) return self.upload_path(path, artifact, limits=limits) def describe(self, ref: str) -> dict[str, Any]: """Return validated object metadata without exposing MinIO credentials.""" key = self._parse_ref(ref) _stat, metadata = self._validated_stat(key) return { "artifact_ref": ref, "digest": metadata["sha256"], "row_count": int(metadata["row-count"]), "schema_hash": metadata["schema-sha256"], "expires_at": metadata["expires-at"], } def describe_optional(self, ref: str) -> dict[str, Any] | None: """Return None only for a confirmed missing object.""" try: return self.describe(ref) except KeyError: return None except S3Error as exc: if exc.code in {"NoSuchKey", "NoSuchObject", "NotFound"}: return None raise def read( self, ref: str, expected_digest: str, *, expected_schema_fields: list[dict[str, Any]] | None = None, limits: dict[str, int] | None = None, ) -> pl.LazyFrame: with self.stage( ref, expected_digest, expected_schema_fields=expected_schema_fields, limits=limits, ) as path: effective = self._limits(limits) try: frame = pl.read_parquet( path, n_rows=effective["max_rows"] + 1, memory_map=False, ) except Exception as exc: raise ValueError("artifact is not valid Parquet") from exc if frame.estimated_size() > effective["memory_limit_bytes"]: raise ValueError("artifact frame exceeds the configured memory limit") _validate_frame_schema( frame, _normalized_schema_fields(expected_schema_fields) ) return frame.lazy() @contextmanager def stage( self, ref: str, expected_digest: str, *, expected_schema_fields: list[dict[str, Any]] | None = None, limits: dict[str, int] | None = None, ): effective = self._limits(limits) if _DIGEST.fullmatch(str(expected_digest or "")) is None: raise ValueError("expected artifact digest is invalid") key = self._parse_ref(ref) _stat, metadata = self._validated_stat( key, expected_digest=expected_digest, limits=effective, ) if expected_schema_fields is None: raise ValueError("expected artifact schema fields are required") fields = _normalized_schema_fields(expected_schema_fields) if canonical_schema_hash(fields) != metadata["schema-sha256"]: raise ValueError("artifact schema contract is not expected") response = self.client.get_object(self.bucket, key) digest = hashlib.sha256() path = None size = 0 try: try: with tempfile.NamedTemporaryFile( prefix="dataops-rule-stage-", suffix=".parquet", delete=False, ) as handle: path = handle.name while chunk := response.read(1024 * 1024): size += len(chunk) if size > effective["max_artifact_bytes"]: raise ValueError( "artifact size exceeds the configured limit" ) digest.update(chunk) handle.write(chunk) finally: response.close() release = getattr(response, "release_conn", None) if callable(release): release() if digest.hexdigest() != expected_digest: raise ValueError("artifact digest does not match content") footer_rows, uncompressed = _parquet_footer_bounds(path) if footer_rows != int(metadata["row-count"]): raise ValueError("artifact row count does not match metadata") if footer_rows > effective["max_rows"]: raise ValueError( "artifact row count exceeds the configured limit" ) if uncompressed > effective["memory_limit_bytes"]: raise ValueError( "Parquet footer uncompressed size exceeds the memory limit" ) lazy = pl.scan_parquet(path) _validate_frame_schema(lazy, fields) yield path finally: if path is not None: with suppress(FileNotFoundError): os.unlink(path) def cleanup_expired(self, correlation_id: str) -> int: correlation = _uid(correlation_id, "correlation_id") prefix = f"rules/{correlation}/" removed = 0 for item in self.client.list_objects( self.bucket, prefix=prefix, recursive=True, ): key = str(getattr(item, "object_name", "")) if not key.startswith(prefix): continue try: self._parse_ref(f"minio://{self.bucket}/{key}") stat = self.client.stat_object(self.bucket, key) metadata = _metadata(getattr(stat, "metadata", None)) expired = _parse_timestamp( metadata["expires-at"] ) <= _now_utc(self.clock) except ValueError: continue if expired: self.client.remove_object(self.bucket, key) removed += 1 return removed def delete(self, ref: str) -> None: """Delete one exact store-owned artifact after validating its key.""" key = self._parse_ref(ref) self.client.remove_object(self.bucket, key) class PostgresArtifactResolver: """Resolve a canonical artifact binding without accepting caller paths.""" def __init__(self, engine, artifact_store: ArtifactStore): self.engine = engine self.artifact_store = artifact_store def resolve( self, *, binding_id: str, correlation_id: str, kind: str, ) -> dict[str, Any]: binding = _uid(binding_id, "artifact binding id") correlation = _uid(correlation_id, "artifact correlation id") if kind not in {"input", "lookup", "output"}: raise ValueError("artifact kind is invalid") statement = text( """ SELECT a.artifact_ref, a.artifact_digest, a.row_count, a.schema_hash, a.schema_fields, a.expires_at, a.binding_hash AS catalog_binding_hash, b.binding_hash AS current_binding_hash FROM public.rule_run_artifacts a JOIN public.dataflow_dataset_bindings b ON b.id = a.binding_id WHERE a.binding_id = CAST(:binding_id AS uuid) AND a.correlation_id = CAST(:correlation_id AS uuid) AND a.artifact_kind = :artifact_kind AND a.handoff_status = 'ready' AND a.expires_at > CURRENT_TIMESTAMP AND a.binding_hash = b.binding_hash AND b.object_kind = 'parquet_artifact' AND b.access_mode IN ('read', 'read_write') ORDER BY a.created_at DESC LIMIT 1 """ ) with self.engine.connect() as connection: row = connection.execute( statement, { "binding_id": binding, "correlation_id": correlation, "artifact_kind": kind, }, ).mappings().one_or_none() if row is None: raise ValueError("canonical artifact binding was not found") artifact_ref = str(row["artifact_ref"]) key = self.artifact_store._parse_ref(artifact_ref) if not key.startswith(f"rules/{correlation}/"): raise ValueError( "catalog artifact does not match the execution correlation" ) described = self.artifact_store.describe(artifact_ref) row_fields = row["schema_fields"] if isinstance(row_fields, str): row_fields = json.loads(row_fields) if ( described["digest"] != str(row["artifact_digest"]) or described["row_count"] != int(row["row_count"]) or described["schema_hash"] != str(row["schema_hash"]) ): raise ValueError("catalog artifact metadata does not match storage") return { **described, "schema_fields": _normalized_schema_fields(row_fields), "binding_hash": str(row["catalog_binding_hash"]), } def attest_binding( self, *, binding_id: str, binding_hash: str, access_mode: str, ) -> dict[str, str]: binding = _uid(binding_id, "artifact binding id") if _DIGEST.fullmatch(str(binding_hash or "")) is None: raise ValueError("artifact binding hash is invalid") allowed = { "read": {"read", "read_write"}, "write": {"write", "read_write"}, }.get(access_mode) if allowed is None: raise ValueError("artifact access mode is invalid") with self.engine.connect() as connection: row = connection.execute( text( """ SELECT binding_hash, access_mode, object_kind FROM public.dataflow_dataset_bindings WHERE id = CAST(:binding_id AS uuid) """ ), {"binding_id": binding}, ).mappings().one_or_none() if ( row is None or row["object_kind"] != "parquet_artifact" or row["access_mode"] not in allowed or str(row["binding_hash"]) != binding_hash ): raise ValueError("canonical artifact binding no longer matches") return {"binding_hash": str(row["binding_hash"])} @staticmethod def _catalog_artifact(row: Mapping[str, Any]) -> dict[str, Any]: fields = row["schema_fields"] if isinstance(fields, str): fields = json.loads(fields) expires_at = row["expires_at"] return { "artifact_ref": str(row["artifact_ref"]), "digest": str(row["artifact_digest"]), "row_count": int(row["row_count"]), "schema_hash": str(row["schema_hash"]), "schema_fields": _normalized_schema_fields(fields), "expires_at": ( _timestamp(expires_at) if isinstance(expires_at, datetime) else str(expires_at) ), } @staticmethod def _attest_binding_locked( connection, *, binding_id: str, binding_hash: str, kind: str, ) -> None: row = connection.execute( text( """ SELECT binding_hash, access_mode, object_kind FROM public.dataflow_dataset_bindings WHERE id = CAST(:binding_id AS uuid) FOR SHARE """ ), {"binding_id": binding_id}, ).mappings().one_or_none() allowed = ( {"write", "read_write"} if kind == "output" else {"read", "read_write"} ) if ( row is None or row["object_kind"] != "parquet_artifact" or row["access_mode"] not in allowed or str(row["binding_hash"]) != binding_hash ): raise ValueError("canonical artifact binding no longer matches") def _lookup_handoff( self, *, correlation_id: str, binding_id: str, kind: str, ) -> dict[str, Any] | None: with self.engine.connect() as connection: row = connection.execute( text( """ SELECT id::text AS id, correlation_id::text, binding_id::text, artifact_ref, artifact_digest, row_count, schema_hash, schema_fields, artifact_kind, binding_hash, expires_at, handoff_status FROM public.rule_run_artifacts WHERE correlation_id = CAST(:correlation_id AS uuid) AND binding_id = CAST(:binding_id AS uuid) AND artifact_kind = :artifact_kind """ ), { "correlation_id": correlation_id, "binding_id": binding_id, "artifact_kind": kind, }, ).mappings().one_or_none() return dict(row) if row is not None else None def reserve( self, *, binding_id: str, correlation_id: str, artifact: dict[str, Any], kind: str, binding_hash: str, ) -> dict[str, Any]: """Atomically attest the binding and reserve one pending handoff.""" binding = _uid(binding_id, "artifact binding id") correlation = _uid(correlation_id, "artifact correlation id") if kind not in {"input", "lookup", "output"}: raise ValueError("artifact kind is invalid") if _DIGEST.fullmatch(str(binding_hash or "")) is None: raise ValueError("artifact binding hash is invalid") if not isinstance(artifact, dict): raise ValueError("artifact metadata is invalid") artifact_ref = artifact.get("artifact_ref") key = self.artifact_store._parse_ref(artifact_ref) if not key.startswith(f"rules/{correlation}/"): raise ValueError( "artifact does not match the execution correlation" ) fields = _normalized_schema_fields(artifact.get("schema_fields")) if canonical_schema_hash(fields) != artifact.get("schema_hash"): raise ValueError("artifact schema contract does not match") if _DIGEST.fullmatch(str(artifact.get("digest") or "")) is None: raise ValueError("artifact digest is invalid") reservation_id = new_governance_uid() parameters = { "id": reservation_id, "correlation_id": correlation, "binding_id": binding, "artifact_ref": artifact_ref, "artifact_digest": artifact["digest"], "row_count": int(artifact["row_count"]), "schema_hash": artifact["schema_hash"], "schema_fields": json.dumps( fields, sort_keys=True, separators=(",", ":"), ), "artifact_kind": kind, "binding_hash": binding_hash, "expires_at": artifact["expires_at"], } selected = None inserted = False try: with self.engine.begin() as connection: self._attest_binding_locked( connection, binding_id=binding, binding_hash=binding_hash, kind=kind, ) selected = connection.execute( text( """ INSERT INTO public.rule_run_artifacts ( id, correlation_id, binding_id, artifact_ref, artifact_digest, row_count, schema_hash, schema_fields, artifact_kind, binding_hash, handoff_status, expires_at ) VALUES ( CAST(:id AS uuid), CAST(:correlation_id AS uuid), CAST(:binding_id AS uuid), :artifact_ref, :artifact_digest, :row_count, :schema_hash, CAST(:schema_fields AS jsonb), :artifact_kind, :binding_hash, 'pending', CAST(:expires_at AS timestamptz) ) ON CONFLICT ( correlation_id, binding_id, artifact_kind ) DO NOTHING RETURNING id::text AS id, correlation_id::text, binding_id::text, artifact_ref, artifact_digest, row_count, schema_hash, schema_fields, artifact_kind, binding_hash, expires_at, handoff_status """ ), parameters, ).mappings().one_or_none() inserted = selected is not None if selected is None: selected = connection.execute( text( """ SELECT id::text AS id, correlation_id::text, binding_id::text, artifact_ref, artifact_digest, row_count, schema_hash, schema_fields, artifact_kind, binding_hash, expires_at, handoff_status FROM public.rule_run_artifacts WHERE correlation_id = CAST(:correlation_id AS uuid) AND binding_id = CAST(:binding_id AS uuid) AND artifact_kind = :artifact_kind FOR UPDATE """ ), parameters, ).mappings().one_or_none() except ValueError: raise except Exception as exc: try: selected = self._lookup_handoff( correlation_id=correlation, binding_id=binding, kind=kind, ) except Exception as recheck_exc: raise ArtifactCommitUnknown( "artifact reservation commit outcome is unknown" ) from recheck_exc if ( selected is None or str(selected["artifact_ref"]) != artifact_ref or str(selected["artifact_digest"]) != artifact["digest"] or str(selected["binding_hash"]) != binding_hash ): raise ArtifactCommitUnknown( "artifact reservation commit outcome is unknown" ) from exc inserted = True if selected is None: raise ArtifactCommitUnknown( "artifact reservation commit outcome is unknown" ) row = dict(selected) if str(row["binding_hash"]) != binding_hash: raise ValueError("artifact reservation binding hash conflicts") if str(row["artifact_digest"]) != artifact["digest"]: raise ValueError( "immutable artifact catalog digest conflicts with retry" ) status = str(row["handoff_status"]) if not inserted: if status == "ready": return { **self._catalog_artifact(row), "correlation_id": correlation, "reservation_id": str(row["id"]), "handoff_status": "ready", "upload_required": False, } if status == "pending": raise ArtifactHandoffPending( "artifact handoff is already pending" ) raise ValueError("artifact handoff has failed") return { **self._catalog_artifact(row), "correlation_id": correlation, "reservation_id": str(row["id"]), "handoff_status": status, "upload_required": status == "pending", } def _abort_pending(self, reservation_id: str) -> None: with self.engine.begin() as connection: connection.execute( text( """ DELETE FROM public.rule_run_artifacts WHERE id = CAST(:id AS uuid) AND handoff_status = 'pending' """ ), {"id": reservation_id}, ) def finalize( self, *, reservation: dict[str, Any], binding_id: str, binding_hash: str, kind: str, ) -> dict[str, Any]: reservation_id = _uid( reservation.get("reservation_id"), "artifact reservation id" ) binding = _uid(binding_id, "artifact binding id") selected = None try: with self.engine.begin() as connection: self._attest_binding_locked( connection, binding_id=binding, binding_hash=binding_hash, kind=kind, ) selected = connection.execute( text( """ UPDATE public.rule_run_artifacts SET handoff_status = 'ready', ready_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP, failure_code = NULL, failed_at = NULL WHERE id = CAST(:id AS uuid) AND binding_id = CAST(:binding_id AS uuid) AND binding_hash = :binding_hash AND artifact_digest = :artifact_digest AND handoff_status = 'pending' RETURNING id::text AS id, correlation_id::text, binding_id::text, artifact_ref, artifact_digest, row_count, schema_hash, schema_fields, artifact_kind, binding_hash, expires_at, handoff_status """ ), { "id": reservation_id, "binding_id": binding, "binding_hash": binding_hash, "artifact_digest": reservation["digest"], }, ).mappings().one_or_none() if selected is None: raise ValueError( "pending artifact handoff no longer matches" ) except ValueError: raise except Exception as exc: try: selected = self._lookup_handoff( correlation_id=_uid( reservation["correlation_id"], "artifact correlation id", ), binding_id=binding, kind=kind, ) except Exception as recheck_exc: raise ArtifactCommitUnknown( "artifact finalize commit outcome is unknown" ) from recheck_exc if ( selected is None or str(selected["id"]) != reservation_id or str(selected["artifact_digest"]) != reservation["digest"] or str(selected["binding_hash"]) != binding_hash or str(selected["handoff_status"]) != "ready" ): raise ArtifactCommitUnknown( "artifact finalize commit outcome is unknown" ) from exc return self._catalog_artifact(selected) def publish_path( self, path: str, *, binding_id: str, binding_hash: str, correlation_id: str, kind: str, ttl_seconds: int, schema_fields: list[dict[str, Any]], limits: dict[str, int] | None = None, ) -> dict[str, Any]: """Reserve, upload, and finalize one durable artifact handoff.""" prepared = self.artifact_store.prepare_path( path, correlation_id, ttl_seconds, schema_fields=schema_fields, limits=limits, ) reservation = self.reserve( binding_id=binding_id, correlation_id=correlation_id, artifact=prepared, kind=kind, binding_hash=binding_hash, ) if not reservation["upload_required"]: stored = self.artifact_store.describe( reservation["artifact_ref"] ) if stored["digest"] != reservation["digest"]: raise ValueError( "ready artifact catalog does not match storage" ) return { key: reservation[key] for key in ( "artifact_ref", "digest", "row_count", "schema_hash", "schema_fields", "expires_at", ) } reserved_artifact = { key: reservation[key] for key in ( "artifact_ref", "digest", "row_count", "schema_hash", "schema_fields", "expires_at", ) } try: self.artifact_store.upload_path( path, reserved_artifact, limits=limits, ) except Exception: with suppress(Exception): self.artifact_store.delete( reserved_artifact["artifact_ref"] ) with suppress(Exception): self._abort_pending(reservation["reservation_id"]) raise return self.finalize( reservation=reservation, binding_id=binding_id, binding_hash=binding_hash, kind=kind, ) def _mark_failed( self, *, row_id: str, failure_code: str, ) -> None: with self.engine.begin() as connection: connection.execute( text( """ UPDATE public.rule_run_artifacts SET handoff_status = 'failed', failed_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP, failure_code = :failure_code WHERE id = CAST(:id AS uuid) AND handoff_status IN ('pending','ready') """ ), {"id": row_id, "failure_code": failure_code}, ) def reconcile( self, *, limit: int = 100, grace_seconds: int = 300, ) -> dict[str, int]: """Repair bounded catalog/store drift after the grace period.""" if ( isinstance(limit, bool) or not isinstance(limit, int) or limit < 1 or limit > 1_000 ): raise ValueError("artifact reconciliation limit is invalid") if ( isinstance(grace_seconds, bool) or not isinstance(grace_seconds, int) or grace_seconds < 30 or grace_seconds > 86_400 ): raise ValueError("artifact reconciliation grace is invalid") result = { "pending_finalized": 0, "pending_deleted": 0, "ready_failed": 0, "orphans_deleted": 0, } with self.engine.connect() as connection: rows = connection.execute( text( """ SELECT id::text AS id, correlation_id::text, binding_id::text, artifact_ref, artifact_digest, row_count, schema_hash, schema_fields, artifact_kind, binding_hash, expires_at, handoff_status FROM public.rule_run_artifacts WHERE handoff_status IN ('pending','ready') AND updated_at <= CURRENT_TIMESTAMP - make_interval(secs => :grace_seconds) ORDER BY updated_at, id LIMIT :limit """ ), { "grace_seconds": grace_seconds, "limit": limit, }, ).mappings().all() for raw_row in rows: row = dict(raw_row) row_id = str(row["id"]) status = str(row["handoff_status"]) invalid_object = False try: self.artifact_store._parse_ref(row["artifact_ref"]) stored = self.artifact_store.describe_optional( row["artifact_ref"] ) except Exception: invalid_object = True stored = None if invalid_object: if status == "pending": with suppress(Exception): self.artifact_store.delete(row["artifact_ref"]) self._mark_failed( row_id=row_id, failure_code="pending_object_invalid", ) else: self._mark_failed( row_id=row_id, failure_code="ready_object_invalid", ) result["ready_failed"] += 1 continue if status == "ready": if stored is None: self._mark_failed( row_id=row_id, failure_code="ready_object_missing", ) result["ready_failed"] += 1 continue if stored is None: with self.engine.begin() as connection: deleted = connection.execute( text( """ DELETE FROM public.rule_run_artifacts WHERE id = CAST(:id AS uuid) AND handoff_status = 'pending' """ ), {"id": row_id}, ) if int(deleted.rowcount or 0) == 1: result["pending_deleted"] += 1 continue if any( ( stored["digest"] != str(row["artifact_digest"]), stored["row_count"] != int(row["row_count"]), stored["schema_hash"] != str(row["schema_hash"]), ) ): with suppress(Exception): self.artifact_store.delete(row["artifact_ref"]) self._mark_failed( row_id=row_id, failure_code="pending_object_invalid", ) continue with self.engine.begin() as connection: finalized = connection.execute( text( """ UPDATE public.rule_run_artifacts a SET handoff_status = 'ready', ready_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP, failure_code = NULL, failed_at = NULL WHERE a.id = CAST(:id AS uuid) AND a.handoff_status = 'pending' AND EXISTS ( SELECT 1 FROM public.dataflow_dataset_bindings b WHERE b.id = a.binding_id AND b.binding_hash = a.binding_hash AND b.object_kind = 'parquet_artifact' AND b.access_mode IN ( 'read','write','read_write' ) ) RETURNING a.id """ ), {"id": row_id}, ) if int(finalized.rowcount or 0) == 1: result["pending_finalized"] += 1 else: self._mark_failed( row_id=row_id, failure_code="pending_binding_changed", ) remaining = limit - len(rows) if remaining <= 0: return result now = _now_utc(self.artifact_store.clock) candidates = [] scanned = 0 for item in self.artifact_store.client.list_objects( self.artifact_store.bucket, prefix="rules/", recursive=True, ): scanned += 1 if scanned > limit * 10 or len(candidates) >= remaining: break key = str(getattr(item, "object_name", "")) ref = f"minio://{self.artifact_store.bucket}/{key}" try: self.artifact_store._parse_ref(ref) except ValueError: continue modified = getattr(item, "last_modified", None) if not isinstance(modified, datetime): continue if modified.tzinfo is None: modified = modified.replace(tzinfo=UTC) if modified.astimezone(UTC) > now - timedelta( seconds=grace_seconds ): continue candidates.append(ref) if not candidates: return result with self.engine.connect() as connection: referenced = { str(row["artifact_ref"]) for row in connection.execute( text( """ SELECT artifact_ref FROM public.rule_run_artifacts WHERE artifact_ref = ANY(CAST(:artifact_refs AS text[])) """ ), {"artifact_refs": candidates}, ).mappings().all() } for ref in candidates: if ref in referenced: continue self.artifact_store.delete(ref) result["orphans_deleted"] += 1 return result def cleanup_expired(self, *, limit: int = 100) -> int: if ( isinstance(limit, bool) or not isinstance(limit, int) or limit < 1 or limit > 1_000 ): raise ValueError("artifact cleanup limit is invalid") removed = 0 with self.engine.begin() as connection: rows = connection.execute( text( """ SELECT id::text AS id, artifact_ref FROM public.rule_run_artifacts WHERE expires_at <= CURRENT_TIMESTAMP ORDER BY expires_at, id LIMIT :limit FOR UPDATE SKIP LOCKED """ ), {"limit": limit}, ).mappings().all() for row in rows: self.artifact_store.delete(str(row["artifact_ref"])) connection.execute( text( """ DELETE FROM public.rule_run_artifacts WHERE id = CAST(:id AS uuid) """ ), {"id": str(row["id"])}, ) removed += 1 return removed