fix_title_version.py 6.98 KB
#!/usr/bin/env python3
"""修复已导入的 crawler 歌曲表 title/version 字段。

原导入逻辑从 hk_songs.name(歌曲名)解析 title/version,
现需改为从 hk_music_record.record_name(录音名)解析。

用法:
    # dry-run 预览差异(不写库)
    .venv/bin/python fix_title_version.py

    # 实际执行更新
    .venv/bin/python fix_title_version.py --apply
"""
import argparse
import logging

from etl_to_crawler.config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE
from etl_to_crawler.connections import get_source_conn, get_test_crawler_conn, close_all_pools
from etl_to_crawler.utils import split_title_version

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
log = logging.getLogger(__name__)

# ── 平台配置 ──────────────────────────────────────────────────────────────
# 每个平台的: PG表名, MySQL平台码, PG platform_song_id 对应 hk_music_record 的哪个字段
PLATFORM_CONFIGS = [
    {
        'name': 'QQ',
        'pg_table': 'crawler_qqmusic_songs',
        'mysql_platform': PLATFORM_QQ,
        'match_field': 'platform_mid',  # PG.platform_song_id ↔ mr.platform_mid
    },
    {
        'name': 'Kugou',
        'pg_table': 'crawler_kugou_songs',
        'mysql_platform': PLATFORM_KUGOU,
        'match_field': 'platform_unique_key',  # PG.platform_song_id ↔ mr.platform_unique_key
    },
    {
        'name': 'Netease',
        'pg_table': 'crawler_netease_songs',
        'mysql_platform': PLATFORM_NETEASE,
        'match_field': 'platform_unique_key',  # PG.platform_song_id ↔ mr.platform_unique_key
    },
]

BATCH_SIZE = 2000


def fetch_record_name_map(mysql_conn, platform: str, match_field: str) -> dict[str, str]:
    """从 hk_music_record 查询 platform_unique_key/platform_mid → record_name 映射。"""
    sql = f"""
        SELECT {match_field} AS match_key, record_name
        FROM hk_music_record
        WHERE platform = %s
          AND deleted = 0
          AND {match_field} IS NOT NULL
          AND TRIM({match_field}) != ''
          AND record_name IS NOT NULL
          AND TRIM(record_name) != ''
    """
    with mysql_conn.cursor() as cur:
        cur.execute(sql, (platform,))
        rows = cur.fetchall()
    mapping = {}
    for row in rows:
        key = str(row['match_key']).strip()
        if key:
            mapping[key] = row['record_name']
    log.info("  MySQL %s: 获取 %d 条 record_name 映射", platform, len(mapping))
    return mapping


def fetch_pg_songs(pg_conn, table: str) -> list[dict]:
    """从 PG 歌曲表查询 id, platform_song_id, title, version。"""
    sql = f"SELECT id, platform_song_id, title, version FROM {table}"
    with pg_conn.cursor() as cur:
        cur.execute(sql)
        rows = cur.fetchall()
    # pg8000 返回 tuple,转为 dict
    return [{'id': str(r[0]), 'platform_song_id': r[1], 'title': r[2], 'version': r[3]} for r in rows]


def compute_updates(songs: list[dict], name_map: dict[str, str]) -> list[tuple]:
    """对比已有 title/version 与 record_name 解析结果,返回需更新的 (new_title, new_version, id) 列表。"""
    updates = []
    skipped = 0
    for song in songs:
        psid = str(song['platform_song_id']).strip() if song['platform_song_id'] is not None else ''
        record_name = name_map.get(psid)
        if not record_name:
            skipped += 1
            continue
        new_title, new_version = split_title_version(record_name)
        old_title = song.get('title') or ''
        old_version = song.get('version') or ''
        if new_title != old_title or new_version != old_version:
            updates.append((new_title, new_version, str(song['id'])))
    if skipped:
        log.info("  %d 条记录在 hk_music_record 中未找到匹配,跳过", skipped)
    return updates


def batch_update(pg_conn, table: str, updates: list[tuple]) -> int:
    """批量 UPDATE title, version,使用参数化 VALUES 构造单条 UPDATE。"""
    if not updates:
        return 0
    total = 0
    for i in range(0, len(updates), BATCH_SIZE):
        chunk = updates[i:i + BATCH_SIZE]
        placeholders = ", ".join(["(%s, %s, %s::uuid)"] * len(chunk))
        sql = f"""
            UPDATE {table} AS t
            SET title = v.new_title,
                version = v.new_version,
                updated_at = NOW()
            FROM (VALUES {placeholders}) AS v(new_title, new_version, id)
            WHERE t.id = v.id
        """
        params = []
        for new_title, new_version, song_id in chunk:
            params.extend([new_title, new_version, song_id])
        with pg_conn.cursor() as cur:
            cur.execute(sql, params)
        pg_conn.commit()
        total += len(chunk)
        log.info("  已更新 %d / %d", total, len(updates))
    return total


def main():
    parser = argparse.ArgumentParser(description='修复 crawler 歌曲表 title/version 字段')
    parser.add_argument('--apply', action='store_true', help='实际执行更新(默认 dry-run)')
    args = parser.parse_args()

    mysql_conn = get_source_conn()
    pg_conn = get_test_crawler_conn()

    total_updated = 0
    total_diff = 0

    for cfg in PLATFORM_CONFIGS:
        log.info("=== 处理 %s (%s) ===", cfg['name'], cfg['pg_table'])

        # 1. 从 MySQL 获取 record_name 映射
        name_map = fetch_record_name_map(mysql_conn, cfg['mysql_platform'], cfg['match_field'])

        # 2. 从 PG 获取现有歌曲
        songs = fetch_pg_songs(pg_conn, cfg['pg_table'])
        log.info("  PG %s: %d 条歌曲记录", cfg['name'], len(songs))

        # 3. 计算差异
        updates = compute_updates(songs, name_map)
        log.info("  差异记录: %d 条", len(updates))
        total_diff += len(updates)

        # 4. dry-run 预览前 20 条
        if updates:
            log.info("  差异示例(前20条):")
            for new_title, new_version, song_id in updates[:20]:
                # 找到对应的旧值
                old = next((s for s in songs if str(s['id']) == song_id), None)
                if old:
                    log.info("    id=%s  title: '%s' → '%s'  version: '%s' → '%s'",
                             song_id,
                             old.get('title', ''), new_title,
                             old.get('version', ''), new_version)

        # 5. 执行更新
        if args.apply and updates:
            updated = batch_update(pg_conn, cfg['pg_table'], updates)
            total_updated += updated
            log.info("  %s 更新完成: %d 条", cfg['name'], updated)
        elif not args.apply:
            log.info("  [dry-run] 跳过更新,添加 --apply 执行")

    log.info("=== 汇总 ===")
    log.info("差异总数: %d", total_diff)
    if args.apply:
        log.info("已更新总数: %d", total_updated)
    else:
        log.info("[dry-run] 未执行更新,添加 --apply 参数以实际执行")

    close_all_pools()


if __name__ == '__main__':
    main()