Commit dcddb309 dcddb309a8260585eb0584d92a7739fc32793938 by 沈秋雨

feat(etl): 支持断点续传功能提升数据导入可靠性

- run_etl.py中新增resume相关参数支持断点续传和手动起始位置
- reader.py中修改iter_hk_songs_batches函数,改用id游标替代offset分页
- runner.py新增断点状态文件管理功能,支持状态读写与断点继续
- run函数接受断点续传参数并根据状态文件恢复进度
- 处理每批次导入成功后保存断点状态,确保可中断恢复
- 新增单元测试覆盖断点续传逻辑,验证正确使用start_after_id参数
- 调整接口参数传递与兼容性,保持向后兼容性和灵活性
1 parent 0ab1bdc0
......@@ -8,11 +8,12 @@ SELECT id, name, lyricist, composer, audio_url, lyrics_url,
cover_url, singer, issue_time, source_song_id, song_time
FROM hk_songs_test
WHERE deleted = '0'
AND id > %s
AND name IS NOT NULL AND name != ''
AND audio_url IS NOT NULL AND audio_url != ''
AND singer IS NOT NULL AND singer != ''
ORDER BY id
LIMIT %s OFFSET %s
LIMIT %s
"""
_PLATFORM_QUERY = """
......@@ -42,18 +43,22 @@ ORDER BY sar.song_id,
"""
def iter_hk_songs_batches(conn: pymysql.Connection, batch_size: int) -> Iterator[list[dict]]:
offset = 0
def iter_hk_songs_batches(
conn: pymysql.Connection,
batch_size: int,
start_after_id: int = 0,
) -> Iterator[list[dict]]:
last_id = start_after_id
with conn.cursor() as cur:
while True:
cur.execute(_HK_SONGS_QUERY, (batch_size, offset))
cur.execute(_HK_SONGS_QUERY, (last_id, batch_size))
rows = cur.fetchall()
if not rows:
break
yield rows
last_id = int(rows[-1]['id'])
if len(rows) < batch_size:
break
offset += batch_size
def fetch_platform_records(source_conn: pymysql.Connection, song_ids: list[int]) -> list[dict]:
......
import uuid
import json
import logging
from pathlib import Path
from tqdm import tqdm
from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, OSS_CONFIG
......@@ -308,19 +309,59 @@ _PROCESSORS = {
}
def run(platforms: list[str], max_batches: int | None = None) -> None:
DEFAULT_STATE_FILE = Path('output/etl_to_crawler_state.json')
def _load_state(path: Path) -> dict:
if not path.exists():
return {}
with path.open('r', encoding='utf-8') as f:
return json.load(f)
def _save_state(path: Path, state: dict) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open('w', encoding='utf-8') as f:
json.dump(state, f, ensure_ascii=False, indent=2, sort_keys=True)
f.write('\n')
def _state_last_id(state: dict, key: str) -> int:
value = state.get(key, {}).get('last_hk_songs_id', 0)
return int(value or 0)
def run(
platforms: list[str],
max_batches: int | None = None,
resume: bool = False,
state_file: str | Path = DEFAULT_STATE_FILE,
state_key: str = 'all',
start_after_id: int | None = None,
reset_state: bool = False,
) -> 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']
state_path = Path(state_file)
state = _load_state(state_path) if resume and not reset_state else {}
resume_start_id = 0
if resume:
resume_start_id = _state_last_id(state, state_key)
if start_after_id is not None:
resume_start_id = start_after_id
total_ok = total_err = 0
imported: list[dict] = []
try:
for i, batch in enumerate(tqdm(iter_hk_songs_batches(hk_conn, BATCH_SIZE), desc='batches')):
for i, batch in enumerate(tqdm(
iter_hk_songs_batches(hk_conn, BATCH_SIZE, start_after_id=resume_start_id),
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')]
......@@ -360,6 +401,9 @@ def run(platforms: list[str], max_batches: int | None = None) -> None:
hk_row.get('name'), pr['platform'], e)
total_err += 1
pg_conn.commit()
if resume:
state.setdefault(state_key, {})['last_hk_songs_id'] = int(batch[-1]['id'])
_save_state(state_path, state)
finally:
hk_conn.close()
......
......@@ -16,6 +16,14 @@ if __name__ == '__main__':
help='要导入的平台(默认 all)')
parser.add_argument('--max-batches', type=int, default=None,
help='最多处理多少批次(冒烟测试用)')
parser.add_argument('--resume', action='store_true',
help='启用断点续传,按 hk_songs_test.id 从上次成功提交的批次后继续')
parser.add_argument('--state-file', default='output/etl_to_crawler_state.json',
help='断点状态文件路径')
parser.add_argument('--start-after-id', type=int, default=None,
help='手动指定从哪个 hk_songs_test.id 之后开始读取')
parser.add_argument('--reset-state', action='store_true',
help='忽略已有断点状态,从头或 --start-after-id 指定位置重新开始')
args = parser.parse_args()
if args.platform == 'all':
......@@ -24,4 +32,12 @@ if __name__ == '__main__':
platforms = [PLATFORM_MAP[args.platform]]
print(f"Starting ETL for platforms: {platforms}")
run(platforms, max_batches=args.max_batches)
run(
platforms,
max_batches=args.max_batches,
resume=args.resume,
state_file=args.state_file,
state_key=args.platform,
start_after_id=args.start_after_id,
reset_state=args.reset_state,
)
......
from unittest.mock import MagicMock
from etl_to_crawler.reader import fetch_platform_records, select_primary_record
from etl_to_crawler.reader import (
fetch_platform_records,
iter_hk_songs_batches,
select_primary_record,
)
def _make_cursor(rows):
......@@ -122,3 +126,22 @@ def test_select_primary_record_uses_earliest_pub_time_when_priority_ties():
]
assert select_primary_record(rows)['record_id'] == 101
def test_iter_hk_songs_batches_uses_id_cursor_instead_of_offset():
cur = _make_cursor([])
cur.fetchall.side_effect = [
[{'id': 10, 'source_song_id': 100}, {'id': 20, 'source_song_id': 200}],
[{'id': 25, 'source_song_id': 250}],
]
conn = MagicMock()
conn.cursor.return_value = cur
batches = list(iter_hk_songs_batches(conn, batch_size=2, start_after_id=5))
assert batches == [
[{'id': 10, 'source_song_id': 100}, {'id': 20, 'source_song_id': 200}],
[{'id': 25, 'source_song_id': 250}],
]
assert cur.execute.call_args_list[0][0][1] == (5, 2)
assert cur.execute.call_args_list[1][0][1] == (20, 2)
......
......@@ -48,7 +48,7 @@ def test_run_imports_all_platform_records_but_writes_one_primary_yinyan_relation
monkeypatch.setattr(runner, 'get_spider_conn', lambda: _Connection())
monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
monkeypatch.setattr(runner, 'get_oss_bucket', lambda: object())
monkeypatch.setattr(runner, 'iter_hk_songs_batches', lambda conn, batch_size: [[{
monkeypatch.setattr(runner, 'iter_hk_songs_batches', lambda conn, batch_size, start_after_id=0: [[{
'source_song_id': 10,
'name': '歌',
'audio_url': 'https://example.com/a.mp3',
......@@ -87,3 +87,32 @@ def test_run_imports_all_platform_records_but_writes_one_primary_yinyan_relation
processors['2'].assert_called_once()
yinyan_writer.assert_called_once_with(pg_conn.cur, [(10, 200)])
assert pg_conn.commits == 1
def test_run_resume_starts_after_saved_id_and_updates_state_after_commit(monkeypatch, tmp_path):
pg_conn = _PgConnection()
state_file = tmp_path / 'etl_state.json'
state_file.write_text('{"all": {"last_hk_songs_id": 40}}', encoding='utf-8')
seen_start_ids = []
monkeypatch.setattr(runner, 'get_hk_songs_conn', lambda: _Connection())
monkeypatch.setattr(runner, 'get_source_conn', lambda: _Connection())
monkeypatch.setattr(runner, 'get_spider_conn', lambda: _Connection())
monkeypatch.setattr(runner, 'get_pg_conn', lambda: pg_conn)
monkeypatch.setattr(runner, 'get_oss_bucket', lambda: object())
def fake_batches(conn, batch_size, start_after_id=0):
seen_start_ids.append(start_after_id)
return iter([[
{'id': 50, 'source_song_id': 10, 'name': '歌', 'audio_url': 'https://example.com/a.mp3', 'singer': '歌手'},
{'id': 60, 'source_song_id': 20, 'name': '歌2', 'audio_url': 'https://example.com/b.mp3', 'singer': '歌手2'},
]])
monkeypatch.setattr(runner, 'iter_hk_songs_batches', fake_batches)
monkeypatch.setattr(runner, 'fetch_platform_records', lambda conn, song_ids: [])
runner.run(['1', '2', '4'], resume=True, state_file=state_file, state_key='all')
assert seen_start_ids == [40]
assert pg_conn.commits == 1
assert '"last_hk_songs_id": 60' in state_file.read_text(encoding='utf-8')
......