check_record_titles.py 8.51 KB
#!/usr/bin/env python3
"""核对 spider 平台原始标题与正式 crawler 歌曲标题,并输出 Excel。

数据链路:
    archive_crawler.yinyan_song_records.(platform, platform_song_id)
      -> hikoon-data-spider.media_*_songs.id -> spider_title
      -> archive_crawler.crawler_*_songs.platform_song_id -> crawler_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_archive_crawler_conn, get_spider_conn


log = logging.getLogger(__name__)

DEFAULT_OUTPUT = "record_title_check.xlsx"
DEFAULT_BATCH_SIZE = 2000
ALLOWED_IMPORT_TABLES = {"yinyan_song_records", "yinyan_song_records2"}
SPIDER_TABLES = {
    "1": "media_tencent_songs",
    "2": "media_ku_gou_songs",
    "4": "media_netease_songs",
}

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],
            "crawler_title": row[4],
        }
        for row in rows
    ]


def fetch_spider_titles(mysql_conn, crawler_rows: list[dict], batch_size: int) -> dict[tuple[str, int], str]:
    """按 platform_song_id 分平台批量读取 spider 原始 title。"""
    titles: dict[tuple[str, int], str] = {}
    for platform, table in SPIDER_TABLES.items():
        ids = sorted({
            int(row["platform_song_id"])
            for row in crawler_rows
            if str(row["platform"]) == platform and row["platform_song_id"] is not None
        })
        for batch in chunked(ids, batch_size):
            placeholders = ", ".join(["%s"] * len(batch))
            with mysql_conn.cursor() as cursor:
                cursor.execute(
                    f"SELECT id, title FROM {table} WHERE id IN ({placeholders})",
                    tuple(batch),
                )
                for row in cursor.fetchall():
                    if row["title"] is not None:
                        titles[(platform, int(row["id"]))] = str(row["title"])
    return titles


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


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

    workbook = Workbook()
    worksheet = workbook.active
    worksheet.title = "Spider与Crawler标题核对"
    worksheet.freeze_panes = "A2"

    headers = (
        "record_id",
        "spider_title",
        "crawler_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["spider_title"],
            row["crawler_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="核对 spider title 与正式 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="待核对的 ARCHIVE_CRAWLER_DB 关联表",
    )
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE, help="spider 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_archive_crawler_conn()
        mysql_conn = get_spider_conn()

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

        spider_titles = fetch_spider_titles(mysql_conn, crawler_rows, args.batch_size)
        log.info("从 hikoon-data-spider 读取 %d 个平台标题", len(spider_titles))

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

        matched = sum(row["matched"] for row in rows)
        missing_spider = sum(row["spider_title"] is None for row in rows)
        missing_crawler = sum(row["crawler_title"] is None for row in rows)
        log.info(
            "核对完成:总数=%d,相同=%d,不同=%d,缺少 spider_title=%d,缺少 crawler_title=%d",
            len(rows), matched, len(rows) - matched, missing_spider, missing_crawler,
        )
        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()