|
|
@@ -14,6 +14,7 @@ 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 (
|
|
|
@@ -26,6 +27,14 @@ PARQUET_CONTENT_TYPE = "application/x-parquet"
|
|
|
_DIGEST = re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
|
|
|
|
|
+class ArtifactCommitUnknown(RuntimeError):
|
|
|
+ """A catalog transaction may have committed but cannot be confirmed."""
|
|
|
+
|
|
|
+
|
|
|
+class ArtifactHandoffPending(RuntimeError):
|
|
|
+ """Another publisher owns the durable pending handoff."""
|
|
|
+
|
|
|
+
|
|
|
def _parquet_footer_bounds(path: str) -> tuple[int, int]:
|
|
|
try:
|
|
|
metadata = pq.ParquetFile(path).metadata
|
|
|
@@ -436,7 +445,7 @@ class ArtifactStore:
|
|
|
"expires_at": expires_at,
|
|
|
}
|
|
|
|
|
|
- def write_path(
|
|
|
+ def prepare_path(
|
|
|
self,
|
|
|
path: str,
|
|
|
correlation_id: str,
|
|
|
@@ -445,7 +454,7 @@ class ArtifactStore:
|
|
|
schema_fields: list[dict[str, Any]],
|
|
|
limits: dict[str, int] | None = None,
|
|
|
) -> dict[str, Any]:
|
|
|
- """Upload a worker-produced Parquet file without collecting it."""
|
|
|
+ """Validate a local Parquet file and reserve its server-owned key."""
|
|
|
|
|
|
effective = self._limits(limits)
|
|
|
correlation = _uid(correlation_id, "correlation_id")
|
|
|
@@ -480,6 +489,63 @@ class ArtifactStore:
|
|
|
)
|
|
|
artifact_id = new_governance_uid()
|
|
|
key = f"rules/{correlation}/{artifact_id}.parquet"
|
|
|
+ return {
|
|
|
+ "artifact_ref": f"minio://{self.bucket}/{key}",
|
|
|
+ "digest": digest_hex,
|
|
|
+ "row_count": rows,
|
|
|
+ "schema_hash": schema_digest,
|
|
|
+ "schema_fields": fields,
|
|
|
+ "expires_at": expires_at,
|
|
|
+ }
|
|
|
+
|
|
|
+ def upload_path(
|
|
|
+ self,
|
|
|
+ path: str,
|
|
|
+ artifact: dict[str, Any],
|
|
|
+ *,
|
|
|
+ limits: dict[str, int] | None = None,
|
|
|
+ ) -> dict[str, Any]:
|
|
|
+ """Upload to an already reserved exact key and verify the object."""
|
|
|
+
|
|
|
+ if not isinstance(artifact, dict):
|
|
|
+ raise ValueError("prepared artifact metadata is invalid")
|
|
|
+ required = {
|
|
|
+ "artifact_ref",
|
|
|
+ "digest",
|
|
|
+ "row_count",
|
|
|
+ "schema_hash",
|
|
|
+ "schema_fields",
|
|
|
+ "expires_at",
|
|
|
+ }
|
|
|
+ if set(artifact) != required:
|
|
|
+ raise ValueError("prepared artifact metadata has a closed shape")
|
|
|
+ key = self._parse_ref(artifact["artifact_ref"])
|
|
|
+ correlation = key.split("/", 2)[1]
|
|
|
+ ttl_seconds = int(
|
|
|
+ (
|
|
|
+ _parse_timestamp(artifact["expires_at"])
|
|
|
+ - _now_utc(self.clock)
|
|
|
+ ).total_seconds()
|
|
|
+ )
|
|
|
+ if ttl_seconds < 1 or ttl_seconds > self.max_ttl_seconds:
|
|
|
+ raise ValueError("prepared artifact TTL is invalid")
|
|
|
+ expected = self.prepare_path(
|
|
|
+ path,
|
|
|
+ correlation,
|
|
|
+ ttl_seconds,
|
|
|
+ schema_fields=artifact["schema_fields"],
|
|
|
+ limits=limits,
|
|
|
+ )
|
|
|
+ for name in (
|
|
|
+ "digest",
|
|
|
+ "row_count",
|
|
|
+ "schema_hash",
|
|
|
+ "schema_fields",
|
|
|
+ ):
|
|
|
+ if expected[name] != artifact[name]:
|
|
|
+ raise ValueError("prepared artifact no longer matches its path")
|
|
|
+ effective = self._limits(limits)
|
|
|
+ size = os.path.getsize(path)
|
|
|
uploaded = False
|
|
|
try:
|
|
|
with open(path, "rb") as handle:
|
|
|
@@ -490,19 +556,18 @@ class ArtifactStore:
|
|
|
size,
|
|
|
content_type=PARQUET_CONTENT_TYPE,
|
|
|
metadata={
|
|
|
- "sha256": digest_hex,
|
|
|
- "row-count": str(rows),
|
|
|
- "schema-sha256": schema_digest,
|
|
|
- "expires-at": expires_at,
|
|
|
+ "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
|
|
|
- artifact_ref = f"minio://{self.bucket}/{key}"
|
|
|
with self.stage(
|
|
|
- artifact_ref,
|
|
|
- digest_hex,
|
|
|
- expected_schema_fields=fields,
|
|
|
+ artifact["artifact_ref"],
|
|
|
+ artifact["digest"],
|
|
|
+ expected_schema_fields=artifact["schema_fields"],
|
|
|
limits=effective,
|
|
|
):
|
|
|
pass
|
|
|
@@ -511,14 +576,27 @@ class ArtifactStore:
|
|
|
with suppress(Exception):
|
|
|
self.client.remove_object(self.bucket, key)
|
|
|
raise
|
|
|
- return {
|
|
|
- "artifact_ref": artifact_ref,
|
|
|
- "digest": digest_hex,
|
|
|
- "row_count": rows,
|
|
|
- "schema_hash": schema_digest,
|
|
|
- "schema_fields": fields,
|
|
|
- "expires_at": expires_at,
|
|
|
- }
|
|
|
+ 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."""
|
|
|
@@ -533,6 +611,18 @@ class ArtifactStore:
|
|
|
"expires_at": metadata["expires-at"],
|
|
|
}
|
|
|
|
|
|
+ def describe_optional(self, ref: str) -> dict[str, Any] | None:
|
|
|
+ """Return None only for a confirmed missing object."""
|
|
|
+
|
|
|
+ try:
|
|
|
+ return self.describe(ref)
|
|
|
+ except KeyError:
|
|
|
+ return None
|
|
|
+ except S3Error as exc:
|
|
|
+ if exc.code in {"NoSuchKey", "NoSuchObject", "NotFound"}:
|
|
|
+ return None
|
|
|
+ raise
|
|
|
+
|
|
|
def read(
|
|
|
self,
|
|
|
ref: str,
|
|
|
@@ -692,14 +782,17 @@ class PostgresArtifactResolver:
|
|
|
a.schema_hash,
|
|
|
a.schema_fields,
|
|
|
a.expires_at,
|
|
|
- b.binding_hash
|
|
|
+ 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
|
|
|
@@ -736,7 +829,7 @@ class PostgresArtifactResolver:
|
|
|
return {
|
|
|
**described,
|
|
|
"schema_fields": _normalized_schema_fields(row_fields),
|
|
|
- "binding_hash": str(row["binding_hash"]),
|
|
|
+ "binding_hash": str(row["catalog_binding_hash"]),
|
|
|
}
|
|
|
|
|
|
def attest_binding(
|
|
|
@@ -775,7 +868,88 @@ class PostgresArtifactResolver:
|
|
|
raise ValueError("canonical artifact binding no longer matches")
|
|
|
return {"binding_hash": str(row["binding_hash"])}
|
|
|
|
|
|
- def register(
|
|
|
+ @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,
|
|
|
@@ -784,15 +958,14 @@ class PostgresArtifactResolver:
|
|
|
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")
|
|
|
- self.attest_binding(
|
|
|
- binding_id=binding,
|
|
|
- binding_hash=binding_hash,
|
|
|
- access_mode="write" if kind == "output" else "read",
|
|
|
- )
|
|
|
+ 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")
|
|
|
@@ -801,118 +974,544 @@ class PostgresArtifactResolver:
|
|
|
raise ValueError(
|
|
|
"artifact does not match the execution correlation"
|
|
|
)
|
|
|
- described = self.artifact_store.describe(artifact_ref)
|
|
|
- for key in (
|
|
|
- "artifact_ref",
|
|
|
- "digest",
|
|
|
- "row_count",
|
|
|
- "schema_hash",
|
|
|
- "expires_at",
|
|
|
- ):
|
|
|
- if described[key] != artifact.get(key):
|
|
|
- raise ValueError("artifact metadata does not match storage")
|
|
|
fields = _normalized_schema_fields(artifact.get("schema_fields"))
|
|
|
- if canonical_schema_hash(fields) != described["schema_hash"]:
|
|
|
- raise ValueError("artifact schema contract does not match storage")
|
|
|
- inserted = None
|
|
|
+ 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:
|
|
|
- inserted = connection.execute(
|
|
|
+ connection.execute(
|
|
|
text(
|
|
|
"""
|
|
|
- INSERT INTO public.rule_run_artifacts (
|
|
|
- id, correlation_id, binding_id, artifact_ref,
|
|
|
- artifact_digest, row_count, schema_hash, schema_fields,
|
|
|
- artifact_kind, expires_at
|
|
|
- ) VALUES (
|
|
|
- CAST(:id AS uuid), CAST(:correlation_id AS uuid),
|
|
|
- CAST(:binding_id AS uuid), :artifact_ref,
|
|
|
- :artifact_digest, :row_count, :schema_hash,
|
|
|
- CAST(:schema_fields AS jsonb), :artifact_kind,
|
|
|
- CAST(:expires_at AS timestamptz)
|
|
|
- )
|
|
|
- ON CONFLICT (
|
|
|
- correlation_id, binding_id, artifact_kind
|
|
|
- ) DO NOTHING
|
|
|
- RETURNING artifact_ref, artifact_digest, row_count,
|
|
|
- schema_hash, schema_fields, expires_at
|
|
|
+ DELETE FROM public.rule_run_artifacts
|
|
|
+ WHERE id = CAST(:id AS uuid)
|
|
|
+ AND handoff_status = 'pending'
|
|
|
"""
|
|
|
),
|
|
|
- {
|
|
|
- "id": new_governance_uid(),
|
|
|
- "correlation_id": correlation,
|
|
|
- "binding_id": binding,
|
|
|
- "artifact_ref": described["artifact_ref"],
|
|
|
- "artifact_digest": described["digest"],
|
|
|
- "row_count": described["row_count"],
|
|
|
- "schema_hash": described["schema_hash"],
|
|
|
- "schema_fields": json.dumps(
|
|
|
- fields,
|
|
|
- sort_keys=True,
|
|
|
- separators=(",", ":"),
|
|
|
- ),
|
|
|
- "artifact_kind": kind,
|
|
|
- "expires_at": described["expires_at"],
|
|
|
- },
|
|
|
- ).mappings().one_or_none()
|
|
|
- if inserted is None:
|
|
|
- inserted = connection.execute(
|
|
|
+ {"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(
|
|
|
"""
|
|
|
- SELECT artifact_ref, artifact_digest, row_count,
|
|
|
- schema_hash, schema_fields, expires_at
|
|
|
- FROM public.rule_run_artifacts
|
|
|
- WHERE correlation_id = CAST(:correlation_id AS uuid)
|
|
|
+ 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 artifact_kind = :artifact_kind
|
|
|
- FOR UPDATE
|
|
|
+ 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
|
|
|
"""
|
|
|
),
|
|
|
{
|
|
|
- "correlation_id": correlation,
|
|
|
+ "id": reservation_id,
|
|
|
"binding_id": binding,
|
|
|
- "artifact_kind": kind,
|
|
|
+ "binding_hash": binding_hash,
|
|
|
+ "artifact_digest": reservation["digest"],
|
|
|
},
|
|
|
).mappings().one_or_none()
|
|
|
- if inserted is None:
|
|
|
- self.artifact_store.delete(artifact_ref)
|
|
|
- raise ValueError("immutable artifact catalog handoff was not found")
|
|
|
- existing = dict(inserted)
|
|
|
- existing_fields = existing["schema_fields"]
|
|
|
- if isinstance(existing_fields, str):
|
|
|
- existing_fields = json.loads(existing_fields)
|
|
|
- registered = {
|
|
|
- "artifact_ref": str(existing["artifact_ref"]),
|
|
|
- "digest": str(existing["artifact_digest"]),
|
|
|
- "row_count": int(existing["row_count"]),
|
|
|
- "schema_hash": str(existing["schema_hash"]),
|
|
|
- "schema_fields": _normalized_schema_fields(existing_fields),
|
|
|
- "expires_at": (
|
|
|
- _timestamp(existing["expires_at"])
|
|
|
- if isinstance(existing["expires_at"], datetime)
|
|
|
- else str(existing["expires_at"])
|
|
|
- ),
|
|
|
- }
|
|
|
- if registered["digest"] != described["digest"]:
|
|
|
- self.artifact_store.delete(artifact_ref)
|
|
|
- raise ValueError(
|
|
|
- "immutable artifact catalog digest conflicts with retry"
|
|
|
- )
|
|
|
- if registered["artifact_ref"] != artifact_ref:
|
|
|
- self.artifact_store.delete(artifact_ref)
|
|
|
+ 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(
|
|
|
- registered["artifact_ref"]
|
|
|
+ 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",
|
|
|
- ):
|
|
|
- if stored[key] != registered[key]:
|
|
|
- raise ValueError(
|
|
|
- "immutable artifact catalog does not match storage"
|
|
|
+ )
|
|
|
+ }
|
|
|
+ try:
|
|
|
+ self.artifact_store.upload_path(
|
|
|
+ path,
|
|
|
+ reserved_artifact,
|
|
|
+ limits=limits,
|
|
|
+ )
|
|
|
+ except Exception:
|
|
|
+ with suppress(Exception):
|
|
|
+ self.artifact_store.delete(
|
|
|
+ reserved_artifact["artifact_ref"]
|
|
|
+ )
|
|
|
+ with suppress(Exception):
|
|
|
+ self._abort_pending(reservation["reservation_id"])
|
|
|
+ raise
|
|
|
+ return self.finalize(
|
|
|
+ reservation=reservation,
|
|
|
+ binding_id=binding_id,
|
|
|
+ binding_hash=binding_hash,
|
|
|
+ kind=kind,
|
|
|
+ )
|
|
|
+
|
|
|
+ def _mark_failed(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ row_id: str,
|
|
|
+ failure_code: str,
|
|
|
+ ) -> None:
|
|
|
+ with self.engine.begin() as connection:
|
|
|
+ connection.execute(
|
|
|
+ text(
|
|
|
+ """
|
|
|
+ UPDATE public.rule_run_artifacts
|
|
|
+ SET handoff_status = 'failed',
|
|
|
+ failed_at = CURRENT_TIMESTAMP,
|
|
|
+ updated_at = CURRENT_TIMESTAMP,
|
|
|
+ failure_code = :failure_code
|
|
|
+ WHERE id = CAST(:id AS uuid)
|
|
|
+ AND handoff_status IN ('pending','ready')
|
|
|
+ """
|
|
|
+ ),
|
|
|
+ {"id": row_id, "failure_code": failure_code},
|
|
|
+ )
|
|
|
+
|
|
|
+ def reconcile(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ limit: int = 100,
|
|
|
+ grace_seconds: int = 300,
|
|
|
+ ) -> dict[str, int]:
|
|
|
+ """Repair bounded catalog/store drift after the grace period."""
|
|
|
+
|
|
|
+ if (
|
|
|
+ isinstance(limit, bool)
|
|
|
+ or not isinstance(limit, int)
|
|
|
+ or limit < 1
|
|
|
+ or limit > 1_000
|
|
|
+ ):
|
|
|
+ raise ValueError("artifact reconciliation limit is invalid")
|
|
|
+ if (
|
|
|
+ isinstance(grace_seconds, bool)
|
|
|
+ or not isinstance(grace_seconds, int)
|
|
|
+ or grace_seconds < 30
|
|
|
+ or grace_seconds > 86_400
|
|
|
+ ):
|
|
|
+ raise ValueError("artifact reconciliation grace is invalid")
|
|
|
+ result = {
|
|
|
+ "pending_finalized": 0,
|
|
|
+ "pending_deleted": 0,
|
|
|
+ "ready_failed": 0,
|
|
|
+ "orphans_deleted": 0,
|
|
|
+ }
|
|
|
+ with self.engine.connect() as connection:
|
|
|
+ rows = connection.execute(
|
|
|
+ text(
|
|
|
+ """
|
|
|
+ SELECT id::text AS id, correlation_id::text,
|
|
|
+ binding_id::text, artifact_ref, artifact_digest,
|
|
|
+ row_count, schema_hash, schema_fields,
|
|
|
+ artifact_kind, binding_hash, expires_at,
|
|
|
+ handoff_status
|
|
|
+ FROM public.rule_run_artifacts
|
|
|
+ WHERE handoff_status IN ('pending','ready')
|
|
|
+ AND updated_at <= CURRENT_TIMESTAMP
|
|
|
+ - make_interval(secs => :grace_seconds)
|
|
|
+ ORDER BY updated_at, id
|
|
|
+ LIMIT :limit
|
|
|
+ """
|
|
|
+ ),
|
|
|
+ {
|
|
|
+ "grace_seconds": grace_seconds,
|
|
|
+ "limit": limit,
|
|
|
+ },
|
|
|
+ ).mappings().all()
|
|
|
+ for raw_row in rows:
|
|
|
+ row = dict(raw_row)
|
|
|
+ row_id = str(row["id"])
|
|
|
+ status = str(row["handoff_status"])
|
|
|
+ invalid_object = False
|
|
|
+ try:
|
|
|
+ self.artifact_store._parse_ref(row["artifact_ref"])
|
|
|
+ stored = self.artifact_store.describe_optional(
|
|
|
+ row["artifact_ref"]
|
|
|
+ )
|
|
|
+ except Exception:
|
|
|
+ invalid_object = True
|
|
|
+ stored = None
|
|
|
+ if invalid_object:
|
|
|
+ if status == "pending":
|
|
|
+ with suppress(Exception):
|
|
|
+ self.artifact_store.delete(row["artifact_ref"])
|
|
|
+ self._mark_failed(
|
|
|
+ row_id=row_id,
|
|
|
+ failure_code="pending_object_invalid",
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ self._mark_failed(
|
|
|
+ row_id=row_id,
|
|
|
+ failure_code="ready_object_invalid",
|
|
|
)
|
|
|
- return registered
|
|
|
+ result["ready_failed"] += 1
|
|
|
+ continue
|
|
|
+ if status == "ready":
|
|
|
+ if stored is None:
|
|
|
+ self._mark_failed(
|
|
|
+ row_id=row_id,
|
|
|
+ failure_code="ready_object_missing",
|
|
|
+ )
|
|
|
+ result["ready_failed"] += 1
|
|
|
+ continue
|
|
|
+ if stored is None:
|
|
|
+ with self.engine.begin() as connection:
|
|
|
+ deleted = connection.execute(
|
|
|
+ text(
|
|
|
+ """
|
|
|
+ DELETE FROM public.rule_run_artifacts
|
|
|
+ WHERE id = CAST(:id AS uuid)
|
|
|
+ AND handoff_status = 'pending'
|
|
|
+ """
|
|
|
+ ),
|
|
|
+ {"id": row_id},
|
|
|
+ )
|
|
|
+ if int(deleted.rowcount or 0) == 1:
|
|
|
+ result["pending_deleted"] += 1
|
|
|
+ continue
|
|
|
+ if any(
|
|
|
+ (
|
|
|
+ stored["digest"] != str(row["artifact_digest"]),
|
|
|
+ stored["row_count"] != int(row["row_count"]),
|
|
|
+ stored["schema_hash"] != str(row["schema_hash"]),
|
|
|
+ )
|
|
|
+ ):
|
|
|
+ with suppress(Exception):
|
|
|
+ self.artifact_store.delete(row["artifact_ref"])
|
|
|
+ self._mark_failed(
|
|
|
+ row_id=row_id,
|
|
|
+ failure_code="pending_object_invalid",
|
|
|
+ )
|
|
|
+ continue
|
|
|
+ with self.engine.begin() as connection:
|
|
|
+ finalized = connection.execute(
|
|
|
+ text(
|
|
|
+ """
|
|
|
+ UPDATE public.rule_run_artifacts a
|
|
|
+ SET handoff_status = 'ready',
|
|
|
+ ready_at = CURRENT_TIMESTAMP,
|
|
|
+ updated_at = CURRENT_TIMESTAMP,
|
|
|
+ failure_code = NULL,
|
|
|
+ failed_at = NULL
|
|
|
+ WHERE a.id = CAST(:id AS uuid)
|
|
|
+ AND a.handoff_status = 'pending'
|
|
|
+ AND EXISTS (
|
|
|
+ SELECT 1
|
|
|
+ FROM public.dataflow_dataset_bindings b
|
|
|
+ WHERE b.id = a.binding_id
|
|
|
+ AND b.binding_hash = a.binding_hash
|
|
|
+ AND b.object_kind = 'parquet_artifact'
|
|
|
+ AND b.access_mode IN (
|
|
|
+ 'read','write','read_write'
|
|
|
+ )
|
|
|
+ )
|
|
|
+ RETURNING a.id
|
|
|
+ """
|
|
|
+ ),
|
|
|
+ {"id": row_id},
|
|
|
+ )
|
|
|
+ if int(finalized.rowcount or 0) == 1:
|
|
|
+ result["pending_finalized"] += 1
|
|
|
+ else:
|
|
|
+ self._mark_failed(
|
|
|
+ row_id=row_id,
|
|
|
+ failure_code="pending_binding_changed",
|
|
|
+ )
|
|
|
+
|
|
|
+ remaining = limit - len(rows)
|
|
|
+ if remaining <= 0:
|
|
|
+ return result
|
|
|
+ now = _now_utc(self.artifact_store.clock)
|
|
|
+ candidates = []
|
|
|
+ scanned = 0
|
|
|
+ for item in self.artifact_store.client.list_objects(
|
|
|
+ self.artifact_store.bucket,
|
|
|
+ prefix="rules/",
|
|
|
+ recursive=True,
|
|
|
+ ):
|
|
|
+ scanned += 1
|
|
|
+ if scanned > limit * 10 or len(candidates) >= remaining:
|
|
|
+ break
|
|
|
+ key = str(getattr(item, "object_name", ""))
|
|
|
+ ref = f"minio://{self.artifact_store.bucket}/{key}"
|
|
|
+ try:
|
|
|
+ self.artifact_store._parse_ref(ref)
|
|
|
+ except ValueError:
|
|
|
+ continue
|
|
|
+ modified = getattr(item, "last_modified", None)
|
|
|
+ if not isinstance(modified, datetime):
|
|
|
+ continue
|
|
|
+ if modified.tzinfo is None:
|
|
|
+ modified = modified.replace(tzinfo=UTC)
|
|
|
+ if modified.astimezone(UTC) > now - timedelta(
|
|
|
+ seconds=grace_seconds
|
|
|
+ ):
|
|
|
+ continue
|
|
|
+ candidates.append(ref)
|
|
|
+ if not candidates:
|
|
|
+ return result
|
|
|
+ with self.engine.connect() as connection:
|
|
|
+ referenced = {
|
|
|
+ str(row["artifact_ref"])
|
|
|
+ for row in connection.execute(
|
|
|
+ text(
|
|
|
+ """
|
|
|
+ SELECT artifact_ref
|
|
|
+ FROM public.rule_run_artifacts
|
|
|
+ WHERE artifact_ref =
|
|
|
+ ANY(CAST(:artifact_refs AS text[]))
|
|
|
+ """
|
|
|
+ ),
|
|
|
+ {"artifact_refs": candidates},
|
|
|
+ ).mappings().all()
|
|
|
+ }
|
|
|
+ for ref in candidates:
|
|
|
+ if ref in referenced:
|
|
|
+ continue
|
|
|
+ self.artifact_store.delete(ref)
|
|
|
+ result["orphans_deleted"] += 1
|
|
|
+ return result
|
|
|
|
|
|
def cleanup_expired(self, *, limit: int = 100) -> int:
|
|
|
if (
|
马小龙