282 lines
10 KiB
Python
282 lines
10 KiB
Python
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()
|