Commit d3cca312 d3cca3128cd3ca2eedd0ba5c44f953e276e093fb by 沈秋雨

feat(etl): 支持重试失败导入并修复相关逻辑

1 parent 443980dc
...@@ -56,6 +56,8 @@ if YINYAN_IMPORT_TABLE not in _ALLOWED_YINYAN_IMPORT_TABLES: ...@@ -56,6 +56,8 @@ if YINYAN_IMPORT_TABLE not in _ALLOWED_YINYAN_IMPORT_TABLES:
56 f"got {YINYAN_IMPORT_TABLE!r}" 56 f"got {YINYAN_IMPORT_TABLE!r}"
57 ) 57 )
58 58
59 TARGET_TABLE_NAME = os.environ.get('TARGET_TABLE_NAME', 'hk_songs').strip()
60
59 BATCH_SIZE = 10 61 BATCH_SIZE = 10
60 BACKFILL_BATCH_SIZE = int(os.environ.get('BACKFILL_BATCH_SIZE', '5000')) 62 BACKFILL_BATCH_SIZE = int(os.environ.get('BACKFILL_BATCH_SIZE', '5000'))
61 HTTP_POOL_MAXSIZE = int(os.environ.get('HTTP_POOL_MAXSIZE', '128')) 63 HTTP_POOL_MAXSIZE = int(os.environ.get('HTTP_POOL_MAXSIZE', '128'))
......
...@@ -328,6 +328,4 @@ def refresh_conn(conn, pool_type: str): ...@@ -328,6 +328,4 @@ def refresh_conn(conn, pool_type: str):
328 if pool is None: 328 if pool is None:
329 raise ValueError(f"unknown pool_type: {pool_type}") 329 raise ValueError(f"unknown pool_type: {pool_type}")
330 330
331 if pool_type == 'pg':
332 return pool().check(conn)
333 return pool().check(conn) 331 return pool().check(conn)
......
1 import re
2
3 _TIMESTAMP_RE = re.compile(r'\[\d{2}:\d{2}\.\d{2,3}\]')
4 _META_TAG_RE = re.compile(r'\[[a-zA-Z]+:[^\]]*\]')
5
6
7 def ensure_newlines(text: str | None) -> str: 1 def ensure_newlines(text: str | None) -> str:
8 """确保歌词每行之间有换行符,保留时间戳等原始内容""" 2 """确保歌词每行之间有换行符,保留时间戳等原始内容"""
9 if not text: 3 if not text:
...@@ -13,11 +7,3 @@ def ensure_newlines(text: str | None) -> str: ...@@ -13,11 +7,3 @@ def ensure_newlines(text: str | None) -> str:
13 lines = [line.strip() for line in text.splitlines() if line.strip()] 7 lines = [line.strip() for line in text.splitlines() if line.strip()]
14 return '\n'.join(lines) 8 return '\n'.join(lines)
15 9
16
17 def strip_timestamps(text: str | None) -> str:
18 if not text:
19 return ''
20 text = _META_TAG_RE.sub('', text)
21 text = _TIMESTAMP_RE.sub('', text)
22 lines = [line.strip() for line in text.splitlines() if line.strip()]
23 return '\n'.join(lines)
......
1 from typing import Iterator 1 from typing import Iterator
2 import pymysql 2 import pymysql
3 from .config import PLATFORMS 3 from .config import PLATFORMS, TARGET_TABLE_NAME
4 4
5 5
6 _HK_SONGS_QUERY = """ 6 _HK_SONGS_QUERY = f"""
7 SELECT id, name, lyricist, composer, audio_url, lyrics_url, 7 SELECT id, name, lyricist, composer, audio_url, lyrics_url,
8 cover_url, singer, issue_time, source_song_id, song_time 8 cover_url, singer, issue_time, source_song_id, song_time
9 FROM hk_songs_test 9 FROM {TARGET_TABLE_NAME}
10 WHERE deleted = '0' 10 WHERE deleted = '0'
11 AND id > %s 11 AND id > %s
12 AND name IS NOT NULL AND name != '' 12 AND name IS NOT NULL AND name != ''
...@@ -16,34 +16,18 @@ ORDER BY id ...@@ -16,34 +16,18 @@ ORDER BY id
16 LIMIT %s 16 LIMIT %s
17 """ 17 """
18 18
19 _HK_SONGS_BY_SOURCE_IDS_QUERY = """ 19 _HK_SONGS_BY_SOURCE_IDS_QUERY = f"""
20 SELECT id, name, lyricist, composer, audio_url, lyrics_url, 20 SELECT id, name, lyricist, composer, audio_url, lyrics_url,
21 cover_url, singer, issue_time, source_song_id, song_time 21 cover_url, singer, issue_time, source_song_id, song_time
22 FROM hk_songs_test 22 FROM {TARGET_TABLE_NAME}
23 WHERE deleted = '0' 23 WHERE deleted = '0'
24 AND source_song_id IN ({placeholders}) 24 AND source_song_id IN ({{placeholders}})
25 AND name IS NOT NULL AND name != '' 25 AND name IS NOT NULL AND name != ''
26 AND audio_url IS NOT NULL AND audio_url != '' 26 AND audio_url IS NOT NULL AND audio_url != ''
27 AND singer IS NOT NULL AND singer != '' 27 AND singer IS NOT NULL AND singer != ''
28 ORDER BY id 28 ORDER BY id
29 """ 29 """
30 30
31 # records2 初始化不要求封面或歌词;正式导入时这两个字段允许为空。
32 _HK_SONGS_RECORDS2_QUERY = """
33 SELECT id, source_song_id
34 FROM hk_songs_test
35 WHERE deleted = '0'
36 AND id > %s
37 AND source_song_id IS NOT NULL
38 AND name IS NOT NULL AND TRIM(name) != ''
39 AND song_time IS NOT NULL AND song_time > 0
40 AND singer IS NOT NULL AND TRIM(singer) != ''
41 AND audio_url IS NOT NULL AND TRIM(audio_url) != ''
42 AND issue_time IS NOT NULL
43 ORDER BY id
44 LIMIT %s
45 """
46
47 _PLATFORM_QUERY = """ 31 _PLATFORM_QUERY = """
48 SELECT 32 SELECT
49 sar.song_id AS source_song_id, 33 sar.song_id AS source_song_id,
...@@ -143,30 +127,17 @@ def iter_hk_songs_batches( ...@@ -143,30 +127,17 @@ def iter_hk_songs_batches(
143 break 127 break
144 128
145 129
146 def iter_hk_songs_records2_batches( 130 def fetch_hk_songs_by_source_ids(
147 conn: pymysql.Connection, 131 conn: pymysql.Connection,
148 batch_size: int, 132 source_song_ids: list[int],
149 start_after_id: int = 0, 133 include_deleted: bool = False,
150 ) -> Iterator[list[dict]]: 134 ) -> dict[int, dict]:
151 """遍历具备 records2 基础入库字段的歌曲,封面和歌词可为空。"""
152 last_id = start_after_id
153 with conn.cursor() as cur:
154 while True:
155 cur.execute(_HK_SONGS_RECORDS2_QUERY, (last_id, batch_size))
156 rows = cur.fetchall()
157 if not rows:
158 break
159 yield rows
160 last_id = int(rows[-1]['id'])
161 if len(rows) < batch_size:
162 break
163
164
165 def fetch_hk_songs_by_source_ids(conn: pymysql.Connection, source_song_ids: list[int]) -> dict[int, dict]:
166 if not source_song_ids: 135 if not source_song_ids:
167 return {} 136 return {}
168 placeholders = ','.join(['%s'] * len(source_song_ids)) 137 placeholders = ','.join(['%s'] * len(source_song_ids))
169 query = _HK_SONGS_BY_SOURCE_IDS_QUERY.format(placeholders=placeholders) 138 query = _HK_SONGS_BY_SOURCE_IDS_QUERY.format(placeholders=placeholders)
139 if include_deleted:
140 query = query.replace("WHERE deleted = '0'", 'WHERE 1 = 1', 1)
170 with conn.cursor() as cur: 141 with conn.cursor() as cur:
171 cur.execute(query, source_song_ids) 142 cur.execute(query, source_song_ids)
172 rows = cur.fetchall() 143 rows = cur.fetchall()
...@@ -178,20 +149,35 @@ def fetch_hk_songs_by_source_ids(conn: pymysql.Connection, source_song_ids: list ...@@ -178,20 +149,35 @@ def fetch_hk_songs_by_source_ids(conn: pymysql.Connection, source_song_ids: list
178 149
179 150
180 def mark_hk_songs_deleted(conn: pymysql.Connection, source_song_ids: list[int]) -> int: 151 def mark_hk_songs_deleted(conn: pymysql.Connection, source_song_ids: list[int]) -> int:
181 """将无法导入的源歌曲在 hk_songs_test 中软删除。""" 152 """将无法导入的源歌曲在目标表中软删除。"""
182 if not source_song_ids: 153 if not source_song_ids:
183 return 0 154 return 0
184 unique_ids = list(dict.fromkeys(int(song_id) for song_id in source_song_ids)) 155 unique_ids = list(dict.fromkeys(int(song_id) for song_id in source_song_ids))
185 placeholders = ','.join(['%s'] * len(unique_ids)) 156 placeholders = ','.join(['%s'] * len(unique_ids))
186 with conn.cursor() as cur: 157 with conn.cursor() as cur:
187 cur.execute( 158 cur.execute(
188 f"UPDATE hk_songs_test SET deleted = '1' " 159 f"UPDATE {TARGET_TABLE_NAME} SET deleted = '1' "
189 f"WHERE source_song_id IN ({placeholders}) AND deleted = '0'", 160 f"WHERE source_song_id IN ({placeholders}) AND deleted = '0'",
190 unique_ids, 161 unique_ids,
191 ) 162 )
192 return cur.rowcount 163 return cur.rowcount
193 164
194 165
166 def mark_hk_songs_active(conn: pymysql.Connection, source_song_ids: list[int]) -> int:
167 """恢复需要重新导入的 HK 源歌曲。"""
168 if not source_song_ids:
169 return 0
170 unique_ids = list(dict.fromkeys(int(song_id) for song_id in source_song_ids))
171 placeholders = ','.join(['%s'] * len(unique_ids))
172 with conn.cursor() as cur:
173 cur.execute(
174 f"UPDATE {TARGET_TABLE_NAME} SET deleted = '0' "
175 f"WHERE source_song_id IN ({placeholders}) AND deleted = '1'",
176 unique_ids,
177 )
178 return cur.rowcount
179
180
195 def fetch_platform_records(source_conn: pymysql.Connection, song_ids: list[int]) -> list[dict]: 181 def fetch_platform_records(source_conn: pymysql.Connection, song_ids: list[int]) -> list[dict]:
196 if not song_ids: 182 if not song_ids:
197 return [] 183 return []
......
...@@ -13,6 +13,7 @@ from .reader import ( ...@@ -13,6 +13,7 @@ from .reader import (
13 iter_hk_songs_batches, 13 iter_hk_songs_batches,
14 fetch_hk_songs_by_source_ids, 14 fetch_hk_songs_by_source_ids,
15 mark_hk_songs_deleted, 15 mark_hk_songs_deleted,
16 mark_hk_songs_active,
16 fetch_platform_records, 17 fetch_platform_records,
17 fetch_all_platform_records, 18 fetch_all_platform_records,
18 fetch_all_song_record_relations, 19 fetch_all_song_record_relations,
...@@ -28,6 +29,7 @@ from .spider import ( ...@@ -28,6 +29,7 @@ from .spider import (
28 ) 29 )
29 from .writer import ( 30 from .writer import (
30 fetch_pending_yinyan_song_records, 31 fetch_pending_yinyan_song_records,
32 reset_failed_yinyan_song_records,
31 fetch_yinyan_records_missing_platform, 33 fetch_yinyan_records_missing_platform,
32 insert_yinyan_song_records, 34 insert_yinyan_song_records,
33 insert_yinyan_song_records2, 35 insert_yinyan_song_records2,
...@@ -1384,7 +1386,7 @@ def initialize_yinyan_song_records(platforms: list[str], max_batches: int | None ...@@ -1384,7 +1386,7 @@ def initialize_yinyan_song_records(platforms: list[str], max_batches: int | None
1384 pg_conn.commit() 1386 pg_conn.commit()
1385 total += len(init_rows) 1387 total += len(init_rows)
1386 1388
1387 # 本批次中无法匹配到任何录音的歌曲 → 软删 hk_songs_test 1389 # 本批次中无法匹配到任何录音的歌曲 → 软删目标表
1388 matched_src_ids = {int(hk_row['source_song_id']) for hk_row in batch 1390 matched_src_ids = {int(hk_row['source_song_id']) for hk_row in batch
1389 if hk_row.get('source_song_id') and int(hk_row['source_song_id']) in pr_by_song} 1391 if hk_row.get('source_song_id') and int(hk_row['source_song_id']) in pr_by_song}
1390 unmatched_ids = [ 1392 unmatched_ids = [
...@@ -1804,7 +1806,7 @@ def backfill_qq_invalid_covers(max_batches: int | None = None) -> None: ...@@ -1804,7 +1806,7 @@ def backfill_qq_invalid_covers(max_batches: int | None = None) -> None:
1804 removals.append(row) 1806 removals.append(row)
1805 continue 1807 continue
1806 if not hk_row: 1808 if not hk_row:
1807 log.warning('QQ cover backfill cannot find hk_songs_test row: source_song_id=%s', row['song_id']) 1809 log.warning('QQ cover backfill cannot find target table row: source_song_id=%s', row['song_id'])
1808 retry_later += 1 1810 retry_later += 1
1809 continue 1811 continue
1810 1812
...@@ -1863,6 +1865,7 @@ def backfill_qq_invalid_covers(max_batches: int | None = None) -> None: ...@@ -1863,6 +1865,7 @@ def backfill_qq_invalid_covers(max_batches: int | None = None) -> None:
1863 def run( 1865 def run(
1864 platforms: list[str], 1866 platforms: list[str],
1865 max_batches: int | None = None, 1867 max_batches: int | None = None,
1868 retry_failed: bool = False,
1866 ) -> None: 1869 ) -> None:
1867 hk_conn = get_hk_songs_conn() 1870 hk_conn = get_hk_songs_conn()
1868 src_conn = get_source_conn() 1871 src_conn = get_source_conn()
...@@ -1874,6 +1877,18 @@ def run( ...@@ -1874,6 +1877,18 @@ def run(
1874 imported: list[dict] = [] 1877 imported: list[dict] = []
1875 1878
1876 try: 1879 try:
1880 if retry_failed:
1881 with pg_conn.cursor() as pg_cur:
1882 reset_count, retry_song_ids = reset_failed_yinyan_song_records(pg_cur, platforms)
1883 pg_conn.commit()
1884 if retry_song_ids:
1885 restored = mark_hk_songs_active(hk_conn, retry_song_ids)
1886 hk_conn.commit()
1887 log.info(
1888 'Reset failed imports for retry: rows=%d songs=%d restored_hk_rows=%d',
1889 reset_count, len(retry_song_ids), restored,
1890 )
1891
1877 batch_index = 0 1892 batch_index = 0
1878 pbar = tqdm(desc='batches') 1893 pbar = tqdm(desc='batches')
1879 while max_batches is None or batch_index < max_batches: 1894 while max_batches is None or batch_index < max_batches:
...@@ -1884,12 +1899,15 @@ def run( ...@@ -1884,12 +1899,15 @@ def run(
1884 pg_conn = refresh_conn(pg_conn, 'pg') 1899 pg_conn = refresh_conn(pg_conn, 'pg')
1885 1900
1886 with pg_conn.cursor() as pg_cur: 1901 with pg_conn.cursor() as pg_cur:
1887 pending_records = fetch_pending_yinyan_song_records(pg_cur, BATCH_SIZE) 1902 pending_records = fetch_pending_yinyan_song_records(pg_cur, BATCH_SIZE, platforms)
1888 if not pending_records: 1903 if not pending_records:
1889 break 1904 break
1890 1905
1891 song_ids = [int(r['song_id']) for r in pending_records] 1906 song_ids = [int(r['song_id']) for r in pending_records]
1892 record_ids = [int(r['record_id']) for r in pending_records] 1907 record_ids = [int(r['record_id']) for r in pending_records]
1908 if retry_failed:
1909 hk_by_song = fetch_hk_songs_by_source_ids(hk_conn, song_ids, include_deleted=True)
1910 else:
1893 hk_by_song = fetch_hk_songs_by_source_ids(hk_conn, song_ids) 1911 hk_by_song = fetch_hk_songs_by_source_ids(hk_conn, song_ids)
1894 platform_records = fetch_platform_records_by_record_ids(src_conn, record_ids) 1912 platform_records = fetch_platform_records_by_record_ids(src_conn, record_ids)
1895 pr_by_state_key: dict[tuple[int, int, str], dict] = {} 1913 pr_by_state_key: dict[tuple[int, int, str], dict] = {}
......
1 """工具函数模块""" 1 """工具函数模块"""
2 import hashlib
3 import re
4 from urllib.parse import urlparse
5
6
7 def extract_plain_lyric(lyric: str) -> str:
8 """去除 LRC 歌词中的时间戳和标签,保留纯文本
9
10 Args:
11 lyric: 原始歌词内容(可能含 [mm:ss.xx] 时间戳和 [ti:xxx] 等标签)
12
13 Returns:
14 纯文本歌词,行之间用换行符连接
15 """
16 if not lyric:
17 return ""
18 lines = []
19 for line in lyric.splitlines():
20 # 去除 LRC 时间戳 [mm:ss.xx] 或 [mm:ss.xxx]
21 cleaned = re.sub(r"\[\d{2}:\d{2}(?:\.\d{2,3})?\]", "", line)
22 # 去除标签 [xx:yy]
23 cleaned = re.sub(r"\[[a-zA-Z]+:[^\]]+\]", "", cleaned)
24 cleaned = cleaned.strip()
25 if cleaned:
26 lines.append(cleaned)
27 return "\n".join(lines)
28
29
30 def split_title_version(title: str | None) -> tuple[str, str]:
31 """拆分歌名末尾括号版本信息。
32
33 例如:化风行万里 (DJ默涵版) -> (化风行万里, DJ默涵版)
34 """
35 if not title:
36 return "", ""
37 text = title.strip()
38 match = re.match(r"^(?P<title>.+?)\s*[\((](?P<version>[^()()]+)[\))]\s*$", text)
39 if not match:
40 return text, ""
41 clean_title = match.group("title").strip()
42 version = match.group("version").strip()
43 return clean_title or text, version
44
45
46 def upload_plain_lyric_to_bucket(platform: str, unique_id: str, lyric: str, bucket, base_url: str) -> str:
47 """将歌词去时间戳后上传到当前 ETL 使用的 OSS bucket"""
48 plain_lyric = extract_plain_lyric(lyric)
49 if not plain_lyric:
50 return ""
51 oss_key = f"crawler/lyric/{platform}/{unique_id}.txt"
52 bucket.put_object(oss_key, plain_lyric.encode("utf-8"), headers={'Content-Type': 'text/plain; charset=utf-8'})
53 return f"{base_url.rstrip('/')}/{oss_key}"
54
55
56 def compute_audio_md5(audio_bytes: bytes) -> str:
57 """计算音频字节流的 MD5 值(32 位小写十六进制字符串)
58
59 Args:
60 audio_bytes: 音频文件字节流
61
62 Returns:
63 MD5 字符串,输入为空时返回空字符串
64 """
65 if not audio_bytes:
66 return ""
67 return hashlib.md5(audio_bytes).hexdigest()
68 """工具函数模块"""
2 import asyncio 69 import asyncio
3 import hashlib 70 import hashlib
4 import logging 71 import logging
......
...@@ -112,7 +112,18 @@ def fetch_existing_yinyan_song_ids(cur) -> set[int]: ...@@ -112,7 +112,18 @@ def fetch_existing_yinyan_song_ids(cur) -> set[int]:
112 return {row[0] for row in cur.fetchall()} 112 return {row[0] for row in cur.fetchall()}
113 113
114 114
115 def fetch_pending_yinyan_song_records(cur, limit: int) -> list[dict]: 115 def fetch_pending_yinyan_song_records(
116 cur,
117 limit: int,
118 platforms: list[str] | None = None,
119 ) -> list[dict]:
120 platform_filter = ''
121 params: list = []
122 if platforms:
123 placeholders = ', '.join(['%s'] * len(platforms))
124 platform_filter = f'AND platform IN ({placeholders})'
125 params.extend(platforms)
126 params.append(limit)
116 cur.execute( 127 cur.execute(
117 f""" 128 f"""
118 SELECT song_id, record_id, platform 129 SELECT song_id, record_id, platform
...@@ -120,15 +131,36 @@ def fetch_pending_yinyan_song_records(cur, limit: int) -> list[dict]: ...@@ -120,15 +131,36 @@ def fetch_pending_yinyan_song_records(cur, limit: int) -> list[dict]:
120 WHERE is_yinyan_push = FALSE 131 WHERE is_yinyan_push = FALSE
121 AND platform IS NOT NULL 132 AND platform IS NOT NULL
122 AND platform_song_id IS NULL 133 AND platform_song_id IS NULL
134 {platform_filter}
123 ORDER BY song_id 135 ORDER BY song_id
124 LIMIT %s 136 LIMIT %s
125 """, 137 """,
126 (limit,), 138 tuple(params),
127 ) 139 )
128 rows = cur.fetchall() 140 rows = cur.fetchall()
129 return [{'song_id': row[0], 'record_id': row[1], 'platform': row[2]} for row in rows] 141 return [{'song_id': row[0], 'record_id': row[1], 'platform': row[2]} for row in rows]
130 142
131 143
144 def reset_failed_yinyan_song_records(cur, platforms: list[str]) -> tuple[int, list[int]]:
145 """将指定平台的失败行恢复为待导入,返回 (行数, HK song_ids)。"""
146 if not platforms:
147 return 0, []
148 placeholders = ', '.join(['%s'] * len(platforms))
149 cur.execute(
150 f"""
151 UPDATE {YINYAN_IMPORT_TABLE}
152 SET platform_song_id = NULL
153 WHERE is_yinyan_push = FALSE
154 AND platform_song_id < 0
155 AND platform IN ({placeholders})
156 RETURNING song_id
157 """,
158 tuple(platforms),
159 )
160 rows = cur.fetchall()
161 return len(rows), list(dict.fromkeys(int(row[0]) for row in rows))
162
163
132 def fetch_yinyan_records_missing_platform(cur, limit: int) -> list[dict]: 164 def fetch_yinyan_records_missing_platform(cur, limit: int) -> list[dict]:
133 cur.execute( 165 cur.execute(
134 """ 166 """
......
...@@ -16,6 +16,8 @@ if __name__ == '__main__': ...@@ -16,6 +16,8 @@ if __name__ == '__main__':
16 help='要导入的平台(默认 all)') 16 help='要导入的平台(默认 all)')
17 parser.add_argument('--max-batches', type=int, default=None, 17 parser.add_argument('--max-batches', type=int, default=None,
18 help='最多处理多少批次(冒烟测试用)') 18 help='最多处理多少批次(冒烟测试用)')
19 parser.add_argument('--retry-failed', action='store_true',
20 help='将当前导入状态表中的失败记录恢复为待处理并重新导入')
19 parser.add_argument('--init-yinyan-records', action='store_true', 21 parser.add_argument('--init-yinyan-records', action='store_true',
20 help='只初始化 yinyan_song_records 待导入状态表,不执行 crawler 导入') 22 help='只初始化 yinyan_song_records 待导入状态表,不执行 crawler 导入')
21 parser.add_argument('--init-yinyan-records2', action='store_true', 23 parser.add_argument('--init-yinyan-records2', action='store_true',
...@@ -45,4 +47,4 @@ if __name__ == '__main__': ...@@ -45,4 +47,4 @@ if __name__ == '__main__':
45 elif args.backfill_qq_invalid_covers: 47 elif args.backfill_qq_invalid_covers:
46 backfill_qq_invalid_covers(max_batches=args.max_batches) 48 backfill_qq_invalid_covers(max_batches=args.max_batches)
47 else: 49 else:
48 run(platforms, max_batches=args.max_batches) 50 run(platforms, max_batches=args.max_batches, retry_failed=args.retry_failed)
......
1 from etl_to_crawler.lyric import strip_timestamps
2
3 def test_strips_standard_timestamps():
4 lrc = "[00:12.34]第一行歌词\n[01:23.45]第二行歌词"
5 assert strip_timestamps(lrc) == "第一行歌词\n第二行歌词"
6
7 def test_strips_millisecond_timestamps():
8 lrc = "[00:12.345]带三位毫秒"
9 assert strip_timestamps(lrc) == "带三位毫秒"
10
11 def test_strips_meta_tags():
12 lrc = "[ti:歌曲名]\n[ar:歌手]\n[00:01.00]歌词"
13 result = strip_timestamps(lrc)
14 assert "歌词" in result
15 assert "[ti:" not in result
16
17 def test_empty_and_none():
18 assert strip_timestamps(None) == ''
19 assert strip_timestamps('') == ''
20
21 def test_plain_text_unchanged():
22 assert strip_timestamps("没有时间戳的歌词") == "没有时间戳的歌词"
1 from unittest.mock import MagicMock 1 from unittest.mock import MagicMock
2 import etl_to_crawler.reader as reader
2 3
3 from etl_to_crawler.reader import ( 4 from etl_to_crawler.reader import (
4 fetch_all_song_record_relations, 5 fetch_all_song_record_relations,
...@@ -87,6 +88,32 @@ def test_fetch_platform_records_by_record_ids_queries_exact_records_once(): ...@@ -87,6 +88,32 @@ def test_fetch_platform_records_by_record_ids_queries_exact_records_once():
87 assert result[0]['platform'] == '2' 88 assert result[0]['platform'] == '2'
88 89
89 90
91 def test_retry_fetch_can_include_soft_deleted_hk_songs():
92 conn = MagicMock()
93 conn.cursor.return_value = _make_cursor([])
94
95 reader.fetch_hk_songs_by_source_ids(conn, [10], include_deleted=True)
96
97 sql, params = conn.cursor.return_value.execute.call_args.args
98 assert "deleted = '0'" not in sql
99 assert 'WHERE 1 = 1' in sql
100 assert params == [10]
101
102
103 def test_mark_hk_songs_active_restores_retry_sources():
104 conn = MagicMock()
105 conn.cursor.return_value = _make_cursor([])
106 conn.cursor.return_value.rowcount = 2
107
108 restored = reader.mark_hk_songs_active(conn, [10, 10, 11])
109
110 sql, params = conn.cursor.return_value.execute.call_args.args
111 assert "SET deleted = '0'" in sql
112 assert "AND deleted = '1'" in sql
113 assert params == [10, 11]
114 assert restored == 2
115
116
90 def test_fetch_all_song_record_relations_keeps_all_records_for_a_song(): 117 def test_fetch_all_song_record_relations_keeps_all_records_for_a_song():
91 conn = MagicMock() 118 conn = MagicMock()
92 conn.cursor.return_value = _make_cursor([ 119 conn.cursor.return_value = _make_cursor([
......
...@@ -340,7 +340,7 @@ def test_run_imports_only_pending_yinyan_platform_record(monkeypatch): ...@@ -340,7 +340,7 @@ def test_run_imports_only_pending_yinyan_platform_record(monkeypatch):
340 monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn) 340 monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
341 monkeypatch.setattr(runner, 'get_oss_bucket', lambda: object()) 341 monkeypatch.setattr(runner, 'get_oss_bucket', lambda: object())
342 monkeypatch.setattr(runner, 'refresh_conn', lambda conn, name: conn) 342 monkeypatch.setattr(runner, 'refresh_conn', lambda conn, name: conn)
343 monkeypatch.setattr(runner, 'fetch_pending_yinyan_song_records', lambda cur, batch_size: [ 343 monkeypatch.setattr(runner, 'fetch_pending_yinyan_song_records', lambda cur, batch_size, platforms: [
344 {'song_id': 10, 'record_id': 200, 'platform': '2'}, 344 {'song_id': 10, 'record_id': 200, 'platform': '2'},
345 ] if pg_conn.commits == 0 else []) 345 ] if pg_conn.commits == 0 else [])
346 monkeypatch.setattr(runner, 'fetch_hk_songs_by_source_ids', lambda conn, song_ids: {10: { 346 monkeypatch.setattr(runner, 'fetch_hk_songs_by_source_ids', lambda conn, song_ids: {10: {
...@@ -389,6 +389,29 @@ def test_run_imports_only_pending_yinyan_platform_record(monkeypatch): ...@@ -389,6 +389,29 @@ def test_run_imports_only_pending_yinyan_platform_record(monkeypatch):
389 assert pg_conn.commits == 1 389 assert pg_conn.commits == 1
390 390
391 391
392 def test_run_retry_failed_resets_failures_and_restores_hk_sources(monkeypatch):
393 pg_conn = _PgConnection()
394 hk_conn = _Connection()
395 reset = MagicMock(return_value=(2, [10, 11]))
396 restore = MagicMock(return_value=2)
397
398 monkeypatch.setattr(runner, 'get_hk_songs_conn', lambda: hk_conn)
399 monkeypatch.setattr(runner, 'get_source_conn', lambda: _Connection())
400 monkeypatch.setattr(runner, 'get_spider_conn', lambda: _Connection())
401 monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
402 monkeypatch.setattr(runner, 'get_oss_bucket', lambda: object())
403 monkeypatch.setattr(runner, 'refresh_conn', lambda conn, name: conn)
404 monkeypatch.setattr(runner, 'reset_failed_yinyan_song_records', reset)
405 monkeypatch.setattr(runner, 'mark_hk_songs_active', restore)
406 monkeypatch.setattr(runner, 'fetch_pending_yinyan_song_records', lambda cur, batch_size, platforms: [])
407
408 runner.run(['1'], retry_failed=True)
409
410 reset.assert_called_once_with(pg_conn.cur, ['1'])
411 restore.assert_called_once_with(hk_conn, [10, 11])
412 assert pg_conn.commits == 1
413
414
392 def test_initialize_yinyan_song_records_inserts_primary_records(monkeypatch): 415 def test_initialize_yinyan_song_records_inserts_primary_records(monkeypatch):
393 pg_conn = _PgConnection() 416 pg_conn = _PgConnection()
394 inserted = [] 417 inserted = []
......
...@@ -208,6 +208,33 @@ def test_fetch_pending_yinyan_song_records_reads_platform_for_single_platform_im ...@@ -208,6 +208,33 @@ def test_fetch_pending_yinyan_song_records_reads_platform_for_single_platform_im
208 ] 208 ]
209 209
210 210
211 def test_fetch_pending_yinyan_song_records_filters_requested_platforms():
212 cur = MagicMock()
213 cur.fetchall.return_value = []
214
215 fetch_pending_yinyan_song_records(cur, 10, ['4'])
216
217 sql, params = cur.execute.call_args.args
218 assert 'AND platform IN (%s)' in sql
219 assert params == ('4', 10)
220
221
222 def test_reset_failed_yinyan_song_records_resets_selected_platforms_once():
223 cur = MagicMock()
224 cur.fetchall.return_value = [(10,), (10,), (11,)]
225
226 reset_count, song_ids = writer.reset_failed_yinyan_song_records(cur, ['1', '2'])
227
228 sql, params = cur.execute.call_args.args
229 assert f'UPDATE {writer.YINYAN_IMPORT_TABLE}' in sql
230 assert 'platform_song_id = NULL' in sql
231 assert 'platform_song_id < 0' in sql
232 assert 'RETURNING song_id' in sql
233 assert params == ('1', '2')
234 assert reset_count == 3
235 assert song_ids == [10, 11]
236
237
211 def test_import_state_queries_can_target_records2(monkeypatch): 238 def test_import_state_queries_can_target_records2(monkeypatch):
212 monkeypatch.setattr(writer, 'YINYAN_IMPORT_TABLE', 'yinyan_song_records2') 239 monkeypatch.setattr(writer, 'YINYAN_IMPORT_TABLE', 'yinyan_song_records2')
213 cur = MagicMock() 240 cur = MagicMock()
......