Files
govoplan-dataflow/src/govoplan_dataflow/backend/run_worker.py
T

758 lines
23 KiB
Python

from __future__ import annotations
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
import logging
import socket
from sqlalchemy import func, or_, select
from sqlalchemy.orm import Session
from govoplan_core.auth import ApiPrincipal
from govoplan_core.core.automation import (
AutomationPrincipalRequest,
automation_principal_provider,
)
from govoplan_core.core.notifications import (
NotificationDispatchRequest,
notification_dispatch_provider,
)
from govoplan_core.db.base import utcnow
from govoplan_dataflow.backend.db.models import (
DataflowPipeline,
DataflowPipelineRevision,
DataflowRun,
DataflowTrigger,
)
from govoplan_dataflow.backend.governance import require_definition_action
from govoplan_dataflow.backend.recovery import (
DataflowRecoveryError,
DataflowRunRecovery,
begin_dataflow_run_recovery,
claim_stale_dataflow_recovery,
dataflow_run_recovery_state,
)
from govoplan_dataflow.backend.service import (
_execute_pipeline_run,
_require_deployed_revision,
pipeline_run_request,
)
MAX_ACTIVE_RUNS_PER_TENANT = 4
DEFAULT_LEASE_SECONDS = 120
logger = logging.getLogger(__name__)
class DataflowWorkerError(RuntimeError):
pass
def dispatch_pending_runs(
session: Session,
*,
registry: object | None,
now: datetime | None = None,
limit: int = 10,
worker_id: str | None = None,
) -> dict[str, object]:
current = _as_utc(now or utcnow())
resolved_worker_id = (
str(worker_id).strip()
if worker_id and str(worker_id).strip()
else socket.gethostname()
)
recovered, recovered_outcome_unknown = _recover_expired_leases(
session,
now=current,
)
session.commit()
summary: dict[str, object] = {
"claimed": 0,
"succeeded": 0,
"retrying": 0,
"failed": 0,
"cancelled": 0,
"outcome_unknown": recovered_outcome_unknown,
"recovered": recovered,
"runs": [],
}
for _ in range(max(1, min(int(limit), 100))):
run_id = _claim_next_run(
session,
now=current,
worker_id=resolved_worker_id,
)
if run_id is None:
session.rollback()
break
session.commit()
summary["claimed"] = int(summary["claimed"]) + 1
outcome = _execute_claimed_run(
session,
run_id=run_id,
registry=registry,
now=current,
)
session.commit()
summary[outcome] = int(summary[outcome]) + 1
runs = list(summary["runs"])
runs.append({"ref": f"dataflow-run:{run_id}", "status": outcome})
summary["runs"] = runs
return summary
def purge_expired_runs(
session: Session,
*,
now: datetime | None = None,
limit: int = 500,
) -> dict[str, object]:
current = _as_utc(now or utcnow())
runs = list(
session.scalars(
select(DataflowRun)
.where(
DataflowRun.status.in_(
("succeeded", "failed", "cancelled")
),
DataflowRun.retention_until.is_not(None),
DataflowRun.retention_until <= current,
DataflowRun.purged_at.is_(None),
)
.order_by(DataflowRun.retention_until, DataflowRun.id)
.limit(max(1, min(int(limit), 5_000)))
.with_for_update(skip_locked=True)
)
)
for run in runs:
run.request_ = {}
run.source_fingerprints = []
run.diagnostics = []
authorization = dict(run.authorization_)
authorization.pop("submitted_principal", None)
authorization["personal_data_purged"] = True
run.authorization_ = authorization
run.purged_at = current
session.flush()
return {
"purged": len(runs),
"run_refs": [f"dataflow-run:{run.id}" for run in runs],
}
def run_metrics(
session: Session,
*,
tenant_id: str,
now: datetime | None = None,
) -> dict[str, object]:
current = _as_utc(now or utcnow())
rows = session.execute(
select(DataflowRun.status, func.count())
.where(DataflowRun.tenant_id == tenant_id)
.group_by(DataflowRun.status)
).all()
statuses = {str(status): int(count) for status, count in rows}
recent = list(
session.execute(
select(
DataflowRun.status,
DataflowRun.started_at,
DataflowRun.finished_at,
).where(
DataflowRun.tenant_id == tenant_id,
DataflowRun.finished_at.is_not(None),
DataflowRun.finished_at >= current - timedelta(hours=24),
)
)
)
durations = [
(
_as_utc(finished_at) - _as_utc(started_at)
).total_seconds()
for _status, started_at, finished_at in recent
if finished_at is not None and started_at is not None
]
oldest_queued_at = session.scalar(
select(func.min(DataflowRun.queued_at)).where(
DataflowRun.tenant_id == tenant_id,
DataflowRun.status.in_(("queued", "retrying")),
)
)
return {
"statuses": statuses,
"active": sum(
statuses.get(status, 0)
for status in ("queued", "retrying", "running")
),
"completed_last_24_hours": len(recent),
"failed_last_24_hours": sum(
1 for status, _started_at, _finished_at in recent
if status == "failed"
),
"average_duration_seconds": (
round(sum(durations) / len(durations), 3)
if durations
else None
),
"oldest_queued_at": oldest_queued_at,
}
def _recover_expired_leases(
session: Session,
*,
now: datetime,
) -> tuple[int, int]:
runs = list(
session.scalars(
select(DataflowRun)
.where(
DataflowRun.status == "running",
DataflowRun.lease_expires_at.is_not(None),
DataflowRun.lease_expires_at < now,
)
.with_for_update(skip_locked=True)
)
)
recovered = 0
outcome_unknown = 0
for run in runs:
recovery = dataflow_run_recovery_state(session, run_id=run.id)
recovery_status: str | None = None
recovery_metadata = dict(run.authorization_).get("recovery")
effect_started = bool(
isinstance(recovery_metadata, Mapping)
and recovery_metadata.get("boundary") == "output-publication"
)
if recovery is not None and recovery.get("operation_id"):
try:
recovery_status = claim_stale_dataflow_recovery(
session,
operation_id=str(recovery["operation_id"]),
lease_ttl_seconds=DEFAULT_LEASE_SECONDS,
effect_started=effect_started,
)
except DataflowRecoveryError:
logger.warning(
"Could not claim stale Dataflow recovery operation %s",
recovery["operation_id"],
exc_info=True,
)
continue
run.worker_id = None
run.claimed_at = None
run.lease_expires_at = None
run.heartbeat_at = None
if run.cancellation_requested_at is not None:
_cancel_run(run, now=now)
elif recovery_status == "outcome_unknown" or effect_started:
run.status = "outcome_unknown"
run.finished_at = now
run.progress_phase = "outcome_unknown"
run.error = (
"The worker lease expired after output publication may have "
"started. Verify the sink before retrying."
)
outcome_unknown += 1
elif run.attempts < run.max_attempts:
run.status = "retrying"
run.available_at = now
run.progress_phase = "retrying_after_worker_loss"
run.error = "The previous worker lease expired."
else:
_fail_run(
run,
now=now,
message="The worker lease expired after the final attempt.",
)
recovered += 1
return recovered, outcome_unknown
def _claim_next_run(
session: Session,
*,
now: datetime,
worker_id: str,
) -> str | None:
candidates = list(
session.scalars(
select(DataflowRun)
.where(
DataflowRun.status.in_(("queued", "retrying")),
or_(
DataflowRun.available_at.is_(None),
DataflowRun.available_at <= now,
),
)
.order_by(DataflowRun.available_at, DataflowRun.created_at)
.limit(40)
.with_for_update(skip_locked=True)
)
)
for run in candidates:
active = int(
session.scalar(
select(func.count())
.select_from(DataflowRun)
.where(
DataflowRun.tenant_id == run.tenant_id,
DataflowRun.status == "running",
)
)
or 0
)
if active >= MAX_ACTIVE_RUNS_PER_TENANT:
continue
run.status = "running"
run.attempts += 1
run.worker_id = worker_id[:255]
run.claimed_at = now
run.heartbeat_at = now
run.lease_expires_at = now + timedelta(
seconds=_lease_seconds(run)
)
run.started_at = run.started_at or now
run.finished_at = None
run.progress_percent = max(run.progress_percent, 5)
run.progress_phase = "authorizing"
session.flush()
return run.id
return None
def _execute_claimed_run(
session: Session,
*,
run_id: str,
registry: object | None,
now: datetime,
) -> str:
run = session.get(DataflowRun, run_id)
if run is None:
return "failed"
if run.cancellation_requested_at is not None:
_cancel_run(run, now=now)
return "cancelled"
recovery: DataflowRunRecovery | None = None
try:
principal, provenance = _resolve_principal(
session,
run=run,
registry=registry,
)
run.authorization_ = {
**dict(run.authorization_),
"last_resolution": provenance,
"resolved_at": now.isoformat(),
}
pipeline = session.get(DataflowPipeline, run.pipeline_id)
revision = session.get(
DataflowPipelineRevision,
run.pipeline_revision_id,
)
if pipeline is None or revision is None:
raise DataflowWorkerError(
"The pipeline or pinned revision no longer exists."
)
action = (
"run"
if run.invocation_kind in {"manual", "api", "backfill"}
else "automate"
)
require_definition_action(
pipeline,
principal=principal,
registry=registry,
action=action,
)
_require_deployed_revision(
session,
tenant_id=run.tenant_id,
pipeline=pipeline,
revision=revision,
environment=run.environment,
)
session.commit()
recovery = begin_dataflow_run_recovery(
session,
run=run,
lease_ttl_seconds=_lease_seconds(run),
)
if recovery.replayed:
session.rollback()
run = session.get(DataflowRun, run_id)
return run.status if run is not None else "failed"
_notify_run(
session,
registry=registry,
run=run,
event_kind="dataflow.run.started",
subject=f"Dataflow run started: {pipeline.name}",
)
retryable = _execute_pipeline_run(
session,
run=run,
pipeline=pipeline,
revision=revision,
request=pipeline_run_request(run),
principal=principal,
registry=registry,
recovery=recovery,
)
if (
run.status == "failed"
and retryable
and run.attempts < run.max_attempts
):
_retry_run(run, now=now)
_finish_claim(run)
_update_trigger_status(session, run)
_notify_terminal_run(
session,
registry=registry,
run=run,
pipeline_name=pipeline.name,
)
recovery.finish(session, run=run)
return run.status
except DataflowRecoveryError:
logger.exception(
"Dataflow recovery finalization failed for run %s",
run_id,
)
session.rollback()
run = session.get(DataflowRun, run_id)
if run is not None:
if recovery is not None and recovery.publication_started:
run.status = "outcome_unknown"
run.finished_at = now
run.progress_phase = "outcome_unknown"
run.error = (
"Output publication completed without verifiable terminal "
"recovery evidence; reconcile the sink."
)
_finish_claim(run)
return "outcome_unknown"
if recovery is None:
_fail_run(
run,
now=now,
message=(
"The recovery ledger was unavailable; no Dataflow "
"execution started."
),
)
_finish_claim(run)
else:
_fail_run(
run,
now=now,
message=(
"The database-only run could not commit verifiable "
"recovery evidence."
),
)
_finish_claim(run)
return "failed"
except (DataflowWorkerError, PermissionError, ValueError) as exc:
_fail_run(run, now=now, message=str(exc))
_finish_claim(run)
_update_trigger_status(session, run)
_notify_terminal_run(
session,
registry=registry,
run=run,
pipeline_name="Dataflow",
)
if recovery is not None:
recovery.finish(session, run=run)
return "failed"
except Exception as exc:
logger.exception(
"Unexpected Dataflow worker failure for run %s",
run_id,
)
session.rollback()
run = session.get(DataflowRun, run_id)
if run is None:
return "failed"
if recovery is not None and recovery.publication_started:
run.status = "outcome_unknown"
run.finished_at = now
run.error = (
"The output provider failed after dispatch began; verify the "
"sink before retrying."
)
run.progress_phase = "outcome_unknown"
_finish_claim(run)
try:
recovery.finish(session, run=run)
except DataflowRecoveryError:
logger.exception(
"Could not persist the uncertain output state for run %s",
run_id,
)
session.rollback()
return "outcome_unknown"
_fail_run(
run,
now=now,
message=(
"Unexpected worker failure "
f"({type(exc).__name__})."
),
)
_finish_claim(run)
if recovery is not None:
try:
recovery.finish(session, run=run)
except DataflowRecoveryError:
logger.exception(
"Could not finalize failed Dataflow recovery for run %s",
run_id,
)
session.rollback()
return "failed"
def _resolve_principal(
session: Session,
*,
run: DataflowRun,
registry: object | None,
) -> tuple[ApiPrincipal, Mapping[str, object]]:
provider = automation_principal_provider(registry)
if provider is None:
raise DataflowWorkerError(
"Queued Dataflow execution requires the automation-principal "
"capability."
)
authorization = dict(run.authorization_)
subject_kind = str(
authorization.get("subject_kind") or "delegated_user"
)
common = {
"tenant_id": run.tenant_id,
"authorization_ref": str(
authorization.get("authorization_ref")
or f"dataflow-run:{run.id}"
),
"grant_scopes": tuple(
str(scope)
for scope in authorization.get("grant_scopes") or ()
),
"context": {
"run_ref": f"dataflow-run:{run.id}",
"pipeline_ref": f"pipeline:{run.pipeline_id}",
"environment": run.environment,
"attempt": run.attempts,
},
}
if subject_kind == "service_account":
service_account_id = str(
authorization.get("service_account_id") or ""
).strip()
request = AutomationPrincipalRequest.service_account(
service_account_id=service_account_id,
**common,
)
else:
request = AutomationPrincipalRequest.delegated_user(
account_id=str(authorization.get("account_id") or ""),
membership_id=str(
authorization.get("membership_id") or ""
),
**common,
)
resolution = provider.resolve_automation_principal(
session,
request=request,
)
if not resolution.allowed or not isinstance(
resolution.principal,
ApiPrincipal,
):
raise DataflowWorkerError(
resolution.reason or "Dataflow run authorization was denied."
)
return resolution.principal, resolution.provenance
def _retry_run(run: DataflowRun, *, now: datetime) -> None:
delay_seconds = min(15 * (2 ** max(0, run.attempts - 1)), 900)
run.status = "retrying"
run.available_at = now + timedelta(seconds=delay_seconds)
run.finished_at = None
run.progress_percent = 0
run.progress_phase = "retrying"
def _finish_claim(run: DataflowRun) -> None:
run.worker_id = None
run.claimed_at = None
run.lease_expires_at = None
run.heartbeat_at = None
def _cancel_run(run: DataflowRun, *, now: datetime) -> None:
run.status = "cancelled"
run.finished_at = now
run.error = "Cancelled by request."
run.progress_phase = "cancelled"
_finish_claim(run)
def _fail_run(
run: DataflowRun,
*,
now: datetime,
message: str,
) -> None:
run.status = "failed"
run.finished_at = now
run.error = message
run.progress_phase = "failed"
def _update_trigger_status(session: Session, run: DataflowRun) -> None:
if not run.trigger_id:
return
trigger = session.get(DataflowTrigger, run.trigger_id)
if trigger is None:
return
trigger.last_status = run.status
trigger.last_error = run.error
def _notify_terminal_run(
session: Session,
*,
registry: object | None,
run: DataflowRun,
pipeline_name: str,
) -> None:
if run.status == "retrying":
event_kind = "dataflow.run.retrying"
subject = f"Dataflow run will retry: {pipeline_name}"
elif run.status == "succeeded":
event_kind = "dataflow.run.completed"
subject = f"Dataflow run completed: {pipeline_name}"
elif run.status == "cancelled":
event_kind = "dataflow.run.cancelled"
subject = f"Dataflow run cancelled: {pipeline_name}"
else:
event_kind = "dataflow.run.failed"
subject = f"Dataflow run failed: {pipeline_name}"
_notify_run(
session,
registry=registry,
run=run,
event_kind=event_kind,
subject=subject,
)
def _notify_run(
session: Session,
*,
registry: object | None,
run: DataflowRun,
event_kind: str,
subject: str,
) -> None:
provider = notification_dispatch_provider(registry)
authorization = dict(run.authorization_)
recipient_id = str(authorization.get("account_id") or "").strip()
if provider is None or not recipient_id:
return
try:
provider.enqueue_notification(
session,
NotificationDispatchRequest(
tenant_id=run.tenant_id,
source_module="dataflow",
source_resource_type="dataflow_run",
source_resource_id=run.id,
event_kind=event_kind,
recipient_type="account",
recipient_id=recipient_id,
subject=subject,
body_text=run.error,
action_url="/dataflow",
payload={
"run_ref": f"dataflow-run:{run.id}",
"pipeline_ref": f"pipeline:{run.pipeline_id}",
"status": run.status,
"environment": run.environment,
"attempt": run.attempts,
},
),
)
except Exception:
# Notification delivery is optional and cannot change run outcome.
logger.warning(
"Could not enqueue Dataflow notification for run %s",
run.id,
exc_info=True,
)
return
def _lease_seconds(run: DataflowRun) -> int:
budget = dict(run.resource_budget)
wall_seconds = float(budget.get("max_wall_seconds") or 30.0)
return max(DEFAULT_LEASE_SECONDS, int(wall_seconds) + 60)
def _as_utc(value: datetime | None) -> datetime:
if value is None:
return datetime.now(timezone.utc)
if value.tzinfo is None:
return value.replace(tzinfo=timezone.utc)
return value.astimezone(timezone.utc)
class SqlDataflowRunWorker:
def __init__(self, *, registry: object | None = None) -> None:
self._registry = registry
def dispatch_pending(
self,
session: object,
*,
now: datetime | None = None,
limit: int = 10,
worker_id: str | None = None,
) -> Mapping[str, object]:
if not isinstance(session, Session):
raise TypeError("Dataflow run dispatch requires a Session.")
return dispatch_pending_runs(
session,
registry=self._registry,
now=now,
limit=limit,
worker_id=worker_id,
)
def purge_expired(
self,
session: object,
*,
now: datetime | None = None,
limit: int = 500,
) -> Mapping[str, object]:
if not isinstance(session, Session):
raise TypeError("Dataflow retention requires a Session.")
return purge_expired_runs(session, now=now, limit=limit)
__all__ = [
"SqlDataflowRunWorker",
"dispatch_pending_runs",
"purge_expired_runs",
"run_metrics",
]