sync_singer_hot_songs.py 10.2 KB
#!/usr/bin/env python3
"""将采集库 crawler_*_singers.hot_songs 同步到档案库 archive_subject_account.hot_songs。

支持测试 / 正式环境,采集和档案同时切换。

数据链路:
    测试: crawler_dev  -> data_dev
    正式: archive_crawler -> archive_data
    ON archive_subject_account.crawler_singer_id = crawler_*_singers.id
    UPDATE archive_subject_account.hot_songs

仅在 crawler_singer_id 能对应上时才更新,且只更新值有变化的行。
默认 dry-run,加 --apply 才会提交。

用法:
    .venv/bin/python sync_singer_hot_songs.py                        # 测试环境 dry-run
    .venv/bin/python sync_singer_hot_songs.py --apply                # 测试环境执行
    .venv/bin/python sync_singer_hot_songs.py --env prod             # 正式环境 dry-run
    .venv/bin/python sync_singer_hot_songs.py --env prod --apply     # 正式环境执行
"""

import argparse
import json
import logging
import sys

from etl_to_crawler.connections import (
    close_all_pools,
    get_test_crawler_conn,
    get_test_archive_conn,
    get_archive_crawler_conn,
    get_archive_conn,
)

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

CRAWLER_SINGER_TABLES = [
    "crawler_qqmusic_singers",
    "crawler_kugou_singers",
    "crawler_netease_singers",
]

BATCH_SIZE = 500


def fetch_crawler_singer_hot_songs(pg_conn) -> list[dict]:
    """从采集库读取所有有 hot_songs 的歌手记录。"""
    all_rows: list[dict] = []
    for table in CRAWLER_SINGER_TABLES:
        sql = f"""
            SELECT id, name, hot_songs
            FROM {table}
            WHERE hot_songs IS NOT NULL
        """
        with pg_conn.cursor() as cur:
            cur.execute(sql)
            rows = cur.fetchall()
        for row in rows:
            all_rows.append({
                "id": str(row[0]),  # 统一转 str,与档案库 varchar 类型对齐
                "name": row[1],
                "hot_songs": row[2],
                "source_table": table,
            })
        log.info("  %s: %d 条有 hot_songs 的歌手", table, len(rows))
    return all_rows


def fetch_archive_current_hot_songs(archive_conn, crawler_ids: set) -> dict:
    """查询档案库中匹配记录的当前 hot_songs 值。

    返回的 dict 的 key 即为已匹配的 crawler_singer_id,
    未出现在 dict 中的 ID 表示档案库中无对应记录。
    """
    if not crawler_ids:
        return {}
    current: dict = {}
    id_list = list(crawler_ids)
    for i in range(0, len(id_list), BATCH_SIZE):
        batch = id_list[i : i + BATCH_SIZE]
        sql = """
            SELECT crawler_singer_id, hot_songs
            FROM archive_subject_account
            WHERE crawler_singer_id = ANY(%s)
        """
        with archive_conn.cursor() as cur:
            cur.execute(sql, (batch,))
            for row in cur.fetchall():
                current[row[0]] = row[1]
    return current


def batch_update_archive(archive_conn, updates: list[tuple], apply: bool) -> int:
    """批量更新 archive_subject_account.hot_songs。"""
    if not updates:
        return 0
    total = 0
    for i in range(0, len(updates), BATCH_SIZE):
        batch = updates[i : i + BATCH_SIZE]
        values_sql_parts = []
        params: list = []
        for idx, (crawler_singer_id, hot_songs) in enumerate(batch):
            values_sql_parts.append(f"(%s, %s::jsonb)")
            params.extend([crawler_singer_id, json.dumps(hot_songs, ensure_ascii=False)])
        values_sql = ", ".join(values_sql_parts)
        sql = f"""
            UPDATE archive_subject_account AS a
            SET hot_songs = v.hot_songs
            FROM (VALUES {values_sql}) AS v(crawler_singer_id, hot_songs)
            WHERE a.crawler_singer_id = v.crawler_singer_id
        """
        with archive_conn.cursor() as cur:
            cur.execute(sql, params)
        total += len(batch)
        log.info("  已处理 %d / %d", total, len(updates))
    if apply:
        archive_conn.commit()
    return total


def _json_equal(a, b) -> bool:
    """比较两个 JSON 值是否逻辑相等。"""
    if a is None and b is None:
        return True
    if a is None or b is None:
        return False
    # 统一用 Python 对象比较(jsonb 反序列化后 key 序无关)
    if isinstance(a, str):
        a = json.loads(a)
    if isinstance(b, str):
        b = json.loads(b)
    return a == b


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="同步采集库 crawler_*_singers.hot_songs 到档案库 archive_subject_account"
    )
    parser.add_argument(
        "--env",
        choices=["test", "prod"],
        default="test",
        help="运行环境: test=测试(crawler_dev+data_dev), prod=正式(archive_crawler+archive_data)(默认 test)",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="实际执行更新(默认 dry-run 只预览不写库)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    mode = "APPLY" if args.apply else "DRY-RUN"
    env_label = "正式环境" if args.env == "prod" else "测试环境"

    # 根据环境选择连接函数
    if args.env == "prod":
        get_crawler_conn = get_archive_crawler_conn   # 正式采集库: archive_crawler
        get_archive_conn_fn = get_archive_conn          # 正式档案库: archive_data
        crawler_db = "archive_crawler"
        archive_db = "archive_data"
    else:
        get_crawler_conn = get_test_crawler_conn        # 测试采集库: crawler_dev
        get_archive_conn_fn = get_test_archive_conn      # 测试档案库: data_dev
        crawler_db = "crawler_dev"
        archive_db = "data_dev"

    log.info("=== sync_singer_hot_songs [%s] [%s] ===", env_label, mode)
    log.info("采集库: %s  |  档案库: %s", crawler_db, archive_db)

    # 1. 从采集库读取歌手 hot_songs
    log.info("步骤 1: 读取采集库 crawler_*_singers.hot_songs (%s)", crawler_db)
    crawler_conn = get_crawler_conn()
    crawler_rows = fetch_crawler_singer_hot_songs(crawler_conn)
    if not crawler_rows:
        log.info("采集库无任何 hot_songs 数据,退出")
        return
    log.info("共读取 %d 条歌手记录", len(crawler_rows))

    # 2+3. 一次查询档案库:匹配 + 获取当前值 + 比对变化
    crawler_ids = {row["id"] for row in crawler_rows}
    log.info("步骤 2: 查询档案库并比对变化 (%s,共 %d 个候选)", archive_db, len(crawler_ids))
    archive_conn = get_archive_conn_fn()
    current_hot_songs = fetch_archive_current_hot_songs(archive_conn, crawler_ids)
    matched_ids = set(current_hot_songs.keys())  # dict 的 key 即为匹配到的 ID
    log.info("匹配到 %d 个 crawler_singer_id", len(matched_ids))

    if not matched_ids:
        log.info("无可匹配的记录,退出")
        return

    updates: list[tuple] = []       # (crawler_singer_id, new_hot_songs)
    changes: list[dict] = []        # 变更明细
    updated_ids: set = set()        # 已产生变更的 ID(去重)
    for row in crawler_rows:
        sid = row["id"]
        if sid not in matched_ids or sid in updated_ids:
            continue
        new_val = row["hot_songs"]
        old_val = current_hot_songs.get(sid)
        if not _json_equal(old_val, new_val):
            updated_ids.add(sid)
            updates.append((sid, new_val))
            changes.append({
                "crawler_singer_id": sid,
                "name": row["name"],
                "source_table": row["source_table"],
                "old_hot_songs": old_val,
                "new_hot_songs": new_val,
            })

    skipped = len(matched_ids) - len(updated_ids)
    log.info("需更新: %d 条,无变化跳过: %d 条,无法匹配跳过: %d 条",
             len(updates), skipped, len(crawler_ids) - len(matched_ids))

    if not updates:
        print("\n" + "=" * 80)
        print("统计摘要")
        print("=" * 80)
        print(f"  环境           : {env_label}")
        print(f"  采集库         : {crawler_db}")
        print(f"  档案库         : {archive_db}")
        print(f"  采集库歌手总数 : {len(crawler_rows)}")
        print(f"  匹配档案数     : {len(matched_ids)}")
        print(f"  需更新         : 0")
        print(f"  无变化跳过     : {len(matched_ids)}")
        print(f"  无法匹配跳过   : {len(crawler_ids) - len(matched_ids)}")
        print("=" * 80)
        log.info("无变化需要更新,退出")
        return

    # 4. 执行更新
    if args.apply:
        log.info("步骤 4: 执行更新 (--apply)")
    else:
        log.info("步骤 4: DRY-RUN 预览(加 --apply 实际执行)")

    if args.apply:
        batch_update_archive(archive_conn, updates, apply=True)
        log.info("更新完成,已提交")
    else:
        log.info("DRY-RUN 模式,未写库")

    # 5. 输出变更明细
    print("\n" + "=" * 80)
    print(f"变更明细(共 {len(changes)} 条)")
    print("=" * 80)
    for c in changes:
        print(f"\n  crawler_singer_id : {c['crawler_singer_id']}")
        print(f"  name              : {c['name']}")
        print(f"  source_table      : {c['source_table']}")
        old_display = json.dumps(c["old_hot_songs"], ensure_ascii=False, indent=2) if c["old_hot_songs"] else "NULL"
        new_display = json.dumps(c["new_hot_songs"], ensure_ascii=False, indent=2) if c["new_hot_songs"] else "NULL"
        print(f"  old_hot_songs     : {old_display}")
        print(f"  new_hot_songs     : {new_display}")

    # 6. 统计摘要
    print("\n" + "=" * 80)
    print("统计摘要")
    print("=" * 80)
    print(f"  环境           : {env_label}")
    print(f"  采集库         : {crawler_db}")
    print(f"  档案库         : {archive_db}")
    print(f"  采集库歌手总数 : {len(crawler_rows)}")
    print(f"  匹配档案数     : {len(matched_ids)}")
    print(f"  需更新         : {len(updates)}")
    print(f"  无变化跳过     : {skipped}")
    print(f"  无法匹配跳过   : {len(crawler_ids) - len(matched_ids)}")
    print(f"  执行模式       : {'已提交 (--apply)' if args.apply else 'DRY-RUN(未写库)'}")
    print("=" * 80)


if __name__ == "__main__":
    try:
        main()
    finally:
        close_all_pools()