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 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 = []

    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 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 row is not None:
                row["handoff_status"] = "ready"
            return _Result(copy.deepcopy(row), rowcount=int(row is not None))
        if (
            "UPDATE public.rule_run_artifacts" in sql
            and "handoff_status = 'failed'" in sql
        ):
            if row is not None:
                row["handoff_status"] = "failed"
                row["failure_code"] = parameters["failure_code"]
            return _Result(rowcount=int(row is not None))
        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
    )