167 lines
5.3 KiB
Python
167 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Exercise the connector -> datasource -> dataflow capability path."""
|
|
|
|
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.datasources import datasource_catalogue, datasource_lifecycle
|
|
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,
|
|
PipelinePreviewRequest,
|
|
)
|
|
from govoplan_dataflow.backend.service import preview_pipeline
|
|
from govoplan_datasources.backend.db.models import (
|
|
DatasourceMaterializationRecord,
|
|
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__,
|
|
],
|
|
)
|
|
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)
|
|
if writer is None or lifecycle is None or catalogue 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.")
|
|
engine.dispose()
|
|
print("Connector -> Datasources -> Dataflow 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())
|