feat(etl): 实现批次编排与 CLI 入口,支持 --max-batches 冒烟测试参数
Showing
2 changed files
with
385 additions
and
0 deletions
etl_to_crawler/runner.py
0 → 100644
| 1 | import uuid | ||
| 2 | import json | ||
| 3 | import logging | ||
| 4 | from tqdm import tqdm | ||
| 5 | |||
| 6 | from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, OSS_CONFIG | ||
| 7 | from .connections import get_hk_songs_conn, get_source_conn, get_spider_conn, get_pg_conn, get_oss_bucket | ||
| 8 | from .reader import iter_hk_songs_batches, fetch_platform_records | ||
| 9 | from .spider import ( | ||
| 10 | fetch_qq_songs, fetch_qq_singers, | ||
| 11 | fetch_kugou_songs, fetch_kugou_singers, | ||
| 12 | fetch_netease_songs, fetch_netease_singers, | ||
| 13 | ) | ||
| 14 | from .writer import ( | ||
| 15 | upsert_qq_singers, upsert_qq_albums, upsert_qq_songs, | ||
| 16 | upsert_qq_singer_songs, upsert_qq_singer_albums, | ||
| 17 | upsert_kugou_singers, upsert_kugou_albums, upsert_kugou_songs, | ||
| 18 | upsert_kugou_singer_songs, upsert_kugou_singer_albums, | ||
| 19 | upsert_netease_singers, upsert_netease_albums, upsert_netease_songs, | ||
| 20 | upsert_netease_singer_songs, upsert_netease_singer_albums, | ||
| 21 | ) | ||
| 22 | from .oss import transfer_url, build_oss_key | ||
| 23 | from .lyric import strip_timestamps | ||
| 24 | |||
| 25 | logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s') | ||
| 26 | log = logging.getLogger(__name__) | ||
| 27 | |||
| 28 | |||
| 29 | def _safe_transfer(url, oss_key, bucket, base_url): | ||
| 30 | try: | ||
| 31 | return transfer_url(url, oss_key, bucket, base_url) | ||
| 32 | except Exception as e: | ||
| 33 | log.warning("OSS transfer failed for %s: %s", url, e) | ||
| 34 | return url # 失败时保留原 URL,不阻断流程 | ||
| 35 | |||
| 36 | |||
| 37 | def _process_qq(hk_row: dict, pr: dict, spider_conn, pg_cur, bucket, base_url): | ||
| 38 | mid = pr['platform_unique_key'] | ||
| 39 | songs_map = fetch_qq_songs(spider_conn, [mid]) | ||
| 40 | if mid not in songs_map: | ||
| 41 | return | ||
| 42 | sp = songs_map[mid] | ||
| 43 | song_id_int = sp['id'] | ||
| 44 | |||
| 45 | singers_map = fetch_qq_singers(spider_conn, [song_id_int]) | ||
| 46 | singer_list = singers_map.get(song_id_int, []) | ||
| 47 | |||
| 48 | # OSS 转移 | ||
| 49 | audio_url = _safe_transfer( | ||
| 50 | hk_row['audio_url'], | ||
| 51 | build_oss_key('qq', 'audio', mid + '.mp3'), | ||
| 52 | bucket, base_url | ||
| 53 | ) | ||
| 54 | cover_url = _safe_transfer( | ||
| 55 | hk_row.get('cover_url') or sp.get('cover', ''), | ||
| 56 | build_oss_key('qq', 'cover', str(song_id_int) + '.jpg'), | ||
| 57 | bucket, base_url | ||
| 58 | ) | ||
| 59 | album_cover = '' | ||
| 60 | if sp.get('album_id') and sp.get('album_cover'): | ||
| 61 | album_cover = _safe_transfer( | ||
| 62 | sp['album_cover'], | ||
| 63 | build_oss_key('qq', 'album', str(sp['album_id']) + '.jpg'), | ||
| 64 | bucket, base_url | ||
| 65 | ) | ||
| 66 | |||
| 67 | # 歌手头像转移 + 写入 singers | ||
| 68 | singer_rows = [] | ||
| 69 | for sg in singer_list: | ||
| 70 | avatar = _safe_transfer( | ||
| 71 | sg.get('avatar', ''), | ||
| 72 | build_oss_key('qq', 'singer', sg['mid'] + '.jpg'), | ||
| 73 | bucket, base_url | ||
| 74 | ) | ||
| 75 | singer_rows.append({**sg, 'id': sg['singer_id'], 'avatar': avatar}) | ||
| 76 | upsert_qq_singers(pg_cur, singer_rows) | ||
| 77 | |||
| 78 | # singers JSONB | ||
| 79 | singers_json = json.dumps([{ | ||
| 80 | 'name': sg['name'], | ||
| 81 | 'singer_id': sg['singer_id'], | ||
| 82 | 'platform_singer_id': sg['mid'], | ||
| 83 | } for sg in singer_list], ensure_ascii=False) | ||
| 84 | |||
| 85 | # 写入 album | ||
| 86 | album_id = None | ||
| 87 | if sp.get('album_id'): | ||
| 88 | album_id = sp['album_id'] | ||
| 89 | upsert_qq_albums(pg_cur, [{ | ||
| 90 | 'id': sp['album_id'], 'mid': sp.get('album_mid') or '', | ||
| 91 | 'cover': album_cover, 'title': sp.get('album_title') or '', | ||
| 92 | 'intro': sp.get('album_intro'), 'type': sp.get('album_type') or '', | ||
| 93 | 'company_id': sp.get('company_id') or 0, 'company': sp.get('company') or '', | ||
| 94 | 'is_owner': sp.get('is_owner') or 0, 'published_at': sp.get('album_published_at'), | ||
| 95 | }]) | ||
| 96 | |||
| 97 | # 写入 song | ||
| 98 | platform_song_id = int(pr['platform_mid']) if pr.get('platform_mid') else song_id_int | ||
| 99 | song_uuid = str(uuid.uuid4()) | ||
| 100 | upsert_qq_songs(pg_cur, [{ | ||
| 101 | 'song_uuid': song_uuid, | ||
| 102 | 'platform_song_id': platform_song_id, | ||
| 103 | 'mid': mid, | ||
| 104 | 'album_id': album_id, | ||
| 105 | 'cover': cover_url, | ||
| 106 | 'title': sp.get('title', hk_row['name']), | ||
| 107 | 'name': hk_row['name'], | ||
| 108 | 'duration': sp.get('duration', 0) or 0, | ||
| 109 | 'lyric': strip_timestamps(sp.get('lyric')), | ||
| 110 | 'composer_name': sp.get('composer_name') or hk_row.get('composer'), | ||
| 111 | 'lyricist_name': sp.get('lyricist_name') or hk_row.get('lyricist'), | ||
| 112 | 'url': audio_url, | ||
| 113 | 'lyric_url': hk_row.get('lyrics_url'), | ||
| 114 | 'platform_index_url': sp.get('platform_index_url'), | ||
| 115 | 'published_at': sp.get('published_at') or hk_row.get('issue_time'), | ||
| 116 | 'singers_json': singers_json, | ||
| 117 | }]) | ||
| 118 | # 查询实际 UUID(ON CONFLICT DO NOTHING 时使用已有 UUID) | ||
| 119 | pg_cur.execute('SELECT id FROM crawler_qqmusic_songs WHERE platform_song_id = %s', (platform_song_id,)) | ||
| 120 | row = pg_cur.fetchone() | ||
| 121 | if row: | ||
| 122 | song_uuid = str(row[0]) | ||
| 123 | |||
| 124 | # singer_songs / singer_albums | ||
| 125 | upsert_qq_singer_songs(pg_cur, [(sg['singer_id'], song_uuid) for sg in singer_list]) | ||
| 126 | if album_id: | ||
| 127 | upsert_qq_singer_albums(pg_cur, [(sg['singer_id'], album_id) for sg in singer_list]) | ||
| 128 | |||
| 129 | |||
| 130 | def _process_kugou(hk_row: dict, pr: dict, spider_conn, pg_cur, bucket, base_url): | ||
| 131 | song_id = int(pr['platform_unique_key']) | ||
| 132 | songs_map = fetch_kugou_songs(spider_conn, [song_id]) | ||
| 133 | if song_id not in songs_map: | ||
| 134 | return | ||
| 135 | sp = songs_map[song_id] | ||
| 136 | |||
| 137 | singers_map = fetch_kugou_singers(spider_conn, [song_id]) | ||
| 138 | singer_list = singers_map.get(song_id, []) | ||
| 139 | |||
| 140 | audio_url = _safe_transfer( | ||
| 141 | hk_row['audio_url'], | ||
| 142 | build_oss_key('kugou', 'audio', str(song_id) + '.mp3'), | ||
| 143 | bucket, base_url | ||
| 144 | ) | ||
| 145 | cover_url = _safe_transfer( | ||
| 146 | hk_row.get('cover_url') or sp.get('cover', ''), | ||
| 147 | build_oss_key('kugou', 'cover', str(song_id) + '.jpg'), | ||
| 148 | bucket, base_url | ||
| 149 | ) | ||
| 150 | album_cover = '' | ||
| 151 | if sp.get('album_id') and sp.get('album_cover'): | ||
| 152 | album_cover = _safe_transfer( | ||
| 153 | sp['album_cover'], | ||
| 154 | build_oss_key('kugou', 'album', str(sp['album_id']) + '.jpg'), | ||
| 155 | bucket, base_url | ||
| 156 | ) | ||
| 157 | |||
| 158 | singer_rows = [] | ||
| 159 | for sg in singer_list: | ||
| 160 | avatar = _safe_transfer( | ||
| 161 | sg.get('avatar', ''), | ||
| 162 | build_oss_key('kugou', 'singer', str(sg['singer_id']) + '.jpg'), | ||
| 163 | bucket, base_url | ||
| 164 | ) | ||
| 165 | singer_rows.append({**sg, 'id': sg['singer_id'], 'avatar': avatar}) | ||
| 166 | upsert_kugou_singers(pg_cur, singer_rows) | ||
| 167 | |||
| 168 | singers_json = json.dumps([{ | ||
| 169 | 'name': sg['name'], | ||
| 170 | 'singer_id': sg['singer_id'], | ||
| 171 | 'platform_singer_id': str(sg['singer_id']), | ||
| 172 | } for sg in singer_list], ensure_ascii=False) | ||
| 173 | |||
| 174 | album_id = None | ||
| 175 | if sp.get('album_id'): | ||
| 176 | album_id = sp['album_id'] | ||
| 177 | upsert_kugou_albums(pg_cur, [{ | ||
| 178 | 'id': sp['album_id'], 'cover': album_cover, | ||
| 179 | 'title': sp.get('album_title') or '', 'intro': sp.get('album_intro'), | ||
| 180 | 'type': sp.get('album_type') or '', 'company_id': sp.get('company_id') or 0, | ||
| 181 | 'company': sp.get('company') or '', 'is_owner': sp.get('is_owner') or 0, | ||
| 182 | 'published_at': sp.get('album_published_at'), | ||
| 183 | }]) | ||
| 184 | |||
| 185 | song_uuid = str(uuid.uuid4()) | ||
| 186 | upsert_kugou_songs(pg_cur, [{ | ||
| 187 | 'song_uuid': song_uuid, | ||
| 188 | 'platform_song_id': song_id, | ||
| 189 | 'hash': sp.get('hid', pr.get('platform_mid', '')), | ||
| 190 | 'album_audio_id': sp.get('album_audio_id') or pr.get('album_audio_id') or 0, | ||
| 191 | 'album_id': album_id, | ||
| 192 | 'cover': cover_url, | ||
| 193 | 'title': sp.get('title', hk_row['name']), | ||
| 194 | 'name': hk_row['name'], | ||
| 195 | 'duration': sp.get('duration', 0) or 0, | ||
| 196 | 'lyric': strip_timestamps(sp.get('lyric')), | ||
| 197 | 'composer_name': sp.get('composer_name') or hk_row.get('composer'), | ||
| 198 | 'lyricist_name': sp.get('lyricist_name') or hk_row.get('lyricist'), | ||
| 199 | 'url': audio_url, | ||
| 200 | 'lyric_url': hk_row.get('lyrics_url'), | ||
| 201 | 'platform_index_url': sp.get('platform_index_url'), | ||
| 202 | 'published_at': sp.get('published_at') or hk_row.get('issue_time'), | ||
| 203 | 'singers_json': singers_json, | ||
| 204 | }]) | ||
| 205 | # 查询实际 UUID(ON CONFLICT DO NOTHING 时使用已有 UUID) | ||
| 206 | pg_cur.execute('SELECT id FROM crawler_kugou_songs WHERE platform_song_id = %s', (song_id,)) | ||
| 207 | row = pg_cur.fetchone() | ||
| 208 | if row: | ||
| 209 | song_uuid = str(row[0]) | ||
| 210 | |||
| 211 | upsert_kugou_singer_songs(pg_cur, [(sg['singer_id'], song_uuid) for sg in singer_list]) | ||
| 212 | if album_id: | ||
| 213 | upsert_kugou_singer_albums(pg_cur, [(sg['singer_id'], album_id) for sg in singer_list]) | ||
| 214 | |||
| 215 | |||
| 216 | def _process_netease(hk_row: dict, pr: dict, spider_conn, pg_cur, bucket, base_url): | ||
| 217 | song_id = int(pr['platform_unique_key']) | ||
| 218 | songs_map = fetch_netease_songs(spider_conn, [song_id]) | ||
| 219 | if song_id not in songs_map: | ||
| 220 | return | ||
| 221 | sp = songs_map[song_id] | ||
| 222 | |||
| 223 | singers_map = fetch_netease_singers(spider_conn, [song_id]) | ||
| 224 | singer_list = singers_map.get(song_id, []) | ||
| 225 | |||
| 226 | audio_url = _safe_transfer( | ||
| 227 | hk_row['audio_url'], | ||
| 228 | build_oss_key('netease', 'audio', str(song_id) + '.mp3'), | ||
| 229 | bucket, base_url | ||
| 230 | ) | ||
| 231 | cover_url = _safe_transfer( | ||
| 232 | hk_row.get('cover_url') or sp.get('cover', ''), | ||
| 233 | build_oss_key('netease', 'cover', str(song_id) + '.jpg'), | ||
| 234 | bucket, base_url | ||
| 235 | ) | ||
| 236 | album_cover = '' | ||
| 237 | if sp.get('album_id') and sp.get('album_cover'): | ||
| 238 | album_cover = _safe_transfer( | ||
| 239 | sp['album_cover'], | ||
| 240 | build_oss_key('netease', 'album', str(sp['album_id']) + '.jpg'), | ||
| 241 | bucket, base_url | ||
| 242 | ) | ||
| 243 | |||
| 244 | singer_rows = [] | ||
| 245 | for sg in singer_list: | ||
| 246 | avatar = _safe_transfer( | ||
| 247 | sg.get('avatar', ''), | ||
| 248 | build_oss_key('netease', 'singer', str(sg['singer_id']) + '.jpg'), | ||
| 249 | bucket, base_url | ||
| 250 | ) | ||
| 251 | singer_rows.append({**sg, 'id': sg['singer_id'], 'avatar': avatar}) | ||
| 252 | upsert_netease_singers(pg_cur, singer_rows) | ||
| 253 | |||
| 254 | singers_json = json.dumps([{ | ||
| 255 | 'name': sg['name'], | ||
| 256 | 'singer_id': sg['singer_id'], | ||
| 257 | 'platform_singer_id': str(sg['singer_id']), | ||
| 258 | } for sg in singer_list], ensure_ascii=False) | ||
| 259 | |||
| 260 | album_id = None | ||
| 261 | if sp.get('album_id'): | ||
| 262 | album_id = sp['album_id'] | ||
| 263 | upsert_netease_albums(pg_cur, [{ | ||
| 264 | 'id': sp['album_id'], 'cover': album_cover, | ||
| 265 | 'title': sp.get('album_title') or '', 'intro': sp.get('album_intro'), | ||
| 266 | 'type': sp.get('album_type') or '', 'company_id': sp.get('company_id') or 0, | ||
| 267 | 'company': sp.get('company') or '', 'is_owner': sp.get('is_owner') or 0, | ||
| 268 | 'published_at': sp.get('album_published_at'), | ||
| 269 | }]) | ||
| 270 | |||
| 271 | song_uuid = str(uuid.uuid4()) | ||
| 272 | upsert_netease_songs(pg_cur, [{ | ||
| 273 | 'song_uuid': song_uuid, | ||
| 274 | 'platform_song_id': song_id, | ||
| 275 | 'album_id': album_id, | ||
| 276 | 'cover': cover_url, | ||
| 277 | 'title': sp.get('title', hk_row['name']), | ||
| 278 | 'name': hk_row['name'], | ||
| 279 | 'duration': sp.get('duration', 0) or 0, | ||
| 280 | 'lyric': strip_timestamps(sp.get('lyric')), | ||
| 281 | 'composer_name': sp.get('composer_name') or hk_row.get('composer'), | ||
| 282 | 'lyricist_name': sp.get('lyricist_name') or hk_row.get('lyricist'), | ||
| 283 | 'url': audio_url, | ||
| 284 | 'lyric_url': hk_row.get('lyrics_url'), | ||
| 285 | 'platform_index_url': sp.get('platform_index_url'), | ||
| 286 | 'published_at': sp.get('published_at') or hk_row.get('issue_time'), | ||
| 287 | 'singers_json': singers_json, | ||
| 288 | }]) | ||
| 289 | # 查询实际 UUID(ON CONFLICT DO NOTHING 时使用已有 UUID) | ||
| 290 | pg_cur.execute('SELECT id FROM crawler_netease_songs WHERE platform_song_id = %s', (song_id,)) | ||
| 291 | row = pg_cur.fetchone() | ||
| 292 | if row: | ||
| 293 | song_uuid = str(row[0]) | ||
| 294 | |||
| 295 | upsert_netease_singer_songs(pg_cur, [(sg['singer_id'], song_uuid) for sg in singer_list]) | ||
| 296 | if album_id: | ||
| 297 | upsert_netease_singer_albums(pg_cur, [(sg['singer_id'], album_id) for sg in singer_list]) | ||
| 298 | |||
| 299 | |||
| 300 | _PROCESSORS = { | ||
| 301 | PLATFORM_QQ: _process_qq, | ||
| 302 | PLATFORM_KUGOU: _process_kugou, | ||
| 303 | PLATFORM_NETEASE: _process_netease, | ||
| 304 | } | ||
| 305 | |||
| 306 | |||
| 307 | def run(platforms: list[str], max_batches: int | None = None) -> None: | ||
| 308 | hk_conn = get_hk_songs_conn() | ||
| 309 | src_conn = get_source_conn() | ||
| 310 | spider_conn = get_spider_conn() | ||
| 311 | pg_conn = get_pg_conn() | ||
| 312 | bucket = get_oss_bucket() | ||
| 313 | base_url = OSS_CONFIG['base_url'] | ||
| 314 | |||
| 315 | total_ok = total_err = 0 | ||
| 316 | |||
| 317 | try: | ||
| 318 | for i, batch in enumerate(tqdm(iter_hk_songs_batches(hk_conn, BATCH_SIZE), desc='batches')): | ||
| 319 | if max_batches is not None and i >= max_batches: | ||
| 320 | break | ||
| 321 | song_ids = [int(r['source_song_id']) for r in batch if r.get('source_song_id')] | ||
| 322 | platform_records = fetch_platform_records(src_conn, song_ids) | ||
| 323 | |||
| 324 | # index platform records by source_song_id | ||
| 325 | pr_by_song: dict[int, list] = {} | ||
| 326 | for pr in platform_records: | ||
| 327 | if pr['platform'] in platforms: | ||
| 328 | pr_by_song.setdefault(int(pr['source_song_id']), []).append(pr) | ||
| 329 | |||
| 330 | with pg_conn.cursor() as pg_cur: | ||
| 331 | for hk_row in batch: | ||
| 332 | src_id = int(hk_row['source_song_id']) if hk_row.get('source_song_id') else None | ||
| 333 | if not src_id or src_id not in pr_by_song: | ||
| 334 | continue | ||
| 335 | for pr in pr_by_song[src_id]: | ||
| 336 | processor = _PROCESSORS.get(pr['platform']) | ||
| 337 | if not processor: | ||
| 338 | continue | ||
| 339 | pg_cur.execute('SAVEPOINT sp_song') | ||
| 340 | try: | ||
| 341 | processor(hk_row, pr, spider_conn, pg_cur, bucket, base_url) | ||
| 342 | pg_cur.execute('RELEASE SAVEPOINT sp_song') | ||
| 343 | total_ok += 1 | ||
| 344 | except Exception as e: | ||
| 345 | pg_cur.execute('ROLLBACK TO SAVEPOINT sp_song') | ||
| 346 | pg_cur.execute('RELEASE SAVEPOINT sp_song') | ||
| 347 | log.error("Error processing song %s platform %s: %s", | ||
| 348 | hk_row.get('name'), pr['platform'], e) | ||
| 349 | total_err += 1 | ||
| 350 | pg_conn.commit() | ||
| 351 | |||
| 352 | finally: | ||
| 353 | hk_conn.close() | ||
| 354 | src_conn.close() | ||
| 355 | spider_conn.close() | ||
| 356 | pg_conn.close() | ||
| 357 | |||
| 358 | log.info("Done. OK=%d ERR=%d", total_ok, total_err) |
run_etl.py
0 → 100644
| 1 | #!/usr/bin/env python3 | ||
| 2 | import argparse | ||
| 3 | from etl_to_crawler.config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, PLATFORMS | ||
| 4 | from etl_to_crawler.runner import run | ||
| 5 | |||
| 6 | PLATFORM_MAP = { | ||
| 7 | 'qq': PLATFORM_QQ, | ||
| 8 | 'kugou': PLATFORM_KUGOU, | ||
| 9 | 'netease': PLATFORM_NETEASE, | ||
| 10 | } | ||
| 11 | |||
| 12 | if __name__ == '__main__': | ||
| 13 | parser = argparse.ArgumentParser(description='ETL: hk_songs_test → crawler_dev') | ||
| 14 | parser.add_argument('--platform', default='all', | ||
| 15 | choices=['qq', 'kugou', 'netease', 'all'], | ||
| 16 | help='要导入的平台(默认 all)') | ||
| 17 | parser.add_argument('--max-batches', type=int, default=None, | ||
| 18 | help='最多处理多少批次(冒烟测试用)') | ||
| 19 | args = parser.parse_args() | ||
| 20 | |||
| 21 | if args.platform == 'all': | ||
| 22 | platforms = PLATFORMS | ||
| 23 | else: | ||
| 24 | platforms = [PLATFORM_MAP[args.platform]] | ||
| 25 | |||
| 26 | print(f"Starting ETL for platforms: {platforms}") | ||
| 27 | run(platforms, max_batches=args.max_batches) |
-
Please register or sign in to post a comment