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 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,
)

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 _schema_hash(frame: pl.DataFrame | pl.LazyFrame) -> str:
    schema = (
        frame.collect_schema()
        if isinstance(frame, pl.LazyFrame)
        else frame.schema
    )
    canonical = [
        {"name": name, "dtype": str(dtype)}
        for name, dtype in schema.items()
    ]
    encoded = json.dumps(
        canonical,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )
    return hashlib.sha256(encoded.encode("utf-8")).hexdigest()


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")
    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 _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,
    ) -> tuple[Any, dict[str, str]]:
        stat = self.client.stat_object(self.bucket, key)
        size = int(getattr(stat, "size", -1))
        if size < 1 or size > self.max_artifact_bytes:
            raise ValueError("artifact size exceeds the configured 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 > self.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,
    ) -> dict[str, Any]:
        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(self.max_rows + 1).collect(engine="streaming")
        if collected.height > self.max_rows:
            raise ValueError("artifact row count exceeds the configured limit")
        if collected.estimated_size() > self.memory_limit_bytes:
            raise ValueError("artifact frame exceeds the configured memory limit")
        schema_digest = _schema_hash(collected)
        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
        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 > self.max_artifact_bytes:
                raise ValueError("artifact size exceeds the configured 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),
                    },
                )
            self._validated_stat(key, expected_digest=digest_hex)
            artifact_ref = f"minio://{self.bucket}/{key}"
            self.read(artifact_ref, digest_hex)
        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,
            "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)
        return {
            "artifact_ref": ref,
            "digest": metadata["sha256"],
            "row_count": int(metadata["row-count"]),
            "schema_hash": metadata["schema-sha256"],
            "expires_at": metadata["expires-at"],
        }

    def read(self, ref: str, expected_digest: str) -> pl.LazyFrame:
        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
        )
        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 > self.max_artifact_bytes:
                    raise ValueError(
                        "artifact size exceeds the configured 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)
        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 > self.max_rows:
            raise ValueError("artifact row count exceeds the configured limit")
        if frame.estimated_size() > self.memory_limit_bytes:
            raise ValueError("artifact frame exceeds the configured memory limit")
        if _schema_hash(frame) != metadata["schema-sha256"]:
            raise ValueError("artifact schema does not match metadata")
        return frame.lazy()


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 object_ref, binding_hash
            FROM public.dataflow_dataset_bindings
            WHERE id = CAST(:binding_id AS uuid)
              AND object_kind = 'parquet_artifact'
              AND access_mode IN ('read', 'read_write')
            """
        )
        with self.engine.connect() as connection:
            row = connection.execute(
                statement, {"binding_id": binding}
            ).mappings().one_or_none()
        if row is None:
            raise ValueError("canonical artifact binding was not found")
        artifact_ref = str(row["object_ref"])
        if f"/rules/{correlation}/" not in artifact_ref:
            raise ValueError(
                "artifact binding does not match the execution correlation"
            )
        return {
            **self.artifact_store.describe(artifact_ref),
            "binding_hash": str(row["binding_hash"]),
        }