utils.py 10.4 KB
"""工具函数模块"""
import asyncio
import hashlib
import logging
import re
from datetime import datetime
from urllib.parse import urlparse
from zoneinfo import ZoneInfo

try:
    import httpx
except ModuleNotFoundError:
    httpx = None

try:
    from app.core.config import settings
    from app.core.oss_client import oss_client
except ModuleNotFoundError:
    settings = None
    oss_client = None

CHINA_TZ = ZoneInfo("Asia/Shanghai")
logger = logging.getLogger(__name__)

# OSS 上传默认超时(秒),优先使用全局配置
_OSS_UPLOAD_TIMEOUT = settings.OSS_UPLOAD_TIMEOUT_SECONDS if settings else 30
# 音频下载默认超时(秒),优先使用全局配置
_AUDIO_DOWNLOAD_TIMEOUT = settings.AUDIO_DOWNLOAD_TIMEOUT_SECONDS if settings else 30
# 音频下载最大限制(字节),优先使用全局配置
_MAX_AUDIO_SIZE = settings.MAX_AUDIO_DOWNLOAD_SIZE if settings else 50 * 1024 * 1024

# 歌词中常见的词曲作者匹配规则
_LYRICIST_PATTERNS = [
    r"作\s*词\s*[::]\s*([^\r\n//]+)",
    r"词\s*[::]\s*([^\r\n//]+)",
    r"Lyrics?\s*[::]\s*([^\r\n//]+)",
]
_COMPOSER_PATTERNS = [
    r"作\s*曲\s*[::]\s*([^\r\n//]+)",
    r"曲\s*[::]\s*([^\r\n//]+)",
    r"Composer\s*[::]\s*([^\r\n//]+)",
]


def _search_with_patterns(patterns: list, text: str) -> str:
    """根据多个正则匹配文本,返回首个命中分组"""
    if not text:
        return ""
    for pattern in patterns:
        match = re.search(pattern, text)
        if match:
            return match.group(1).strip()
    return ""


def extract_lyricist_composer(lyric: str) -> tuple[str, str]:
    """从歌词文本中提取词作者和曲作者

    Returns:
        (lyricist_name, composer_name)
    """
    lyricist = _search_with_patterns(_LYRICIST_PATTERNS, lyric)
    composer = _search_with_patterns(_COMPOSER_PATTERNS, lyric)
    return lyricist, composer


def now_cn() -> datetime:
    """返回中国时区的当前时间(无时区信息),保留微秒"""
    return datetime.now(CHINA_TZ).replace(tzinfo=None)


def extract_plain_lyric(lyric: str) -> str:
    """去除 LRC 歌词中的时间戳和标签,保留纯文本

    Args:
        lyric: 原始歌词内容(可能含 [mm:ss.xx] 时间戳和 [ti:xxx] 等标签)

    Returns:
        纯文本歌词,行之间用换行符连接
    """
    if not lyric:
        return ""
    lines = []
    for line in lyric.splitlines():
        # 去除 LRC 时间戳 [mm:ss.xx] 或 [mm:ss.xxx]
        cleaned = re.sub(r"\[\d{2}:\d{2}(?:\.\d{2,3})?\]", "", line)
        # 去除标签 [xx:yy]
        cleaned = re.sub(r"\[[a-zA-Z]+:[^\]]+\]", "", cleaned)
        cleaned = cleaned.strip()
        if cleaned:
            lines.append(cleaned)
    return "\n".join(lines)


def split_title_version(title: str | None) -> tuple[str, str]:
    """拆分歌名末尾括号版本信息。

    例如:化风行万里 (DJ默涵版) -> (化风行万里, DJ默涵版)
    """
    if not title:
        return "", ""
    text = title.strip()
    match = re.match(r"^(?P<title>.+?)\s*[\((](?P<version>[^()()]+)[\))]\s*$", text)
    if not match:
        return text, ""
    clean_title = match.group("title").strip()
    version = match.group("version").strip()
    return clean_title or text, version


def upload_plain_lyric_to_bucket(platform: str, unique_id: str, lyric: str, bucket, base_url: str) -> str:
    """将歌词去时间戳后上传到当前 ETL 使用的 OSS bucket"""
    plain_lyric = extract_plain_lyric(lyric)
    if not plain_lyric:
        return ""
    oss_key = f"crawler/lyric/{platform}/{unique_id}.txt"
    bucket.put_object(oss_key, plain_lyric.encode("utf-8"), headers={'Content-Type': 'text/plain; charset=utf-8'})
    return f"{base_url.rstrip('/')}/{oss_key}"


async def upload_lyric_to_oss(platform: str, unique_id: str, lyric: str) -> str:
    """将歌词去除时间戳后上传为纯文本文件到 OSS

    Args:
        platform: 平台标识,如 qqmusic/kugou/kuwo/netease/migu
        unique_id: 平台唯一标识,如 song_mid/hash/rid 等
        lyric: 原始歌词内容

    Returns:
        OSS 文件访问 URL,上传失败或歌词为空返回空字符串
    """
    if not lyric:
        return ""
    if oss_client is None:
        logger.error("OSS client is not configured")
        return ""
    plain_lyric = extract_plain_lyric(lyric)
    if not plain_lyric:
        return ""
    oss_key = f"lyrics/{platform}/{unique_id}.txt"
    try:
        file_url = await asyncio.wait_for(
            asyncio.to_thread(
                oss_client.upload_bytes,
                plain_lyric.encode("utf-8"),
                oss_key,
                "public-read",
                "text/plain; charset=utf-8",
            ),
            timeout=_OSS_UPLOAD_TIMEOUT,
        )
        return file_url or ""
    except asyncio.TimeoutError:
        logger.exception(f"上传歌词到 OSS 超时 {_OSS_UPLOAD_TIMEOUT}s {platform}/{unique_id}")
        return ""
    except Exception as e:
        logger.exception(f"上传歌词到 OSS 失败 {platform}/{unique_id}: {e}")
        return ""


async def download_and_upload_audio(
    audio_url: str,
    oss_key: str,
    *,
    headers: dict | None = None,
    allow_external_oss_url: bool = False,
) -> tuple[str, str]:
    """流式下载音频并上传到 OSS

    内置大小限制、下载超时、上传超时保护,防止大文件或慢网络拖垮 worker。

    Args:
        audio_url: 音频下载地址
        oss_key: OSS 对象键,如 songs/qqmusic/xxxx.mp3
        headers: 可选的下载请求头
        allow_external_oss_url: 若 audio_url 已是当前 OSS 域名下的地址,是否直接透传

    Returns:
        (OSS 文件 URL, 音频 MD5),失败返回 ("", "")
    """
    if not audio_url:
        return "", ""

    if allow_external_oss_url and settings.OSS_FILE_BASE_NAME and audio_url.startswith(settings.OSS_FILE_BASE_NAME):
        return audio_url, ""

    audio_bytes = bytearray()
    try:
        async with httpx.AsyncClient(timeout=_AUDIO_DOWNLOAD_TIMEOUT) as client:
            async with client.stream("GET", audio_url, headers=headers or {}) as response:
                response.raise_for_status()
                content_length = response.headers.get("Content-Length")
                if content_length and int(content_length) > _MAX_AUDIO_SIZE:
                    logger.warning(
                        f"音频文件过大,跳过: {oss_key}, "
                        f"size={int(content_length) / 1024 / 1024:.2f}MB"
                    )
                    return "", ""
                async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
                    audio_bytes.extend(chunk)
                    if len(audio_bytes) > _MAX_AUDIO_SIZE:
                        logger.warning(
                            f"音频下载超过大小限制,跳过: {oss_key}, "
                            f"size>{_MAX_AUDIO_SIZE / 1024 / 1024:.2f}MB"
                        )
                        return "", ""
    except Exception as e:
        logger.error(f"下载音频失败 {oss_key}: {e}")
        return "", ""

    audio_bytes = bytes(audio_bytes)
    audio_md5 = compute_audio_md5(audio_bytes)
    try:
        file_url = await asyncio.wait_for(
            asyncio.to_thread(oss_client.upload_bytes, audio_bytes, oss_key),
            timeout=_OSS_UPLOAD_TIMEOUT,
        )
    except asyncio.TimeoutError:
        logger.error(f"上传音频到 OSS 超时 {_OSS_UPLOAD_TIMEOUT}s: {oss_key}")
        return "", ""
    except Exception as e:
        logger.error(f"上传音频到 OSS 失败 {oss_key}: {e}")
        return "", ""
    return file_url or "", audio_md5


_CONTENT_TYPE_EXT_MAP = {
    "image/jpeg": "jpg",
    "image/jpg": "jpg",
    "image/png": "png",
    "image/webp": "webp",
    "image/gif": "gif",
}


def _guess_image_ext(content_type: str, url: str) -> str:
    """根据 Content-Type 或 URL 路径推断图片扩展名,默认 jpg"""
    if content_type:
        mime = content_type.split(";")[0].strip().lower()
        ext = _CONTENT_TYPE_EXT_MAP.get(mime)
        if ext:
            return ext
    path = urlparse(url).path
    suffix = path.rsplit(".", 1)[-1].lower()
    if suffix in {"jpg", "jpeg", "png", "webp", "gif"}:
        return "jpg" if suffix == "jpeg" else suffix
    return "jpg"


async def upload_cover_to_oss(platform: str, unique_id: str, image_url: str) -> str:
    """下载封面图片并上传到 OSS

    Args:
        platform: 平台标识,如 qqmusic/kugou/kuwo/netease/migu
        unique_id: 平台唯一标识,如 song_mid/hash/rid 等
        image_url: 封面图片原始 URL

    Returns:
        OSS 文件访问 URL,上传失败或 URL 为空返回空字符串
    """
    if not image_url:
        return ""
    try:
        async with httpx.AsyncClient(timeout=10, follow_redirects=True) as client:
            resp = await client.get(image_url, headers={
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
                "Referer": "https://www.kugou.com/",
            })
            resp.raise_for_status()
            image_bytes = resp.content
            content_type = resp.headers.get("content-type", "")
    except Exception as e:
        logger.exception(f"下载封面图片失败 {platform}/{unique_id} {image_url}: {e}")
        return ""

    ext = _guess_image_ext(content_type, image_url)
    oss_key = f"covers/{platform}/{unique_id}.{ext}"
    mime = _CONTENT_TYPE_EXT_MAP.get(content_type.split(";")[0].strip().lower(), f"image/{ext}")
    try:
        file_url = await asyncio.wait_for(
            asyncio.to_thread(
                oss_client.upload_bytes,
                image_bytes,
                oss_key,
                "public-read",
                mime,
            ),
            timeout=_OSS_UPLOAD_TIMEOUT,
        )
        return file_url or ""
    except asyncio.TimeoutError:
        logger.exception(f"上传封面到 OSS 超时 {_OSS_UPLOAD_TIMEOUT}s {platform}/{unique_id}")
        return ""
    except Exception as e:
        logger.exception(f"上传封面到 OSS 失败 {platform}/{unique_id}: {e}")
        return ""


def compute_audio_md5(audio_bytes: bytes) -> str:
    """计算音频字节流的 MD5 值(32 位小写十六进制字符串)

    Args:
        audio_bytes: 音频文件字节流

    Returns:
        MD5 字符串,输入为空时返回空字符串
    """
    if not audio_bytes:
        return ""
    return hashlib.md5(audio_bytes).hexdigest()