xiaoxia-saas/apps/worker/video_processing/oss_helpers.py

"""OSS 工具函数 — 从 generation.py / edit_plan_generation.py 提取的共享 OSS 操作.

提供 OSS 配置读取、Bucket 创建、素材上传/下载、asset_id → 本地路径解析
等能力，供 render_edit_plan 和 generate_video 共同复用。
"""

from __future__ import annotations

import hashlib
import logging
import os
import threading
from pathlib import Path
from urllib.parse import urlparse

import oss2

logger = logging.getLogger(__name__)

# OSS 上传配置
OSS_CONNECT_TIMEOUT = 10  # 连接超时（秒），防止 TCP 握手挂死
OSS_UPLOAD_TOTAL_TIMEOUT = 300  # 单文件上传总超时（秒），防止网络慢时无限卡住
OSS_MULTIPART_THRESHOLD = 100 * 1024 * 1024  # 分片上传阈值：100MB 以上走分片
OSS_PART_SIZE = 8 * 1024 * 1024  # 分片大小：8MB
OSS_MULTIPART_NUM_THREADS = 3  # 分片上传并发数


# ── OSS 配置 ──────────────────────────────────────────────────────────────────


def oss_settings() -> tuple[str, str, str, str] | None:
    """获取 OSS 配置。

    Returns:
        (access_key_id, access_key_secret, endpoint, bucket_name) 元组，
        配置缺失时返回 None。
    """
    access_key_id = os.getenv("OSS_ACCESS_KEY_ID")
    access_key_secret = os.getenv("OSS_ACCESS_KEY_SECRET")
    endpoint = os.getenv("OSS_ENDPOINT")
    bucket_name = os.getenv("OSS_BUCKET_NAME")
    if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
        return None
    return access_key_id, access_key_secret, endpoint, bucket_name


def oss_bucket() -> oss2.Bucket | None:
    """获取 OSS Bucket 实例。

    P0-2 修复：endpoint 不带 scheme 时自动补 https:// 前缀，
    确保 sign_url 等依赖 scheme 的方法返回 HTTPS URL。

    P0-staging 修复：增加 connect_timeout=10s，防止网络抖动时
    TCP 握手阶段无限挂死，导致 worker 进程卡死。

    Returns:
        oss2.Bucket 实例，配置缺失时返回 None。
    """
    settings = oss_settings()
    if settings is None:
        return None
    access_key_id, access_key_secret, endpoint, bucket_name = settings
    # endpoint 无 scheme 时补 https://，与 API 端 storage.py 保持一致
    if not endpoint.startswith(("http://", "https://")):
        endpoint = f"https://{endpoint}"
    return oss2.Bucket(
        oss2.Auth(access_key_id, access_key_secret),
        endpoint,
        bucket_name,
        connect_timeout=OSS_CONNECT_TIMEOUT,
    )


def normalize_storage_key(storage_key_or_url: str) -> str:
    """标准化存储键 — 如果是完整 URL 则提取 path 部分。

    Examples:
        "https://bucket.oss-cn-hangzhou.aliyuncs.com/path/to/file.mp4"
        → "path/to/file.mp4"
        "path/to/file.mp4" → "path/to/file.mp4"
    """
    if storage_key_or_url.startswith(("http://", "https://")):
        return urlparse(storage_key_or_url).path.lstrip("/")
    return storage_key_or_url.lstrip("/")


# ── 上传 / 下载 ───────────────────────────────────────────────────────────────


def download_asset(asset_storage_key: str, local_path: Path) -> bool:
    """从 OSS 下载素材文件到本地路径。

    Args:
        asset_storage_key: 素材的存储键（或完整 URL）
        local_path: 本地保存路径

    Returns:
        True 表示下载成功，False 表示失败。
    """
    bucket = oss_bucket()
    if bucket is None:
        return False
    try:
        bucket.get_object_to_file(normalize_storage_key(asset_storage_key), str(local_path))
        return local_path.exists() and local_path.stat().st_size > 0
    except Exception:
        logger.exception("下载素材失败: %s", asset_storage_key)
        return False


def upload_to_oss(local_path: Path, storage_key: str) -> str | None:
    """上传文件到 OSS，返回公开 URL。

    大文件（>100MB）自动走分片上传，降低内存峰值，减少 OOM 风险。
    上传加总超时保护（默认 300s），防止网络异常时无限挂死。

    Args:
        local_path: 本地文件路径
        storage_key: 目标存储键

    Returns:
        公开访问 URL，上传失败或 OSS 未配置时返回 None。
    """
    bucket = oss_bucket()
    if bucket is None:
        return None

    result: dict = {"url": None, "error": None, "file_size": 0}
    done = threading.Event()

    def _do_upload():
        try:
            # 尝试获取文件大小，用于分片判断和日志；stat 失败时 fallback 走普通上传
            try:
                file_size = local_path.stat().st_size
                result["file_size"] = file_size
                use_multipart = file_size >= OSS_MULTIPART_THRESHOLD
            except OSError:
                use_multipart = False
                file_size = 0

            if use_multipart:
                # 分片上传：降低内存峰值，每片 8MB，3 线程并发
                logger.info(
                    "大文件分片上传: storage_key=%s, size=%.1fMB, part_size=%dMB, threads=%d",
                    storage_key[:80],
                    file_size / 1024 / 1024,
                    OSS_PART_SIZE // 1024 // 1024,
                    OSS_MULTIPART_NUM_THREADS,
                )
                oss2.resumable_upload(
                    bucket,
                    storage_key,
                    str(local_path),
                    multipart_threshold=OSS_MULTIPART_THRESHOLD,
                    part_size=OSS_PART_SIZE,
                    num_threads=OSS_MULTIPART_NUM_THREADS,
                )
            else:
                bucket.put_object_from_file(storage_key, str(local_path))

            # 构造返回 URL
            settings = oss_settings()
            if settings:
                _, _, endpoint, bucket_name = settings
                endpoint_clean = endpoint.replace("https://", "").replace("http://", "")
                result["url"] = f"https://{bucket_name}.{endpoint_clean}/{storage_key}"
        except Exception as e:
            result["error"] = e
            logger.exception("上传 OSS 失败: %s", storage_key)
        finally:
            done.set()

    upload_thread = threading.Thread(target=_do_upload, daemon=True)
    upload_thread.start()
    finished = done.wait(timeout=OSS_UPLOAD_TOTAL_TIMEOUT)

    if not finished:
        logger.error(
            "OSS 上传超时（%.0fs），强制中止: storage_key=%s, size=%.1fMB",
            OSS_UPLOAD_TOTAL_TIMEOUT,
            storage_key[:80],
            result["file_size"] / 1024 / 1024 if result["file_size"] else 0,
        )
        return None

    if result["error"]:
        return None

    return result["url"]


def get_signed_download_url(storage_key_or_url: str, expires_seconds: int = 3600) -> str | None:
    """生成预签名下载 URL（用于私有 bucket 的 URL 校验或临时下载）。

    Args:
        storage_key_or_url: 存储键或完整 URL（URL 会自动提取 path）
        expires_seconds: 签名有效期（秒）

    Returns:
        预签名 URL，失败或 OSS 未配置时返回 None。
    """
    bucket = oss_bucket()
    if bucket is None:
        return None
    try:
        storage_key = normalize_storage_key(storage_key_or_url)
        signed = bucket.sign_url("GET", storage_key, expires_seconds)
        logger.info("生成预签名URL: key=%s url_prefix=%s", storage_key[:80], signed[:60])
        return signed
    except Exception:
        logger.exception("生成预签名URL失败: %s", storage_key_or_url[:80])
        return None


# ── Asset 解析 ────────────────────────────────────────────────────────────────


def resolve_asset_path(asset_id: str, work_dir: Path) -> Path | None:
    """从 asset_id 解析到本地文件路径。

    策略（按优先级）：
    1. 如果 asset_id 是本地绝对路径（/var/storage/...）→ 直接返回
    2. 如果 work_dir 下已有缓存文件 → 返回缓存路径
    3. 从 OSS 下载到 work_dir/{hash}.mp4 → 返回下载路径
    4. 下载失败 → 返回 None

    缓存策略：以 asset_id 的 SHA256 前 16 位为文件名，避免重复下载。
    """
    # 1. 本地绝对路径
    if asset_id.startswith("/") and os.path.exists(asset_id):
        return Path(asset_id)

    # 2. 缓存命中
    cache_hash = hashlib.sha256(asset_id.encode()).hexdigest()[:16]
    cached_path = work_dir / f"{cache_hash}.mp4"
    if cached_path.exists() and cached_path.stat().st_size > 0:
        return cached_path

    # 3. 从 OSS 下载
    if download_asset(asset_id, cached_path):
        return cached_path

    return None


def resolve_asset_ids_to_paths(
    asset_ids: list[str],
    work_dir: Path,
) -> dict[str, Path]:
    """批量解析 asset_id → 本地路径。

    Args:
        asset_ids: 素材 ID 列表
        work_dir: 工作目录

    Returns:
        {asset_id: local_path} 映射，仅包含成功解析的条目。
    """
    result: dict[str, Path] = {}
    for aid in asset_ids:
        local_path = resolve_asset_path(aid, work_dir)
        if local_path:
            result[aid] = local_path
    return result