xiaoxia-saas/packages/domain/job.py

"""Job 领域模型 — Phase 8 任务 2.10.

统一的异步任务抽象，为视频合成等耗时操作提供异步任务管理能力。

状态机：
  pending → running → success
                  ↘ failed → pending (重试)
       ↘ cancelled
"""

from __future__ import annotations

import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone

if sys.version_info >= (3, 11):
    from enum import StrEnum
else:
    from enum import Enum

    class StrEnum(str, Enum):
        pass


from uuid import uuid4


class JobType(StrEnum):
    """任务类型枚举。"""

    VIDEO_COMPOSE = "video_compose"
    """视频合成（VideoComposeService）"""

    RENDER_EDIT_PLAN = "render_edit_plan"
    """剪辑计划渲染"""

    ASSET_INGEST = "asset_ingest"
    """素材导入"""

    CLASSIFICATION = "classification"
    """素材分类"""

    VOICE_EXTRACTION = "voice_extraction"
    """语音提取"""

    GENERATION = "generation"
    """通用生成"""


class JobStatus(StrEnum):
    """任务状态枚举。"""

    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    CANCELLED = "cancelled"


# 终态集合
TERMINAL_STATUSES = frozenset({JobStatus.SUCCESS, JobStatus.FAILED, JobStatus.CANCELLED})

# 合法状态转换
_VALID_TRANSITIONS: dict[JobStatus, set[JobStatus]] = {
    JobStatus.PENDING: {JobStatus.RUNNING, JobStatus.SUCCESS, JobStatus.CANCELLED},
    JobStatus.RUNNING: {JobStatus.SUCCESS, JobStatus.FAILED, JobStatus.CANCELLED},
    JobStatus.FAILED: {JobStatus.PENDING},  # 重试回到 pending
}


@dataclass(slots=True)
class Job:
    """统一异步任务实体。

    Attributes:
        id: 任务唯一标识
        project_id: 所属项目
        job_type: 任务类型
        status: 当前状态
        progress: 进度百分比 (0.0 ~ 100.0)
        current_stage: 当前阶段描述（人类可读）
        payload: 任务输入参数（JSON 序列化）
        result: 任务结果（JSON 序列化）
        error_message: 错误信息
        retry_count: 已重试次数
        max_retries: 最大重试次数
        celery_task_id: Celery 异步任务 ID
        source_id: 关联的业务实体 ID（如 edit_plan_id, generation_task_id）
        created_by_user_id: 创建人
        started_at: 开始执行时间
        completed_at: 完成时间
        created_at: 创建时间
        updated_at: 最后更新时间
    """

    id: str
    project_id: str
    job_type: JobType
    status: JobStatus = JobStatus.PENDING
    progress: float = 0.0
    current_stage: str = ""
    payload: dict = field(default_factory=dict)
    result: dict = field(default_factory=dict)
    error_message: str = ""
    retry_count: int = 0
    max_retries: int = 3
    celery_task_id: str = ""
    source_id: str = ""
    created_by_user_id: str = ""
    started_at: datetime | None = None
    completed_at: datetime | None = None
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    @classmethod
    def create(
        cls,
        project_id: str,
        job_type: JobType | str,
        *,
        payload: dict | None = None,
        source_id: str = "",
        created_by_user_id: str = "",
        max_retries: int = 3,
    ) -> Job:
        """创建新任务。

        Args:
            project_id: 项目 ID
            job_type: 任务类型
            payload: 任务输入参数
            source_id: 关联业务实体 ID
            created_by_user_id: 创建人 ID
            max_retries: 最大重试次数

        Returns:
            新建的 Job 实例

        Raises:
            ValueError: 参数校验失败
        """
        if not project_id.strip():
            raise ValueError("project_id 不能为空")

        # 兼容字符串类型
        if isinstance(job_type, str):
            try:
                job_type = JobType(job_type)
            except ValueError as _e:
                raise ValueError(f"不支持的任务类型: {job_type}") from _e

        return cls(
            id=uuid4().hex,
            project_id=project_id.strip(),
            job_type=job_type,
            payload=payload or {},
            source_id=source_id.strip(),
            created_by_user_id=created_by_user_id.strip(),
            max_retries=max_retries,
        )

    @property
    def is_terminal(self) -> bool:
        """是否处于终态。"""
        return self.status in TERMINAL_STATUSES

    @property
    def is_retryable(self) -> bool:
        """是否可重试（失败且未超过重试上限）。"""
        return self.status == JobStatus.FAILED and self.retry_count < self.max_retries

    def transition_to(self, new_status: JobStatus | str) -> None:
        """执行状态转换。

        Args:
            new_status: 目标状态

        Raises:
            ValueError: 非法状态转换
        """
        if isinstance(new_status, str):
            try:
                new_status = JobStatus(new_status)
            except ValueError as _e:
                raise ValueError(f"无效状态: {new_status}") from _e

        allowed = _VALID_TRANSITIONS.get(self.status, set())
        if new_status not in allowed:
            raise ValueError(
                f"非法状态转换: {self.status.value} → {new_status.value}，"
                f"允许: {{{', '.join(s.value for s in allowed)}}}"
            )

        now = datetime.now(timezone.utc)
        self.status = new_status
        self.updated_at = now

        # 自动设置时间戳
        if new_status == JobStatus.RUNNING and self.started_at is None:
            self.started_at = now
        elif new_status in (JobStatus.SUCCESS, JobStatus.FAILED):
            self.completed_at = now

    def mark_running(self, stage: str = "") -> None:
        """标记为执行中。"""
        self.transition_to(JobStatus.RUNNING)
        if stage:
            self.current_stage = stage

    def mark_success(self, result: dict | None = None) -> None:
        """标记为成功。"""
        self.transition_to(JobStatus.SUCCESS)
        self.progress = 100.0
        self.current_stage = "完成"
        if result:
            self.result = result

    def mark_failed(self, error_message: str) -> None:
        """标记为失败。"""
        self.transition_to(JobStatus.FAILED)
        self.error_message = error_message
        self.current_stage = "失败"

    def mark_cancelled(self) -> None:
        """标记为取消。"""
        self.transition_to(JobStatus.CANCELLED)
        self.current_stage = "已取消"

    def update_progress(self, progress: float, stage: str = "") -> None:
        """更新进度。

        Args:
            progress: 进度百分比 (0.0 ~ 100.0)
            stage: 当前阶段描述

        Raises:
            ValueError: 进度值非法
        """
        if not 0.0 <= progress <= 100.0:
            raise ValueError(f"进度必须在 0~100 之间，当前: {progress}")
        self.progress = progress
        if stage:
            self.current_stage = stage
        self.updated_at = datetime.now(timezone.utc)

    def prepare_retry(self) -> None:
        """准备重试：重置状态为 pending。

        Raises:
            ValueError: 不可重试
        """
        if not self.is_retryable:
            raise ValueError(
                f"任务不可重试: status={self.status.value}, "
                f"retry_count={self.retry_count}, max_retries={self.max_retries}"
            )
        self.retry_count += 1
        self.transition_to(JobStatus.PENDING)
        self.progress = 0.0
        self.current_stage = f"第 {self.retry_count} 次重试"
        self.error_message = ""
        self.started_at = None
        self.completed_at = None
        self.celery_task_id = ""

    def to_dict(self) -> dict:
        """序列化为字典。"""
        return {
            "id": self.id,
            "project_id": self.project_id,
            "job_type": self.job_type.value,
            "status": self.status.value,
            "progress": self.progress,
            "current_stage": self.current_stage,
            "payload": self.payload,
            "result": self.result,
            "error_message": self.error_message,
            "retry_count": self.retry_count,
            "max_retries": self.max_retries,
            "celery_task_id": self.celery_task_id,
            "source_id": self.source_id,
            "created_by_user_id": self.created_by_user_id,
            "is_retryable": self.is_retryable,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }