DataOps-platform/app/runner/artifacts.py

"""Digest-bound, bounded Parquet artifacts owned by the DataOps Runner."""

from __future__ import annotations

import hashlib
import io
import json
import os
import re
import tempfile
from base64 import urlsafe_b64decode, urlsafe_b64encode
from collections.abc import Mapping
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from typing import Any

import polars as pl
from sqlalchemy import text

from app.core.common.identifiers import (
    ensure_governance_uid,
    new_governance_uid,
)
from app.core.data_rules.execution_contracts import canonical_schema_hash

PARQUET_CONTENT_TYPE = "application/x-parquet"
_DIGEST = re.compile(r"^[0-9a-f]{64}$")


def _now_utc(clock) -> datetime:
    value = clock()
    if not isinstance(value, datetime):
        raise ValueError("artifact clock must return a datetime")
    if value.tzinfo is None:
        value = value.replace(tzinfo=UTC)
    return value.astimezone(UTC)


def _timestamp(value: datetime) -> str:
    return value.astimezone(UTC).isoformat().replace("+00:00", "Z")


def _parse_timestamp(value: Any) -> datetime:
    if not isinstance(value, str) or not value.endswith("Z"):
        raise ValueError("artifact expiry metadata is invalid")
    try:
        parsed = datetime.fromisoformat(value[:-1] + "+00:00")
    except ValueError as exc:
        raise ValueError("artifact expiry metadata is invalid") from exc
    return parsed.astimezone(UTC)


def _uid(value: Any, label: str) -> str:
    try:
        return ensure_governance_uid({"uid": str(value)})
    except ValueError as exc:
        raise ValueError(f"{label} must be a valid UUIDv7") from exc


def _inferred_schema_fields(
    frame: pl.DataFrame | pl.LazyFrame,
) -> list[dict[str, Any]]:
    schema = (
        frame.collect_schema()
        if isinstance(frame, pl.LazyFrame)
        else frame.schema
    )
    fields = []
    for name, dtype in schema.items():
        field: dict[str, Any] = {
            "name": name,
            "nullable": True,
        }
        if dtype == pl.Boolean:
            field["type"] = "boolean"
        elif dtype == pl.Date:
            field["type"] = "date"
        elif dtype == pl.String:
            field["type"] = "string"
        elif dtype.is_integer():
            field["type"] = "integer"
        elif dtype == pl.Float32:
            field["type"] = "float"
        elif dtype == pl.Float64:
            field["type"] = "double"
        elif dtype.is_decimal():
            field.update(
                {
                    "type": "decimal",
                    "precision": dtype.precision,
                    "scale": dtype.scale,
                }
            )
        elif isinstance(dtype, pl.Datetime):
            field["type"] = (
                "timestamptz" if dtype.time_zone else "timestamp"
            )
            if dtype.time_zone:
                field["timezone"] = dtype.time_zone
        else:
            raise ValueError(f"unsupported artifact dtype for {name}")
        fields.append(field)
    return sorted(fields, key=lambda item: item["name"])


def _normalized_schema_fields(value: Any) -> list[dict[str, Any]]:
    canonical_schema_hash(value)
    return sorted(
        [dict(field) for field in value],
        key=lambda item: item["name"],
    )


def _schema_contract(value: Any) -> tuple[list[dict[str, Any]], str, str]:
    fields = _normalized_schema_fields(value)
    encoded = json.dumps(
        fields,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    ).encode("utf-8")
    return (
        fields,
        canonical_schema_hash(fields),
        urlsafe_b64encode(encoded).decode("ascii"),
    )


def _decode_schema_contract(value: str) -> list[dict[str, Any]]:
    try:
        decoded = urlsafe_b64decode(value.encode("ascii"))
        fields = json.loads(decoded.decode("utf-8"))
    except Exception as exc:
        raise ValueError("artifact schema contract metadata is invalid") from exc
    return _normalized_schema_fields(fields)


def _validate_frame_schema(
    frame: pl.DataFrame | pl.LazyFrame,
    fields: list[dict[str, Any]],
) -> None:
    schema = (
        frame.collect_schema()
        if isinstance(frame, pl.LazyFrame)
        else frame.schema
    )
    expected_names = {field["name"] for field in fields}
    if set(schema.names()) != expected_names:
        raise ValueError("artifact schema fields do not match")
    for field in fields:
        dtype = schema[field["name"]]
        field_type = field["type"]
        matches = (
            (field_type == "boolean" and dtype == pl.Boolean)
            or (field_type == "date" and dtype == pl.Date)
            or (field_type == "string" and dtype == pl.String)
            or (field_type == "integer" and dtype.is_integer())
            or (field_type == "float" and dtype == pl.Float32)
            or (field_type == "double" and dtype == pl.Float64)
            or (
                field_type == "decimal"
                and dtype.is_decimal()
                and dtype.precision == field.get("precision")
                and dtype.scale == field.get("scale")
            )
            or (
                field_type == "timestamp"
                and isinstance(dtype, pl.Datetime)
                and dtype.time_zone is None
            )
            or (
                field_type == "timestamptz"
                and isinstance(dtype, pl.Datetime)
                and dtype.time_zone == field.get("timezone")
            )
        )
        if not matches:
            raise ValueError(
                f"artifact schema type for {field['name']} does not match"
            )
    if isinstance(frame, pl.DataFrame):
        for field in fields:
            if not field["nullable"] and frame[field["name"]].null_count():
                raise ValueError(
                    f"artifact nullable contract for {field['name']} does not match"
                )


def _metadata(value: Any) -> dict[str, str]:
    if not isinstance(value, Mapping):
        raise ValueError("artifact content metadata is missing")
    normalized = {}
    for key, item in value.items():
        name = str(key).lower()
        if name.startswith("x-amz-meta-"):
            name = name[len("x-amz-meta-") :]
        if name in {
            "sha256",
            "row-count",
            "schema-sha256",
            "expires-at",
            "artifact-bytes",
            "schema-contract",
        }:
            normalized[name] = str(item)
    required = {
        "sha256",
        "row-count",
        "schema-sha256",
        "expires-at",
        "artifact-bytes",
        "schema-contract",
    }
    if set(normalized) != required:
        raise ValueError("artifact content metadata is incomplete")
    return normalized


class ArtifactStore:
    """Read and write only server-owned, bounded Parquet artifacts."""

    def __init__(
        self,
        client,
        *,
        bucket: str,
        max_artifact_bytes: int,
        max_rows: int,
        memory_limit_bytes: int,
        max_ttl_seconds: int = 86400,
        clock=None,
    ):
        if not re.fullmatch(r"[a-z0-9][a-z0-9.-]{1,61}[a-z0-9]", bucket):
            raise ValueError("artifact bucket name is invalid")
        self.client = client
        self.bucket = bucket
        self.max_artifact_bytes = int(max_artifact_bytes)
        self.max_rows = int(max_rows)
        self.memory_limit_bytes = int(memory_limit_bytes)
        self.max_ttl_seconds = int(max_ttl_seconds)
        self.clock = clock or (lambda: datetime.now(UTC))
        if (
            self.max_artifact_bytes < 1024
            or self.max_rows < 1
            or self.memory_limit_bytes < self.max_artifact_bytes
            or self.max_ttl_seconds < 1
        ):
            raise ValueError("artifact resource limits are invalid")
        if not self.client.bucket_exists(self.bucket):
            raise ValueError("artifact bucket does not exist")

    def _limits(self, value: Any = None) -> dict[str, int]:
        configured = {
            "max_rows": self.max_rows,
            "max_artifact_bytes": self.max_artifact_bytes,
            "memory_limit_bytes": self.memory_limit_bytes,
        }
        if value is None:
            return configured
        if not isinstance(value, dict) or set(value) != set(configured):
            raise ValueError("artifact limits must have a closed shape")
        result = {}
        for key, ceiling in configured.items():
            item = value[key]
            if (
                isinstance(item, bool)
                or not isinstance(item, int)
                or item < 1
                or item > ceiling
            ):
                raise ValueError(
                    f"artifact {key} exceeds the configured ceiling"
                )
            result[key] = item
        return result

    def _parse_ref(self, ref: Any) -> str:
        prefix = f"minio://{self.bucket}/"
        if not isinstance(ref, str) or not ref.startswith(prefix):
            raise ValueError("artifact reference is not owned by this store")
        key = ref[len(prefix) :]
        match = re.fullmatch(
            r"rules/([0-9a-f-]{36})/([0-9a-f-]{36})\.parquet",
            key,
        )
        if match is None:
            raise ValueError("artifact reference is invalid")
        _uid(match.group(1), "artifact correlation id")
        _uid(match.group(2), "artifact id")
        return key

    def _validated_stat(
        self,
        key: str,
        *,
        expected_digest: str | None = None,
        limits: dict[str, int] | None = None,
    ) -> tuple[Any, dict[str, str]]:
        effective = self._limits(limits)
        stat = self.client.stat_object(self.bucket, key)
        size = int(getattr(stat, "size", -1))
        if size < 1 or size > effective["max_artifact_bytes"]:
            raise ValueError("artifact size exceeds the configured limit")
        if size > effective["memory_limit_bytes"]:
            raise ValueError("artifact download exceeds the memory limit")
        if str(getattr(stat, "content_type", "")).lower() != PARQUET_CONTENT_TYPE:
            raise ValueError("artifact content type is invalid")
        metadata = _metadata(getattr(stat, "metadata", None))
        digest = metadata["sha256"]
        if _DIGEST.fullmatch(digest) is None:
            raise ValueError("artifact digest metadata is invalid")
        if expected_digest is not None and digest != expected_digest:
            raise ValueError("artifact digest does not match")
        try:
            row_count = int(metadata["row-count"])
            metadata_size = int(metadata["artifact-bytes"])
        except (TypeError, ValueError) as exc:
            raise ValueError("artifact count metadata is invalid") from exc
        if row_count < 0 or row_count > effective["max_rows"]:
            raise ValueError("artifact row count exceeds the configured limit")
        if metadata_size != size:
            raise ValueError("artifact size metadata does not match")
        if _DIGEST.fullmatch(metadata["schema-sha256"]) is None:
            raise ValueError("artifact schema metadata is invalid")
        if _parse_timestamp(metadata["expires-at"]) <= _now_utc(self.clock):
            raise ValueError("artifact has expired")
        return stat, metadata

    def write(
        self,
        frame: pl.LazyFrame | pl.DataFrame,
        correlation_id: str,
        ttl_seconds: int,
        *,
        schema_fields: list[dict[str, Any]] | None = None,
        limits: dict[str, int] | None = None,
    ) -> dict[str, Any]:
        effective = self._limits(limits)
        correlation = _uid(correlation_id, "correlation_id")
        if (
            isinstance(ttl_seconds, bool)
            or not isinstance(ttl_seconds, int)
            or ttl_seconds < 1
            or ttl_seconds > self.max_ttl_seconds
        ):
            raise ValueError("artifact TTL is outside the configured limit")
        if isinstance(frame, pl.DataFrame):
            lazy = frame.lazy()
        elif isinstance(frame, pl.LazyFrame):
            lazy = frame
        else:
            raise ValueError("artifact frame must be a Polars frame")
        collected = lazy.head(effective["max_rows"] + 1).collect(
            engine="streaming"
        )
        if collected.height > effective["max_rows"]:
            raise ValueError("artifact row count exceeds the configured limit")
        if collected.estimated_size() > effective["memory_limit_bytes"]:
            raise ValueError("artifact frame exceeds the configured memory limit")
        fields, schema_digest, schema_encoded = _schema_contract(
            schema_fields or _inferred_schema_fields(collected)
        )
        _validate_frame_schema(collected, fields)
        expires_at = _timestamp(
            _now_utc(self.clock) + timedelta(seconds=ttl_seconds)
        )
        artifact_id = new_governance_uid()
        key = f"rules/{correlation}/{artifact_id}.parquet"
        path = None
        uploaded = False
        try:
            with tempfile.NamedTemporaryFile(
                prefix="dataops-rule-artifact-",
                suffix=".parquet",
                delete=False,
            ) as handle:
                path = handle.name
            collected.write_parquet(path)
            size = os.path.getsize(path)
            if size < 1 or size > effective["max_artifact_bytes"]:
                raise ValueError("artifact size exceeds the configured limit")
            if size > effective["memory_limit_bytes"]:
                raise ValueError("serialized artifact exceeds the memory limit")
            if (
                size + collected.estimated_size()
                > effective["memory_limit_bytes"]
            ):
                raise ValueError(
                    "artifact serialization exceeds the memory limit"
                )
            digest = hashlib.sha256()
            with open(path, "rb") as handle:
                while chunk := handle.read(1024 * 1024):
                    digest.update(chunk)
            digest_hex = digest.hexdigest()
            with open(path, "rb") as handle:
                self.client.put_object(
                    self.bucket,
                    key,
                    handle,
                    size,
                    content_type=PARQUET_CONTENT_TYPE,
                    metadata={
                        "sha256": digest_hex,
                        "row-count": str(collected.height),
                        "schema-sha256": schema_digest,
                        "expires-at": expires_at,
                        "artifact-bytes": str(size),
                        "schema-contract": schema_encoded,
                    },
                )
            uploaded = True
            self._validated_stat(
                key,
                expected_digest=digest_hex,
                limits=effective,
            )
            artifact_ref = f"minio://{self.bucket}/{key}"
            self.read(
                artifact_ref,
                digest_hex,
                expected_schema_fields=fields,
                limits=effective,
            )
        except Exception:
            if uploaded:
                with suppress(Exception):
                    self.client.remove_object(self.bucket, key)
            raise
        finally:
            if path is not None:
                with suppress(FileNotFoundError):
                    os.unlink(path)
        return {
            "artifact_ref": artifact_ref,
            "digest": digest_hex,
            "row_count": collected.height,
            "schema_hash": schema_digest,
            "schema_fields": fields,
            "expires_at": expires_at,
        }

    def describe(self, ref: str) -> dict[str, Any]:
        """Return validated object metadata without exposing MinIO credentials."""

        key = self._parse_ref(ref)
        _stat, metadata = self._validated_stat(key)
        fields = _decode_schema_contract(metadata["schema-contract"])
        if canonical_schema_hash(fields) != metadata["schema-sha256"]:
            raise ValueError("artifact schema contract does not match")
        return {
            "artifact_ref": ref,
            "digest": metadata["sha256"],
            "row_count": int(metadata["row-count"]),
            "schema_hash": metadata["schema-sha256"],
            "schema_fields": fields,
            "expires_at": metadata["expires-at"],
        }

    def read(
        self,
        ref: str,
        expected_digest: str,
        *,
        expected_schema_fields: list[dict[str, Any]] | None = None,
        limits: dict[str, int] | None = None,
    ) -> pl.LazyFrame:
        effective = self._limits(limits)
        if _DIGEST.fullmatch(str(expected_digest or "")) is None:
            raise ValueError("expected artifact digest is invalid")
        key = self._parse_ref(ref)
        _stat, metadata = self._validated_stat(
            key,
            expected_digest=expected_digest,
            limits=effective,
        )
        fields = _decode_schema_contract(metadata["schema-contract"])
        if canonical_schema_hash(fields) != metadata["schema-sha256"]:
            raise ValueError("artifact schema contract does not match")
        if expected_schema_fields is not None:
            expected = _normalized_schema_fields(expected_schema_fields)
            if expected != fields:
                raise ValueError("artifact schema contract is not expected")
        response = self.client.get_object(self.bucket, key)
        digest = hashlib.sha256()
        payload = io.BytesIO()
        size = 0
        try:
            while chunk := response.read(1024 * 1024):
                size += len(chunk)
                if size > effective["max_artifact_bytes"]:
                    raise ValueError(
                        "artifact size exceeds the configured limit"
                    )
                if size > effective["memory_limit_bytes"]:
                    raise ValueError(
                        "artifact download exceeds the memory limit"
                    )
                digest.update(chunk)
                payload.write(chunk)
        finally:
            response.close()
            release = getattr(response, "release_conn", None)
            if callable(release):
                release()
        if digest.hexdigest() != expected_digest:
            raise ValueError("artifact digest does not match content")
        payload.seek(0)
        try:
            frame = pl.read_parquet(
                payload,
                n_rows=effective["max_rows"] + 1,
                memory_map=False,
            )
        except Exception as exc:
            raise ValueError("artifact is not valid Parquet") from exc
        if frame.height != int(metadata["row-count"]):
            raise ValueError("artifact row count does not match metadata")
        if frame.height > effective["max_rows"]:
            raise ValueError("artifact row count exceeds the configured limit")
        if frame.estimated_size() > effective["memory_limit_bytes"]:
            raise ValueError("artifact frame exceeds the configured memory limit")
        if size + frame.estimated_size() > effective["memory_limit_bytes"]:
            raise ValueError("artifact decompression exceeds the memory limit")
        _validate_frame_schema(frame, fields)
        return frame.lazy()

    def cleanup_expired(self, correlation_id: str) -> int:
        correlation = _uid(correlation_id, "correlation_id")
        prefix = f"rules/{correlation}/"
        removed = 0
        for item in self.client.list_objects(
            self.bucket,
            prefix=prefix,
            recursive=True,
        ):
            key = str(getattr(item, "object_name", ""))
            if not key.startswith(prefix):
                continue
            try:
                self._parse_ref(f"minio://{self.bucket}/{key}")
                stat = self.client.stat_object(self.bucket, key)
                metadata = _metadata(getattr(stat, "metadata", None))
                expired = _parse_timestamp(
                    metadata["expires-at"]
                ) <= _now_utc(self.clock)
            except ValueError:
                continue
            if expired:
                self.client.remove_object(self.bucket, key)
                removed += 1
        return removed

    def delete(self, ref: str) -> None:
        """Delete one exact store-owned artifact after validating its key."""

        key = self._parse_ref(ref)
        self.client.remove_object(self.bucket, key)


class PostgresArtifactResolver:
    """Resolve a canonical artifact binding without accepting caller paths."""

    def __init__(self, engine, artifact_store: ArtifactStore):
        self.engine = engine
        self.artifact_store = artifact_store

    def resolve(self, *, binding_id: str, correlation_id: str) -> dict[str, Any]:
        binding = _uid(binding_id, "artifact binding id")
        correlation = _uid(correlation_id, "artifact correlation id")
        statement = text(
            """
            SELECT
                a.artifact_ref,
                a.artifact_digest,
                a.row_count,
                a.schema_hash,
                a.schema_fields,
                a.expires_at,
                b.binding_hash
            FROM public.rule_run_artifacts a
            JOIN public.dataflow_dataset_bindings b
              ON b.id = a.binding_id
            WHERE a.binding_id = CAST(:binding_id AS uuid)
              AND a.correlation_id = CAST(:correlation_id AS uuid)
              AND a.expires_at > CURRENT_TIMESTAMP
              AND b.object_kind = 'parquet_artifact'
              AND b.access_mode IN ('read', 'read_write')
            ORDER BY a.created_at DESC
            LIMIT 1
            """
        )
        with self.engine.connect() as connection:
            row = connection.execute(
                statement,
                {
                    "binding_id": binding,
                    "correlation_id": correlation,
                },
            ).mappings().one_or_none()
        if row is None:
            raise ValueError("canonical artifact binding was not found")
        artifact_ref = str(row["artifact_ref"])
        key = self.artifact_store._parse_ref(artifact_ref)
        if not key.startswith(f"rules/{correlation}/"):
            raise ValueError(
                "catalog artifact does not match the execution correlation"
            )
        described = self.artifact_store.describe(artifact_ref)
        row_fields = row["schema_fields"]
        if isinstance(row_fields, str):
            row_fields = json.loads(row_fields)
        if (
            described["digest"] != str(row["artifact_digest"])
            or described["row_count"] != int(row["row_count"])
            or described["schema_hash"] != str(row["schema_hash"])
            or described["schema_fields"]
            != _normalized_schema_fields(row_fields)
        ):
            raise ValueError("catalog artifact metadata does not match storage")
        return {
            **described,
            "binding_hash": str(row["binding_hash"]),
        }

    def attest_binding(
        self,
        *,
        binding_id: str,
        binding_hash: str,
        access_mode: str,
    ) -> dict[str, str]:
        binding = _uid(binding_id, "artifact binding id")
        if _DIGEST.fullmatch(str(binding_hash or "")) is None:
            raise ValueError("artifact binding hash is invalid")
        allowed = {
            "read": {"read", "read_write"},
            "write": {"write", "read_write"},
        }.get(access_mode)
        if allowed is None:
            raise ValueError("artifact access mode is invalid")
        with self.engine.connect() as connection:
            row = connection.execute(
                text(
                    """
                    SELECT binding_hash, access_mode, object_kind
                    FROM public.dataflow_dataset_bindings
                    WHERE id = CAST(:binding_id AS uuid)
                    """
                ),
                {"binding_id": binding},
            ).mappings().one_or_none()
        if (
            row is None
            or row["object_kind"] != "parquet_artifact"
            or row["access_mode"] not in allowed
            or str(row["binding_hash"]) != binding_hash
        ):
            raise ValueError("canonical artifact binding no longer matches")
        return {"binding_hash": str(row["binding_hash"])}

    def register(
        self,
        *,
        binding_id: str,
        correlation_id: str,
        artifact: dict[str, Any],
        kind: str,
        binding_hash: str,
    ) -> None:
        binding = _uid(binding_id, "artifact binding id")
        correlation = _uid(correlation_id, "artifact correlation id")
        if kind not in {"input", "lookup", "output"}:
            raise ValueError("artifact kind is invalid")
        self.attest_binding(
            binding_id=binding,
            binding_hash=binding_hash,
            access_mode="write" if kind == "output" else "read",
        )
        if not isinstance(artifact, dict):
            raise ValueError("artifact metadata is invalid")
        artifact_ref = artifact.get("artifact_ref")
        key = self.artifact_store._parse_ref(artifact_ref)
        if not key.startswith(f"rules/{correlation}/"):
            raise ValueError(
                "artifact does not match the execution correlation"
            )
        described = self.artifact_store.describe(artifact_ref)
        for key in (
            "artifact_ref",
            "digest",
            "row_count",
            "schema_hash",
            "schema_fields",
            "expires_at",
        ):
            if described[key] != artifact.get(key):
                raise ValueError("artifact metadata does not match storage")
        with self.engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO public.rule_run_artifacts (
                        id, correlation_id, binding_id, artifact_ref,
                        artifact_digest, row_count, schema_hash, schema_fields,
                        artifact_kind, expires_at
                    ) VALUES (
                        CAST(:id AS uuid), CAST(:correlation_id AS uuid),
                        CAST(:binding_id AS uuid), :artifact_ref,
                        :artifact_digest, :row_count, :schema_hash,
                        CAST(:schema_fields AS jsonb), :artifact_kind,
                        CAST(:expires_at AS timestamptz)
                    )
                    ON CONFLICT (
                        correlation_id, binding_id, artifact_digest
                    ) DO NOTHING
                    """
                ),
                {
                    "id": new_governance_uid(),
                    "correlation_id": correlation,
                    "binding_id": binding,
                    "artifact_ref": described["artifact_ref"],
                    "artifact_digest": described["digest"],
                    "row_count": described["row_count"],
                    "schema_hash": described["schema_hash"],
                    "schema_fields": json.dumps(
                        described["schema_fields"],
                        sort_keys=True,
                        separators=(",", ":"),
                    ),
                    "artifact_kind": kind,
                    "expires_at": described["expires_at"],
                },
            )