from __future__ import annotations from datetime import UTC, datetime import unittest from govoplan_core.auth import ApiPrincipal from govoplan_core.core.access import PrincipalRef from govoplan_core.core.dataflows import DataflowDatasetRequest, DataflowRunConflictError from govoplan_core.core.datasources import ( CAPABILITY_DATASOURCE_CATALOGUE, DatasourceDescriptor, DatasourceMaterialization, DatasourceReadResult, ) from govoplan_core.db.base import Base from govoplan_core.db.session import configure_database, reset_database from govoplan_dataflow.backend.dataset_output import SqlDataflowDatasetOutputProvider from govoplan_dataflow.backend.db.models import ( DataflowPipeline, DataflowPipelineRevision, DataflowRun, ) from govoplan_dataflow.backend.schemas import PipelineGraph from govoplan_dataflow.backend.service import definition_hash def principal(tenant_id: str = "tenant-1") -> ApiPrincipal: return ApiPrincipal( principal=PrincipalRef( account_id="account-1", membership_id="membership-1", tenant_id=tenant_id, scopes=frozenset({"dataflow:pipeline:read", "dataflow:pipeline:run"}), ), account=object(), user=object(), ) class PublishedOutputCatalogue: rows = ( {"recipient_key": "published", "email": "published@example.test"}, ) def list_datasources(self, *_args, **_kwargs): return () def get_datasource(self, *_args, **_kwargs): return None def list_materializations(self, *_args, **_kwargs): return () def read_datasource(self, _session, _principal, *, request): materialization = DatasourceMaterialization( ref="materialization:published-1", datasource_ref="datasource:published-1", revision=1, state="published", fingerprint="f" * 64, row_count=len(self.rows), created_at=datetime(2026, 8, 4, 9, 0, tzinfo=UTC), ) rows = self.rows[request.offset : request.offset + request.limit] return DatasourceReadResult( datasource=DatasourceDescriptor( ref="datasource:published-1", source_name="published_output", name="Published output", kind="custom", mode="static", shape="tabular", fingerprint=materialization.fingerprint, current_materialization_ref=materialization.ref, row_count=len(self.rows), ), materialization=materialization, rows=rows, total_rows=len(self.rows), truncated=request.offset + len(rows) < len(self.rows), ) class CapabilityRegistry: def __init__(self, catalogue) -> None: self.catalogue = catalogue def has_capability(self, name: str) -> bool: return name == CAPABILITY_DATASOURCE_CATALOGUE def capability(self, name: str): return self.catalogue if self.has_capability(name) else None class DataflowDatasetOutputTests(unittest.TestCase): def setUp(self) -> None: self.database = configure_database("sqlite:///:memory:") Base.metadata.create_all( self.database.engine, tables=[ DataflowPipeline.__table__, DataflowPipelineRevision.__table__, DataflowRun.__table__, ], ) def tearDown(self) -> None: reset_database(dispose=True) def test_pinned_revision_executes_through_typed_registry(self) -> None: graph = PipelineGraph.model_validate( { "nodes": [ { "id": "source", "type": "source.inline", "label": "Source", "position": {"x": 0, "y": 0}, "config": { "source_name": "audience", "rows": [ {"recipient_key": "one", "email": "one@example.test"}, {"recipient_key": "two", "email": "two@example.test"}, ] }, }, { "id": "output", "type": "output", "label": "Output", "position": {"x": 200, "y": 0}, "config": {}, }, ], "edges": [{"id": "edge-1", "source": "source", "target": "output"}], } ) content_hash = definition_hash(graph, None) with self.database.session() as session: pipeline = DataflowPipeline( id="pipeline-1", tenant_id="tenant-1", scope_type="tenant", scope_id="tenant-1", definition_kind="flow", name="Audience", status="active", current_revision=1, allow_run=True, metadata_={}, ) revision = DataflowPipelineRevision( id="revision-1", tenant_id="tenant-1", pipeline_id=pipeline.id, revision=1, graph=graph.model_dump(mode="json"), sql_text=None, editor_mode="graph", content_hash=content_hash, ) session.add_all((pipeline, revision)) session.flush() provider = SqlDataflowDatasetOutputProvider() descriptors = provider.list_outputs(session, principal()) result = provider.read_output( session, principal(), request=DataflowDatasetRequest( pipeline_ref=pipeline.id, revision=1, expected_definition_hash=content_hash, ), ) self.assertEqual(1, len(descriptors)) self.assertEqual(2, result.total_rows) self.assertEqual("one", result.rows[0]["recipient_key"]) self.assertEqual(64, len(result.output_hash)) replay = provider.read_output( session, principal(), request=DataflowDatasetRequest( pipeline_ref=pipeline.id, revision=1, expected_definition_hash=content_hash, expected_source_fingerprints=result.source_fingerprints, ), ) self.assertEqual(result.output_hash, replay.output_hash) with self.assertRaises(DataflowRunConflictError): provider.read_output( session, principal(), request=DataflowDatasetRequest( pipeline_ref=pipeline.id, revision=1, expected_definition_hash="wrong", ), ) run = DataflowRun( id="run-published-1", tenant_id="tenant-1", pipeline_id=pipeline.id, pipeline_revision_id=revision.id, run_type="published", status="succeeded", execution_backend="duckdb", environment="production", executor_version="duckdb-v1", definition_hash=content_hash, request_={}, source_fingerprints=[ {"node_id": "source", "fingerprint": "source-v1"} ], result_schema=[], diagnostics=[], output_row_count=1, output_publication_ref="publication:published-1", output_datasource_ref="datasource:published-1", output_materialization_ref="materialization:published-1", finished_at=datetime(2026, 8, 4, 9, 0, tzinfo=UTC), ) session.add(run) session.flush() published_provider = SqlDataflowDatasetOutputProvider( CapabilityRegistry(PublishedOutputCatalogue()) ) published = published_provider.read_output( session, principal(), request=DataflowDatasetRequest( pipeline_ref=pipeline.id, revision=1, run_ref="dataflow-run:run-published-1", expected_definition_hash=content_hash, expected_source_fingerprints=( {"node_id": "source", "fingerprint": "source-v1"}, ), ), ) self.assertEqual("published", published.rows[0]["recipient_key"]) self.assertEqual("dataflow-run:run-published-1", published.run_ref) self.assertTrue(published.provenance["immutable_run"]) self.assertEqual( "materialization:published-1", published.provenance["materialization_ref"], ) with self.assertRaises(DataflowRunConflictError): published_provider.read_output( session, principal(), request=DataflowDatasetRequest( pipeline_ref=pipeline.id, revision=1, run_ref="dataflow-run:run-published-1", parameters={"changed": True}, ), ) with self.assertRaises(DataflowRunConflictError): provider.read_output( session, principal(), request=DataflowDatasetRequest( pipeline_ref=pipeline.id, revision=1, expected_definition_hash=content_hash, expected_source_fingerprints=( {"node_id": "source", "fingerprint": "changed"}, ), ), ) if __name__ == "__main__": unittest.main()