DataOps-platform/tests/runner/test_artifact_handoff.py

from __future__ import annotations

import copy
import json
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace

import polars as pl
import pytest
from minio.error import S3Error

from app.core.common.identifiers import new_governance_uid
from tests.runner.test_artifacts import FakeMinio, _store


class _Result:
    def __init__(self, row=None, *, rowcount=0):
        self.row = row
        self.rowcount = rowcount

    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 _HandoffConnection:
    def __init__(self, engine, *, transactional):
        self.engine = engine
        self.transactional = transactional
        self.backup = None
        self.finalized = False

    def __enter__(self):
        self.backup = copy.deepcopy(self.engine.row)
        return self

    def __exit__(self, exc_type, *_args):
        if exc_type is not None:
            self.engine.row = self.backup
            return False
        if self.finalized and self.engine.finalize_mode:
            if self.engine.finalize_mode == "rollback_unknown":
                self.engine.row = self.backup
                self.engine.fail_recheck = True
            raise RuntimeError("database commit acknowledgement lost")
        return False

    def execute(self, statement, parameters):
        sql = str(statement)
        self.engine.events.append(sql)
        if (
            "FROM public.dataflow_dataset_bindings" in sql
            and "JOIN" not in sql
        ):
            if self.engine.fail_reserve:
                raise RuntimeError("database unavailable")
            return _Result(copy.deepcopy(self.engine.binding))
        if "INSERT INTO public.rule_run_artifacts" in sql:
            if self.engine.row is not None:
                return _Result(None)
            self.engine.row = {
                "id": parameters["id"],
                "correlation_id": parameters["correlation_id"],
                "binding_id": parameters["binding_id"],
                "artifact_ref": parameters["artifact_ref"],
                "artifact_digest": parameters["artifact_digest"],
                "row_count": parameters["row_count"],
                "schema_hash": parameters["schema_hash"],
                "schema_fields": json.loads(parameters["schema_fields"]),
                "artifact_kind": parameters["artifact_kind"],
                "binding_hash": parameters["binding_hash"],
                "expires_at": parameters["expires_at"],
                "handoff_status": "pending",
            }
            return _Result(copy.deepcopy(self.engine.row), rowcount=1)
        if (
            "UPDATE public.rule_run_artifacts" in sql
            and "handoff_status = 'ready'" in sql
        ):
            if self.engine.row is not None:
                self.engine.row["handoff_status"] = "ready"
                self.finalized = True
                return _Result(copy.deepcopy(self.engine.row), rowcount=1)
            return _Result(None)
        if (
            "DELETE FROM public.rule_run_artifacts" in sql
            and "handoff_status = 'pending'" in sql
        ):
            deleted = self.engine.row is not None
            self.engine.row = None
            return _Result(rowcount=int(deleted))
        if "FROM public.rule_run_artifacts" in sql:
            if self.engine.fail_recheck:
                raise RuntimeError("database recheck unavailable")
            row = copy.deepcopy(self.engine.row)
            if "JOIN public.dataflow_dataset_bindings" in sql:
                if (
                    row is None
                    or row["handoff_status"] != "ready"
                    or row["binding_hash"] != self.engine.binding[
                        "binding_hash"
                    ]
                ):
                    row = None
                elif row is not None:
                    row["current_binding_hash"] = self.engine.binding[
                        "binding_hash"
                    ]
            return _Result(row)
        raise AssertionError(sql)


class HandoffEngine:
    def __init__(self, binding):
        self.binding = dict(binding)
        self.row = None
        self.events = []
        self.fail_reserve = False
        self.fail_recheck = False
        self.finalize_mode = None

    def begin(self):
        return _HandoffConnection(self, transactional=True)

    def connect(self):
        return _HandoffConnection(self, transactional=False)


class EventMinio(FakeMinio):
    def __init__(self, events):
        super().__init__()
        self.events = events
        self.fail_upload = False

    def put_object(self, *args, **kwargs):
        self.events.append("MINIO PUT")
        if self.fail_upload:
            raise RuntimeError("object upload failed")
        return super().put_object(*args, **kwargs)


def _binding(binding_hash):
    return {
        "binding_hash": binding_hash,
        "access_mode": "read_write",
        "object_kind": "parquet_artifact",
    }


def _publish_fixture(tmp_path, *, engine=None, client=None):
    from app.runner.artifacts import PostgresArtifactResolver

    binding_id = new_governance_uid()
    binding_hash = "b" * 64
    engine = engine or HandoffEngine(_binding(binding_hash))
    client = client or EventMinio(engine.events)
    store = _store(client)
    path = tmp_path / "handoff.parquet"
    pl.DataFrame({"id": [1]}).write_parquet(path)
    resolver = PostgresArtifactResolver(engine, store)
    return {
        "binding_id": binding_id,
        "binding_hash": binding_hash,
        "client": client,
        "engine": engine,
        "path": path,
        "resolver": resolver,
        "store": store,
    }


def _publish(fixture):
    return fixture["resolver"].publish_path(
        str(fixture["path"]),
        binding_id=fixture["binding_id"],
        binding_hash=fixture["binding_hash"],
        correlation_id=new_governance_uid(),
        kind="output",
        ttl_seconds=60,
        schema_fields=[
            {"name": "id", "type": "integer", "nullable": False}
        ],
    )


def test_publish_reserves_pending_before_upload_and_finalizes_ready(tmp_path):
    fixture = _publish_fixture(tmp_path)

    result = _publish(fixture)

    events = fixture["engine"].events
    assert next(
        index
        for index, event in enumerate(events)
        if "INSERT INTO public.rule_run_artifacts" in event
    ) < events.index("MINIO PUT")
    assert events.index("MINIO PUT") < next(
        index
        for index, event in enumerate(events)
        if "handoff_status = 'ready'" in event
    )
    assert fixture["engine"].row["handoff_status"] == "ready"
    assert fixture["engine"].row["binding_hash"] == fixture["binding_hash"]
    assert result["artifact_ref"] == fixture["engine"].row["artifact_ref"]


def test_reserve_locks_and_rejects_stale_binding_before_upload(tmp_path):
    fixture = _publish_fixture(tmp_path)
    fixture["engine"].binding["binding_hash"] = "c" * 64

    with pytest.raises(ValueError, match="binding"):
        _publish(fixture)

    assert any(
        "FOR SHARE" in event
        for event in fixture["engine"].events
        if "dataflow_dataset_bindings" in event
    )
    assert "MINIO PUT" not in fixture["engine"].events
    assert fixture["engine"].row is None


def test_reserve_crash_never_uploads_an_object(tmp_path):
    from app.runner.artifacts import ArtifactCommitUnknown

    fixture = _publish_fixture(tmp_path)
    fixture["engine"].fail_reserve = True

    with pytest.raises(ArtifactCommitUnknown):
        _publish(fixture)

    assert "MINIO PUT" not in fixture["engine"].events
    assert fixture["client"].objects == {}


def test_upload_failure_removes_pending_reservation(tmp_path):
    fixture = _publish_fixture(tmp_path)
    fixture["client"].fail_upload = True

    with pytest.raises(RuntimeError, match="upload"):
        _publish(fixture)

    assert fixture["engine"].row is None
    assert fixture["client"].objects == {}


def test_finalize_commit_after_client_error_is_rechecked_as_success(tmp_path):
    fixture = _publish_fixture(tmp_path)
    fixture["engine"].finalize_mode = "committed_then_error"

    result = _publish(fixture)

    assert fixture["engine"].row["handoff_status"] == "ready"
    assert result["artifact_ref"] == fixture["engine"].row["artifact_ref"]
    assert fixture["client"].objects


def test_unconfirmed_finalize_returns_unknown_without_deleting_object(tmp_path):
    from app.runner.artifacts import ArtifactCommitUnknown

    fixture = _publish_fixture(tmp_path)
    fixture["engine"].finalize_mode = "rollback_unknown"

    with pytest.raises(ArtifactCommitUnknown):
        _publish(fixture)

    assert fixture["client"].objects


class ReconcileMinio(FakeMinio):
    def __init__(self, now):
        super().__init__()
        self.now = now
        self.ages = {}

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


class ReconcileEngine:
    def __init__(self, rows):
        self.rows = [dict(row) for row in rows]
        self.events = []
        self.finalize_race_to_ready = False

    def begin(self):
        return _ReconcileConnection(self)

    def connect(self):
        return _ReconcileConnection(self)


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

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return False

    def execute(self, statement, parameters):
        sql = str(statement)
        self.engine.events.append(sql)
        if (
            "SELECT id::text AS id" in sql
            and "handoff_status IN" in sql
        ):
            return _Result(
                copy.deepcopy(self.engine.rows[: parameters["limit"]])
            )
        if (
            "SELECT id::text AS id" in sql
            and "WHERE id = CAST(:id AS uuid)" in sql
        ):
            row = next(
                (
                    item
                    for item in self.engine.rows
                    if item["id"] == parameters["id"]
                ),
                None,
            )
            return _Result(copy.deepcopy(row))
        if "SELECT artifact_ref" in sql and "handoff_status" not in sql:
            return _Result(
                [{"artifact_ref": row["artifact_ref"]} for row in self.engine.rows]
            )
        row = next(
            (
                item
                for item in self.engine.rows
                if item["id"] == parameters.get("id")
            ),
            None,
        )
        if (
            "UPDATE public.rule_run_artifacts" in sql
            and "handoff_status = 'ready'" in sql
        ):
            if self.engine.finalize_race_to_ready and row is not None:
                row["handoff_status"] = "ready"
                return _Result(None, rowcount=0)
            if row is not None and row["handoff_status"] == "pending":
                row["handoff_status"] = "ready"
                return _Result(copy.deepcopy(row), rowcount=1)
            return _Result(None, rowcount=0)
        if (
            "UPDATE public.rule_run_artifacts" in sql
            and "handoff_status = 'failed'" in sql
        ):
            if (
                row is not None
                and row["handoff_status"] == parameters["expected_status"]
            ):
                row["handoff_status"] = "failed"
                row["failure_code"] = parameters["failure_code"]
                return _Result(rowcount=1)
            return _Result(rowcount=0)
        if "DELETE FROM public.rule_run_artifacts" in sql:
            before = len(self.engine.rows)
            self.engine.rows = [
                item
                for item in self.engine.rows
                if item["id"] != parameters["id"]
            ]
            return _Result(rowcount=before - len(self.engine.rows))
        raise AssertionError(sql)


def _catalog_row(artifact, *, status, binding_hash="d" * 64):
    return {
        "id": new_governance_uid(),
        "artifact_ref": artifact["artifact_ref"],
        "artifact_digest": artifact["digest"],
        "row_count": artifact["row_count"],
        "schema_hash": artifact["schema_hash"],
        "schema_fields": artifact["schema_fields"],
        "expires_at": artifact["expires_at"],
        "handoff_status": status,
        "binding_hash": binding_hash,
        "binding_id": new_governance_uid(),
        "artifact_kind": "output",
        "correlation_id": artifact["artifact_ref"].split("/")[4],
    }


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

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    valid_pending = store.write(
        pl.DataFrame({"id": [1]}),
        new_governance_uid(),
        600,
    )
    missing_pending = store.write(
        pl.DataFrame({"id": [2]}),
        new_governance_uid(),
        600,
    )
    store.delete(missing_pending["artifact_ref"])
    missing_ready = store.write(
        pl.DataFrame({"id": [3]}),
        new_governance_uid(),
        600,
    )
    store.delete(missing_ready["artifact_ref"])
    invalid_pending = store.write(
        pl.DataFrame({"id": [6]}),
        new_governance_uid(),
        600,
    )
    invalid_key = invalid_pending["artifact_ref"].split("/", 3)[-1]
    client.objects[("dataops-rules", invalid_key)][
        "content_type"
    ] = "text/plain"
    old_orphan = store.write(
        pl.DataFrame({"id": [4]}),
        new_governance_uid(),
        600,
    )
    fresh_orphan = store.write(
        pl.DataFrame({"id": [5]}),
        new_governance_uid(),
        600,
    )
    old_key = old_orphan["artifact_ref"].split("/", 3)[-1]
    fresh_key = fresh_orphan["artifact_ref"].split("/", 3)[-1]
    client.ages[old_key] = now - timedelta(minutes=20)
    client.ages[fresh_key] = now - timedelta(seconds=30)
    client.objects[
        ("dataops-rules", "rules/not-a-safe-catalog-key.parquet")
    ] = copy.deepcopy(next(iter(client.objects.values())))
    client.ages["rules/not-a-safe-catalog-key.parquet"] = now - timedelta(
        hours=1
    )
    engine = ReconcileEngine(
        [
            _catalog_row(valid_pending, status="pending"),
            _catalog_row(missing_pending, status="pending"),
            _catalog_row(missing_ready, status="ready"),
            _catalog_row(invalid_pending, status="pending"),
        ]
    )

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=10,
        grace_seconds=300,
    )

    assert result == {
        "pending_finalized": 1,
        "pending_deleted": 1,
        "ready_failed": 1,
        "orphans_deleted": 1,
    }
    assert ("dataops-rules", old_key) not in client.objects
    assert ("dataops-rules", fresh_key) in client.objects
    assert ("dataops-rules", invalid_key) not in client.objects
    assert (
        "dataops-rules",
        "rules/not-a-safe-catalog-key.parquet",
    ) in client.objects
    assert any(
        row["handoff_status"] == "ready"
        and row["artifact_ref"] == valid_pending["artifact_ref"]
        for row in engine.rows
    )
    assert any(
        row["handoff_status"] == "failed"
        and row["artifact_ref"] == missing_ready["artifact_ref"]
        for row in engine.rows
    )
    assert any(
        row["handoff_status"] == "failed"
        and row["artifact_ref"] == invalid_pending["artifact_ref"]
        and row["failure_code"] == "pending_object_invalid"
        for row in engine.rows
    )


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

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        new_governance_uid(),
        600,
    )
    engine = ReconcileEngine([_catalog_row(artifact, status="pending")])
    engine.finalize_race_to_ready = True

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=1,
        grace_seconds=300,
    )

    assert result["pending_finalized"] == 1
    assert engine.rows[0]["handoff_status"] == "ready"
    assert engine.rows[0].get("failure_code") is None


@pytest.mark.parametrize(
    ("status", "storage_error"),
    [
        ("pending", TimeoutError("MinIO timed out")),
        ("ready", TimeoutError("MinIO timed out")),
        ("pending", ConnectionError("MinIO network unavailable")),
        ("ready", ConnectionError("MinIO network unavailable")),
        (
            "pending",
            S3Error(
                "AccessDenied",
                "authentication unavailable",
                None,
                None,
                None,
                None,
            ),
        ),
        (
            "ready",
            S3Error(
                "AccessDenied",
                "authentication unavailable",
                None,
                None,
                None,
                None,
            ),
        ),
    ],
)
def test_reconcile_preserves_catalog_on_transient_storage_error(
    status, storage_error
):
    from app.runner.artifacts import PostgresArtifactResolver

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        new_governance_uid(),
        600,
    )
    key = artifact["artifact_ref"].split("/", 3)[-1]
    original_stat = client.stat_object

    def fail_stat(bucket, object_key):
        if object_key == key:
            raise storage_error
        return original_stat(bucket, object_key)

    client.stat_object = fail_stat
    engine = ReconcileEngine([_catalog_row(artifact, status=status)])

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=1,
        grace_seconds=300,
    )

    assert result == {
        "pending_finalized": 0,
        "pending_deleted": 0,
        "ready_failed": 0,
        "orphans_deleted": 0,
    }
    assert engine.rows[0]["handoff_status"] == status
    assert ("dataops-rules", key) in client.objects


@pytest.mark.parametrize("status", ["pending", "ready"])
def test_reconcile_preserves_catalog_on_transient_download_error(status):
    from app.runner.artifacts import PostgresArtifactResolver

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        new_governance_uid(),
        600,
    )
    key = artifact["artifact_ref"].split("/", 3)[-1]

    def fail_download(_bucket, object_key):
        assert object_key == key
        raise TimeoutError("MinIO download timed out")

    client.get_object = fail_download
    engine = ReconcileEngine([_catalog_row(artifact, status=status)])

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=1,
        grace_seconds=300,
    )

    assert result["pending_finalized"] == 0
    assert result["ready_failed"] == 0
    assert engine.rows[0]["handoff_status"] == status
    assert ("dataops-rules", key) in client.objects


@pytest.mark.parametrize("status", ["pending", "ready"])
def test_reconcile_handles_confirmed_not_found_during_stream_validation(status):
    from app.runner.artifacts import PostgresArtifactResolver

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1]}),
        new_governance_uid(),
        600,
    )

    def disappear_on_download(bucket, object_key):
        client.objects.pop((bucket, object_key), None)
        raise KeyError(object_key)

    client.get_object = disappear_on_download
    engine = ReconcileEngine([_catalog_row(artifact, status=status)])

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=1,
        grace_seconds=300,
    )

    if status == "pending":
        assert result["pending_deleted"] == 1
        assert engine.rows == []
    else:
        assert result["ready_failed"] == 1
        assert engine.rows[0]["handoff_status"] == "failed"
        assert engine.rows[0]["failure_code"] == "ready_object_missing"


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

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1, 2, 3]}),
        new_governance_uid(),
        600,
    )
    key = artifact["artifact_ref"].split("/", 3)[-1]
    stored = client.objects[("dataops-rules", key)]
    payload = stored["payload"]
    midpoint = len(payload) // 2
    stored["payload"] = (
        payload[:midpoint]
        + bytes([payload[midpoint] ^ 1])
        + payload[midpoint + 1 :]
    )
    client.get_calls.clear()
    engine = ReconcileEngine([_catalog_row(artifact, status="pending")])

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=1,
        grace_seconds=300,
    )

    assert result["pending_finalized"] == 0
    assert engine.rows[0]["handoff_status"] == "failed"
    assert engine.rows[0]["failure_code"] == "pending_object_invalid"
    assert ("dataops-rules", key) not in client.objects
    assert client.get_calls == [("dataops-rules", key)]


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

    now = datetime(2026, 7, 23, 12, 0, tzinfo=UTC)
    client = ReconcileMinio(now)
    store = _store(client, clock=lambda: now)
    artifact = store.write(
        pl.DataFrame({"id": [1, 2]}),
        new_governance_uid(),
        600,
    )
    key = artifact["artifact_ref"].split("/", 3)[-1]
    stored = client.objects[("dataops-rules", key)]
    payload = stored["payload"]
    stored["payload"] = bytes([payload[0] ^ 1]) + payload[1:]
    engine = ReconcileEngine([_catalog_row(artifact, status="ready")])

    result = PostgresArtifactResolver(engine, store).reconcile(
        limit=1,
        grace_seconds=300,
    )

    assert result["ready_failed"] == 1
    assert engine.rows[0]["handoff_status"] == "failed"
    assert engine.rows[0]["failure_code"] == "ready_object_invalid"