261 lines
8.4 KiB
Python
261 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Exercise connector -> datasource -> dataflow publication capabilities."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from govoplan_connectors.backend.db.models import ConnectorTabularSource
|
|
from govoplan_core.auth import ApiPrincipal
|
|
from govoplan_core.core.access import PrincipalRef
|
|
from govoplan_core.core.dataflows import (
|
|
DataflowPublicationTarget,
|
|
DataflowRunRequest,
|
|
dataflow_run_lifecycle,
|
|
)
|
|
from govoplan_core.core.datasources import (
|
|
DatasourceReadRequest,
|
|
datasource_catalogue,
|
|
datasource_lifecycle,
|
|
datasource_publication,
|
|
)
|
|
from govoplan_core.core.modules import ModuleContext
|
|
from govoplan_core.core.tabular_sources import (
|
|
TabularSnapshotInput,
|
|
tabular_snapshot_writer,
|
|
)
|
|
from govoplan_core.db.base import Base
|
|
from govoplan_core.server.registry import build_platform_registry
|
|
from govoplan_dataflow.backend.schemas import (
|
|
GraphEdge,
|
|
GraphNode,
|
|
GraphPosition,
|
|
PipelineGraph,
|
|
PipelineCreateRequest,
|
|
PipelinePreviewRequest,
|
|
)
|
|
from govoplan_dataflow.backend.db.models import (
|
|
DataflowPipeline,
|
|
DataflowPipelineRevision,
|
|
DataflowRun,
|
|
)
|
|
from govoplan_dataflow.backend.service import create_pipeline, preview_pipeline
|
|
from govoplan_datasources.backend.db.models import (
|
|
DatasourceMaterializationRecord,
|
|
DatasourcePublicationRecord,
|
|
DatasourceRecord,
|
|
DatasourceStageRecord,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
registry = build_platform_registry(
|
|
("connectors", "datasources", "dataflow", "workflow")
|
|
)
|
|
registry.configure_capability_context(
|
|
ModuleContext(registry=registry, settings=object())
|
|
)
|
|
engine = create_engine("sqlite:///:memory:")
|
|
Base.metadata.create_all(
|
|
engine,
|
|
tables=[
|
|
ConnectorTabularSource.__table__,
|
|
DatasourceRecord.__table__,
|
|
DatasourceMaterializationRecord.__table__,
|
|
DatasourceStageRecord.__table__,
|
|
DatasourcePublicationRecord.__table__,
|
|
DataflowPipeline.__table__,
|
|
DataflowPipelineRevision.__table__,
|
|
DataflowRun.__table__,
|
|
],
|
|
)
|
|
session_factory = sessionmaker(bind=engine)
|
|
with session_factory() as session:
|
|
principal = _principal()
|
|
writer = tabular_snapshot_writer(registry)
|
|
lifecycle = datasource_lifecycle(registry)
|
|
catalogue = datasource_catalogue(registry)
|
|
publisher = datasource_publication(registry)
|
|
runner = dataflow_run_lifecycle(registry)
|
|
if (
|
|
writer is None
|
|
or lifecycle is None
|
|
or catalogue is None
|
|
or publisher is None
|
|
or runner is None
|
|
):
|
|
raise RuntimeError("Datasource composition capabilities are incomplete.")
|
|
|
|
origin = writer.create_snapshot(
|
|
session,
|
|
principal,
|
|
snapshot=TabularSnapshotInput(
|
|
name="Monthly cases",
|
|
source_name="connector_monthly_cases",
|
|
rows=(
|
|
{"id": 1, "amount": 5},
|
|
{"id": 2, "amount": 15},
|
|
),
|
|
),
|
|
)
|
|
datasource = lifecycle.register_origin(
|
|
session,
|
|
principal,
|
|
origin_ref=origin.ref,
|
|
name="Monthly cases cache",
|
|
source_name="monthly_cases",
|
|
mode="cached",
|
|
)
|
|
result = preview_pipeline(
|
|
session,
|
|
tenant_id="tenant-1",
|
|
actor_id="account-1",
|
|
payload=PipelinePreviewRequest(
|
|
graph=_graph(
|
|
datasource_ref=datasource.ref,
|
|
fingerprint=datasource.fingerprint,
|
|
),
|
|
row_limit=100,
|
|
),
|
|
principal=principal,
|
|
registry=registry,
|
|
)
|
|
expected_rows = [
|
|
{"id": 1, "amount": 5},
|
|
{"id": 2, "amount": 15},
|
|
]
|
|
if result.status != "succeeded":
|
|
raise RuntimeError(f"Dataflow preview failed: {result.diagnostics}")
|
|
if result.rows != expected_rows:
|
|
raise RuntimeError(f"Unexpected Dataflow rows: {result.rows!r}")
|
|
if result.source_fingerprints[0]["source_ref"] != datasource.ref:
|
|
raise RuntimeError("Dataflow lineage did not retain the datasource reference.")
|
|
pipeline = create_pipeline(
|
|
session,
|
|
tenant_id="tenant-1",
|
|
actor_id="account-1",
|
|
payload=PipelineCreateRequest(
|
|
name="Monthly case output",
|
|
status="active",
|
|
graph=_graph(
|
|
datasource_ref=datasource.ref,
|
|
fingerprint=datasource.fingerprint,
|
|
),
|
|
editor_mode="graph",
|
|
),
|
|
)
|
|
run_request = DataflowRunRequest(
|
|
pipeline_ref=f"pipeline:{pipeline.id}",
|
|
revision=1,
|
|
idempotency_key="composition-run-1",
|
|
publication=DataflowPublicationTarget(
|
|
name="Monthly case result",
|
|
source_name="monthly_case_result",
|
|
freeze=True,
|
|
frozen_label="Composition evidence",
|
|
),
|
|
)
|
|
published = runner.start_run(
|
|
session,
|
|
principal,
|
|
request=run_request,
|
|
)
|
|
replayed = runner.start_run(
|
|
session,
|
|
principal,
|
|
request=run_request,
|
|
)
|
|
if published.status != "succeeded":
|
|
raise RuntimeError(f"Dataflow publication failed: {published.error}")
|
|
if replayed.ref != published.ref or not replayed.replayed:
|
|
raise RuntimeError("Dataflow run idempotency did not replay the prior run.")
|
|
if (
|
|
not published.output_datasource_ref
|
|
or not published.output_materialization_ref
|
|
):
|
|
raise RuntimeError("Dataflow publication did not retain output references.")
|
|
output = catalogue.read_datasource(
|
|
session,
|
|
principal,
|
|
request=DatasourceReadRequest(
|
|
datasource_ref=published.output_datasource_ref,
|
|
),
|
|
)
|
|
if list(output.rows) != expected_rows:
|
|
raise RuntimeError(
|
|
f"Unexpected published Dataflow rows: {list(output.rows)!r}"
|
|
)
|
|
if (
|
|
output.materialization is None
|
|
or output.materialization.ref != published.output_materialization_ref
|
|
or output.materialization.frozen_at is None
|
|
):
|
|
raise RuntimeError(
|
|
"Published Datasource materialization is not pinned and frozen."
|
|
)
|
|
engine.dispose()
|
|
print(
|
|
"Connector -> Datasources -> pinned Dataflow publication composition passed."
|
|
)
|
|
return 0
|
|
|
|
|
|
def _principal() -> ApiPrincipal:
|
|
return ApiPrincipal(
|
|
principal=PrincipalRef(
|
|
account_id="account-1",
|
|
membership_id="membership-1",
|
|
tenant_id="tenant-1",
|
|
scopes=frozenset(
|
|
{
|
|
"connectors:source:read",
|
|
"connectors:source:write",
|
|
"datasources:catalogue:read",
|
|
"datasources:source:write",
|
|
"datasources:stage:write",
|
|
"dataflow:pipeline:run",
|
|
}
|
|
),
|
|
),
|
|
account=object(),
|
|
user=object(),
|
|
)
|
|
|
|
|
|
def _graph(*, datasource_ref: str, fingerprint: str) -> PipelineGraph:
|
|
return PipelineGraph(
|
|
nodes=[
|
|
GraphNode(
|
|
id="source",
|
|
type="source.reference",
|
|
label="Cases",
|
|
position=GraphPosition(x=0, y=0),
|
|
config={
|
|
"source_ref": datasource_ref,
|
|
"source_name": "monthly_cases",
|
|
"expected_fingerprint": fingerprint,
|
|
"consistency": "current",
|
|
},
|
|
),
|
|
GraphNode(
|
|
id="output",
|
|
type="output",
|
|
label="Output",
|
|
position=GraphPosition(x=200, y=0),
|
|
config={},
|
|
),
|
|
],
|
|
edges=[
|
|
GraphEdge(
|
|
id="source-output",
|
|
source="source",
|
|
target="output",
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|