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"]) == {
        key: artifact[key]
        for key in (
            "artifact_ref",
            "digest",
            "row_count",
            "schema_hash",
            "expires_at",
        )
    }
    stored = next(iter(client.objects.values()))
    assert "x-amz-meta-schema-contract" not in stored["metadata"]
    assert sum(
        len(key) + len(value)
        for key, value in stored["metadata"].items()
    ) <= 2_048
    frame = store.read(
        artifact["artifact_ref"],
        artifact["digest"],
        expected_schema_fields=artifact["schema_fields"],
    )
    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"],
            expected_schema_fields=artifact["schema_fields"],
        )


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,
            expected_schema_fields=[
                {"name": "id", "type": "integer", "nullable": True}
            ],
        )


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"],
            expected_schema_fields=artifact["schema_fields"],
            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_stage_streams_to_tempfile_and_preflights_parquet_footer(
):
    import inspect

    from app.runner.artifacts import ArtifactStore

    client = FakeMinio()
    store = _store(client)
    correlation_id = new_governance_uid()
    artifact = store.write(
        pl.DataFrame({"id": [1, 2]}),
        correlation_id,
        60,
    )

    assert "BytesIO" not in inspect.getsource(ArtifactStore.stage)

    with store.stage(
        artifact["artifact_ref"],
        artifact["digest"],
        expected_schema_fields=artifact["schema_fields"],
    ) as staged:
        assert pl.scan_parquet(staged).collect().height == 2


def test_artifact_stage_rejects_compression_bomb_from_footer_before_scan():
    from app.runner.artifacts import ArtifactStore

    client = FakeMinio()
    store = ArtifactStore(
        client,
        bucket="dataops-rules",
        max_artifact_bytes=4 * 1024 * 1024,
        max_rows=1_000,
        memory_limit_bytes=64 * 1024 * 1024,
        max_ttl_seconds=3600,
    )
    correlation_id = new_governance_uid()
    artifact = store.write(
        pl.DataFrame(
            {
                "payload": [
                    f"{index}-" + ("compressible-value-" * 10_000)
                    for index in range(100)
                ]
            }
        ),
        correlation_id,
        60,
    )
    calls_before = len(client.get_calls)

    with pytest.raises(
        ValueError, match="uncompressed|footer"
    ), store.stage(
        artifact["artifact_ref"],
        artifact["digest"],
        expected_schema_fields=artifact["schema_fields"],
        limits={
            "max_rows": 1_000,
            "max_artifact_bytes": 4 * 1024 * 1024,
            "memory_limit_bytes": 1024 * 1024,
        },
    ):
        raise AssertionError("compression bomb must not be exposed")
    assert len(client.get_calls) == calls_before + 1


def test_artifact_prepare_reserves_key_before_exact_path_upload(tmp_path):
    client = FakeMinio()
    store = _store(client)
    correlation_id = new_governance_uid()
    path = tmp_path / "prepared.parquet"
    pl.DataFrame({"id": [1, 2]}).write_parquet(path)
    fields = [{"name": "id", "type": "integer", "nullable": False}]

    prepared = store.prepare_path(
        str(path),
        correlation_id,
        60,
        schema_fields=fields,
    )

    assert client.objects == {}
    assert prepared["artifact_ref"].startswith(
        f"minio://dataops-rules/rules/{correlation_id}/"
    )

    uploaded = store.upload_path(str(path), prepared)

    assert uploaded == prepared
    assert store.describe(prepared["artifact_ref"])["digest"] == prepared[
        "digest"
    ]


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

    def all(self):
        if self.row is None:
            return []
        return self.row if isinstance(self.row, list) else [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"],
                "catalog_binding_hash": binding_hash,
                "current_binding_hash": binding_hash,
            }
        raise AssertionError(sql)

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

    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]
    )
    assert "a.handoff_status = 'ready'" in engine.calls[0][0]
    assert "a.binding_hash = b.binding_hash" in engine.calls[0][0]

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


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

    store = _store(FakeMinio())
    correlation_id = new_governance_uid()
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        correlation_id,
        60,
    )
    row_id = new_governance_uid()
    deleted_ids = []

    def handler(sql, parameters):
        if "DELETE FROM public.rule_run_artifacts" in sql:
            deleted_ids.append(parameters["id"])
            return None
        if "FROM public.rule_run_artifacts" in sql:
            assert "FOR UPDATE SKIP LOCKED" in sql
            assert parameters["limit"] == 10
            return [{"id": row_id, "artifact_ref": artifact["artifact_ref"]}]
        raise AssertionError(sql)

    resolver = PostgresArtifactResolver(_Engine(handler), store)

    assert resolver.cleanup_expired(limit=10) == 1
    assert deleted_ids == [row_id]
    assert not store.client.objects