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

"""批量修复素材缩略图 — 为历史视频素材生成缩略图。

使用方式：
  从管理接口或 shell 触发：
    celery_app.send_task("worker.batch_generate_thumbnails")

逻辑：
  1. 查询所有 file_type=video 且 thumbnail_url 为空或为旧格式公开 URL 的素材
  2. 逐个：下载视频 → 抽第一帧 → 上传 OSS → 更新 thumbnail_url 为 storage_key
  3. 每处理 50 条 commit 一次，失败单条跳过不阻塞
"""

from __future__ import annotations

import tempfile
from pathlib import Path

from celery.utils.log import get_task_logger
from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

logger = get_task_logger(__name__)

# 旧格式 URL 前缀（ingest 旧代码生成的公开 URL），需替换为 storage_key
_OLD_URL_PREFIX = "https://xiaoxia-autocut.oss-cn-hangzhou.aliyuncs.com/"
# 从公开 URL 中提取 storage_key 时，去掉域名前缀即可
_DOMAIN_PREFIXES = [
    "https://xiaoxia-autocut.oss-cn-hangzhou.aliyuncs.com/",
    "http://xiaoxia-autocut.oss-cn-hangzhou.aliyuncs.com/",
]


def _url_to_storage_key(url: str) -> str | None:
    """尝试将旧格式公开 URL 转回 storage_key。"""
    for prefix in _DOMAIN_PREFIXES:
        if url.startswith(prefix):
            return url[len(prefix) :]
    return None


@celery_app.task(name="worker.batch_generate_thumbnails")
def batch_generate_thumbnails() -> dict:
    """为所有缺少缩略图的视频素材批量生成缩略图。

    Returns:
        dict: {total, success, skipped, failed, converted_legacy}
    """
    from video_processing.oss_helpers import download_asset, upload_to_oss
    from video_processing.thumbnail_generator import extract_first_frame

    from packages.adapters.sqlalchemy_impl.models import AssetModel

    db = SessionLocal()
    stats = {"total": 0, "success": 0, "skipped": 0, "failed": 0, "converted_legacy": 0}

    try:
        # 查询所有视频素材中缩略图缺失的
        assets = (
            db.query(AssetModel)
            .filter(
                AssetModel.file_type == "video",
                AssetModel.status != "deleted",
            )
            .all()
        )

        # 筛选需要处理的：thumbnail_url 为空 或 旧格式公开 URL
        to_process = []
        for asset in assets:
            thumb = asset.thumbnail_url or ""
            if not thumb:
                to_process.append((asset, None))  # (asset, None=需要新生成)
            elif thumb.startswith("http"):
                # 旧格式公开 URL → 尝试转为 storage_key
                sk = _url_to_storage_key(thumb)
                if sk:
                    to_process.append((asset, sk))  # 已有文件，只需改 DB
                else:
                    # 非预期 URL 格式，跳过
                    stats["skipped"] += 1
            # else: 已经是 storage_key 格式，跳过

        stats["total"] = len(to_process)
        logger.info(
            "批量缩略图修复启动: total=%d (new=%d, legacy_convert=%d)",
            stats["total"],
            sum(1 for _, sk in to_process if sk is None),
            sum(1 for _, sk in to_process if sk is not None),
        )

        batch_count = 0
        for asset, existing_key in to_process:
            try:
                if existing_key is not None:
                    # 旧 URL → storage_key，只需更新 DB
                    asset.thumbnail_url = existing_key
                    stats["converted_legacy"] += 1
                    stats["success"] += 1
                else:
                    # 需要新生成缩略图
                    if not asset.storage_key:
                        stats["skipped"] += 1
                        continue

                    suffix = Path(asset.storage_key).suffix or ".mp4"
                    local_file = None
                    frame_path = None
                    try:
                        # 下载视频
                        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
                            local_file = Path(tmp.name)

                        if not download_asset(asset.storage_key, local_file):
                            logger.warning("下载失败: asset_id=%s key=%s", asset.id, asset.storage_key[:60])
                            stats["failed"] += 1
                            continue

                        # 抽帧
                        frame_path = extract_first_frame(str(local_file), width=640)

                        # 上传
                        thumb_key = f"assets/{asset.project_id}/thumbnails/{asset.id}.jpg"
                        upload_ok = upload_to_oss(frame_path, thumb_key)
                        if upload_ok:
                            asset.thumbnail_url = thumb_key
                            stats["success"] += 1
                        else:
                            logger.warning("上传失败: asset_id=%s", asset.id)
                            stats["failed"] += 1
                    finally:
                        if local_file and local_file.exists():
                            try:
                                local_file.unlink(missing_ok=True)
                            except OSError:
                                pass
                        if frame_path and Path(frame_path).exists():
                            try:
                                Path(frame_path).unlink(missing_ok=True)
                            except OSError:
                                pass

                batch_count += 1
                if batch_count % 50 == 0:
                    db.commit()
                    logger.info("批量缩略图进度: %d/%d", batch_count, stats["total"])

            except Exception as e:
                logger.warning("单条处理失败: asset_id=%s error=%s", asset.id, e)
                stats["failed"] += 1
                try:
                    db.rollback()
                except Exception:
                    pass

        # 最后提交
        db.commit()
        logger.info("批量缩略图修复完成: %s", stats)
        return stats

    except Exception as e:
        logger.error("批量缩略图修复异常: %s", e)
        db.rollback()
        return {**stats, "error": str(e)}
    finally:
        db.close()