runner.py 13.7 KB
import uuid
import json
import logging
from tqdm import tqdm

from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, OSS_CONFIG
from .connections import get_hk_songs_conn, get_source_conn, get_spider_conn, get_pg_conn, get_oss_bucket
from .reader import iter_hk_songs_batches, fetch_platform_records
from .spider import (
    fetch_qq_songs, fetch_qq_singers,
    fetch_kugou_songs, fetch_kugou_singers,
    fetch_netease_songs, fetch_netease_singers,
)
from .writer import (
    upsert_qq_singers, upsert_qq_albums, upsert_qq_songs,
    upsert_qq_singer_songs, upsert_qq_singer_albums,
    upsert_kugou_singers, upsert_kugou_albums, upsert_kugou_songs,
    upsert_kugou_singer_songs, upsert_kugou_singer_albums,
    upsert_netease_singers, upsert_netease_albums, upsert_netease_songs,
    upsert_netease_singer_songs, upsert_netease_singer_albums,
)
from .oss import transfer_url, build_oss_key
from .lyric import strip_timestamps

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


def _safe_transfer(url, oss_key, bucket, base_url):
    try:
        return transfer_url(url, oss_key, bucket, base_url)
    except Exception as e:
        log.warning("OSS transfer failed for %s: %s", url, e)
        return url  # 失败时保留原 URL,不阻断流程


def _process_qq(hk_row: dict, pr: dict, spider_conn, pg_cur, bucket, base_url):
    mid = pr['platform_unique_key']
    songs_map = fetch_qq_songs(spider_conn, [mid])
    if mid not in songs_map:
        return
    sp = songs_map[mid]
    song_id_int = sp['id']

    singers_map = fetch_qq_singers(spider_conn, [song_id_int])
    singer_list = singers_map.get(song_id_int, [])

    # OSS 转移
    audio_url = _safe_transfer(
        hk_row['audio_url'],
        build_oss_key('qq', 'audio', mid + '.mp3'),
        bucket, base_url
    )
    cover_url = _safe_transfer(
        hk_row.get('cover_url') or sp.get('cover', ''),
        build_oss_key('qq', 'cover', str(song_id_int) + '.jpg'),
        bucket, base_url
    )
    album_cover = ''
    if sp.get('album_id') and sp.get('album_cover'):
        album_cover = _safe_transfer(
            sp['album_cover'],
            build_oss_key('qq', 'album', str(sp['album_id']) + '.jpg'),
            bucket, base_url
        )

    # 歌手头像转移 + 写入 singers
    singer_rows = []
    for sg in singer_list:
        avatar = _safe_transfer(
            sg.get('avatar', ''),
            build_oss_key('qq', 'singer', sg['mid'] + '.jpg'),
            bucket, base_url
        )
        singer_rows.append({**sg, 'id': sg['singer_id'], 'avatar': avatar})
    upsert_qq_singers(pg_cur, singer_rows)

    # singers JSONB
    singers_json = json.dumps([{
        'name': sg['name'],
        'singer_id': sg['singer_id'],
        'platform_singer_id': sg['mid'],
    } for sg in singer_list], ensure_ascii=False)

    # 写入 album
    album_id = None
    if sp.get('album_id'):
        album_id = sp['album_id']
        upsert_qq_albums(pg_cur, [{
            'id': sp['album_id'], 'mid': sp.get('album_mid') or '',
            'cover': album_cover, 'title': sp.get('album_title') or '',
            'intro': sp.get('album_intro'), 'type': sp.get('album_type') or '',
            'company_id': sp.get('company_id') or 0, 'company': sp.get('company') or '',
            'is_owner': sp.get('is_owner') or 0, 'published_at': sp.get('album_published_at'),
        }])

    # 写入 song
    platform_song_id = int(pr['platform_mid']) if pr.get('platform_mid') else song_id_int
    song_uuid = str(uuid.uuid4())
    upsert_qq_songs(pg_cur, [{
        'song_uuid': song_uuid,
        'platform_song_id': platform_song_id,
        'mid': mid,
        'album_id': album_id,
        'cover': cover_url,
        'title': sp.get('title', hk_row['name']),
        'name': hk_row['name'],
        'duration': sp.get('duration', 0) or 0,
        'lyric': strip_timestamps(sp.get('lyric')),
        'composer_name': sp.get('composer_name') or hk_row.get('composer'),
        'lyricist_name': sp.get('lyricist_name') or hk_row.get('lyricist'),
        'url': audio_url,
        'lyric_url': hk_row.get('lyrics_url'),
        'platform_index_url': sp.get('platform_index_url'),
        'published_at': sp.get('published_at') or hk_row.get('issue_time'),
        'singers_json': singers_json,
    }])
    # 查询实际 UUID(ON CONFLICT DO NOTHING 时使用已有 UUID)
    pg_cur.execute('SELECT id FROM crawler_qqmusic_songs WHERE platform_song_id = %s', (platform_song_id,))
    row = pg_cur.fetchone()
    if row:
        song_uuid = str(row[0])

    # singer_songs / singer_albums
    upsert_qq_singer_songs(pg_cur, [(sg['singer_id'], song_uuid) for sg in singer_list])
    if album_id:
        upsert_qq_singer_albums(pg_cur, [(sg['singer_id'], album_id) for sg in singer_list])


def _process_kugou(hk_row: dict, pr: dict, spider_conn, pg_cur, bucket, base_url):
    song_id = int(pr['platform_unique_key'])
    songs_map = fetch_kugou_songs(spider_conn, [song_id])
    if song_id not in songs_map:
        return
    sp = songs_map[song_id]

    singers_map = fetch_kugou_singers(spider_conn, [song_id])
    singer_list = singers_map.get(song_id, [])

    audio_url = _safe_transfer(
        hk_row['audio_url'],
        build_oss_key('kugou', 'audio', str(song_id) + '.mp3'),
        bucket, base_url
    )
    cover_url = _safe_transfer(
        hk_row.get('cover_url') or sp.get('cover', ''),
        build_oss_key('kugou', 'cover', str(song_id) + '.jpg'),
        bucket, base_url
    )
    album_cover = ''
    if sp.get('album_id') and sp.get('album_cover'):
        album_cover = _safe_transfer(
            sp['album_cover'],
            build_oss_key('kugou', 'album', str(sp['album_id']) + '.jpg'),
            bucket, base_url
        )

    singer_rows = []
    for sg in singer_list:
        avatar = _safe_transfer(
            sg.get('avatar', ''),
            build_oss_key('kugou', 'singer', str(sg['singer_id']) + '.jpg'),
            bucket, base_url
        )
        singer_rows.append({**sg, 'id': sg['singer_id'], 'avatar': avatar})
    upsert_kugou_singers(pg_cur, singer_rows)

    singers_json = json.dumps([{
        'name': sg['name'],
        'singer_id': sg['singer_id'],
        'platform_singer_id': str(sg['singer_id']),
    } for sg in singer_list], ensure_ascii=False)

    album_id = None
    if sp.get('album_id'):
        album_id = sp['album_id']
        upsert_kugou_albums(pg_cur, [{
            'id': sp['album_id'], 'cover': album_cover,
            'title': sp.get('album_title') or '', 'intro': sp.get('album_intro'),
            'type': sp.get('album_type') or '', 'company_id': sp.get('company_id') or 0,
            'company': sp.get('company') or '', 'is_owner': sp.get('is_owner') or 0,
            'published_at': sp.get('album_published_at'),
        }])

    song_uuid = str(uuid.uuid4())
    upsert_kugou_songs(pg_cur, [{
        'song_uuid': song_uuid,
        'platform_song_id': song_id,
        'hash': sp.get('hid', pr.get('platform_mid', '')),
        'album_audio_id': sp.get('album_audio_id') or pr.get('album_audio_id') or 0,
        'album_id': album_id,
        'cover': cover_url,
        'title': sp.get('title', hk_row['name']),
        'name': hk_row['name'],
        'duration': sp.get('duration', 0) or 0,
        'lyric': strip_timestamps(sp.get('lyric')),
        'composer_name': sp.get('composer_name') or hk_row.get('composer'),
        'lyricist_name': sp.get('lyricist_name') or hk_row.get('lyricist'),
        'url': audio_url,
        'lyric_url': hk_row.get('lyrics_url'),
        'platform_index_url': sp.get('platform_index_url'),
        'published_at': sp.get('published_at') or hk_row.get('issue_time'),
        'singers_json': singers_json,
    }])
    # 查询实际 UUID(ON CONFLICT DO NOTHING 时使用已有 UUID)
    pg_cur.execute('SELECT id FROM crawler_kugou_songs WHERE platform_song_id = %s', (song_id,))
    row = pg_cur.fetchone()
    if row:
        song_uuid = str(row[0])

    upsert_kugou_singer_songs(pg_cur, [(sg['singer_id'], song_uuid) for sg in singer_list])
    if album_id:
        upsert_kugou_singer_albums(pg_cur, [(sg['singer_id'], album_id) for sg in singer_list])


def _process_netease(hk_row: dict, pr: dict, spider_conn, pg_cur, bucket, base_url):
    song_id = int(pr['platform_unique_key'])
    songs_map = fetch_netease_songs(spider_conn, [song_id])
    if song_id not in songs_map:
        return
    sp = songs_map[song_id]

    singers_map = fetch_netease_singers(spider_conn, [song_id])
    singer_list = singers_map.get(song_id, [])

    audio_url = _safe_transfer(
        hk_row['audio_url'],
        build_oss_key('netease', 'audio', str(song_id) + '.mp3'),
        bucket, base_url
    )
    cover_url = _safe_transfer(
        hk_row.get('cover_url') or sp.get('cover', ''),
        build_oss_key('netease', 'cover', str(song_id) + '.jpg'),
        bucket, base_url
    )
    album_cover = ''
    if sp.get('album_id') and sp.get('album_cover'):
        album_cover = _safe_transfer(
            sp['album_cover'],
            build_oss_key('netease', 'album', str(sp['album_id']) + '.jpg'),
            bucket, base_url
        )

    singer_rows = []
    for sg in singer_list:
        avatar = _safe_transfer(
            sg.get('avatar', ''),
            build_oss_key('netease', 'singer', str(sg['singer_id']) + '.jpg'),
            bucket, base_url
        )
        singer_rows.append({**sg, 'id': sg['singer_id'], 'avatar': avatar})
    upsert_netease_singers(pg_cur, singer_rows)

    singers_json = json.dumps([{
        'name': sg['name'],
        'singer_id': sg['singer_id'],
        'platform_singer_id': str(sg['singer_id']),
    } for sg in singer_list], ensure_ascii=False)

    album_id = None
    if sp.get('album_id'):
        album_id = sp['album_id']
        upsert_netease_albums(pg_cur, [{
            'id': sp['album_id'], 'cover': album_cover,
            'title': sp.get('album_title') or '', 'intro': sp.get('album_intro'),
            'type': sp.get('album_type') or '', 'company_id': sp.get('company_id') or 0,
            'company': sp.get('company') or '', 'is_owner': sp.get('is_owner') or 0,
            'published_at': sp.get('album_published_at'),
        }])

    song_uuid = str(uuid.uuid4())
    upsert_netease_songs(pg_cur, [{
        'song_uuid': song_uuid,
        'platform_song_id': song_id,
        'album_id': album_id,
        'cover': cover_url,
        'title': sp.get('title', hk_row['name']),
        'name': hk_row['name'],
        'duration': sp.get('duration', 0) or 0,
        'lyric': strip_timestamps(sp.get('lyric')),
        'composer_name': sp.get('composer_name') or hk_row.get('composer'),
        'lyricist_name': sp.get('lyricist_name') or hk_row.get('lyricist'),
        'url': audio_url,
        'lyric_url': hk_row.get('lyrics_url'),
        'platform_index_url': sp.get('platform_index_url'),
        'published_at': sp.get('published_at') or hk_row.get('issue_time'),
        'singers_json': singers_json,
    }])
    # 查询实际 UUID(ON CONFLICT DO NOTHING 时使用已有 UUID)
    pg_cur.execute('SELECT id FROM crawler_netease_songs WHERE platform_song_id = %s', (song_id,))
    row = pg_cur.fetchone()
    if row:
        song_uuid = str(row[0])

    upsert_netease_singer_songs(pg_cur, [(sg['singer_id'], song_uuid) for sg in singer_list])
    if album_id:
        upsert_netease_singer_albums(pg_cur, [(sg['singer_id'], album_id) for sg in singer_list])


_PROCESSORS = {
    PLATFORM_QQ: _process_qq,
    PLATFORM_KUGOU: _process_kugou,
    PLATFORM_NETEASE: _process_netease,
}


def run(platforms: list[str], max_batches: int | None = None) -> None:
    hk_conn = get_hk_songs_conn()
    src_conn = get_source_conn()
    spider_conn = get_spider_conn()
    pg_conn = get_pg_conn()
    bucket = get_oss_bucket()
    base_url = OSS_CONFIG['base_url']

    total_ok = total_err = 0

    try:
        for i, batch in enumerate(tqdm(iter_hk_songs_batches(hk_conn, BATCH_SIZE), desc='batches')):
            if max_batches is not None and i >= max_batches:
                break
            song_ids = [int(r['source_song_id']) for r in batch if r.get('source_song_id')]
            platform_records = fetch_platform_records(src_conn, song_ids)

            # index platform records by source_song_id
            pr_by_song: dict[int, list] = {}
            for pr in platform_records:
                if pr['platform'] in platforms:
                    pr_by_song.setdefault(int(pr['source_song_id']), []).append(pr)

            with pg_conn.cursor() as pg_cur:
                for hk_row in batch:
                    src_id = int(hk_row['source_song_id']) if hk_row.get('source_song_id') else None
                    if not src_id or src_id not in pr_by_song:
                        continue
                    for pr in pr_by_song[src_id]:
                        processor = _PROCESSORS.get(pr['platform'])
                        if not processor:
                            continue
                        pg_cur.execute('SAVEPOINT sp_song')
                        try:
                            processor(hk_row, pr, spider_conn, pg_cur, bucket, base_url)
                            pg_cur.execute('RELEASE SAVEPOINT sp_song')
                            total_ok += 1
                        except Exception as e:
                            pg_cur.execute('ROLLBACK TO SAVEPOINT sp_song')
                            pg_cur.execute('RELEASE SAVEPOINT sp_song')
                            log.error("Error processing song %s platform %s: %s",
                                      hk_row.get('name'), pr['platform'], e)
                            total_err += 1
                pg_conn.commit()

    finally:
        hk_conn.close()
        src_conn.close()
        spider_conn.close()
        pg_conn.close()

    log.info("Done. OK=%d ERR=%d", total_ok, total_err)