company-ai-platform/app/modules/workflows/service.py

from typing import Any
from uuid import uuid4

from fastapi import HTTPException, status
from sqlalchemy import func, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.core.constants import ActorValue
from app.core.http.pagination import bounded_limit
from app.core.utils.time import utc_now
from app.modules.business.service import serialize_model
from app.modules.workflows.constants import (
    WORKFLOW_ACTION_CODE_PREFIX,
    WORKFLOW_CODE_PREFIX,
    WorkflowErrorDetail,
    WorkflowStatus,
)
from app.modules.workflows.models import WorkflowAction, WorkflowInstance


class WorkflowService:
    """Track V3 workflow instances and append-only workflow actions."""

    def __init__(self, db: Session):
        self.db = db

    def start_or_update(
        self,
        workflow_type: str,
        aggregate_type: str,
        aggregate_id: str | int | None,
        status_value: str,
        action: str,
        actor: str = ActorValue.SYSTEM,
        payload: dict[str, Any] | None = None,
        commit: bool = True,
        source_event_id: str | None = None,
    ) -> WorkflowInstance:
        existing_workflow = self._workflow_for_source_event(source_event_id)
        if existing_workflow is not None:
            return existing_workflow

        aggregate_id_text = str(aggregate_id) if aggregate_id is not None else None
        workflow_query = (
            select(WorkflowInstance).where(
                WorkflowInstance.workflow_type == workflow_type,
                WorkflowInstance.aggregate_type == aggregate_type,
                WorkflowInstance.aggregate_id == aggregate_id_text,
            )
            .with_for_update()
        )
        record = self.db.execute(workflow_query).scalar_one_or_none()
        previous_status = None
        now = utc_now()
        if record is None:
            candidate = WorkflowInstance(
                code=(
                    f"{WORKFLOW_CODE_PREFIX}-{now:%Y%m%d%H%M%S%f}-"
                    f"{uuid4().hex[:8]}"
                ),
                workflow_type=workflow_type,
                aggregate_type=aggregate_type,
                aggregate_id=aggregate_id_text,
                status=status_value,
                actor=actor,
                current_step=action,
                payload=payload or {},
            )
            try:
                with self.db.begin_nested():
                    self.db.add(candidate)
                    self.db.flush()
                record = candidate
            except IntegrityError:
                record = self.db.execute(workflow_query).scalar_one()
                previous_status = record.status
        else:
            previous_status = record.status

        existing_workflow = self._workflow_for_source_event(source_event_id)
        if existing_workflow is not None:
            if existing_workflow.code != record.code:
                raise ValueError(
                    "Source event is already attached to a different workflow"
                )
            return existing_workflow

        if previous_status is not None:
            record.status = status_value
            record.actor = actor
            record.current_step = action
            record.payload = payload or {}
            record.updated_at = now

        if status_value in {
            WorkflowStatus.BLOCKED,
            WorkflowStatus.COMPLETED,
            WorkflowStatus.FAILED,
        }:
            record.completed_at = now
        elif previous_status in {
            WorkflowStatus.BLOCKED,
            WorkflowStatus.COMPLETED,
            WorkflowStatus.FAILED,
        }:
            record.completed_at = None

        self.db.add(
            WorkflowAction(
                code=(
                    f"{WORKFLOW_ACTION_CODE_PREFIX}-{now:%Y%m%d%H%M%S%f}-"
                    f"{uuid4().hex[:8]}"
                ),
                workflow_code=record.code,
                source_event_id=source_event_id,
                action=action,
                actor=actor,
                from_status=previous_status,
                to_status=status_value,
                payload=payload or {},
            )
        )
        if commit:
            self.db.commit()
            self.db.refresh(record)
        else:
            self.db.flush()
        return record

    def _workflow_for_source_event(
        self,
        source_event_id: str | None,
    ) -> WorkflowInstance | None:
        if not source_event_id:
            return None
        action = self.db.execute(
            select(WorkflowAction).where(
                WorkflowAction.source_event_id == source_event_id
            )
        ).scalar_one_or_none()
        if action is None:
            return None
        return self.db.execute(
            select(WorkflowInstance).where(WorkflowInstance.code == action.workflow_code)
        ).scalar_one()

    def list_workflows(
        self,
        status_filter: str | None = None,
        workflow_type: str | None = None,
        limit: int = 100,
    ) -> list[dict[str, Any]]:
        stmt = (
            select(WorkflowInstance)
            .order_by(WorkflowInstance.id.desc())
            .limit(bounded_limit(limit))
        )
        if status_filter:
            stmt = stmt.where(WorkflowInstance.status == status_filter)
        if workflow_type:
            stmt = stmt.where(WorkflowInstance.workflow_type == workflow_type)
        return [serialize_model(item) for item in self.db.execute(stmt).scalars()]

    def get_workflow(self, code: str) -> dict[str, Any]:
        record = self.db.execute(
            select(WorkflowInstance).where(WorkflowInstance.code == code)
        ).scalar_one_or_none()
        if record is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=WorkflowErrorDetail.WORKFLOW_NOT_FOUND,
            )
        actions = self.db.execute(
            select(WorkflowAction)
            .where(WorkflowAction.workflow_code == code)
            .order_by(WorkflowAction.id.asc())
        ).scalars()
        data = serialize_model(record)
        data["actions"] = [serialize_model(item) for item in actions]
        return data

    def count_by_status(self) -> dict[str, int]:
        rows = self.db.execute(
            select(WorkflowInstance.status, func.count()).group_by(WorkflowInstance.status)
        ).all()
        return {str(status_value): int(count) for status_value, count in rows}