"""Digest-bound, bounded Parquet artifacts owned by the DataOps Runner.""" from __future__ import annotations import hashlib import io import json import os import re import tempfile from base64 import urlsafe_b64decode, urlsafe_b64encode from collections.abc import Mapping from contextlib import suppress from datetime import UTC, datetime, timedelta from typing import Any import polars as pl 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}$") 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, str]: fields = _normalized_schema_fields(value) encoded = json.dumps( fields, sort_keys=True, separators=(",", ":"), ensure_ascii=False, ).encode("utf-8") return ( fields, canonical_schema_hash(fields), urlsafe_b64encode(encoded).decode("ascii"), ) def _decode_schema_contract(value: str) -> list[dict[str, Any]]: try: decoded = urlsafe_b64decode(value.encode("ascii")) fields = json.loads(decoded.decode("utf-8")) except Exception as exc: raise ValueError("artifact schema contract metadata is invalid") from exc return _normalized_schema_fields(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", "schema-contract", }: normalized[name] = str(item) required = { "sha256", "row-count", "schema-sha256", "expires-at", "artifact-bytes", "schema-contract", } if set(normalized) != required: raise ValueError("artifact content metadata is incomplete") 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_encoded = _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), "schema-contract": schema_encoded, }, ) uploaded = True self._validated_stat( key, expected_digest=digest_hex, limits=effective, ) artifact_ref = f"minio://{self.bucket}/{key}" self.read( artifact_ref, digest_hex, expected_schema_fields=fields, limits=effective, ) 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 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) fields = _decode_schema_contract(metadata["schema-contract"]) if canonical_schema_hash(fields) != metadata["schema-sha256"]: raise ValueError("artifact schema contract does not match") return { "artifact_ref": ref, "digest": metadata["sha256"], "row_count": int(metadata["row-count"]), "schema_hash": metadata["schema-sha256"], "schema_fields": fields, "expires_at": metadata["expires-at"], } 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: 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, ) fields = _decode_schema_contract(metadata["schema-contract"]) if canonical_schema_hash(fields) != metadata["schema-sha256"]: raise ValueError("artifact schema contract does not match") if expected_schema_fields is not None: expected = _normalized_schema_fields(expected_schema_fields) if expected != fields: raise ValueError("artifact schema contract is not expected") response = self.client.get_object(self.bucket, key) digest = hashlib.sha256() payload = io.BytesIO() size = 0 try: while chunk := response.read(1024 * 1024): size += len(chunk) if 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" ) digest.update(chunk) payload.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") payload.seek(0) try: frame = pl.read_parquet( payload, n_rows=effective["max_rows"] + 1, memory_map=False, ) except Exception as exc: raise ValueError("artifact is not valid Parquet") from exc if frame.height != int(metadata["row-count"]): raise ValueError("artifact row count does not match metadata") if frame.height > effective["max_rows"]: raise ValueError("artifact row count exceeds the configured limit") if frame.estimated_size() > effective["memory_limit_bytes"]: raise ValueError("artifact frame exceeds the configured memory limit") if size + frame.estimated_size() > effective["memory_limit_bytes"]: raise ValueError("artifact decompression exceeds the memory limit") _validate_frame_schema(frame, fields) return frame.lazy() 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) -> dict[str, Any]: binding = _uid(binding_id, "artifact binding id") correlation = _uid(correlation_id, "artifact correlation id") statement = text( """ SELECT a.artifact_ref, a.artifact_digest, a.row_count, a.schema_hash, a.schema_fields, a.expires_at, b.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.expires_at > CURRENT_TIMESTAMP 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, }, ).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"]) or described["schema_fields"] != _normalized_schema_fields(row_fields) ): raise ValueError("catalog artifact metadata does not match storage") return { **described, "binding_hash": str(row["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"])} def register( self, *, binding_id: str, correlation_id: str, artifact: dict[str, Any], kind: str, binding_hash: str, ) -> None: 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") self.attest_binding( binding_id=binding, binding_hash=binding_hash, access_mode="write" if kind == "output" else "read", ) 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" ) described = self.artifact_store.describe(artifact_ref) for key in ( "artifact_ref", "digest", "row_count", "schema_hash", "schema_fields", "expires_at", ): if described[key] != artifact.get(key): raise ValueError("artifact metadata does not match storage") with self.engine.begin() as connection: 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, 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, CAST(:expires_at AS timestamptz) ) ON CONFLICT ( correlation_id, binding_id, artifact_digest ) DO NOTHING """ ), { "id": new_governance_uid(), "correlation_id": correlation, "binding_id": binding, "artifact_ref": described["artifact_ref"], "artifact_digest": described["digest"], "row_count": described["row_count"], "schema_hash": described["schema_hash"], "schema_fields": json.dumps( described["schema_fields"], sort_keys=True, separators=(",", ":"), ), "artifact_kind": kind, "expires_at": described["expires_at"], }, )