Commit cdf0d9e2 cdf0d9e2b55b0eab9141153e9eab1104ef59f076 by 沈秋雨

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 插入及去重逻辑,保障代码质量和功能稳定性
1 parent 7fdc8f0d
...@@ -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:
......
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)]
......