feat(etl): 新增 yinyan_song_records2 初始化功能及空字符清理支持
- 扩展 run_etl.py,添加 --init-yinyan-records2 参数支持初始化 yinyan_song_records2 - 提升 BATCH_SIZE 从 100 到 1000,提升批处理效率 - 实现 _NulSafePgConnection 及 _NulSafePgCursor,自动清理 PostgreSQL 参数中的 NUL 字符,防止插入错误 - 新增查询接口 iter_hk_songs_records2_batches 和 fetch_all_song_record_relations,用于提取满足 records2 条件的歌曲关联数据 - 在 runner.py 中实现 initialize_yinyan_song_records2,写入所有合格歌曲关系并统一去除 records1 已存在的重复关系 - 在 writer.py 中新增 insert_yinyan_song_records2 和 delete_yinyan_song_records2_existing_relations 函数,实现批量插入和基于 records1 的重复关系删除 - 新增单元测试覆盖 NUL 字符清理、records2 插入及去重逻辑,保障代码质量和功能稳定性
Showing
10 changed files
with
330 additions
and
4 deletions
| ... | @@ -48,7 +48,7 @@ PLATFORM_KUGOU = '2' | ... | @@ -48,7 +48,7 @@ PLATFORM_KUGOU = '2' |
| 48 | PLATFORM_NETEASE = '4' | 48 | PLATFORM_NETEASE = '4' |
| 49 | PLATFORMS = [PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE] | 49 | PLATFORMS = [PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE] |
| 50 | 50 | ||
| 51 | BATCH_SIZE = 100 | 51 | BATCH_SIZE = 1000 |
| 52 | BACKFILL_BATCH_SIZE = int(os.environ.get('BACKFILL_BATCH_SIZE', '5000')) | 52 | BACKFILL_BATCH_SIZE = int(os.environ.get('BACKFILL_BATCH_SIZE', '5000')) |
| 53 | HTTP_POOL_MAXSIZE = int(os.environ.get('HTTP_POOL_MAXSIZE', '128')) | 53 | HTTP_POOL_MAXSIZE = int(os.environ.get('HTTP_POOL_MAXSIZE', '128')) |
| 54 | HTTP_TRANSFER_RETRIES = int(os.environ.get('HTTP_TRANSFER_RETRIES', '3')) | 54 | HTTP_TRANSFER_RETRIES = int(os.environ.get('HTTP_TRANSFER_RETRIES', '3')) | ... | ... |
| ... | @@ -5,6 +5,58 @@ import oss2 | ... | @@ -5,6 +5,58 @@ import oss2 |
| 5 | from .config import SOURCE_DB, HK_SONGS_DB, CRAWLER_DB, OSS_CONFIG, OSS_CONNECTION_POOL_SIZE | 5 | from .config import SOURCE_DB, HK_SONGS_DB, CRAWLER_DB, OSS_CONFIG, OSS_CONNECTION_POOL_SIZE |
| 6 | 6 | ||
| 7 | 7 | ||
| 8 | def _strip_nul(value): | ||
| 9 | """PostgreSQL text/json 参数不接受 NUL 字符,递归清理全部绑定参数。""" | ||
| 10 | if isinstance(value, str): | ||
| 11 | return value.replace('\x00', '') | ||
| 12 | if isinstance(value, tuple): | ||
| 13 | return tuple(_strip_nul(item) for item in value) | ||
| 14 | if isinstance(value, list): | ||
| 15 | return [_strip_nul(item) for item in value] | ||
| 16 | if isinstance(value, dict): | ||
| 17 | return {key: _strip_nul(item) for key, item in value.items()} | ||
| 18 | return value | ||
| 19 | |||
| 20 | |||
| 21 | class _NulSafePgCursor: | ||
| 22 | """在所有 PostgreSQL 参数化查询前统一移除文本中的 NUL 字符。""" | ||
| 23 | |||
| 24 | def __init__(self, cursor): | ||
| 25 | self._cursor = cursor | ||
| 26 | |||
| 27 | def __enter__(self): | ||
| 28 | self._cursor.__enter__() | ||
| 29 | return self | ||
| 30 | |||
| 31 | def __exit__(self, exc_type, exc_value, traceback): | ||
| 32 | return self._cursor.__exit__(exc_type, exc_value, traceback) | ||
| 33 | |||
| 34 | def execute(self, operation, parameters=None): | ||
| 35 | if parameters is None: | ||
| 36 | return self._cursor.execute(operation) | ||
| 37 | return self._cursor.execute(operation, _strip_nul(parameters)) | ||
| 38 | |||
| 39 | def executemany(self, operation, parameter_sets): | ||
| 40 | return self._cursor.executemany( | ||
| 41 | operation, | ||
| 42 | [_strip_nul(parameters) for parameters in parameter_sets], | ||
| 43 | ) | ||
| 44 | |||
| 45 | def __getattr__(self, name): | ||
| 46 | return getattr(self._cursor, name) | ||
| 47 | |||
| 48 | |||
| 49 | class _NulSafePgConnection: | ||
| 50 | def __init__(self, connection): | ||
| 51 | self._connection = connection | ||
| 52 | |||
| 53 | def cursor(self, *args, **kwargs): | ||
| 54 | return _NulSafePgCursor(self._connection.cursor(*args, **kwargs)) | ||
| 55 | |||
| 56 | def __getattr__(self, name): | ||
| 57 | return getattr(self._connection, name) | ||
| 58 | |||
| 59 | |||
| 8 | def get_source_conn() -> pymysql.Connection: | 60 | def get_source_conn() -> pymysql.Connection: |
| 9 | return pymysql.connect(**SOURCE_DB, cursorclass=pymysql.cursors.DictCursor) | 61 | return pymysql.connect(**SOURCE_DB, cursorclass=pymysql.cursors.DictCursor) |
| 10 | 62 | ||
| ... | @@ -19,10 +71,10 @@ def get_hk_songs_conn() -> pymysql.Connection: | ... | @@ -19,10 +71,10 @@ def get_hk_songs_conn() -> pymysql.Connection: |
| 19 | return pymysql.connect(**HK_SONGS_DB, cursorclass=pymysql.cursors.DictCursor) | 71 | return pymysql.connect(**HK_SONGS_DB, cursorclass=pymysql.cursors.DictCursor) |
| 20 | 72 | ||
| 21 | 73 | ||
| 22 | def get_pg_conn() -> pg8000.Connection: | 74 | def get_pg_conn() -> _NulSafePgConnection: |
| 23 | conn = pg8000.connect(**CRAWLER_DB) | 75 | conn = pg8000.connect(**CRAWLER_DB) |
| 24 | conn.run("SET TimeZone = 'Asia/Shanghai'") | 76 | conn.run("SET TimeZone = 'Asia/Shanghai'") |
| 25 | return conn | 77 | return _NulSafePgConnection(conn) |
| 26 | 78 | ||
| 27 | 79 | ||
| 28 | def get_oss_bucket() -> oss2.Bucket: | 80 | def get_oss_bucket() -> oss2.Bucket: | ... | ... |
| ... | @@ -28,6 +28,24 @@ WHERE deleted = '0' | ... | @@ -28,6 +28,24 @@ WHERE deleted = '0' |
| 28 | ORDER BY id | 28 | ORDER BY id |
| 29 | """ | 29 | """ |
| 30 | 30 | ||
| 31 | # 用于补全 yinyan_song_records2 的歌曲必须具备所有下游所需字段。 | ||
| 32 | # 当前 hk_songs_test 中的 lyrics_url 即歌曲链接字段。 | ||
| 33 | _HK_SONGS_RECORDS2_QUERY = """ | ||
| 34 | SELECT id, source_song_id | ||
| 35 | FROM hk_songs_test | ||
| 36 | WHERE deleted = '0' | ||
| 37 | AND id > %s | ||
| 38 | AND source_song_id IS NOT NULL | ||
| 39 | AND name IS NOT NULL AND TRIM(name) != '' | ||
| 40 | AND song_time IS NOT NULL AND song_time > 0 | ||
| 41 | AND singer IS NOT NULL AND TRIM(singer) != '' | ||
| 42 | AND lyrics_url IS NOT NULL AND TRIM(lyrics_url) != '' | ||
| 43 | AND audio_url IS NOT NULL AND TRIM(audio_url) != '' | ||
| 44 | AND issue_time IS NOT NULL | ||
| 45 | ORDER BY id | ||
| 46 | LIMIT %s | ||
| 47 | """ | ||
| 48 | |||
| 31 | _PLATFORM_QUERY = """ | 49 | _PLATFORM_QUERY = """ |
| 32 | SELECT | 50 | SELECT |
| 33 | sar.song_id AS source_song_id, | 51 | sar.song_id AS source_song_id, |
| ... | @@ -96,6 +114,25 @@ def iter_hk_songs_batches( | ... | @@ -96,6 +114,25 @@ def iter_hk_songs_batches( |
| 96 | break | 114 | break |
| 97 | 115 | ||
| 98 | 116 | ||
| 117 | def iter_hk_songs_records2_batches( | ||
| 118 | conn: pymysql.Connection, | ||
| 119 | batch_size: int, | ||
| 120 | start_after_id: int = 0, | ||
| 121 | ) -> Iterator[list[dict]]: | ||
| 122 | """遍历具备 records2 入库必填字段的歌曲。""" | ||
| 123 | last_id = start_after_id | ||
| 124 | with conn.cursor() as cur: | ||
| 125 | while True: | ||
| 126 | cur.execute(_HK_SONGS_RECORDS2_QUERY, (last_id, batch_size)) | ||
| 127 | rows = cur.fetchall() | ||
| 128 | if not rows: | ||
| 129 | break | ||
| 130 | yield rows | ||
| 131 | last_id = int(rows[-1]['id']) | ||
| 132 | if len(rows) < batch_size: | ||
| 133 | break | ||
| 134 | |||
| 135 | |||
| 99 | def fetch_hk_songs_by_source_ids(conn: pymysql.Connection, source_song_ids: list[int]) -> dict[int, dict]: | 136 | def fetch_hk_songs_by_source_ids(conn: pymysql.Connection, source_song_ids: list[int]) -> dict[int, dict]: |
| 100 | if not source_song_ids: | 137 | if not source_song_ids: |
| 101 | return {} | 138 | return {} |
| ... | @@ -158,6 +195,27 @@ def fetch_all_platform_records(source_conn: pymysql.Connection, song_ids: list[i | ... | @@ -158,6 +195,27 @@ def fetch_all_platform_records(source_conn: pymysql.Connection, song_ids: list[i |
| 158 | return result | 195 | return result |
| 159 | 196 | ||
| 160 | 197 | ||
| 198 | def fetch_all_song_record_relations(source_conn: pymysql.Connection, song_ids: list[int]) -> list[dict]: | ||
| 199 | """返回指定 song_id 的全部有效 (song_id, record_id, platform) 关联。""" | ||
| 200 | if not song_ids: | ||
| 201 | return [] | ||
| 202 | placeholders = ','.join(['%s'] * len(song_ids)) | ||
| 203 | query = _ALL_PLATFORM_RECORDS_QUERY.format(placeholders=placeholders) | ||
| 204 | with source_conn.cursor() as cur: | ||
| 205 | cur.execute(query, song_ids) | ||
| 206 | rows = cur.fetchall() | ||
| 207 | |||
| 208 | # 同一关联可能因源表脏数据重复出现;保留一条即可。 | ||
| 209 | seen_relations: set[tuple[int, int]] = set() | ||
| 210 | result = [] | ||
| 211 | for row in rows: | ||
| 212 | key = (int(row['source_song_id']), int(row['record_id'])) | ||
| 213 | if key not in seen_relations: | ||
| 214 | seen_relations.add(key) | ||
| 215 | result.append(row) | ||
| 216 | return result | ||
| 217 | |||
| 218 | |||
| 161 | def fetch_platform_records_by_record_ids(source_conn: pymysql.Connection, record_ids: list[int]) -> list[dict]: | 219 | def fetch_platform_records_by_record_ids(source_conn: pymysql.Connection, record_ids: list[int]) -> list[dict]: |
| 162 | """按状态表指定的录音 id 批量查询录音,不再按 song_id 拉取全平台候选。""" | 220 | """按状态表指定的录音 id 批量查询录音,不再按 song_id 拉取全平台候选。""" |
| 163 | if not record_ids: | 221 | if not record_ids: | ... | ... |
| ... | @@ -8,10 +8,12 @@ from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, B | ... | @@ -8,10 +8,12 @@ from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, B |
| 8 | from .connections import get_hk_songs_conn, get_source_conn, get_spider_conn, get_pg_conn, get_oss_bucket | 8 | from .connections import get_hk_songs_conn, get_source_conn, get_spider_conn, get_pg_conn, get_oss_bucket |
| 9 | from .reader import ( | 9 | from .reader import ( |
| 10 | iter_hk_songs_batches, | 10 | iter_hk_songs_batches, |
| 11 | iter_hk_songs_records2_batches, | ||
| 11 | fetch_hk_songs_by_source_ids, | 12 | fetch_hk_songs_by_source_ids, |
| 12 | mark_hk_songs_deleted, | 13 | mark_hk_songs_deleted, |
| 13 | fetch_platform_records, | 14 | fetch_platform_records, |
| 14 | fetch_all_platform_records, | 15 | fetch_all_platform_records, |
| 16 | fetch_all_song_record_relations, | ||
| 15 | fetch_platform_records_by_record_ids, | 17 | fetch_platform_records_by_record_ids, |
| 16 | fetch_record_platforms, | 18 | fetch_record_platforms, |
| 17 | select_primary_record, | 19 | select_primary_record, |
| ... | @@ -26,6 +28,8 @@ from .writer import ( | ... | @@ -26,6 +28,8 @@ from .writer import ( |
| 26 | fetch_pending_yinyan_song_records, | 28 | fetch_pending_yinyan_song_records, |
| 27 | fetch_yinyan_records_missing_platform, | 29 | fetch_yinyan_records_missing_platform, |
| 28 | insert_yinyan_song_records, | 30 | insert_yinyan_song_records, |
| 31 | insert_yinyan_song_records2, | ||
| 32 | delete_yinyan_song_records2_existing_relations, | ||
| 29 | update_yinyan_record_platforms, | 33 | update_yinyan_record_platforms, |
| 30 | upsert_yinyan_song_records, | 34 | upsert_yinyan_song_records, |
| 31 | fetch_existing_yinyan_song_ids, | 35 | fetch_existing_yinyan_song_ids, |
| ... | @@ -1239,6 +1243,55 @@ def initialize_yinyan_song_records(platforms: list[str], max_batches: int | None | ... | @@ -1239,6 +1243,55 @@ def initialize_yinyan_song_records(platforms: list[str], max_batches: int | None |
| 1239 | log.info("Initialized yinyan_song_records candidates=%d, skipped_batches=%d", total, skipped_batches) | 1243 | log.info("Initialized yinyan_song_records candidates=%d, skipped_batches=%d", total, skipped_batches) |
| 1240 | 1244 | ||
| 1241 | 1245 | ||
| 1246 | def initialize_yinyan_song_records2(platforms: list[str], max_batches: int | None = None) -> None: | ||
| 1247 | """写入所有合格歌曲关联,再删除 records1 已存在的关联。""" | ||
| 1248 | hk_conn = get_hk_songs_conn() | ||
| 1249 | src_conn = get_source_conn() | ||
| 1250 | pg_conn = get_pg_conn() | ||
| 1251 | total_inserted = 0 | ||
| 1252 | |||
| 1253 | try: | ||
| 1254 | for index, batch in enumerate( | ||
| 1255 | tqdm(iter_hk_songs_records2_batches(hk_conn, BATCH_SIZE), desc='init-yinyan-records2') | ||
| 1256 | ): | ||
| 1257 | if max_batches is not None and index >= max_batches: | ||
| 1258 | break | ||
| 1259 | |||
| 1260 | song_ids = list({int(row['source_song_id']) for row in batch}) | ||
| 1261 | relations = fetch_all_song_record_relations(src_conn, song_ids) | ||
| 1262 | rows = [ | ||
| 1263 | { | ||
| 1264 | 'song_id': int(relation['source_song_id']), | ||
| 1265 | 'record_id': int(relation['record_id']), | ||
| 1266 | 'platform': str(relation['platform']), | ||
| 1267 | } | ||
| 1268 | for relation in relations | ||
| 1269 | if str(relation['platform']) in platforms | ||
| 1270 | ] | ||
| 1271 | if not rows: | ||
| 1272 | continue | ||
| 1273 | |||
| 1274 | with pg_conn.cursor() as pg_cur: | ||
| 1275 | insert_yinyan_song_records2(pg_cur, rows) | ||
| 1276 | pg_conn.commit() | ||
| 1277 | total_inserted += len(rows) | ||
| 1278 | |||
| 1279 | # 必须在全部候选关系写入后,统一根据 records1 做去重。 | ||
| 1280 | with pg_conn.cursor() as pg_cur: | ||
| 1281 | deleted = delete_yinyan_song_records2_existing_relations(pg_cur) | ||
| 1282 | pg_conn.commit() | ||
| 1283 | finally: | ||
| 1284 | hk_conn.close() | ||
| 1285 | src_conn.close() | ||
| 1286 | pg_conn.close() | ||
| 1287 | |||
| 1288 | log.info( | ||
| 1289 | "Initialized yinyan_song_records2 candidates=%d, removed_existing_records1_relations=%d", | ||
| 1290 | total_inserted, | ||
| 1291 | deleted, | ||
| 1292 | ) | ||
| 1293 | |||
| 1294 | |||
| 1242 | def backfill_yinyan_record_platforms(max_batches: int | None = None) -> None: | 1295 | def backfill_yinyan_record_platforms(max_batches: int | None = None) -> None: |
| 1243 | src_conn = get_source_conn() | 1296 | src_conn = get_source_conn() |
| 1244 | pg_conn = get_pg_conn() | 1297 | pg_conn = get_pg_conn() | ... | ... |
| ... | @@ -38,6 +38,38 @@ def insert_yinyan_song_records(cur, records: list[dict]) -> None: | ... | @@ -38,6 +38,38 @@ def insert_yinyan_song_records(cur, records: list[dict]) -> None: |
| 38 | ) | 38 | ) |
| 39 | 39 | ||
| 40 | 40 | ||
| 41 | def insert_yinyan_song_records2(cur, records: list[dict]) -> None: | ||
| 42 | """写入全部候选关联;与 records1 的去重由后续独立步骤完成。""" | ||
| 43 | if not records: | ||
| 44 | return | ||
| 45 | values_sql = ', '.join(['(%s::bigint, %s::bigint, %s::varchar)'] * len(records)) | ||
| 46 | params = [] | ||
| 47 | for record in records: | ||
| 48 | params.extend([record['song_id'], record['record_id'], record['platform']]) | ||
| 49 | cur.execute( | ||
| 50 | f""" | ||
| 51 | INSERT INTO yinyan_song_records2 (song_id, record_id, platform) | ||
| 52 | VALUES {values_sql} | ||
| 53 | ON CONFLICT (song_id, record_id) DO UPDATE | ||
| 54 | SET platform = EXCLUDED.platform | ||
| 55 | """, | ||
| 56 | tuple(params), | ||
| 57 | ) | ||
| 58 | |||
| 59 | |||
| 60 | def delete_yinyan_song_records2_existing_relations(cur) -> int: | ||
| 61 | """删除 records2 中已存在于 records1 的 song-record 关联。""" | ||
| 62 | cur.execute( | ||
| 63 | """ | ||
| 64 | DELETE FROM yinyan_song_records2 AS ysr2 | ||
| 65 | USING yinyan_song_records AS ysr1 | ||
| 66 | WHERE ysr2.song_id = ysr1.song_id | ||
| 67 | AND ysr2.record_id = ysr1.record_id | ||
| 68 | """ | ||
| 69 | ) | ||
| 70 | return cur.rowcount | ||
| 71 | |||
| 72 | |||
| 41 | def fetch_existing_yinyan_song_ids(cur) -> set[int]: | 73 | def fetch_existing_yinyan_song_ids(cur) -> set[int]: |
| 42 | """返回已完成初始化 platform 的 song_id 集合。""" | 74 | """返回已完成初始化 platform 的 song_id 集合。""" |
| 43 | cur.execute("SELECT DISTINCT song_id FROM yinyan_song_records WHERE platform IS NOT NULL") | 75 | cur.execute("SELECT DISTINCT song_id FROM yinyan_song_records WHERE platform IS NOT NULL") | ... | ... |
| 1 | #!/usr/bin/env python3 | 1 | #!/usr/bin/env python3 |
| 2 | import argparse | 2 | import argparse |
| 3 | from etl_to_crawler.config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, PLATFORMS | 3 | from etl_to_crawler.config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, PLATFORMS |
| 4 | from etl_to_crawler.runner import backfill_yinyan_record_platforms, initialize_yinyan_song_records, run, backfill_empty_singers, backfill_qq_invalid_covers | 4 | from etl_to_crawler.runner import backfill_yinyan_record_platforms, initialize_yinyan_song_records, initialize_yinyan_song_records2, run, backfill_empty_singers, backfill_qq_invalid_covers |
| 5 | 5 | ||
| 6 | PLATFORM_MAP = { | 6 | PLATFORM_MAP = { |
| 7 | 'qq': PLATFORM_QQ, | 7 | 'qq': PLATFORM_QQ, |
| ... | @@ -18,6 +18,8 @@ if __name__ == '__main__': | ... | @@ -18,6 +18,8 @@ if __name__ == '__main__': |
| 18 | help='最多处理多少批次(冒烟测试用)') | 18 | help='最多处理多少批次(冒烟测试用)') |
| 19 | parser.add_argument('--init-yinyan-records', action='store_true', | 19 | parser.add_argument('--init-yinyan-records', action='store_true', |
| 20 | help='只初始化 yinyan_song_records 待导入状态表,不执行 crawler 导入') | 20 | help='只初始化 yinyan_song_records 待导入状态表,不执行 crawler 导入') |
| 21 | parser.add_argument('--init-yinyan-records2', action='store_true', | ||
| 22 | help='写入全部合格 song-record 关联至 yinyan_song_records2,再排除 yinyan_song_records 已有关系') | ||
| 21 | parser.add_argument('--backfill-yinyan-platforms', action='store_true', | 23 | parser.add_argument('--backfill-yinyan-platforms', action='store_true', |
| 22 | help='只回填 yinyan_song_records 中为空的 platform 平台代码') | 24 | help='只回填 yinyan_song_records 中为空的 platform 平台代码') |
| 23 | parser.add_argument('--backfill-empty-singers', action='store_true', | 25 | parser.add_argument('--backfill-empty-singers', action='store_true', |
| ... | @@ -36,6 +38,8 @@ if __name__ == '__main__': | ... | @@ -36,6 +38,8 @@ if __name__ == '__main__': |
| 36 | backfill_yinyan_record_platforms(max_batches=args.max_batches) | 38 | backfill_yinyan_record_platforms(max_batches=args.max_batches) |
| 37 | elif args.init_yinyan_records: | 39 | elif args.init_yinyan_records: |
| 38 | initialize_yinyan_song_records(platforms, max_batches=args.max_batches) | 40 | initialize_yinyan_song_records(platforms, max_batches=args.max_batches) |
| 41 | elif args.init_yinyan_records2: | ||
| 42 | initialize_yinyan_song_records2(platforms, max_batches=args.max_batches) | ||
| 39 | elif args.backfill_empty_singers: | 43 | elif args.backfill_empty_singers: |
| 40 | backfill_empty_singers(platforms, max_batches=args.max_batches) | 44 | backfill_empty_singers(platforms, max_batches=args.max_batches) |
| 41 | elif args.backfill_qq_invalid_covers: | 45 | elif args.backfill_qq_invalid_covers: | ... | ... |
tests/test_connections.py
0 → 100644
| 1 | from etl_to_crawler.connections import _NulSafePgConnection | ||
| 2 | |||
| 3 | |||
| 4 | class _Cursor: | ||
| 5 | def __init__(self): | ||
| 6 | self.executed = [] | ||
| 7 | self.executed_many = [] | ||
| 8 | |||
| 9 | def __enter__(self): | ||
| 10 | return self | ||
| 11 | |||
| 12 | def __exit__(self, exc_type, exc_value, traceback): | ||
| 13 | return False | ||
| 14 | |||
| 15 | def execute(self, operation, parameters=None): | ||
| 16 | self.executed.append((operation, parameters)) | ||
| 17 | |||
| 18 | def executemany(self, operation, parameter_sets): | ||
| 19 | self.executed_many.append((operation, parameter_sets)) | ||
| 20 | |||
| 21 | |||
| 22 | class _Connection: | ||
| 23 | def __init__(self): | ||
| 24 | self.raw_cursor = _Cursor() | ||
| 25 | |||
| 26 | def cursor(self): | ||
| 27 | return self.raw_cursor | ||
| 28 | |||
| 29 | |||
| 30 | def test_pg_connection_removes_nul_from_execute_parameters(): | ||
| 31 | raw_connection = _Connection() | ||
| 32 | connection = _NulSafePgConnection(raw_connection) | ||
| 33 | |||
| 34 | with connection.cursor() as cur: | ||
| 35 | cur.execute('INSERT INTO example VALUES (%s, %s)', ('歌\x00词', {'text': '介\x00绍'})) | ||
| 36 | |||
| 37 | assert raw_connection.raw_cursor.executed == [ | ||
| 38 | ('INSERT INTO example VALUES (%s, %s)', ('歌词', {'text': '介绍'})), | ||
| 39 | ] | ||
| 40 | |||
| 41 | |||
| 42 | def test_pg_connection_removes_nul_from_every_executemany_row(): | ||
| 43 | raw_connection = _Connection() | ||
| 44 | connection = _NulSafePgConnection(raw_connection) | ||
| 45 | |||
| 46 | with connection.cursor() as cur: | ||
| 47 | cur.executemany('INSERT INTO example VALUES (%s)', [('A\x00',), ('B\x00',)]) | ||
| 48 | |||
| 49 | assert raw_connection.raw_cursor.executed_many == [ | ||
| 50 | ('INSERT INTO example VALUES (%s)', [('A',), ('B',)]), | ||
| 51 | ] |
| 1 | from unittest.mock import MagicMock | 1 | from unittest.mock import MagicMock |
| 2 | 2 | ||
| 3 | from etl_to_crawler.reader import ( | 3 | from etl_to_crawler.reader import ( |
| 4 | fetch_all_song_record_relations, | ||
| 4 | fetch_platform_records_by_record_ids, | 5 | fetch_platform_records_by_record_ids, |
| 5 | fetch_platform_records, | 6 | fetch_platform_records, |
| 6 | iter_hk_songs_batches, | 7 | iter_hk_songs_batches, |
| ... | @@ -86,6 +87,22 @@ def test_fetch_platform_records_by_record_ids_queries_exact_records_once(): | ... | @@ -86,6 +87,22 @@ def test_fetch_platform_records_by_record_ids_queries_exact_records_once(): |
| 86 | assert result[0]['platform'] == '2' | 87 | assert result[0]['platform'] == '2' |
| 87 | 88 | ||
| 88 | 89 | ||
| 90 | def test_fetch_all_song_record_relations_keeps_all_records_for_a_song(): | ||
| 91 | conn = MagicMock() | ||
| 92 | conn.cursor.return_value = _make_cursor([ | ||
| 93 | {'source_song_id': 10, 'record_id': 100, 'platform': '1'}, | ||
| 94 | {'source_song_id': 10, 'record_id': 101, 'platform': '1'}, | ||
| 95 | {'source_song_id': 10, 'record_id': 200, 'platform': '2'}, | ||
| 96 | {'source_song_id': 10, 'record_id': 101, 'platform': '1'}, | ||
| 97 | ]) | ||
| 98 | |||
| 99 | result = fetch_all_song_record_relations(conn, [10]) | ||
| 100 | |||
| 101 | assert [(row['record_id'], row['platform']) for row in result] == [ | ||
| 102 | (100, '1'), (101, '1'), (200, '2'), | ||
| 103 | ] | ||
| 104 | |||
| 105 | |||
| 89 | def test_select_primary_record_prefers_main_version_then_is_high(): | 106 | def test_select_primary_record_prefers_main_version_then_is_high(): |
| 90 | rows = [ | 107 | rows = [ |
| 91 | { | 108 | { | ... | ... |
| ... | @@ -258,6 +258,38 @@ def test_initialize_yinyan_song_records_inserts_primary_records(monkeypatch): | ... | @@ -258,6 +258,38 @@ def test_initialize_yinyan_song_records_inserts_primary_records(monkeypatch): |
| 258 | assert pg_conn.commits == 1 | 258 | assert pg_conn.commits == 1 |
| 259 | 259 | ||
| 260 | 260 | ||
| 261 | def test_initialize_yinyan_song_records2_writes_all_relations_then_deduplicates(monkeypatch): | ||
| 262 | pg_conn = _PgConnection() | ||
| 263 | inserted = [] | ||
| 264 | dedupe_calls = [] | ||
| 265 | |||
| 266 | monkeypatch.setattr(runner, 'get_hk_songs_conn', lambda: _Connection()) | ||
| 267 | monkeypatch.setattr(runner, 'get_source_conn', lambda: _Connection()) | ||
| 268 | monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn) | ||
| 269 | monkeypatch.setattr(runner, 'iter_hk_songs_records2_batches', lambda conn, batch_size: [[ | ||
| 270 | {'id': 50, 'source_song_id': 10}, | ||
| 271 | ]]) | ||
| 272 | monkeypatch.setattr(runner, 'fetch_all_song_record_relations', lambda conn, song_ids: [ | ||
| 273 | {'source_song_id': 10, 'record_id': 100, 'platform': '1'}, | ||
| 274 | {'source_song_id': 10, 'record_id': 200, 'platform': '2'}, | ||
| 275 | ]) | ||
| 276 | monkeypatch.setattr(runner, 'insert_yinyan_song_records2', lambda cur, rows: inserted.extend(rows)) | ||
| 277 | monkeypatch.setattr( | ||
| 278 | runner, | ||
| 279 | 'delete_yinyan_song_records2_existing_relations', | ||
| 280 | lambda cur: dedupe_calls.append(True) or 1, | ||
| 281 | ) | ||
| 282 | |||
| 283 | runner.initialize_yinyan_song_records2(['1', '2']) | ||
| 284 | |||
| 285 | assert inserted == [ | ||
| 286 | {'song_id': 10, 'record_id': 100, 'platform': '1'}, | ||
| 287 | {'song_id': 10, 'record_id': 200, 'platform': '2'}, | ||
| 288 | ] | ||
| 289 | assert dedupe_calls == [True] | ||
| 290 | assert pg_conn.commits == 2 | ||
| 291 | |||
| 292 | |||
| 261 | def test_backfill_yinyan_record_platforms_updates_missing_platform_rows(monkeypatch): | 293 | def test_backfill_yinyan_record_platforms_updates_missing_platform_rows(monkeypatch): |
| 262 | pg_conn = _PgConnection() | 294 | pg_conn = _PgConnection() |
| 263 | updated = [] | 295 | updated = [] | ... | ... |
| ... | @@ -3,6 +3,8 @@ from etl_to_crawler.writer import ( | ... | @@ -3,6 +3,8 @@ from etl_to_crawler.writer import ( |
| 3 | fetch_pending_yinyan_song_records, | 3 | fetch_pending_yinyan_song_records, |
| 4 | fetch_yinyan_records_missing_platform, | 4 | fetch_yinyan_records_missing_platform, |
| 5 | insert_yinyan_song_records, | 5 | insert_yinyan_song_records, |
| 6 | insert_yinyan_song_records2, | ||
| 7 | delete_yinyan_song_records2_existing_relations, | ||
| 6 | update_yinyan_record_platforms, | 8 | update_yinyan_record_platforms, |
| 7 | upsert_kugou_albums, | 9 | upsert_kugou_albums, |
| 8 | upsert_kugou_singers, | 10 | upsert_kugou_singers, |
| ... | @@ -123,6 +125,31 @@ def test_insert_yinyan_song_records_initializes_unpushed_rows(): | ... | @@ -123,6 +125,31 @@ def test_insert_yinyan_song_records_initializes_unpushed_rows(): |
| 123 | assert rows == [(10, 100, '1'), (11, 101, '2')] | 125 | assert rows == [(10, 100, '1'), (11, 101, '2')] |
| 124 | 126 | ||
| 125 | 127 | ||
| 128 | def test_insert_yinyan_song_records2_writes_all_candidates_before_deduplication(): | ||
| 129 | cur = MagicMock() | ||
| 130 | insert_yinyan_song_records2(cur, [ | ||
| 131 | {'song_id': 10, 'record_id': 100, 'platform': '1'}, | ||
| 132 | {'song_id': 10, 'record_id': 200, 'platform': '2'}, | ||
| 133 | ]) | ||
| 134 | |||
| 135 | sql, params = cur.execute.call_args[0] | ||
| 136 | assert 'INSERT INTO yinyan_song_records2' in sql | ||
| 137 | assert 'VALUES (%s::bigint, %s::bigint, %s::varchar), (%s::bigint, %s::bigint, %s::varchar)' in sql | ||
| 138 | assert 'ON CONFLICT (song_id, record_id) DO UPDATE' in sql | ||
| 139 | assert params == (10, 100, '1', 10, 200, '2') | ||
| 140 | |||
| 141 | |||
| 142 | def test_delete_yinyan_song_records2_existing_relations_uses_records1_as_dedup_source(): | ||
| 143 | cur = MagicMock() | ||
| 144 | delete_yinyan_song_records2_existing_relations(cur) | ||
| 145 | |||
| 146 | sql = cur.execute.call_args[0][0] | ||
| 147 | assert 'DELETE FROM yinyan_song_records2 AS ysr2' in sql | ||
| 148 | assert 'USING yinyan_song_records AS ysr1' in sql | ||
| 149 | assert 'ysr2.song_id = ysr1.song_id' in sql | ||
| 150 | assert 'ysr2.record_id = ysr1.record_id' in sql | ||
| 151 | |||
| 152 | |||
| 126 | def test_fetch_yinyan_records_missing_platform_reads_null_platform_rows(): | 153 | def test_fetch_yinyan_records_missing_platform_reads_null_platform_rows(): |
| 127 | cur = MagicMock() | 154 | cur = MagicMock() |
| 128 | cur.fetchall.return_value = [(10, 100), (11, 101)] | 155 | cur.fetchall.return_value = [(10, 100), (11, 101)] | ... | ... |
-
Please register or sign in to post a comment