migrate_test_to_prod.py 13.1 KB
#!/usr/bin/env python3
"""从 hk_songs_test 迁移数据到 hk_songs,同时将所有资源 URL 转存到 COS。

用法:
    python migrate_test_to_prod.py --limit 10
    python migrate_test_to_prod.py --limit 10 --dry-run
    python migrate_test_to_prod.py --limit 10 --skip-cos   # 跳过文件转存,URL 透传
    python migrate_test_to_prod.py --offset 50 --limit 50  # 分批导入
"""

from __future__ import annotations

import argparse
import logging
import os
import time
from pathlib import Path

import pymysql
import requests
from dotenv import load_dotenv
from qcloud_cos import CosConfig, CosS3Client

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

# ==================== 数据库配置 ====================

DB_CONFIG = {
    "host": os.getenv("TARGET_DB_HOST"),
    "port": int(os.getenv("TARGET_DB_PORT", 3306)),
    "user": os.getenv("TARGET_DB_USER"),
    "password": os.getenv("TARGET_DB_PASSWORD"),
    "database": os.getenv("TARGET_DB_NAME"),
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}

SOURCE_TABLE = "hk_songs_test"
TARGET_TABLE = "hk_songs"

SOURCE_DB_CONFIG = {
    "host": os.getenv("SOURCE_DB_HOST"),
    "port": int(os.getenv("SOURCE_DB_PORT", 3306)),
    "user": os.getenv("SOURCE_DB_USER"),
    "password": os.getenv("SOURCE_DB_PASSWORD"),
    "database": os.getenv("SOURCE_DB_NAME"),
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}

# ==================== COS 配置 ====================

COS_CONFIG = {
    "secret_id": os.getenv("COS_SECRET_ID"),
    "secret_key": os.getenv("COS_SECRET_KEY"),
    "bucket": os.getenv("COS_BUCKET"),
    "region": os.getenv("COS_REGION"),
    "endpoint": os.getenv("COS_ENDPOINT"),
    "cdn_domain": os.getenv("COS_CDN_DOMAIN"),
}

# 需要转存到 COS 的 URL 字段
URL_FIELDS = [
    ("audio_url", "audio"),
    ("accompany_url", "accompany"),
    ("lyrics_url", "lyric"),
    ("lrc_url", "lyric"),
    ("creation_url", "creation"),
    ("opern_url", "opern"),
    ("cover_url", "cover"),
]

# 字段默认扩展名
FIELD_DEFAULT_EXT = {
    "audio": ".mp3",
    "accompany": ".mp3",
    "lyric": ".txt",
    "creation": ".mp3",
    "opern": ".mid",
    "cover": ".jpg",
}

# ==================== COS 客户端(全局单例)====================

def _make_cos_client() -> CosS3Client:
    config = CosConfig(
        Region=COS_CONFIG["region"],
        SecretId=COS_CONFIG["secret_id"],
        SecretKey=COS_CONFIG["secret_key"],
    )
    return CosS3Client(config)


_cos_client: CosS3Client | None = None


def get_cos_client() -> CosS3Client:
    global _cos_client
    if _cos_client is None:
        _cos_client = _make_cos_client()
    return _cos_client


def upload_to_cos(content: bytes, cos_key: str, content_type: str, max_retries: int = 5) -> str | None:
    """上传内容到 COS,返回访问 URL;失败返回 None。"""
    import io

    ct = content_type or "application/octet-stream"
    for attempt in range(1, max_retries + 1):
        try:
            get_cos_client().put_object(
                Bucket=COS_CONFIG["bucket"],
                Body=io.BytesIO(content),
                Key=cos_key,
                ContentType=ct,
            )
            return f"/{cos_key}"
        except Exception as e:
            logger.warning(f"COS 上传失败 ({attempt}/{max_retries}): {cos_key}, {e}")
            if attempt < max_retries:
                time.sleep(2 ** attempt)
    return None


# ==================== 文件下载 ====================

def download_file(url: str, timeout: int = 30, max_retries: int = 3) -> tuple[bytes | None, str | None]:
    if not url or not url.startswith(("http://", "https://")):
        return None, None
    last_err = None
    for attempt in range(1, max_retries + 1):
        try:
            resp = requests.get(url, timeout=timeout, stream=True)
            if resp.status_code == 200:
                return resp.content, resp.headers.get("Content-Type", "").split(";")[0].strip()
            logger.debug(f"下载 HTTP {resp.status_code}: {url}")
        except Exception as e:
            last_err = e
            if attempt < max_retries:
                time.sleep(2 ** attempt)
    logger.warning(f"下载失败(已重试 {max_retries} 次): {url}, 错误: {last_err}")
    return None, None


def _guess_ext(url: str, content_type: str | None, field_type: str) -> str:
    if "." in url.split("/")[-1]:
        ext = "." + url.split("/")[-1].split(".")[-1].split("?")[0]
        if 1 < len(ext) <= 6:
            return ext
    type_map = {
        "audio/mpeg": ".mp3", "audio/wav": ".wav", "audio/x-wav": ".wav",
        "audio/ogg": ".ogg", "audio/flac": ".flac", "audio/aac": ".aac",
        "image/jpeg": ".jpg", "image/png": ".png", "image/webp": ".webp",
        "image/gif": ".gif", "text/plain": ".txt",
        "audio/midi": ".mid", "application/x-midi": ".mid",
    }
    for ct, ext in type_map.items():
        if ct in (content_type or ""):
            return ext
    return FIELD_DEFAULT_EXT.get(field_type, "")


# ==================== BPM 查询 ====================

_BPM_SQL = """
SELECT sp.id AS song_id, rs.bpm
FROM hk_song_platform sp
INNER JOIN (
    SELECT song_id,
           COALESCE(MIN(CASE WHEN is_main_version = 1 THEN record_id END), MIN(record_id)) AS record_id
    FROM hk_song_and_record GROUP BY song_id
) sr ON sr.song_id = sp.id
LEFT JOIN hk_music_record r ON r.id = sr.record_id AND r.deleted = 0x00
LEFT JOIN hk_music_record_state rs ON rs.record_id = r.id AND rs.deleted = 0x00
WHERE sp.id IN ({placeholders})
"""


def _bpm_to_class(bpm_str: str | None) -> int | None:
    if not bpm_str:
        return None
    try:
        bpm = int(float(bpm_str))
    except (ValueError, TypeError):
        return None
    if bpm <= 0:
        return None
    return 1 if bpm < 80 else (2 if bpm <= 120 else 3)


def fetch_bpm_class_map(src_conn, song_ids: list) -> dict[str, int | None]:
    if not song_ids:
        return {}
    placeholders = ",".join(["%s"] * len(song_ids))
    with src_conn.cursor() as cur:
        cur.execute(_BPM_SQL.format(placeholders=placeholders), song_ids)
        return {str(r["song_id"]): _bpm_to_class(r["bpm"]) for r in cur.fetchall()}

def _is_lyrics_text(value: str | None) -> bool:
    """判断 lyrics_url/lrc_url 字段存的是文本内容而非 URL 或路径。"""
    if not value:
        return False
    # http(s):// 开头是完整 URL,/ 开头是路径,都不是文本内容
    return not value.startswith(("http://", "https://", "/"))


def process_url(url: str | None, field_type: str, record_id: int, skip_cos: bool) -> str | None:
    if not url:
        return None
    if skip_cos:
        return url

    # 已经是路径格式(如 /music_library/...),直接透传
    if url.startswith("/"):
        return url

    # lyrics_url/lrc_url 可能存的是文本内容,直接上传文本
    if field_type == "lyric" and _is_lyrics_text(url):
        cos_key = f"music_library/lyric/{record_id}.txt"
        return upload_to_cos(url.encode("utf-8"), cos_key, "text/plain; charset=utf-8") or url

    content, content_type = download_file(url)
    if not content:
        return url  # 下载失败透传

    ext = _guess_ext(url, content_type, field_type)
    cos_key = f"music_library/{field_type}/{record_id}{ext}"
    return upload_to_cos(content, cos_key, content_type or "application/octet-stream") or url


# ==================== 主迁移逻辑 ====================

# hk_songs 目标字段(去掉 staging 专用字段)
TARGET_COLUMNS = [
    "id", "name", "lyricist", "composer", "issue_status", "intro",
    "audio_url", "accompany_url", "lyrics_url", "lrc_url",
    "song_time", "song_start", "song_end",
    "creation_url", "opern_url", "cover_version", "issue_time",
    "cover_url", "animation_type", "bpm_class", "review_status",
    "in_status", "song_status", "commit_time", "review_time",
    "shelf_time", "review_remark", "create_time", "creator",
    "modify_time", "modifier", "deleted", "cooperate_type", "singer",
    "off_shelf_remark", "musician_id", "commit_id", "sheet_music",
    "commit_desc", "price", "source_table_name", "source_song_id",
    "lyric_archive_element_id", "melody_archive_element_id", "audio_fingerprint",
]

INSERT_SQL = (
    f"INSERT INTO `{TARGET_TABLE}` ({', '.join(f'`{c}`' for c in TARGET_COLUMNS)}) "
    f"VALUES ({', '.join(['%s'] * len(TARGET_COLUMNS))})"
)


def migrate(limit: int, offset: int, dry_run: bool, skip_cos: bool) -> None:
    conn = pymysql.connect(**DB_CONFIG)
    src_conn = pymysql.connect(**SOURCE_DB_CONFIG)
    try:
        with conn.cursor() as cur:
            # 查询源数据(只取未软删除的)
            cur.execute(
                f"SELECT * FROM `{SOURCE_TABLE}` WHERE deleted = '0' LIMIT %s OFFSET %s",
                (limit, offset),
            )
            rows = cur.fetchall()

        logger.info(f"读取 {SOURCE_TABLE} {len(rows)} 条(offset={offset})")

        # 批量预查 bpm_class
        song_ids = [r["source_song_id"] for r in rows if r.get("source_song_id")]
        bpm_class_map = fetch_bpm_class_map(src_conn, song_ids)

        inserted_list: list[dict] = []
        skipped_list: list[dict] = []

        for row in rows:
            record_id = row["id"]

            # 检查目标表是否已存在相同 id
            with conn.cursor() as cur:
                cur.execute(f"SELECT id FROM `{TARGET_TABLE}` WHERE id = %s", (record_id,))
                if cur.fetchone():
                    skipped_list.append({"id": record_id, "name": row.get("name"), "reason": "已存在"})
                    continue

            # 转存 URL 字段
            converted = dict(row)

            # 补全 bpm_class(优先用已有值,否则从源库补)
            if not converted.get("bpm_class") and row.get("source_song_id"):
                converted["bpm_class"] = bpm_class_map.get(str(row["source_song_id"]))
            for field, field_type in URL_FIELDS:
                original = row.get(field)
                if original:
                    new_url = process_url(original, field_type, record_id, skip_cos)
                    converted[field] = new_url
                    if new_url != original:
                        logger.info(f"  [{field}] id={record_id}: {original[:60]}... -> COS")

            # 补充 source_table_name(标记来源)
            if not converted.get("source_table_name"):
                converted["source_table_name"] = SOURCE_TABLE

            values = tuple(converted.get(col) for col in TARGET_COLUMNS)

            if dry_run:
                inserted_list.append({"id": record_id, "name": row.get("name"), "lyricist": row.get("lyricist"), "composer": row.get("composer")})
                continue

            with conn.cursor() as cur:
                try:
                    cur.execute(INSERT_SQL, values)
                    conn.commit()
                    inserted_list.append({"id": record_id, "name": row.get("name"), "lyricist": row.get("lyricist"), "composer": row.get("composer")})
                except pymysql.err.IntegrityError as e:
                    conn.rollback()
                    skipped_list.append({"id": record_id, "name": row.get("name"), "reason": str(e)})

        # ── 结尾汇总 ──
        prefix = "[DRY-RUN] " if dry_run else ""
        print(f"\n{'=' * 60}")
        print(f"{prefix}导入结果:成功 {len(inserted_list)} 条,跳过 {len(skipped_list)} 条,共读取 {len(rows)} 条")
        print("=" * 60)

        if inserted_list:
            action = "将导入" if dry_run else "已导入"
            print(f"\n{prefix}{action} ({len(inserted_list)} 条):")
            print(f"  {'ID':<12} {'歌名':<30} {'作词':<15} {'作曲'}")
            print(f"  {'-'*12} {'-'*30} {'-'*15} {'-'*15}")
            for r in inserted_list:
                print(f"  {str(r['id']):<12} {str(r['name'] or ''):<30} {str(r['lyricist'] or ''):<15} {str(r['composer'] or '')}")

        if skipped_list:
            print(f"\n跳过 ({len(skipped_list)} 条):")
            print(f"  {'ID':<12} {'歌名':<30} 原因")
            print(f"  {'-'*12} {'-'*30} {'-'*20}")
            for r in skipped_list:
                print(f"  {str(r['id']):<12} {str(r['name'] or ''):<30} {r['reason']}")
        print()
    finally:
        conn.close()
        src_conn.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="从 hk_songs_test 迁移数据到 hk_songs,资源文件转存 COS")
    parser.add_argument("--limit", type=int, default=10, help="迁移条数(默认 10)")
    parser.add_argument("--offset", type=int, default=0, help="从第几条开始(默认 0)")
    parser.add_argument("--dry-run", action="store_true", help="仅打印计划,不写库也不上传")
    parser.add_argument("--skip-cos", action="store_true", help="跳过 COS 转存,URL 直接透传")
    args = parser.parse_args()

    logger.info(f"开始迁移: {SOURCE_TABLE} -> {TARGET_TABLE}, limit={args.limit}, offset={args.offset}, "
                f"dry_run={args.dry_run}, skip_cos={args.skip_cos}")
    migrate(args.limit, args.offset, args.dry_run, args.skip_cos)


if __name__ == "__main__":
    main()