Model/orchestration/task_orchestrator.py

from typing import Any, Callable, Optional
from uuid import UUID

from domain.tasks.subtasks import SubTask
from domain.tasks.tasks import Source, Task, TaskStatus
from repositories.tasks.subtask_repository import SubTaskRepository
from repositories.tasks.task_repository import TaskRepository
from utilities.private import private


class TaskOrchestrator:
    """Coordinates Task + SubTask lifecycle.

    Exposes primitives (start/complete/fail_subtask) for handlers that want
    fine-grained control, and a high-level run_subtask wrapper that owns the
    try/except so it can replace the body of the legacy subtask_handler
    decorator in backend/utils/subtasks.py.

    Each primitive saves the SubTask, then recomputes the parent Task's
    status from all its children.
    """

    def __init__(
        self,
        task_repo: TaskRepository,
        subtask_repo: SubTaskRepository,
    ) -> None:
        self._tasks = task_repo
        self._subtasks = subtask_repo

    def create_task_with_subtask(
        self,
        *,
        task_source: str,
        inputs: Optional[dict[str, Any]] = None,
        service: Optional[str] = None,
        source: Optional[Source] = None,
        source_id: Optional[str] = None,
    ) -> tuple[Task, SubTask]:
        task = Task.create(
            task_source=task_source,
            service=service,
            source=source,
            source_id=source_id,
        )
        self._tasks.create(task)
        subtask = SubTask.create(task_id=task.id, inputs=inputs)
        self._subtasks.create(subtask)
        return task, subtask

    def create_child_subtask(
        self,
        parent_task_id: UUID,
        *,
        inputs: Optional[dict[str, Any]] = None,
    ) -> SubTask:
        subtask = SubTask.create(task_id=parent_task_id, inputs=inputs)
        self._subtasks.create(subtask)
        return subtask

    def start_subtask(
        self, subtask_id: UUID, cloud_logs_url: Optional[str] = None
    ) -> SubTask:
        subtask = self._subtasks.get(subtask_id)
        subtask.start(cloud_logs_url)
        self._subtasks.save(subtask)
        self._cascade(subtask.task_id)
        return subtask

    def complete_subtask(
        self, subtask_id: UUID, result: Any = None
    ) -> SubTask:
        subtask = self._subtasks.get(subtask_id)
        subtask.complete(result)
        self._subtasks.save(subtask)
        self._cascade(subtask.task_id)
        return subtask

    def fail_subtask(self, subtask_id: UUID, error: BaseException) -> SubTask:
        subtask = self._subtasks.get(subtask_id)
        subtask.fail(error)
        self._subtasks.save(subtask)
        self._cascade(subtask.task_id)
        return subtask

    def run_subtask(
        self,
        subtask_id: UUID,
        work: Callable[[], Any],
        cloud_logs_url: Optional[str] = None,
    ) -> Any:
        self.start_subtask(subtask_id, cloud_logs_url)
        try:
            result = work()
        except Exception as e:
            self.fail_subtask(subtask_id, e)
            raise
        self.complete_subtask(subtask_id, result)
        return result

    def run_subtasks(
        self,
        parent_task_id: UUID,
        inputs_per_subtask: list[dict[str, Any]],
        work: Callable[[SubTask], Any],
        cloud_logs_url: Optional[str] = None,
    ) -> list[Any]:
        """Fan a parent Task out into one child SubTask per item, run ``work`` for
        each (failures isolated per child — a raising item is marked failed and
        its siblings still run), and persist the whole batch in **two** writes
        plus **one** cascade.

        This is the batched form of ``run_subtask``: instead of ~5 writes and a
        full parent re-roll-up *per child* (``create`` + ``start`` + ``complete``
        each cascading — an O(N²) cost on the parent's children), it does one bulk
        ``create_many``, runs every item recording its terminal state in memory,
        then one bulk ``save_many`` and a single ``_cascade``. Children move
        straight from WAITING to their terminal state — the transient IN_PROGRESS
        row is never written, since for a fast batch it only adds DB churn.

        Returns one entry per item in order: the work's result, or ``None`` for an
        item whose work raised.
        """
        subtasks = [
            SubTask.create(task_id=parent_task_id, inputs=inputs)
            for inputs in inputs_per_subtask
        ]
        self._subtasks.create_many(subtasks)

        results: list[Any] = []
        for subtask in subtasks:
            subtask.start(cloud_logs_url)
            try:
                result = work(subtask)
            except Exception as e:  # noqa: BLE001 — isolate per child; siblings continue
                subtask.fail(e)
                results.append(None)
            else:
                subtask.complete(result)
                results.append(result)

        self._subtasks.save_many(subtasks)
        self._cascade(parent_task_id)
        return results

    @private
    def _cascade(self, task_id: UUID) -> None:
        task = self._tasks.get(task_id)
        # FAILED is terminal: once any SubTask has failed the Task is failed and
        # stays failed, so skip the (potentially large) sibling roll-up entirely —
        # no need to list and re-check the SubTasks.
        if task.status is TaskStatus.FAILED:
            return
        statuses = [s.status for s in self._subtasks.list_by_task(task_id)]
        task.recalculate_from_subtasks(statuses)
        self._tasks.save(task)