test: dispatch queued dataflow composition runs
This commit is contained in:
@@ -14,6 +14,10 @@ from govoplan_core.core.dataflows import (
|
||||
DataflowRunRequest,
|
||||
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 (
|
||||
DatasourceReadRequest,
|
||||
datasource_catalogue,
|
||||
@@ -41,6 +45,7 @@ from govoplan_dataflow.backend.db.models import (
|
||||
DataflowRun,
|
||||
)
|
||||
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 (
|
||||
DatasourceMaterializationRecord,
|
||||
DatasourcePayloadRecord,
|
||||
@@ -170,6 +175,29 @@ def main() -> int:
|
||||
principal,
|
||||
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":
|
||||
raise RuntimeError(f"Dataflow publication failed: {published.error}")
|
||||
if replayed.ref != published.ref or not replayed.replayed:
|
||||
@@ -205,6 +233,36 @@ def main() -> int:
|
||||
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:
|
||||
return ApiPrincipal(
|
||||
principal=PrincipalRef(
|
||||
|
||||
Reference in New Issue
Block a user