54 lines
2.6 KiB
Python
54 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from app_runtime.workflow.persistence.checkpoint_repository import CheckpointRepository
|
|
from app_runtime.workflow.persistence.snapshot_sanitizer import WorkflowSnapshotSanitizer
|
|
from app_runtime.workflow.persistence.workflow_repository import WorkflowRepository
|
|
|
|
|
|
class WorkflowPersistence:
|
|
def __init__(self, workflow_repository, checkpoint_repository) -> None:
|
|
self._workflow_repository = workflow_repository
|
|
self._checkpoint_repository = checkpoint_repository
|
|
self._snapshot_sanitizer = WorkflowSnapshotSanitizer()
|
|
|
|
@classmethod
|
|
def create_default(cls, connection_factory=None) -> "WorkflowPersistence":
|
|
return cls(
|
|
workflow_repository=WorkflowRepository(connection_factory),
|
|
checkpoint_repository=CheckpointRepository(connection_factory),
|
|
)
|
|
|
|
def start_run(self, workflow_name: str, start_at: str, snapshot: dict[str, object]) -> int:
|
|
sanitized = self._snapshot_sanitizer.sanitize(snapshot)
|
|
run_id = self._workflow_repository.create_run(workflow_name, sanitized)
|
|
self._checkpoint_repository.save(run_id, start_at, "workflow_started", sanitized)
|
|
return run_id
|
|
|
|
def start_step(self, run_id: int, node_name: str, snapshot: dict[str, object]) -> None:
|
|
self._checkpoint_repository.save(run_id, node_name, "step_started", self._snapshot_sanitizer.sanitize(snapshot))
|
|
|
|
def complete_step(
|
|
self,
|
|
run_id: int,
|
|
node_name: str,
|
|
status: str,
|
|
transition: str,
|
|
snapshot: dict[str, object],
|
|
) -> None:
|
|
sanitized = self._snapshot_sanitizer.sanitize(snapshot)
|
|
self._workflow_repository.record_step(run_id, node_name, status, transition, sanitized)
|
|
self._checkpoint_repository.save(run_id, node_name, "step_completed", sanitized)
|
|
|
|
def fail_step(self, run_id: int, node_name: str, snapshot: dict[str, object], error: Exception) -> None:
|
|
sanitized = self._snapshot_sanitizer.sanitize(snapshot)
|
|
self._workflow_repository.fail_step(run_id, node_name, sanitized, error)
|
|
self._checkpoint_repository.save(run_id, node_name, "step_failed", sanitized)
|
|
|
|
def complete_run(self, run_id: int, snapshot: dict[str, object]) -> None:
|
|
sanitized = self._snapshot_sanitizer.sanitize(snapshot)
|
|
self._workflow_repository.complete_run(run_id, sanitized)
|
|
self._checkpoint_repository.save(run_id, "workflow_done", "workflow_finished", sanitized)
|
|
|
|
def fail_run(self, run_id: int, snapshot: dict[str, object]) -> None:
|
|
self._workflow_repository.fail_run(run_id, self._snapshot_sanitizer.sanitize(snapshot))
|