xiaoxia-saas/apps/worker/tasks.py

import os
import tempfile
from datetime import datetime, timezone

from app.config import get_settings
from app.core.storage import get_minio_service

from packages.adapters.sqlalchemy_impl.asset_repository import SQLAlchemyAssetRepository
from packages.adapters.sqlalchemy_impl.generated_video_repository import (
    SQLAlchemyGeneratedVideoRepository,
)
from packages.adapters.sqlalchemy_impl.generation_task_repository import (
    SQLAlchemyGenerationTaskRepository,
)
from packages.adapters.sqlalchemy_impl.session import (
    SessionLocal,
    build_session_factory,
)
from packages.domain import GeneratedVideo, GenerationTaskStatus

from .celery_app import celery_app
from .video_processing import VideoProcessor

settings = get_settings()
if SessionLocal is None:
    build_session_factory(settings.database_url)


@celery_app.task(name="worker.generate_video")
def generate_video(task_id: str) -> dict:
    session = SessionLocal()
    temp_dir = tempfile.mkdtemp()

    try:
        task_repo = SQLAlchemyGenerationTaskRepository(session)
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        asset_repo = SQLAlchemyAssetRepository(session)
        storage_service = get_minio_service()

        # 1. 获取生成任务
        task = task_repo.get(task_id)
        if task is None:
            return {"ok": False, "error": f"generation task {task_id} not found"}

        # 2. 更新任务状态为运行中
        task.status = GenerationTaskStatus.RUNNING
        task.progress = 10.0
        task.started_at = task.started_at or datetime.now(timezone.utc)
        task_repo.update(task)
        session.commit()

        # 3. 获取素材库中的素材
        assets = asset_repo.list_by_library(task.asset_library_id)
        if not assets:
            raise RuntimeError(f"No assets found in library {task.asset_library_id}")

        task.progress = 20.0
        task_repo.update(task)
        session.commit()

        # 4. 下载素材到本地（简化：只处理前3个视频素材）
        video_assets = [a for a in assets if a.mime_type.startswith("video/")][:3]
        if not video_assets:
            raise RuntimeError("No video assets found")

        local_paths = []
        for i, asset in enumerate(video_assets):
            local_path = os.path.join(temp_dir, f"input_{i}.mp4")
            storage_service.download_file(asset.storage_key, local_path)
            local_paths.append(local_path)

            task.progress = 20.0 + (i + 1) * 10.0
            task_repo.update(task)
            session.commit()

        # 5. 使用 VideoProcessor 生成视频
        processor = VideoProcessor(temp_dir=temp_dir)
        output_filename = f"{task.id}.mp4"
        output_path = os.path.join(temp_dir, output_filename)

        task.progress = 50.0
        task_repo.update(task)
        session.commit()

        result = processor.concatenate_videos(
            input_paths=local_paths,
            output_path=output_path,
            resolution=(1920, 1080),
            fps=25,
        )

        task.progress = 80.0
        task_repo.update(task)
        session.commit()

        # 6. 上传到 MinIO
        storage_key = f"workspaces/{task.workspace_id}/projects/{task.project_id}/generated/{task.id}/{output_filename}"
        thumbnail_key = f"workspaces/{task.workspace_id}/projects/{task.project_id}/generated/{task.id}/thumbnail.jpg"

        storage_service.upload_file(result.output_path, storage_key)
        storage_service.upload_file(result.thumbnail_path, thumbnail_key)

        file_url = storage_service.get_url(storage_key)
        thumbnail_url = storage_service.get_url(thumbnail_key)

        task.progress = 90.0
        task_repo.update(task)
        session.commit()

        # 7. 创建 GeneratedVideo 记录
        video = GeneratedVideo.create(
            workspace_id=task.workspace_id,
            project_id=task.project_id,
            generation_task_id=task.id,
            name=output_filename,
            file_url=file_url,
            file_size=result.file_size,
            duration=result.duration,
            thumbnail_url=thumbnail_url,
            width=result.width,
            height=result.height,
            fps=result.fps,
        )
        video_repo.create(video)

        # 8. 更新任务状态为完成
        task.status = GenerationTaskStatus.COMPLETED
        task.progress = 100.0
        task.result_count = 1
        task.completed_at = datetime.now(timezone.utc)
        task_repo.update(task)
        session.commit()

        return {
            "ok": True,
            "task_id": task.id,
            "video_id": video.id,
            "file_url": file_url,
            "duration": result.duration,
            "file_size": result.file_size,
        }

    except Exception as error:
        try:
            task_repo = SQLAlchemyGenerationTaskRepository(session)
            task = task_repo.get(task_id)
            if task is not None:
                task.status = GenerationTaskStatus.FAILED
                task.error_message = str(error)
                task.completed_at = datetime.now(timezone.utc)
                task_repo.update(task)
                session.commit()
        except Exception:
            pass

        return {"ok": False, "task_id": task_id, "error": str(error)}

    finally:
        session.close()
        # 清理临时文件
        try:
            import shutil

            shutil.rmtree(temp_dir, ignore_errors=True)
        except:
            pass