xiaoxia-saas/apps/worker/worker_app/tasks/batch_download.py

"""批量下载任务 — 将多个成片打包为 zip 上传到 OSS。"""

from __future__ import annotations

import logging
import tempfile
import uuid
import zipfile
from pathlib import Path

from worker_app.celery_app import celery_app

logger = logging.getLogger(__name__)


@celery_app.task(bind=True, name="worker.batch_download_videos", max_retries=1)
def batch_download_videos(self, video_ids: list[str], user_id: str = "") -> dict:
    """批量下载视频并打包为 zip。

    Args:
        video_ids: 视频 ID 列表
        user_id: 发起用户 ID

    Returns:
        {"download_url": "...", "file_count": N, "total_size": total_bytes}
    """
    from video_processing.oss_helpers import upload_to_oss
    from worker_app.db import SessionLocal

    from packages.adapters.sqlalchemy_impl.generated_video_repository import (
        SQLAlchemyGeneratedVideoRepository,
    )

    session = SessionLocal()
    try:
        repo = SQLAlchemyGeneratedVideoRepository(session)
        videos = repo.get_by_ids(video_ids)
    finally:
        session.close()

    if not videos:
        raise ValueError("No videos found for batch download")

    # 创建临时工作目录
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir_path = Path(tmpdir)
        zip_filename = f"videos-{len(videos)}-{video_ids[0][:8]}.zip"
        zip_path = tmpdir_path / zip_filename

        # 逐个下载视频并加入 zip
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as zf:
            for idx, video in enumerate(videos, 1):
                logger.info("Batch download: downloading %d/%d %s", idx, len(videos), video.id)
                try:
                    # 下载视频到临时文件
                    local_name = f"{idx:03d}_{video.name}"
                    local_path = tmpdir_path / local_name

                    # 使用 oss_helpers 的 download_asset，或者直接从 URL 下载
                    if video.file_url:
                        _download_video_to_file(video.file_url, str(local_path))

                    if local_path.exists() and local_path.stat().st_size > 0:
                        zf.write(str(local_path), arcname=local_name)
                        local_path.unlink(missing_ok=True)
                    else:
                        logger.warning("Video %s download failed, skipping", video.id)
                except Exception as e:
                    logger.warning("Failed to download video %s: %s", video.id, e)
                    continue

        # 上传 zip 到 OSS
        if not zip_path.exists() or zip_path.stat().st_size == 0:
            raise RuntimeError("Batch download zip file is empty")
        zip_storage_key = f"batch-downloads/{uuid.uuid4().hex}/{zip_filename}"
        download_url = upload_to_oss(str(zip_path), zip_storage_key)

        total_size = zip_path.stat().st_size
        file_count = len(zipfile.ZipFile(str(zip_path), "r").namelist())

        logger.info(
            "Batch download complete: %d files, %d bytes, url=%s",
            file_count,
            total_size,
            download_url,
        )

        return {
            "download_url": download_url,
            "file_count": file_count,
            "total_size": total_size,
            "video_count": len(videos),
        }


def _download_video_to_file(url: str, dest_path: str) -> None:
    """下载视频文件到本地路径。优先用 OSS SDK 走内网，回退到 HTTP 下载。"""
    from video_processing.oss_helpers import download_asset

    try:
        # 尝试走 OSS 下载（如果是 OSS URL 的话）
        success = download_asset(url, dest_path)
        if success:
            return
    except Exception:
        pass

    # 回退到 HTTP 下载（含 SSRF 防护 + 大小限制 + 类型校验）
    from video_processing.url_security import (
        ALLOWED_VIDEO_MIME_TYPES,
        safe_download_file,
    )

    safe_download_file(
        url,
        dest_path,
        purpose="batch_video_download",
        allowed_mime_types=ALLOWED_VIDEO_MIME_TYPES | {"application/octet-stream"},
        timeout=300.0,
    )