| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700 |
- """Digest-bound, bounded Parquet artifacts owned by the DataOps Runner."""
- from __future__ import annotations
- import hashlib
- import json
- import math
- 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}$")
- def _confirmed_object_missing(exc: BaseException) -> bool:
- return isinstance(exc, KeyError) or (
- isinstance(exc, S3Error)
- and exc.code in {"NoSuchKey", "NoSuchObject", "NotFound"}
- )
- 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 = math.ceil(
- (
- _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 Exception as exc:
- if _confirmed_object_missing(exc):
- 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 resolve_handoff(
- self,
- *,
- artifact_ref: str,
- correlation_id: str,
- ) -> dict[str, Any]:
- """Resolve only a ready upstream output from the same execution."""
- correlation = _uid(correlation_id, "artifact correlation id")
- key = self.artifact_store._parse_ref(artifact_ref)
- if not key.startswith(f"rules/{correlation}/"):
- raise ValueError(
- "artifact handoff does not match the execution correlation"
- )
- with self.engine.connect() as connection:
- row = connection.execute(
- text(
- """
- SELECT a.artifact_digest, a.row_count, a.schema_hash,
- a.schema_fields, a.expires_at
- FROM public.rule_run_artifacts a
- JOIN public.dataflow_dataset_bindings b
- ON b.id = a.binding_id
- WHERE a.artifact_ref = :artifact_ref
- AND a.correlation_id =
- CAST(:correlation_id AS uuid)
- AND a.artifact_kind = 'output'
- AND a.handoff_status = 'ready'
- AND a.expires_at > CURRENT_TIMESTAMP
- AND a.binding_hash = b.binding_hash
- """
- ),
- {
- "artifact_ref": artifact_ref,
- "correlation_id": correlation,
- },
- ).mappings().one_or_none()
- if row is None:
- raise ValueError("ready upstream artifact handoff was not found")
- fields = row["schema_fields"]
- if isinstance(fields, str):
- fields = json.loads(fields)
- described = self.artifact_store.describe(artifact_ref)
- 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(
- "upstream artifact handoff metadata does not match"
- )
- return {
- **described,
- "schema_fields": _normalized_schema_fields(fields),
- }
- 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,
- expected_status: str,
- failure_code: str,
- ) -> bool:
- if expected_status not in {"pending", "ready"}:
- raise ValueError("artifact expected handoff status is invalid")
- with self.engine.begin() as connection:
- updated = 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 = :expected_status
- """
- ),
- {
- "id": row_id,
- "expected_status": expected_status,
- "failure_code": failure_code,
- },
- )
- return int(updated.rowcount or 0) == 1
- def _lookup_handoff_by_id(
- self,
- row_id: 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 id = CAST(:id AS uuid)
- """
- ),
- {"id": row_id},
- ).mappings().one_or_none()
- return dict(row) if row is not None else None
- @staticmethod
- def _matching_ready(
- snapshot: Mapping[str, Any],
- current: Mapping[str, Any] | None,
- ) -> bool:
- if current is None or str(current.get("handoff_status")) != "ready":
- return False
- return all(
- str(current.get(key)) == str(snapshot.get(key))
- for key in (
- "id",
- "correlation_id",
- "binding_id",
- "artifact_ref",
- "artifact_digest",
- "row_count",
- "schema_hash",
- "artifact_kind",
- "binding_hash",
- )
- )
- def _verify_reconcile_object(
- self,
- row: Mapping[str, Any],
- ) -> dict[str, Any] | None:
- artifact_ref = str(row["artifact_ref"])
- self.artifact_store._parse_ref(artifact_ref)
- stored = self.artifact_store.describe_optional(artifact_ref)
- if stored is None:
- return None
- if any(
- (
- stored["digest"] != str(row["artifact_digest"]),
- stored["row_count"] != int(row["row_count"]),
- stored["schema_hash"] != str(row["schema_hash"]),
- )
- ):
- raise ValueError("catalog artifact metadata does not match storage")
- fields = row["schema_fields"]
- if isinstance(fields, str):
- fields = json.loads(fields)
- try:
- with self.artifact_store.stage(
- artifact_ref,
- str(row["artifact_digest"]),
- expected_schema_fields=fields,
- ):
- pass
- except Exception as exc:
- if _confirmed_object_missing(exc):
- return None
- raise
- return stored
- 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"])
- try:
- stored = self._verify_reconcile_object(row)
- except ValueError:
- if status == "pending":
- failed = self._mark_failed(
- row_id=row_id,
- expected_status="pending",
- failure_code="pending_object_invalid",
- )
- if failed:
- with suppress(Exception):
- self.artifact_store.delete(row["artifact_ref"])
- else:
- failed = self._mark_failed(
- row_id=row_id,
- expected_status="ready",
- failure_code="ready_object_invalid",
- )
- if failed:
- result["ready_failed"] += 1
- continue
- except Exception:
- # Authentication, timeout, transport, and server failures are
- # not evidence that a cataloged object is invalid.
- continue
- if status == "ready":
- if stored is None:
- failed = self._mark_failed(
- row_id=row_id,
- expected_status="ready",
- failure_code="ready_object_missing",
- )
- if failed:
- 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
- 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:
- current = self._lookup_handoff_by_id(row_id)
- if self._matching_ready(row, current):
- result["pending_finalized"] += 1
- elif (
- current is not None
- and str(current.get("handoff_status")) == "pending"
- ):
- self._mark_failed(
- row_id=row_id,
- expected_status="pending",
- 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[]))
- UNION
- SELECT artifact_ref
- FROM public.rule_violation_samples
- WHERE handoff_status IN ('pending','ready')
- AND 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
|