DataOps-platform/tests/runner/test_artifacts.py

from __future__ import annotations

import io
from datetime import UTC, datetime
from types import SimpleNamespace

import polars as pl
import pytest

from app.core.common.identifiers import new_governance_uid


class Response(io.BytesIO):
    def release_conn(self):
        return None


class FakeMinio:
    def __init__(self):
        self.buckets = {"dataops-rules"}
        self.objects = {}
        self.get_calls = []
        self.removed = []

    def bucket_exists(self, bucket):
        return bucket in self.buckets

    def make_bucket(self, bucket):
        self.buckets.add(bucket)

    def put_object(
        self,
        bucket,
        key,
        data,
        length,
        *,
        content_type,
        metadata,
    ):
        payload = data.read(length)
        self.objects[(bucket, key)] = {
            "payload": payload,
            "content_type": content_type,
            "metadata": {
                f"x-amz-meta-{name.lower()}": str(value)
                for name, value in metadata.items()
            },
        }

    def stat_object(self, bucket, key):
        item = self.objects[(bucket, key)]
        return SimpleNamespace(
            size=len(item["payload"]),
            content_type=item["content_type"],
            metadata=item["metadata"],
        )

    def get_object(self, bucket, key):
        self.get_calls.append((bucket, key))
        return Response(self.objects[(bucket, key)]["payload"])

    def remove_object(self, bucket, key):
        self.removed.append((bucket, key))
        self.objects.pop((bucket, key), None)

    def list_objects(self, bucket, *, prefix, recursive):
        assert recursive is True
        return [
            SimpleNamespace(object_name=key)
            for object_bucket, key in sorted(self.objects)
            if object_bucket == bucket and key.startswith(prefix)
        ]


def _store(client, *, clock=None, max_rows=100):
    from app.runner.artifacts import ArtifactStore

    return ArtifactStore(
        client,
        bucket="dataops-rules",
        max_artifact_bytes=1024 * 1024,
        max_rows=max_rows,
        memory_limit_bytes=4 * 1024 * 1024,
        max_ttl_seconds=3600,
        clock=clock,
    )


def test_artifact_store_generates_key_and_round_trips_digest_bound_lazyframe():
    client = FakeMinio()
    store = _store(client)
    correlation_id = new_governance_uid()

    artifact = store.write(
        pl.DataFrame(
            {
                "customer_id": [1, 2],
                "name": ["Alice", "Bob"],
            }
        ).lazy(),
        correlation_id,
        300,
    )

    assert artifact["artifact_ref"].startswith(
        f"minio://dataops-rules/rules/{correlation_id}/"
    )
    assert artifact["artifact_ref"].endswith(".parquet")
    assert artifact["digest"]
    assert artifact["row_count"] == 2
    assert artifact["schema_hash"]
    assert artifact["expires_at"].endswith("Z")
    assert "dataops-test" not in repr(artifact)
    assert store.describe(artifact["artifact_ref"]) == artifact
    frame = store.read(artifact["artifact_ref"], artifact["digest"])
    assert isinstance(frame, pl.LazyFrame)
    assert frame.collect().to_dicts() == [
        {"customer_id": 1, "name": "Alice"},
        {"customer_id": 2, "name": "Bob"},
    ]


@pytest.mark.parametrize(
    ("mutation", "message"),
    [
        (
            lambda item: item["metadata"].update(
                {"x-amz-meta-sha256": "0" * 64}
            ),
            "digest",
        ),
        (
            lambda item: item["metadata"].update(
                {"x-amz-meta-schema-sha256": "0" * 64}
            ),
            "schema",
        ),
        (
            lambda item: item["metadata"].update(
                {"x-amz-meta-expires-at": "2000-01-01T00:00:00Z"}
            ),
            "expired",
        ),
        (
            lambda item: item.update({"content_type": "text/plain"}),
            "content type",
        ),
    ],
)
def test_artifact_store_rejects_tampered_digest_schema_ttl_and_content(
    mutation, message
):
    client = FakeMinio()
    now = datetime(2026, 7, 23, 10, 0, tzinfo=UTC)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1]}).lazy(),
        new_governance_uid(),
        300,
    )
    key = artifact["artifact_ref"].split("/", 3)[-1]
    mutation(client.objects[("dataops-rules", key)])

    with pytest.raises(ValueError, match=message):
        store.read(artifact["artifact_ref"], artifact["digest"])


def test_artifact_store_rejects_rows_size_ttl_and_unowned_references():
    client = FakeMinio()
    store = _store(client, max_rows=2)

    with pytest.raises(ValueError, match="row"):
        store.write(
            pl.DataFrame({"id": [1, 2, 3]}).lazy(),
            new_governance_uid(),
            60,
        )
    with pytest.raises(ValueError, match="TTL"):
        store.write(
            pl.DataFrame({"id": [1]}).lazy(),
            new_governance_uid(),
            7200,
        )
    with pytest.raises(ValueError, match="artifact reference"):
        store.read(
            "minio://other-bucket/rules/unsafe/value.parquet",
            "0" * 64,
        )


def test_artifact_store_does_not_return_ref_for_corrupted_server_content():
    class CorruptingMinio(FakeMinio):
        def put_object(self, bucket, key, *args, **kwargs):
            super().put_object(bucket, key, *args, **kwargs)
            payload = self.objects[(bucket, key)]["payload"]
            self.objects[(bucket, key)]["payload"] = bytes(
                [payload[0] ^ 1]
            ) + payload[1:]

    store = _store(CorruptingMinio())

    with pytest.raises(ValueError, match="digest|size"):
        store.write(
            pl.DataFrame({"id": [1]}).lazy(),
            new_governance_uid(),
            60,
        )
    assert not store.client.objects


def test_artifact_store_applies_plan_limits_before_download_and_serialization():
    client = FakeMinio()
    store = _store(client, max_rows=100)
    correlation_id = new_governance_uid()
    artifact = store.write(
        pl.DataFrame({"id": [1, 2]}).lazy(),
        correlation_id,
        60,
    )
    calls_before = len(client.get_calls)

    with pytest.raises(ValueError, match="row count"):
        store.read(
            artifact["artifact_ref"],
            artifact["digest"],
            limits={
                "max_rows": 1,
                "max_artifact_bytes": 1024 * 1024,
                "memory_limit_bytes": 4 * 1024 * 1024,
            },
        )
    assert len(client.get_calls) == calls_before

    with pytest.raises(ValueError, match="row count"):
        store.write(
            pl.DataFrame({"id": [1, 2]}).lazy(),
            correlation_id,
            60,
            limits={
                "max_rows": 1,
                "max_artifact_bytes": 1024 * 1024,
                "memory_limit_bytes": 4 * 1024 * 1024,
            },
        )


def test_artifact_schema_contract_covers_nullability_decimal_and_timezone():
    client = FakeMinio()
    store = _store(client)
    correlation_id = new_governance_uid()
    fields = [
        {
            "name": "amount",
            "type": "decimal",
            "nullable": False,
            "precision": 12,
            "scale": 2,
        },
        {
            "name": "occurred_at",
            "type": "timestamptz",
            "nullable": False,
            "timezone": "Asia/Shanghai",
        },
    ]
    frame = pl.DataFrame(
        {
            "amount": ["12.34"],
            "occurred_at": ["2026-07-23T10:00:00+08:00"],
        }
    ).with_columns(
        pl.col("amount").cast(pl.Decimal(12, 2)),
        pl.col("occurred_at")
        .str.to_datetime(time_zone="Asia/Shanghai")
        .alias("occurred_at"),
    )
    artifact = store.write(
        frame.lazy(),
        correlation_id,
        60,
        schema_fields=fields,
    )

    assert artifact["schema_fields"] == fields
    store.read(
        artifact["artifact_ref"],
        artifact["digest"],
        expected_schema_fields=fields,
    ).collect()

    copy_fields = [dict(field) for field in fields]
    copy_fields[1] = {**copy_fields[1], "timezone": "UTC"}
    with pytest.raises(ValueError, match="timezone|schema"):
        store.read(
            artifact["artifact_ref"],
            artifact["digest"],
            expected_schema_fields=copy_fields,
        )

    with pytest.raises(ValueError, match="nullable"):
        store.write(
            pl.DataFrame(
                {
                    "amount": [None],
                    "occurred_at": [None],
                },
                schema={
                    "amount": pl.Decimal(12, 2),
                    "occurred_at": pl.Datetime(
                        "us", "Asia/Shanghai"
                    ),
                },
            ).lazy(),
            correlation_id,
            60,
            schema_fields=fields,
        )


def test_artifact_cleanup_is_expired_and_correlation_scoped_only():
    client = FakeMinio()
    now = datetime(2026, 7, 23, 10, 0, tzinfo=UTC)
    store = _store(client, clock=lambda: now)
    first = new_governance_uid()
    second = new_governance_uid()
    expired = store.write(pl.DataFrame({"id": [1]}), first, 10)
    active = store.write(pl.DataFrame({"id": [2]}), second, 300)
    now = datetime(2026, 7, 23, 10, 0, 20, tzinfo=UTC)

    assert store.cleanup_expired(first) == 1
    assert not any(
        key == expired["artifact_ref"].split("/", 3)[-1]
        for _bucket, key in client.objects
    )
    assert store.describe(active["artifact_ref"])["row_count"] == 1


class _Rows:
    def __init__(self, row=None):
        self.row = row

    def mappings(self):
        return self

    def one_or_none(self):
        return self.row


class _Connection:
    def __init__(self, engine):
        self.engine = engine

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return None

    def execute(self, statement, parameters):
        sql = str(statement)
        self.engine.calls.append((sql, dict(parameters)))
        return _Rows(self.engine.handler(sql, parameters))


class _Engine:
    def __init__(self, handler):
        self.handler = handler
        self.calls = []

    def connect(self):
        return _Connection(self)

    def begin(self):
        return _Connection(self)


def test_postgres_artifact_resolver_uses_catalog_and_rechecks_binding_hash():
    from app.runner.artifacts import PostgresArtifactResolver

    store = _store(FakeMinio())
    correlation_id = new_governance_uid()
    binding_id = new_governance_uid()
    binding_hash = "a" * 64
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        correlation_id,
        60,
    )

    def handler(sql, _parameters):
        if "FROM public.rule_run_artifacts" in sql:
            return {
                **artifact,
                "artifact_digest": artifact["digest"],
                "binding_hash": binding_hash,
            }
        raise AssertionError(sql)

    engine = _Engine(handler)
    resolved = PostgresArtifactResolver(engine, store).resolve(
        binding_id=binding_id,
        correlation_id=correlation_id,
    )

    assert resolved["artifact_ref"] == artifact["artifact_ref"]
    assert resolved["binding_hash"] == binding_hash
    assert "dataflow_dataset_bindings b" in engine.calls[0][0]
    assert "a.correlation_id = CAST(:correlation_id AS uuid)" in (
        engine.calls[0][0]
    )

    with pytest.raises(ValueError, match="correlation"):
        PostgresArtifactResolver(engine, store).resolve(
            binding_id=binding_id,
            correlation_id=new_governance_uid(),
        )


def test_postgres_artifact_resolver_attests_and_registers_stable_handoff():
    from app.runner.artifacts import PostgresArtifactResolver

    store = _store(FakeMinio())
    correlation_id = new_governance_uid()
    binding_id = new_governance_uid()
    binding_hash = "b" * 64
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        correlation_id,
        60,
    )

    def handler(sql, parameters):
        if "FROM public.dataflow_dataset_bindings" in sql:
            return {
                "binding_hash": binding_hash,
                "access_mode": "read_write",
                "object_kind": "parquet_artifact",
            }
        if "INSERT INTO public.rule_run_artifacts" in sql:
            assert parameters["artifact_digest"] == artifact["digest"]
            assert parameters["schema_fields"]
            return None
        raise AssertionError(sql)

    engine = _Engine(handler)
    resolver = PostgresArtifactResolver(engine, store)
    resolver.attest_binding(
        binding_id=binding_id,
        binding_hash=binding_hash,
        access_mode="write",
    )
    resolver.register(
        binding_id=binding_id,
        correlation_id=correlation_id,
        artifact=artifact,
        kind="output",
        binding_hash=binding_hash,
    )

    assert sum(
        "FROM public.dataflow_dataset_bindings" in sql
        for sql, _parameters in engine.calls
    ) == 2
    assert any(
        "INSERT INTO public.rule_run_artifacts" in sql
        for sql, _parameters in engine.calls
    )