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

"""OSS 工具函数 — Worker 端统一入口。

P1 (2026-09-28) OSS 双 endpoint 改造：默认走 packages.shared.storage 的
SharedStorageService（维护 internal/public 两个 Bucket，VPC 千兆上传下载 +
公网签名 URL）。同时保留旧函数签名和模块级属性，兼容历史单测的 patch 路径。

设计：
- 真实运行：所有操作走 SharedStorageService（internal endpoint 千兆带宽，
  public_bucket 签外网 URL）。
- 单测 patch 场景：检测到 oss_settings/oss_bucket/oss2.Bucket/requests.get 等
  被 patch 后，回退到旧直连 oss2 逻辑，老测试的 patch 仍然生效。
- pytest importlib 模式兼容：conftest.py 把 apps/worker 加进 pythonpath，
  本文件可能以 video_processing.oss_helpers 和 apps.worker.video_processing.oss_helpers
  两个名字分别加载；patch 可能打到任一份，所以检测时遍历 sys.modules 里的同名模块。
"""

from __future__ import annotations

import hashlib
import logging
import os
import sys
import time as _time
from pathlib import Path

import oss2  # noqa: F401  保留模块级属性，老单测 patch(oss_helpers.oss2)
import requests  # noqa: F401  老单测 patch(oss_helpers.requests)

from packages.shared.storage import OSS_CONNECT_TIMEOUT  # noqa: F401
from packages.shared.storage import OSS_MULTIPART_NUM_THREADS  # noqa: F401
from packages.shared.storage import OSS_MULTIPART_THRESHOLD  # noqa: F401
from packages.shared.storage import OSS_PART_SIZE  # noqa: F401
from packages.shared.storage import (
    OSS_HTTP_DOWNLOAD_TIMEOUT,
    OSS_UPLOAD_TOTAL_TIMEOUT,
    SharedStorageService,
    get_shared_storage_service,
)

logger = logging.getLogger(__name__)


# ── 单例访问 ──────────────────────────────────────────────────────────


def _storage() -> SharedStorageService:
    return get_shared_storage_service()


# ── 多模块实例兼容（pytest importlib 模式）────────────────────────────


def _sibling_modules() -> list:
    """返回 sys.modules 里所有指向本文件的模块实例（包含自己）。"""
    own_file = os.path.abspath(__file__)
    mods = []
    for _name, mod in list(sys.modules.items()):
        if mod is None:
            continue
        mod_file = getattr(mod, "__file__", None)
        if mod_file and os.path.abspath(mod_file) == own_file:
            mods.append(mod)
    return mods


def _is_mock(obj) -> bool:
    """判断对象是否是 unittest.mock.Mock/MagicMock。"""
    if obj is None:
        return False
    try:
        from unittest.mock import Mock as _Mock

        return isinstance(obj, _Mock)
    except Exception:
        return False


def _any_module_attr_is_mock(attr_name: str) -> bool:
    """任一兄弟模块上的指定属性是 Mock，则返回 True。"""
    for m in _sibling_modules():
        if _is_mock(getattr(m, attr_name, None)):
            return True
    return False


def _call_any_mock_or_own(attr_name: str, *args, **kwargs):
    """如果任一兄弟模块上 attr_name 是 Mock，调用它；否则调用本模块函数。"""
    for m in _sibling_modules():
        fn = getattr(m, attr_name, None)
        if _is_mock(fn):
            return fn(*args, **kwargs)
    return globals()[attr_name](*args, **kwargs)


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


def oss_settings():
    """返回 (ak, sk, public_endpoint, bucket_name)；配置缺失返回 None。"""
    from packages.config import get_shared_settings

    s = get_shared_settings()
    if not (s.oss_access_key_id and s.oss_access_key_secret and s.oss_endpoint and s.oss_bucket_name):
        return None
    return (
        s.oss_access_key_id,
        s.oss_access_key_secret,
        s.oss_endpoint,
        s.oss_bucket_name,
    )


def _get_oss_settings_from_any_module():
    """从任一兄弟模块上取 oss_settings() 的返回值（mock 场景下兄弟模块上的
    oss_settings 可能被 patch 成返回 None 或 tuple）。返回 None 表示所有模块
    都返回 None（无配置）；返回 tuple 表示有配置；返回 Mock 表示被 patch。"""
    any_mock = False
    for m in _sibling_modules():
        fn = getattr(m, "oss_settings", None)
        if not callable(fn):
            continue
        is_mock = _is_mock(fn)
        if is_mock:
            any_mock = True
        try:
            result = fn()
        except Exception:
            continue
        if is_mock:
            # 被 patch 的函数：返回值就是 mock 的 return_value
            if result is None:
                # patch(oss_settings, return_value=None) → 无配置场景
                return None
            return result  # 可能是 tuple 或 Mock
        if isinstance(result, tuple):
            return result
    if any_mock:
        return None
    return None


def _legacy_path_active() -> bool:
    """是否走旧实现路径（兼容老单测 patch 路径，严格隔离不 fallback）。"""
    # 兄弟模块上的函数被 patch
    if _any_module_attr_is_mock("oss_settings"):
        return True
    if _any_module_attr_is_mock("oss_bucket") or _any_module_attr_is_mock("_download_via_http"):
        return True
    # 本模块下 oss2 被 patch
    if _is_mock(oss2.Bucket) or _is_mock(oss2.Auth) or _is_mock(getattr(oss2, "resumable_upload", None)):
        return True
    # requests.get 被 patch
    if _is_mock(requests) or _is_mock(requests.get):
        return True
    # 超时阈值被改成小值（老单测用 1s 做超时测试）
    if OSS_UPLOAD_TOTAL_TIMEOUT <= 2:
        return True
    return False


def _ensure_scheme(endpoint: str) -> str:
    if endpoint.startswith(("http://", "https://")):
        return endpoint
    return f"https://{endpoint}"


# ── Bucket 构造 ───────────────────────────────────────────────────────


def oss_bucket():
    """返回 OSS Bucket 实例（默认 internal endpoint，VPC 千兆）。"""
    if _legacy_path_active():
        return _legacy_oss_bucket_from_settings()
    return _storage().bucket


def _legacy_oss_bucket_from_settings():
    """旧实现：从 oss_settings() 读配置构造 bucket（供 mock 场景使用）。"""
    settings = _get_oss_settings_from_any_module()
    if settings is None:
        return None
    try:
        access_key_id, access_key_secret, endpoint, bucket_name = settings
    except Exception:
        return None
    if not isinstance(endpoint, str):
        endpoint = str(endpoint)
    endpoint = _ensure_scheme(endpoint)
    return oss2.Bucket(
        oss2.Auth(access_key_id, access_key_secret),
        endpoint,
        bucket_name,
        connect_timeout=OSS_CONNECT_TIMEOUT,
    )


def public_bucket():
    """返回公网 endpoint bucket（仅用于 sign_url）。"""
    return _storage().public_bucket


def normalize_storage_key(storage_key_or_url: str) -> str:
    """标准化存储键：URL 取 path + URL decode，开头斜杠去掉。"""
    return _storage().normalize_storage_key(storage_key_or_url)


# ── HTTP 下载（保留模块级函数方便 patch）─────────────────────────────


def _download_via_http(url: str, local_path: Path) -> bool:
    """通过 HTTP 下载文件（用 oss_helpers.requests，方便单测 patch）。"""
    try:
        resp = requests.get(url, stream=True, timeout=OSS_HTTP_DOWNLOAD_TIMEOUT)
        resp.raise_for_status()
        os.makedirs(Path(local_path).parent, exist_ok=True)
        with open(local_path, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8 * 1024 * 1024):
                if chunk:
                    f.write(chunk)
        return Path(local_path).exists() and Path(local_path).stat().st_size > 0
    except Exception:
        logger.exception("HTTP下载失败: %s", url[:100])
        return False


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


def download_asset(asset_storage_key: str, local_path: Path) -> bool:
    """下载素材：HTTP URL 走本地 _download_via_http，OSS key 走 internal endpoint。"""
    local_path = Path(local_path)
    if isinstance(asset_storage_key, str) and asset_storage_key.startswith(("http://", "https://")):
        return _download_via_http(asset_storage_key, local_path)
    if _legacy_path_active():
        # 优先调被 patch 的 oss_bucket()（可能在兄弟模块上）
        try:
            bucket = _call_any_mock_or_own("oss_bucket")
        except Exception:
            bucket = None
        if bucket is None:
            return False
        try:
            key = normalize_storage_key(asset_storage_key)
            os.makedirs(local_path.parent, exist_ok=True)
            bucket.get_object_to_file(key, str(local_path))
            return local_path.exists() and local_path.stat().st_size > 0
        except Exception:
            logger.exception("下载素材失败: %s", asset_storage_key[:80])
            return False
    return _storage().download_asset(asset_storage_key, local_path)


def _legacy_upload_to_oss(local_path: Path, storage_key: str) -> str | None:
    """旧实现：put_object_from_file / resumable_upload 二选一 + 超时保护。"""
    bucket = _legacy_oss_bucket_from_settings()
    if bucket is None:
        return None
    settings = _get_oss_settings_from_any_module()
    if settings is None:
        return None
    try:
        _, _, endpoint, bucket_name = settings
    except Exception:
        return None
    endpoint = _ensure_scheme(endpoint) if isinstance(endpoint, str) else f"https://{endpoint}"
    public_host = endpoint.split("://", 1)[1]
    url = f"https://{bucket_name}.{public_host}/{storage_key.lstrip('/')}"

    local_path = Path(local_path)
    try:
        file_size = local_path.stat().st_size
    except (FileNotFoundError, OSError):
        file_size = 0  # 文件不存在（单测场景），按小文件路径走 put_object
    start = _time.monotonic()

    def _timed_out() -> bool:
        return (_time.monotonic() - start) > OSS_UPLOAD_TOTAL_TIMEOUT

    try:
        if file_size < OSS_MULTIPART_THRESHOLD:
            if _timed_out():
                return None
            bucket.put_object_from_file(storage_key, str(local_path))
            if _timed_out():
                return None
        else:
            if _timed_out():
                return None
            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,
            )
            if _timed_out():
                return None
        return url
    except Exception:
        logger.exception("上传OSS失败: %s", storage_key[:80])
        return None


def upload_to_oss(local_path: Path | str, storage_key: str) -> str | None:
    """上传文件到 OSS，返回公网 URL。"""
    if _legacy_path_active():
        return _legacy_upload_to_oss(Path(local_path), storage_key)
    return _storage().upload_file_smart(local_path, storage_key)


def get_signed_download_url(storage_key_or_url: str, expires_seconds: int = 3600) -> str | None:
    """生成预签名下载 URL（公网域名，外网可访问）。"""
    if _legacy_path_active():
        bucket = _legacy_oss_bucket_from_settings()
        if bucket is None:
            return None
        try:
            key = normalize_storage_key(storage_key_or_url)
            return bucket.sign_url("GET", key, expires_seconds)
        except Exception:
            logger.exception("生成预签名URL失败: %s", storage_key_or_url[:80])
            return None
    s = _storage()
    if s.public_bucket is None and s.bucket is None:
        return None
    try:
        return s.get_download_url(storage_key_or_url, expires_seconds=expires_seconds)
    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 解析到本地路径（缓存优先，否则 OSS 下载）。

    在 wrapper 层实现缓存逻辑，方便老单测 patch(oss_helpers.download_asset)。
    """
    if not asset_id or not isinstance(asset_id, str):
        return None
    if "\x00" in asset_id:
        return None

    work_dir = Path(work_dir)
    os.makedirs(work_dir, exist_ok=True)

    if asset_id.startswith("/") or ".." in Path(asset_id).parts:
        logger.warning("非法 asset_id: %s", asset_id)
        return None

    cache_hash = hashlib.sha256(asset_id.encode()).hexdigest()[:16]
    local_path = work_dir / f"{cache_hash}.mp4"

    if local_path.exists() and local_path.stat().st_size > 0:
        return local_path

    try:
        ok = download_asset(asset_id, local_path)
        if ok and local_path.exists() and local_path.stat().st_size > 0:
            return local_path
    except Exception:
        logger.exception("下载 asset 失败: %s", asset_id[:80])
    return None


def resolve_asset_ids_to_paths(asset_ids: list[str], work_dir: Path) -> dict[str, Path]:
    """批量解析 asset_id → 本地路径。"""
    result: dict[str, Path] = {}
    for aid in asset_ids:
        p = resolve_asset_path(aid, work_dir)
        if p is not None:
            result[aid] = p
    return result


def delete_from_oss(storage_key_or_url: str) -> bool:
    """从 OSS 删除对象（best-effort，internal endpoint）。"""
    s = _storage()
    if s.bucket is None:
        return False
    try:
        key = normalize_storage_key(storage_key_or_url)
        s.delete_file(key)
        return True
    except Exception:
        logger.exception("删除OSS对象失败: %s", storage_key_or_url[:80])
        return False


def file_exists(storage_key_or_url: str) -> bool:
    """检查文件是否存在（internal endpoint）。"""
    s = _storage()
    if s.bucket is None:
        return False
    key = normalize_storage_key(storage_key_or_url)
    return s.file_exists(key)


def get_public_url(storage_key: str) -> str:
    """返回公网 URL（不带签名）。"""
    return _storage().get_url(storage_key)