|
@@ -3,18 +3,17 @@
|
|
|
from __future__ import annotations
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
import hashlib
|
|
import hashlib
|
|
|
-import io
|
|
|
|
|
import json
|
|
import json
|
|
|
import os
|
|
import os
|
|
|
import re
|
|
import re
|
|
|
import tempfile
|
|
import tempfile
|
|
|
-from base64 import urlsafe_b64decode, urlsafe_b64encode
|
|
|
|
|
from collections.abc import Mapping
|
|
from collections.abc import Mapping
|
|
|
-from contextlib import suppress
|
|
|
|
|
|
|
+from contextlib import contextmanager, suppress
|
|
|
from datetime import UTC, datetime, timedelta
|
|
from datetime import UTC, datetime, timedelta
|
|
|
from typing import Any
|
|
from typing import Any
|
|
|
|
|
|
|
|
import polars as pl
|
|
import polars as pl
|
|
|
|
|
+import pyarrow.parquet as pq
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy import text
|
|
|
|
|
|
|
|
from app.core.common.identifiers import (
|
|
from app.core.common.identifiers import (
|
|
@@ -27,6 +26,22 @@ PARQUET_CONTENT_TYPE = "application/x-parquet"
|
|
|
_DIGEST = re.compile(r"^[0-9a-f]{64}$")
|
|
_DIGEST = re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+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:
|
|
def _now_utc(clock) -> datetime:
|
|
|
value = clock()
|
|
value = clock()
|
|
|
if not isinstance(value, datetime):
|
|
if not isinstance(value, datetime):
|
|
@@ -111,28 +126,9 @@ def _normalized_schema_fields(value: Any) -> list[dict[str, Any]]:
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
-def _schema_contract(value: Any) -> tuple[list[dict[str, Any]], str, str]:
|
|
|
|
|
|
|
+def _schema_contract(value: Any) -> tuple[list[dict[str, Any]], str]:
|
|
|
fields = _normalized_schema_fields(value)
|
|
fields = _normalized_schema_fields(value)
|
|
|
- encoded = json.dumps(
|
|
|
|
|
- fields,
|
|
|
|
|
- sort_keys=True,
|
|
|
|
|
- separators=(",", ":"),
|
|
|
|
|
- ensure_ascii=False,
|
|
|
|
|
- ).encode("utf-8")
|
|
|
|
|
- return (
|
|
|
|
|
- fields,
|
|
|
|
|
- canonical_schema_hash(fields),
|
|
|
|
|
- urlsafe_b64encode(encoded).decode("ascii"),
|
|
|
|
|
- )
|
|
|
|
|
-
|
|
|
|
|
-
|
|
|
|
|
-def _decode_schema_contract(value: str) -> list[dict[str, Any]]:
|
|
|
|
|
- try:
|
|
|
|
|
- decoded = urlsafe_b64decode(value.encode("ascii"))
|
|
|
|
|
- fields = json.loads(decoded.decode("utf-8"))
|
|
|
|
|
- except Exception as exc:
|
|
|
|
|
- raise ValueError("artifact schema contract metadata is invalid") from exc
|
|
|
|
|
- return _normalized_schema_fields(fields)
|
|
|
|
|
|
|
+ return fields, canonical_schema_hash(fields)
|
|
|
|
|
|
|
|
|
|
|
|
|
def _validate_frame_schema(
|
|
def _validate_frame_schema(
|
|
@@ -200,7 +196,6 @@ def _metadata(value: Any) -> dict[str, str]:
|
|
|
"schema-sha256",
|
|
"schema-sha256",
|
|
|
"expires-at",
|
|
"expires-at",
|
|
|
"artifact-bytes",
|
|
"artifact-bytes",
|
|
|
- "schema-contract",
|
|
|
|
|
}:
|
|
}:
|
|
|
normalized[name] = str(item)
|
|
normalized[name] = str(item)
|
|
|
required = {
|
|
required = {
|
|
@@ -209,10 +204,11 @@ def _metadata(value: Any) -> dict[str, str]:
|
|
|
"schema-sha256",
|
|
"schema-sha256",
|
|
|
"expires-at",
|
|
"expires-at",
|
|
|
"artifact-bytes",
|
|
"artifact-bytes",
|
|
|
- "schema-contract",
|
|
|
|
|
}
|
|
}
|
|
|
if set(normalized) != required:
|
|
if set(normalized) != required:
|
|
|
raise ValueError("artifact content metadata is incomplete")
|
|
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
|
|
return normalized
|
|
|
|
|
|
|
|
|
|
|
|
@@ -357,7 +353,7 @@ class ArtifactStore:
|
|
|
raise ValueError("artifact row count exceeds the configured limit")
|
|
raise ValueError("artifact row count exceeds the configured limit")
|
|
|
if collected.estimated_size() > effective["memory_limit_bytes"]:
|
|
if collected.estimated_size() > effective["memory_limit_bytes"]:
|
|
|
raise ValueError("artifact frame exceeds the configured memory limit")
|
|
raise ValueError("artifact frame exceeds the configured memory limit")
|
|
|
- fields, schema_digest, schema_encoded = _schema_contract(
|
|
|
|
|
|
|
+ fields, schema_digest = _schema_contract(
|
|
|
schema_fields or _inferred_schema_fields(collected)
|
|
schema_fields or _inferred_schema_fields(collected)
|
|
|
)
|
|
)
|
|
|
_validate_frame_schema(collected, fields)
|
|
_validate_frame_schema(collected, fields)
|
|
@@ -406,7 +402,6 @@ class ArtifactStore:
|
|
|
"schema-sha256": schema_digest,
|
|
"schema-sha256": schema_digest,
|
|
|
"expires-at": expires_at,
|
|
"expires-at": expires_at,
|
|
|
"artifact-bytes": str(size),
|
|
"artifact-bytes": str(size),
|
|
|
- "schema-contract": schema_encoded,
|
|
|
|
|
},
|
|
},
|
|
|
)
|
|
)
|
|
|
uploaded = True
|
|
uploaded = True
|
|
@@ -416,12 +411,13 @@ class ArtifactStore:
|
|
|
limits=effective,
|
|
limits=effective,
|
|
|
)
|
|
)
|
|
|
artifact_ref = f"minio://{self.bucket}/{key}"
|
|
artifact_ref = f"minio://{self.bucket}/{key}"
|
|
|
- self.read(
|
|
|
|
|
|
|
+ with self.stage(
|
|
|
artifact_ref,
|
|
artifact_ref,
|
|
|
digest_hex,
|
|
digest_hex,
|
|
|
expected_schema_fields=fields,
|
|
expected_schema_fields=fields,
|
|
|
limits=effective,
|
|
limits=effective,
|
|
|
- )
|
|
|
|
|
|
|
+ ):
|
|
|
|
|
+ pass
|
|
|
except Exception:
|
|
except Exception:
|
|
|
if uploaded:
|
|
if uploaded:
|
|
|
with suppress(Exception):
|
|
with suppress(Exception):
|
|
@@ -440,20 +436,100 @@ class ArtifactStore:
|
|
|
"expires_at": expires_at,
|
|
"expires_at": expires_at,
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+ 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]:
|
|
|
|
|
+ """Upload a worker-produced Parquet file without collecting it."""
|
|
|
|
|
+
|
|
|
|
|
+ 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"
|
|
|
|
|
+ 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": digest_hex,
|
|
|
|
|
+ "row-count": str(rows),
|
|
|
|
|
+ "schema-sha256": schema_digest,
|
|
|
|
|
+ "expires-at": 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,
|
|
|
|
|
+ limits=effective,
|
|
|
|
|
+ ):
|
|
|
|
|
+ pass
|
|
|
|
|
+ except Exception:
|
|
|
|
|
+ if uploaded:
|
|
|
|
|
+ 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,
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
def describe(self, ref: str) -> dict[str, Any]:
|
|
def describe(self, ref: str) -> dict[str, Any]:
|
|
|
"""Return validated object metadata without exposing MinIO credentials."""
|
|
"""Return validated object metadata without exposing MinIO credentials."""
|
|
|
|
|
|
|
|
key = self._parse_ref(ref)
|
|
key = self._parse_ref(ref)
|
|
|
_stat, metadata = self._validated_stat(key)
|
|
_stat, metadata = self._validated_stat(key)
|
|
|
- fields = _decode_schema_contract(metadata["schema-contract"])
|
|
|
|
|
- if canonical_schema_hash(fields) != metadata["schema-sha256"]:
|
|
|
|
|
- raise ValueError("artifact schema contract does not match")
|
|
|
|
|
return {
|
|
return {
|
|
|
"artifact_ref": ref,
|
|
"artifact_ref": ref,
|
|
|
"digest": metadata["sha256"],
|
|
"digest": metadata["sha256"],
|
|
|
"row_count": int(metadata["row-count"]),
|
|
"row_count": int(metadata["row-count"]),
|
|
|
"schema_hash": metadata["schema-sha256"],
|
|
"schema_hash": metadata["schema-sha256"],
|
|
|
- "schema_fields": fields,
|
|
|
|
|
"expires_at": metadata["expires-at"],
|
|
"expires_at": metadata["expires-at"],
|
|
|
}
|
|
}
|
|
|
|
|
|
|
@@ -465,6 +541,37 @@ class ArtifactStore:
|
|
|
expected_schema_fields: list[dict[str, Any]] | None = None,
|
|
expected_schema_fields: list[dict[str, Any]] | None = None,
|
|
|
limits: dict[str, int] | None = None,
|
|
limits: dict[str, int] | None = None,
|
|
|
) -> pl.LazyFrame:
|
|
) -> 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)
|
|
effective = self._limits(limits)
|
|
|
if _DIGEST.fullmatch(str(expected_digest or "")) is None:
|
|
if _DIGEST.fullmatch(str(expected_digest or "")) is None:
|
|
|
raise ValueError("expected artifact digest is invalid")
|
|
raise ValueError("expected artifact digest is invalid")
|
|
@@ -474,56 +581,56 @@ class ArtifactStore:
|
|
|
expected_digest=expected_digest,
|
|
expected_digest=expected_digest,
|
|
|
limits=effective,
|
|
limits=effective,
|
|
|
)
|
|
)
|
|
|
- fields = _decode_schema_contract(metadata["schema-contract"])
|
|
|
|
|
|
|
+ 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"]:
|
|
if canonical_schema_hash(fields) != metadata["schema-sha256"]:
|
|
|
- raise ValueError("artifact schema contract does not match")
|
|
|
|
|
- if expected_schema_fields is not None:
|
|
|
|
|
- expected = _normalized_schema_fields(expected_schema_fields)
|
|
|
|
|
- if expected != fields:
|
|
|
|
|
- raise ValueError("artifact schema contract is not expected")
|
|
|
|
|
|
|
+ raise ValueError("artifact schema contract is not expected")
|
|
|
response = self.client.get_object(self.bucket, key)
|
|
response = self.client.get_object(self.bucket, key)
|
|
|
digest = hashlib.sha256()
|
|
digest = hashlib.sha256()
|
|
|
- payload = io.BytesIO()
|
|
|
|
|
|
|
+ path = None
|
|
|
size = 0
|
|
size = 0
|
|
|
try:
|
|
try:
|
|
|
- while chunk := response.read(1024 * 1024):
|
|
|
|
|
- size += len(chunk)
|
|
|
|
|
- if size > effective["max_artifact_bytes"]:
|
|
|
|
|
- raise ValueError(
|
|
|
|
|
- "artifact size exceeds the configured limit"
|
|
|
|
|
- )
|
|
|
|
|
- if size > effective["memory_limit_bytes"]:
|
|
|
|
|
- raise ValueError(
|
|
|
|
|
- "artifact download exceeds the memory limit"
|
|
|
|
|
- )
|
|
|
|
|
- digest.update(chunk)
|
|
|
|
|
- payload.write(chunk)
|
|
|
|
|
|
|
+ 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:
|
|
finally:
|
|
|
- response.close()
|
|
|
|
|
- release = getattr(response, "release_conn", None)
|
|
|
|
|
- if callable(release):
|
|
|
|
|
- release()
|
|
|
|
|
- if digest.hexdigest() != expected_digest:
|
|
|
|
|
- raise ValueError("artifact digest does not match content")
|
|
|
|
|
- payload.seek(0)
|
|
|
|
|
- try:
|
|
|
|
|
- frame = pl.read_parquet(
|
|
|
|
|
- payload,
|
|
|
|
|
- n_rows=effective["max_rows"] + 1,
|
|
|
|
|
- memory_map=False,
|
|
|
|
|
- )
|
|
|
|
|
- except Exception as exc:
|
|
|
|
|
- raise ValueError("artifact is not valid Parquet") from exc
|
|
|
|
|
- if frame.height != int(metadata["row-count"]):
|
|
|
|
|
- raise ValueError("artifact row count does not match metadata")
|
|
|
|
|
- if frame.height > effective["max_rows"]:
|
|
|
|
|
- raise ValueError("artifact row count exceeds the configured limit")
|
|
|
|
|
- if frame.estimated_size() > effective["memory_limit_bytes"]:
|
|
|
|
|
- raise ValueError("artifact frame exceeds the configured memory limit")
|
|
|
|
|
- if size + frame.estimated_size() > effective["memory_limit_bytes"]:
|
|
|
|
|
- raise ValueError("artifact decompression exceeds the memory limit")
|
|
|
|
|
- _validate_frame_schema(frame, fields)
|
|
|
|
|
- return frame.lazy()
|
|
|
|
|
|
|
+ if path is not None:
|
|
|
|
|
+ with suppress(FileNotFoundError):
|
|
|
|
|
+ os.unlink(path)
|
|
|
|
|
|
|
|
def cleanup_expired(self, correlation_id: str) -> int:
|
|
def cleanup_expired(self, correlation_id: str) -> int:
|
|
|
correlation = _uid(correlation_id, "correlation_id")
|
|
correlation = _uid(correlation_id, "correlation_id")
|
|
@@ -565,9 +672,17 @@ class PostgresArtifactResolver:
|
|
|
self.engine = engine
|
|
self.engine = engine
|
|
|
self.artifact_store = artifact_store
|
|
self.artifact_store = artifact_store
|
|
|
|
|
|
|
|
- def resolve(self, *, binding_id: str, correlation_id: str) -> dict[str, Any]:
|
|
|
|
|
|
|
+ def resolve(
|
|
|
|
|
+ self,
|
|
|
|
|
+ *,
|
|
|
|
|
+ binding_id: str,
|
|
|
|
|
+ correlation_id: str,
|
|
|
|
|
+ kind: str,
|
|
|
|
|
+ ) -> dict[str, Any]:
|
|
|
binding = _uid(binding_id, "artifact binding id")
|
|
binding = _uid(binding_id, "artifact binding id")
|
|
|
correlation = _uid(correlation_id, "artifact correlation id")
|
|
correlation = _uid(correlation_id, "artifact correlation id")
|
|
|
|
|
+ if kind not in {"input", "lookup", "output"}:
|
|
|
|
|
+ raise ValueError("artifact kind is invalid")
|
|
|
statement = text(
|
|
statement = text(
|
|
|
"""
|
|
"""
|
|
|
SELECT
|
|
SELECT
|
|
@@ -583,6 +698,7 @@ class PostgresArtifactResolver:
|
|
|
ON b.id = a.binding_id
|
|
ON b.id = a.binding_id
|
|
|
WHERE a.binding_id = CAST(:binding_id AS uuid)
|
|
WHERE a.binding_id = CAST(:binding_id AS uuid)
|
|
|
AND a.correlation_id = CAST(:correlation_id AS uuid)
|
|
AND a.correlation_id = CAST(:correlation_id AS uuid)
|
|
|
|
|
+ AND a.artifact_kind = :artifact_kind
|
|
|
AND a.expires_at > CURRENT_TIMESTAMP
|
|
AND a.expires_at > CURRENT_TIMESTAMP
|
|
|
AND b.object_kind = 'parquet_artifact'
|
|
AND b.object_kind = 'parquet_artifact'
|
|
|
AND b.access_mode IN ('read', 'read_write')
|
|
AND b.access_mode IN ('read', 'read_write')
|
|
@@ -596,6 +712,7 @@ class PostgresArtifactResolver:
|
|
|
{
|
|
{
|
|
|
"binding_id": binding,
|
|
"binding_id": binding,
|
|
|
"correlation_id": correlation,
|
|
"correlation_id": correlation,
|
|
|
|
|
+ "artifact_kind": kind,
|
|
|
},
|
|
},
|
|
|
).mappings().one_or_none()
|
|
).mappings().one_or_none()
|
|
|
if row is None:
|
|
if row is None:
|
|
@@ -614,12 +731,11 @@ class PostgresArtifactResolver:
|
|
|
described["digest"] != str(row["artifact_digest"])
|
|
described["digest"] != str(row["artifact_digest"])
|
|
|
or described["row_count"] != int(row["row_count"])
|
|
or described["row_count"] != int(row["row_count"])
|
|
|
or described["schema_hash"] != str(row["schema_hash"])
|
|
or described["schema_hash"] != str(row["schema_hash"])
|
|
|
- or described["schema_fields"]
|
|
|
|
|
- != _normalized_schema_fields(row_fields)
|
|
|
|
|
):
|
|
):
|
|
|
raise ValueError("catalog artifact metadata does not match storage")
|
|
raise ValueError("catalog artifact metadata does not match storage")
|
|
|
return {
|
|
return {
|
|
|
**described,
|
|
**described,
|
|
|
|
|
+ "schema_fields": _normalized_schema_fields(row_fields),
|
|
|
"binding_hash": str(row["binding_hash"]),
|
|
"binding_hash": str(row["binding_hash"]),
|
|
|
}
|
|
}
|
|
|
|
|
|
|
@@ -667,7 +783,7 @@ class PostgresArtifactResolver:
|
|
|
artifact: dict[str, Any],
|
|
artifact: dict[str, Any],
|
|
|
kind: str,
|
|
kind: str,
|
|
|
binding_hash: str,
|
|
binding_hash: str,
|
|
|
- ) -> None:
|
|
|
|
|
|
|
+ ) -> dict[str, Any]:
|
|
|
binding = _uid(binding_id, "artifact binding id")
|
|
binding = _uid(binding_id, "artifact binding id")
|
|
|
correlation = _uid(correlation_id, "artifact correlation id")
|
|
correlation = _uid(correlation_id, "artifact correlation id")
|
|
|
if kind not in {"input", "lookup", "output"}:
|
|
if kind not in {"input", "lookup", "output"}:
|
|
@@ -691,13 +807,16 @@ class PostgresArtifactResolver:
|
|
|
"digest",
|
|
"digest",
|
|
|
"row_count",
|
|
"row_count",
|
|
|
"schema_hash",
|
|
"schema_hash",
|
|
|
- "schema_fields",
|
|
|
|
|
"expires_at",
|
|
"expires_at",
|
|
|
):
|
|
):
|
|
|
if described[key] != artifact.get(key):
|
|
if described[key] != artifact.get(key):
|
|
|
raise ValueError("artifact metadata does not match storage")
|
|
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
|
|
|
with self.engine.begin() as connection:
|
|
with self.engine.begin() as connection:
|
|
|
- connection.execute(
|
|
|
|
|
|
|
+ inserted = connection.execute(
|
|
|
text(
|
|
text(
|
|
|
"""
|
|
"""
|
|
|
INSERT INTO public.rule_run_artifacts (
|
|
INSERT INTO public.rule_run_artifacts (
|
|
@@ -712,8 +831,10 @@ class PostgresArtifactResolver:
|
|
|
CAST(:expires_at AS timestamptz)
|
|
CAST(:expires_at AS timestamptz)
|
|
|
)
|
|
)
|
|
|
ON CONFLICT (
|
|
ON CONFLICT (
|
|
|
- correlation_id, binding_id, artifact_digest
|
|
|
|
|
|
|
+ correlation_id, binding_id, artifact_kind
|
|
|
) DO NOTHING
|
|
) DO NOTHING
|
|
|
|
|
+ RETURNING artifact_ref, artifact_digest, row_count,
|
|
|
|
|
+ schema_hash, schema_fields, expires_at
|
|
|
"""
|
|
"""
|
|
|
),
|
|
),
|
|
|
{
|
|
{
|
|
@@ -725,11 +846,107 @@ class PostgresArtifactResolver:
|
|
|
"row_count": described["row_count"],
|
|
"row_count": described["row_count"],
|
|
|
"schema_hash": described["schema_hash"],
|
|
"schema_hash": described["schema_hash"],
|
|
|
"schema_fields": json.dumps(
|
|
"schema_fields": json.dumps(
|
|
|
- described["schema_fields"],
|
|
|
|
|
|
|
+ fields,
|
|
|
sort_keys=True,
|
|
sort_keys=True,
|
|
|
separators=(",", ":"),
|
|
separators=(",", ":"),
|
|
|
),
|
|
),
|
|
|
"artifact_kind": kind,
|
|
"artifact_kind": kind,
|
|
|
"expires_at": described["expires_at"],
|
|
"expires_at": described["expires_at"],
|
|
|
},
|
|
},
|
|
|
|
|
+ ).mappings().one_or_none()
|
|
|
|
|
+ if inserted is None:
|
|
|
|
|
+ inserted = 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)
|
|
|
|
|
+ AND binding_id = CAST(:binding_id AS uuid)
|
|
|
|
|
+ AND artifact_kind = :artifact_kind
|
|
|
|
|
+ FOR UPDATE
|
|
|
|
|
+ """
|
|
|
|
|
+ ),
|
|
|
|
|
+ {
|
|
|
|
|
+ "correlation_id": correlation,
|
|
|
|
|
+ "binding_id": binding,
|
|
|
|
|
+ "artifact_kind": kind,
|
|
|
|
|
+ },
|
|
|
|
|
+ ).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)
|
|
|
|
|
+ stored = self.artifact_store.describe(
|
|
|
|
|
+ registered["artifact_ref"]
|
|
|
)
|
|
)
|
|
|
|
|
+ for key in (
|
|
|
|
|
+ "digest",
|
|
|
|
|
+ "row_count",
|
|
|
|
|
+ "schema_hash",
|
|
|
|
|
+ "expires_at",
|
|
|
|
|
+ ):
|
|
|
|
|
+ if stored[key] != registered[key]:
|
|
|
|
|
+ raise ValueError(
|
|
|
|
|
+ "immutable artifact catalog does not match storage"
|
|
|
|
|
+ )
|
|
|
|
|
+ return registered
|
|
|
|
|
+
|
|
|
|
|
+ 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
|