check_record_titles.py 7.91 KB
#!/usr/bin/env python3
"""核对音眼录音名与 crawler 平台歌曲标题,并输出 Excel。

数据链路:
    CRAWLER_DB.yinyan_song_records.record_id
      -> SOURCE_DB.hk_music_record.id -> record_name
    CRAWLER_DB.yinyan_song_records.(platform, platform_song_id)
      -> 对应 crawler_*_songs.platform_song_id -> title

用法:
    .venv/bin/python check_record_titles.py
    .venv/bin/python check_record_titles.py --output output/record_title_check.xlsx
"""

import argparse
import logging
from pathlib import Path
from typing import Iterable

from openpyxl import Workbook
from openpyxl.styles import Alignment, Font, PatternFill
from openpyxl.utils import get_column_letter

from etl_to_crawler.connections import close_all_pools, get_pg_conn, get_source_conn


log = logging.getLogger(__name__)

DEFAULT_OUTPUT = "record_title_check.xlsx"
DEFAULT_BATCH_SIZE = 2000
ALLOWED_IMPORT_TABLES = {"yinyan_song_records", "yinyan_song_records2"}

HEADER_FILL = PatternFill(fill_type="solid", fgColor="4472C4")
MATCH_FILL = PatternFill(fill_type="solid", fgColor="C6EFCE")
HEADER_FONT = Font(color="FFFFFF", bold=True)


def chunked(values: list[int], size: int) -> Iterable[list[int]]:
    """将列表切成固定大小的批次。"""
    for start in range(0, len(values), size):
        yield values[start:start + size]


def fetch_crawler_rows(pg_conn, import_table: str) -> list[dict]:
    """读取关联记录,并按 platform 从对应平台歌曲表取得 title。"""
    if import_table not in ALLOWED_IMPORT_TABLES:
        raise ValueError(f"不支持的导入表: {import_table}")

    sql = f"""
        SELECT
            ysr.record_id,
            ysr.platform,
            ysr.platform_song_id,
            ysr.recording_id,
            CASE ysr.platform
                WHEN '1' THEN qq.title
                WHEN '2' THEN kg.title
                WHEN '4' THEN ne.title
                ELSE NULL
            END AS title
        FROM {import_table} AS ysr
        LEFT JOIN crawler_qqmusic_songs AS qq
          ON ysr.platform = '1'
         AND ysr.platform_song_id = qq.platform_song_id
        LEFT JOIN crawler_kugou_songs AS kg
          ON ysr.platform = '2'
         AND ysr.platform_song_id = kg.platform_song_id
        LEFT JOIN crawler_netease_songs AS ne
          ON ysr.platform = '4'
         AND ysr.platform_song_id = ne.platform_song_id
        WHERE ysr.is_archive_push = TRUE
        ORDER BY ysr.record_id, ysr.platform, ysr.platform_song_id
    """
    with pg_conn.cursor() as cursor:
        cursor.execute(sql)
        rows = cursor.fetchall()

    return [
        {
            "record_id": row[0],
            "platform": row[1],
            "platform_song_id": row[2],
            "recording_id": row[3],
            "title": row[4],
        }
        for row in rows
    ]


def fetch_record_names(mysql_conn, record_ids: Iterable[int], batch_size: int) -> dict[int, str | None]:
    """按 record_id 分批读取 hk_music_record.record_name。"""
    unique_ids = sorted({int(record_id) for record_id in record_ids if record_id is not None})
    names: dict[int, str | None] = {}

    for batch in chunked(unique_ids, batch_size):
        placeholders = ", ".join(["%s"] * len(batch))
        sql = f"""
            SELECT id, record_name
            FROM hk_music_record
            WHERE id IN ({placeholders})
        """
        with mysql_conn.cursor() as cursor:
            cursor.execute(sql, tuple(batch))
            for row in cursor.fetchall():
                names[int(row["id"])] = row["record_name"]

    return names


def combine_rows(crawler_rows: list[dict], record_names: dict[int, str | None]) -> list[dict]:
    """合并两个数据库的查询结果,并计算是否完全相同。"""
    result = []
    for row in crawler_rows:
        record_id = row["record_id"]
        record_name = record_names.get(int(record_id)) if record_id is not None else None
        title = row["title"]
        result.append(
            {
                "record_id": record_id,
                "record_name": record_name,
                "platform": row["platform"],
                "platform_song_id": row["platform_song_id"],
                "recording_id": row["recording_id"],
                "title": title,
                # 两边都必须有值;不做 trim、大小写或繁简转换,按库中原值精确比较。
                "matched": record_name is not None and title is not None and record_name == title,
            }
        )
    return result


def write_excel(rows: list[dict], output_path: Path) -> None:
    """写出三列表格;record_name 与 title 相同时整行标绿。"""
    output_path.parent.mkdir(parents=True, exist_ok=True)

    workbook = Workbook()
    worksheet = workbook.active
    worksheet.title = "录音名与标题核对"
    worksheet.freeze_panes = "A2"

    headers = (
        "record_id",
        "record_name",
        "title",
        "platform",
        "platform_song_id",
        "recording_id",
    )
    for column, header in enumerate(headers, start=1):
        cell = worksheet.cell(row=1, column=column, value=header)
        cell.fill = HEADER_FILL
        cell.font = HEADER_FONT
        cell.alignment = Alignment(horizontal="center", vertical="center")

    for row_number, row in enumerate(rows, start=2):
        values = (
            row["record_id"],
            row["record_name"],
            row["title"],
            row["platform"],
            row["platform_song_id"],
            row["recording_id"],
        )
        for column, value in enumerate(values, start=1):
            cell = worksheet.cell(row=row_number, column=column, value=value)
            cell.alignment = Alignment(vertical="top", wrap_text=True)
            if row["matched"]:
                cell.fill = MATCH_FILL

    for column, width in enumerate((16, 48, 48, 12, 20, 40), start=1):
        worksheet.column_dimensions[get_column_letter(column)].width = width
    worksheet.auto_filter.ref = f"A1:F{max(1, worksheet.max_row)}"

    workbook.save(output_path)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="核对 hk_music_record.record_name 与 crawler title")
    parser.add_argument("--output", default=DEFAULT_OUTPUT, help=f"Excel 输出路径(默认: {DEFAULT_OUTPUT})")
    parser.add_argument(
        "--table",
        choices=sorted(ALLOWED_IMPORT_TABLES),
        default="yinyan_song_records",
        help="待核对的 CRAWLER_DB 关联表",
    )
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE, help="MySQL 单批查询数量")
    args = parser.parse_args()
    if args.batch_size <= 0:
        parser.error("--batch-size 必须大于 0")
    return args


def main() -> None:
    args = parse_args()
    output_path = Path(args.output).expanduser().resolve()

    try:
        pg_conn = get_pg_conn()
        mysql_conn = get_source_conn()

        crawler_rows = fetch_crawler_rows(pg_conn, args.table)
        log.info("从 %s 读取 %d 条关联记录", args.table, len(crawler_rows))

        record_names = fetch_record_names(
            mysql_conn,
            (row["record_id"] for row in crawler_rows),
            args.batch_size,
        )
        log.info("从 hk_music_record 读取 %d 个录音名", len(record_names))

        rows = combine_rows(crawler_rows, record_names)
        write_excel(rows, output_path)

        matched = sum(row["matched"] for row in rows)
        missing_name = sum(row["record_name"] is None for row in rows)
        missing_title = sum(row["title"] is None for row in rows)
        log.info(
            "核对完成:总数=%d,相同=%d,不同=%d,缺少 record_name=%d,缺少 title=%d",
            len(rows), matched, len(rows) - matched, missing_name, missing_title,
        )
        log.info("Excel 已输出至 %s", output_path)
    finally:
        close_all_pools()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    main()