test: dispatch queued dataflow composition runs
This commit is contained in:
@@ -14,6 +14,10 @@ from govoplan_core.core.dataflows import (
|
|||||||
DataflowRunRequest,
|
DataflowRunRequest,
|
||||||
dataflow_run_lifecycle,
|
dataflow_run_lifecycle,
|
||||||
)
|
)
|
||||||
|
from govoplan_core.core.automation import AutomationPrincipalResolution
|
||||||
|
from govoplan_core.core.access import (
|
||||||
|
CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER,
|
||||||
|
)
|
||||||
from govoplan_core.core.datasources import (
|
from govoplan_core.core.datasources import (
|
||||||
DatasourceReadRequest,
|
DatasourceReadRequest,
|
||||||
datasource_catalogue,
|
datasource_catalogue,
|
||||||
@@ -41,6 +45,7 @@ from govoplan_dataflow.backend.db.models import (
|
|||||||
DataflowRun,
|
DataflowRun,
|
||||||
)
|
)
|
||||||
from govoplan_dataflow.backend.service import create_pipeline, preview_pipeline
|
from govoplan_dataflow.backend.service import create_pipeline, preview_pipeline
|
||||||
|
from govoplan_dataflow.backend.run_worker import SqlDataflowRunWorker
|
||||||
from govoplan_datasources.backend.db.models import (
|
from govoplan_datasources.backend.db.models import (
|
||||||
DatasourceMaterializationRecord,
|
DatasourceMaterializationRecord,
|
||||||
DatasourcePayloadRecord,
|
DatasourcePayloadRecord,
|
||||||
@@ -170,6 +175,29 @@ def main() -> int:
|
|||||||
principal,
|
principal,
|
||||||
request=run_request,
|
request=run_request,
|
||||||
)
|
)
|
||||||
|
if published.status != "queued":
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Dataflow run was not queued: {published.status}"
|
||||||
|
)
|
||||||
|
worker = SqlDataflowRunWorker(
|
||||||
|
registry=_AutomationRegistry(registry, principal)
|
||||||
|
)
|
||||||
|
worker_result = worker.dispatch_pending(
|
||||||
|
session,
|
||||||
|
worker_id="composition-worker",
|
||||||
|
)
|
||||||
|
if worker_result["succeeded"] != 1:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Dataflow worker failed: {worker_result!r}"
|
||||||
|
)
|
||||||
|
completed = runner.get_run(
|
||||||
|
session,
|
||||||
|
principal,
|
||||||
|
run_ref=published.ref,
|
||||||
|
)
|
||||||
|
if completed is None:
|
||||||
|
raise RuntimeError("Dataflow run evidence disappeared.")
|
||||||
|
published = completed
|
||||||
if published.status != "succeeded":
|
if published.status != "succeeded":
|
||||||
raise RuntimeError(f"Dataflow publication failed: {published.error}")
|
raise RuntimeError(f"Dataflow publication failed: {published.error}")
|
||||||
if replayed.ref != published.ref or not replayed.replayed:
|
if replayed.ref != published.ref or not replayed.replayed:
|
||||||
@@ -205,6 +233,36 @@ def main() -> int:
|
|||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
class _AutomationProvider:
|
||||||
|
def __init__(self, principal: ApiPrincipal) -> None:
|
||||||
|
self.principal = principal
|
||||||
|
|
||||||
|
def resolve_automation_principal(self, _session, *, request):
|
||||||
|
return AutomationPrincipalResolution(
|
||||||
|
allowed=True,
|
||||||
|
principal=self.principal,
|
||||||
|
granted_scopes=request.grant_scopes,
|
||||||
|
provenance={"status": "composition_recheck"},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _AutomationRegistry:
|
||||||
|
def __init__(self, registry, principal: ApiPrincipal) -> None:
|
||||||
|
self.registry = registry
|
||||||
|
self.provider = _AutomationProvider(principal)
|
||||||
|
|
||||||
|
def has_capability(self, name: str) -> bool:
|
||||||
|
return (
|
||||||
|
name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER
|
||||||
|
or self.registry.has_capability(name)
|
||||||
|
)
|
||||||
|
|
||||||
|
def capability(self, name: str):
|
||||||
|
if name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER:
|
||||||
|
return self.provider
|
||||||
|
return self.registry.capability(name)
|
||||||
|
|
||||||
|
|
||||||
def _principal() -> ApiPrincipal:
|
def _principal() -> ApiPrincipal:
|
||||||
return ApiPrincipal(
|
return ApiPrincipal(
|
||||||
principal=PrincipalRef(
|
principal=PrincipalRef(
|
||||||
|
|||||||
Reference in New Issue
Block a user