| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638 |
- """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 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[]))
- """
- ),
- {"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
|