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 )