49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Literal, Protocol
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
from host_agent.assignment import AssignmentExecutionResult
|
|
from host_agent.client import HostAgentClient
|
|
|
|
|
|
class ActiveAssignmentExecutor(Protocol):
|
|
async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult: ...
|
|
|
|
def request_stop(self) -> None: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssignmentProcessingResult:
|
|
execution: AssignmentExecutionResult
|
|
report_status: Literal["recorded", "already_recorded"]
|
|
|
|
|
|
class AssignmentProcessor:
|
|
def __init__(
|
|
self,
|
|
client: HostAgentClient,
|
|
active_executor: ActiveAssignmentExecutor,
|
|
) -> None:
|
|
self.client = client
|
|
self.active_executor = active_executor
|
|
|
|
def request_stop(self) -> None:
|
|
self.active_executor.request_stop()
|
|
|
|
async def process(self, assignment: AssignmentModel) -> AssignmentProcessingResult:
|
|
execution = await self.active_executor.run(assignment)
|
|
status = "done" if execution.status == "done" else "failed"
|
|
failure_reason = execution.failure_reason if status == "failed" else None
|
|
response = await self.client.report_result(
|
|
assignment,
|
|
status=status,
|
|
failure_reason=failure_reason,
|
|
result=dict(execution.metadata),
|
|
)
|
|
return AssignmentProcessingResult(
|
|
execution=execution,
|
|
report_status=response.status,
|
|
)
|