diff --git a/backend/app/services/index_sync.py b/backend/app/services/index_sync.py index de1c3a1c..ff78e647 100644 --- a/backend/app/services/index_sync.py +++ b/backend/app/services/index_sync.py @@ -189,6 +189,79 @@ def sync_etf_instruments(repo: KlineRepository) -> int: return etf_df.height +def _iter_custom_daily_frames(provider, symbols, asset_type, start_time, end_time, on_chunk_done): + """自定义源日K帧迭代器: iter_daily 流式优先, 仅实现 get_daily 的源包装为单帧。 + + asset_type 按内部契约传 "index"/"etf"; 只服务股票的源 (fuyao/tdxquant/amber + 对 asset_type≠stock 一律空手而归) 由此自然表达「不支持」, 由调用方回退。""" + iter_daily = getattr(provider, "iter_daily", None) + if callable(iter_daily): + yield from iter_daily( + symbols, start_time=start_time, end_time=end_time, + asset_type=asset_type, on_chunk_done=on_chunk_done, + ) + return + df = provider.get_daily( + symbols, start_time=start_time, end_time=end_time, + asset_type=asset_type, on_chunk_done=on_chunk_done, + ) + if not df.is_empty(): + yield df + + +def _sync_custom_daily( + repo: KlineRepository, + symbols: list[str], + *, + asset_type: str, + asset_label: str, + start_time: datetime, + end_time: datetime, + on_chunk_done: Callable[[int, int], None] | None, + persist_frame: Callable[[pl.DataFrame], int], +) -> int: + """日K源为自定义 provider 时接管指数/ETF 日K同步 (#502)。 + + 与个股 sync_and_persist_daily_batch 同一路由语义, 路由判定在 TickFlow 能力 + 检查之前 —— 纯自定义源部署 (未配置 TickFlow) 也能同步指数/ETF。 + + 返回 -1 表示「本次不由自定义源接管」, 调用方继续 TickFlow 原路径: + - 日K源是 tickflow, 或该源未声明 daily 数据集; + - 声明了但对该资产类型空手而归 —— 插件契约本就如此 (tdxquant 注释: + "asset_type 仅接受 stock, 其余空手而归(上层回退)"), 告警后回退, + 与修复前行为一致, 不静默吞掉。 + + persist_frame 接收一帧原始日K并完成 append_*_daily + enriched 落盘, + 返回行数 —— 指数/ETF 各自复用与 TickFlow 路径相同的落盘逻辑, 不另起一套。 + """ + provider_name = preferences.get_daily_data_provider() + if provider_name == "tickflow": + return -1 + from app.data_providers import custom as custom_sources + if not custom_sources.provider_has_dataset(provider_name, "daily"): + return -1 + + provider = custom_sources.get_provider(provider_name) + total_rows = 0 + for raw in _iter_custom_daily_frames( + provider, symbols, asset_type, start_time, end_time, on_chunk_done, + ): + if raw is None or raw.is_empty(): + continue + total_rows += persist_frame(raw) + logger.info("%s daily synced (custom %s): +%d rows", asset_label, provider_name, raw.height) + del raw + gc.collect() + if total_rows > 0: + repo.refresh_index_views() + return total_rows + logger.warning( + "自定义源 %s 未返回%s日K数据 (源仅服务股票时属预期), 本轮回退 TickFlow", + provider_name, asset_label, + ) + return -1 + + def sync_and_persist_index_daily( repo: KlineRepository, capset: CapabilitySet, @@ -204,10 +277,6 @@ def sync_and_persist_index_daily( 否则取 index_instruments 表全量(指数+ETF 合并存储)。 on_chunk_done(current, total) 每个批次完成后回调。 """ - if not capset.has(Cap.KLINE_DAILY_BATCH): - kline_sync._log_cap_unavailable("daily", "指数/ETF 日K同步") - return 0 - if symbols_override: symbols = sorted(set(s for s in symbols_override if s)) if not symbols: @@ -222,12 +291,35 @@ def sync_and_persist_index_daily( if instruments.is_empty() or "symbol" not in instruments.columns: return 0 symbols = sorted(set(instruments["symbol"].to_list())) - limit = resolve_limit(capset, Cap.KLINE_DAILY_BATCH) - batch_size = min_batch(preferences.get_index_daily_batch_size(), limit) end_time = end_date or datetime.now() start_time = start_date or (end_time - timedelta(days=365)) + def _persist_index_frame(raw: pl.DataFrame) -> int: + repo.append_index_daily(raw) + enriched = compute_enriched(raw, factors=None, instruments=None) + repo.append_index_enriched(enriched) + del enriched + return raw.height + + # 自定义日K源优先接管指数日K (#502): 路由判定在 TickFlow 能力检查之前, + # 纯自定义源部署 (未配置 TickFlow) 也能同步; 空手而归时回退下方 TickFlow 路径。 + custom_rows = _sync_custom_daily( + repo, symbols, + asset_type="index", asset_label="指数", + start_time=start_time, end_time=end_time, + on_chunk_done=on_chunk_done, + persist_frame=_persist_index_frame, + ) + if custom_rows >= 0: + return custom_rows + + if not capset.has(Cap.KLINE_DAILY_BATCH): + kline_sync._log_cap_unavailable("daily", "指数/ETF 日K同步") + return 0 + limit = resolve_limit(capset, Cap.KLINE_DAILY_BATCH) + batch_size = min_batch(preferences.get_index_daily_batch_size(), limit) + total_rows = 0 chunks = chunked(symbols, batch_size) for i, chunk in enumerate(chunks): @@ -242,14 +334,11 @@ def sync_and_persist_index_daily( if raw.is_empty(): continue - repo.append_index_daily(raw) - enriched = compute_enriched(raw, factors=None, instruments=None) - repo.append_index_enriched(enriched) - total_rows += raw.height + total_rows += _persist_index_frame(raw) logger.info("index/etf daily synced: %d/%d chunks, +%d rows", i + 1, len(chunks), raw.height) if on_chunk_done: on_chunk_done(i + 1, len(chunks)) - del raw, enriched + del raw gc.collect() repo.refresh_index_views() return total_rows @@ -340,10 +429,6 @@ def sync_and_persist_etf_daily( """同步 ETF 日K到独立 kline_etf_* parquet,并计算 ETF enriched。 on_chunk_done(current, total) 每个批次完成后回调。 """ - if not capset.has(Cap.KLINE_DAILY_BATCH): - kline_sync._log_cap_unavailable("daily", "ETF 日K同步") - return 0 - if symbols_override: symbols = sorted(set(s for s in symbols_override if s)) else: @@ -357,15 +442,47 @@ def sync_and_persist_etf_daily( if not symbols: return 0 - limit = resolve_limit(capset, Cap.KLINE_DAILY_BATCH) - batch_size = min_batch(preferences.get_index_daily_batch_size(), limit) - end_time = end_date or datetime.now() start_time = start_date or (end_time - timedelta(days=365)) + factors = _load_etf_factors(repo) + + def _persist_etf_frame(raw: pl.DataFrame) -> int: + frame_symbols = raw["symbol"].unique().to_list() + repo.append_etf_daily(raw) + batch_factors = ( + factors.filter(pl.col("symbol").is_in(frame_symbols)) + if not factors.is_empty() else factors + ) + # 前复权 ratio = cum/total, 新除权事件会改写全部历史价。只用本次拉取 + # 窗口算 enriched 时, 拆分日前的分区停在未复权价, 日K 留下跳空。 + local = _load_local_etf_daily(repo, frame_symbols) + hist = local if not local.is_empty() else raw + # ETF 使用复权和通用技术指标; 不传 instruments, 避免套用 A股涨跌停/连板逻辑。 + enriched = compute_enriched(hist, factors=batch_factors, instruments=None) + repo.append_etf_enriched(enriched) + del enriched, local, hist + return raw.height + + # 自定义日K源优先接管 ETF 日K (#502), 语义与上方指数分支一致 + custom_rows = _sync_custom_daily( + repo, symbols, + asset_type="etf", asset_label="ETF", + start_time=start_time, end_time=end_time, + on_chunk_done=on_chunk_done, + persist_frame=_persist_etf_frame, + ) + if custom_rows >= 0: + return custom_rows + + if not capset.has(Cap.KLINE_DAILY_BATCH): + kline_sync._log_cap_unavailable("daily", "ETF 日K同步") + return 0 + limit = resolve_limit(capset, Cap.KLINE_DAILY_BATCH) + batch_size = min_batch(preferences.get_index_daily_batch_size(), limit) + total_rows = 0 chunks = chunked(symbols, batch_size) - factors = _load_etf_factors(repo) for i, chunk in enumerate(chunks): sleep_between_batches(i, limit.rpm) raw = kline_sync.sync_daily_batch( @@ -378,20 +495,11 @@ def sync_and_persist_etf_daily( if raw.is_empty(): continue - repo.append_etf_daily(raw) - batch_factors = factors.filter(pl.col("symbol").is_in(chunk)) if not factors.is_empty() else factors - # 前复权 ratio = cum/total, 新除权事件会改写全部历史价。只用本次拉取 - # 窗口算 enriched 时, 拆分日前的分区停在未复权价, 日K 留下跳空。 - local = _load_local_etf_daily(repo, chunk) - hist = local if not local.is_empty() else raw - # ETF 使用复权和通用技术指标; 不传 instruments, 避免套用 A股涨跌停/连板逻辑。 - enriched = compute_enriched(hist, factors=batch_factors, instruments=None) - repo.append_etf_enriched(enriched) - total_rows += raw.height + total_rows += _persist_etf_frame(raw) logger.info("etf daily synced: %d/%d chunks, +%d rows", i + 1, len(chunks), raw.height) if on_chunk_done: on_chunk_done(i + 1, len(chunks)) - del raw, enriched, local, hist + del raw gc.collect() repo.refresh_index_views() return total_rows diff --git a/backend/tests/test_index_sync_custom_daily.py b/backend/tests/test_index_sync_custom_daily.py new file mode 100644 index 00000000..3598c98f --- /dev/null +++ b/backend/tests/test_index_sync_custom_daily.py @@ -0,0 +1,214 @@ +"""指数/ETF 日K同步的自定义数据源路由 (#502)。 + +修复前 sync_and_persist_index_daily / sync_and_persist_etf_daily 经 +kline_sync.sync_daily_batch 硬编码走 TickFlow, 无视日K数据源偏好; +修复后与个股 sync_and_persist_daily_batch 同一路由语义: +声明了 daily 数据集的自定义源接管 (asset_type=index/etf), 空手而归回退 TickFlow。 +""" +from types import SimpleNamespace + +import polars as pl +import pytest + +from app.services import index_sync, kline_sync, preferences + + +class _Caps: + """TickFlow 能力集替身: has() 恒返回配置值。""" + + def __init__(self, ok: bool) -> None: + self.ok = ok + + def has(self, cap) -> bool: + return self.ok + + +class _FakeProvider: + """自定义源替身: iter_daily 记录入参并产出预置帧; frames 为空模拟「仅服务股票」。""" + + def __init__(self, frames: list[pl.DataFrame] | None = None, use_get_daily: bool = False) -> None: + self.frames = frames or [] + self.use_get_daily = use_get_daily + self.iter_calls: list[dict] = [] + + def iter_daily(self, symbols, start_time=None, end_time=None, asset_type="stock", on_chunk_done=None): + self.iter_calls.append({ + "symbols": list(symbols), "asset_type": asset_type, + "start_time": start_time, "end_time": end_time, + }) + yield from self.frames + + def get_daily(self, symbols, start_time=None, end_time=None, asset_type="stock", on_chunk_done=None): + self.iter_calls.append({ + "symbols": list(symbols), "asset_type": asset_type, + "start_time": start_time, "end_time": end_time, + }) + return pl.concat(self.frames, how="diagonal_relaxed") if self.frames else pl.DataFrame() + + +class _FakeRepo: + def __init__(self, tmp_path) -> None: + self.tmp = tmp_path + self.index_rows: list[pl.DataFrame] = [] + self.index_enriched: list[pl.DataFrame] = [] + self.etf_rows: list[pl.DataFrame] = [] + self.etf_enriched: list[pl.DataFrame] = [] + self.refreshed = 0 + + @property + def store(self): + return SimpleNamespace(data_dir=self.tmp) + + def get_index_instruments(self): + return pl.DataFrame({"symbol": ["000001.SH", "399001.SZ"]}) + + def get_etf_instruments(self): + return pl.DataFrame({"symbol": ["510300.SH"]}) + + def append_index_daily(self, df): + self.index_rows.append(df) + + def append_index_enriched(self, df): + self.index_enriched.append(df) + + def append_etf_daily(self, df): + self.etf_rows.append(df) + + def append_etf_enriched(self, df): + self.etf_enriched.append(df) + + def refresh_index_views(self): + self.refreshed += 1 + + +@pytest.fixture +def harness(tmp_path, monkeypatch): + """路由到 fake_src 并隔离 TickFlow 路径副作用; 返回可配置的测试上下文。""" + monkeypatch.setattr(preferences, "get_daily_data_provider", lambda: "fake_src") + # TickFlow 路径副作用隔离: 限速/批量计算与网络调用全部替身化 + tickflow_calls: list[list[str]] = [] + + def _fake_sync_daily_batch(symbols, **kwargs): + tickflow_calls.append(list(symbols)) + return pl.DataFrame() + + monkeypatch.setattr(kline_sync, "sync_daily_batch", _fake_sync_daily_batch) + monkeypatch.setattr(index_sync, "compute_enriched", lambda df, factors=None, instruments=None: df) + monkeypatch.setattr(index_sync, "resolve_limit", lambda capset, cap, **kw: SimpleNamespace(batch=100, rpm=600)) + monkeypatch.setattr(index_sync, "min_batch", lambda size, limit: 50) + monkeypatch.setattr(index_sync, "sleep_between_batches", lambda i, rpm: None) + return SimpleNamespace(tmp_path=tmp_path, monkeypatch=monkeypatch, tickflow_calls=tickflow_calls) + + +def _route(harness, provider): + from app.data_providers import custom as custom_sources + + harness.monkeypatch.setattr(custom_sources, "provider_has_dataset", lambda name, ds: name == "fake_src" and ds == "daily") + harness.monkeypatch.setattr(custom_sources, "get_provider", lambda name: provider) + + +def _index_df() -> pl.DataFrame: + return pl.DataFrame({ + "symbol": ["000001.SH", "399001.SZ"], + "date": ["2026-10-20", "2026-10-20"], + "close": [3813.79, 13118.62], + }) + + +def _etf_df() -> pl.DataFrame: + return pl.DataFrame({"symbol": ["510300.SH"], "date": ["2026-10-20"], "close": [4.236]}) + + +# ---------------------------------------------------------------- 指数日K +def test_custom_provider_takes_over_index_daily(harness): + provider = _FakeProvider(frames=[_index_df()]) + _route(harness, provider) + repo = _FakeRepo(harness.tmp_path) + + # TickFlow 无 KLINE_DAILY_BATCH 能力 (纯自定义源部署) 也要能同步成功 + rows = index_sync.sync_and_persist_index_daily(repo, _Caps(ok=False)) + + assert rows == 2 + assert provider.iter_calls[0]["asset_type"] == "index" + assert provider.iter_calls[0]["symbols"] == ["000001.SH", "399001.SZ"] + assert repo.index_rows and repo.index_rows[0].height == 2 + assert repo.index_enriched and repo.index_enriched[0].height == 2 + assert repo.refreshed == 1 + assert harness.tickflow_calls == [] # 不再走 TickFlow + + +def test_custom_index_daily_via_get_daily_only_provider(harness): + """仅实现 get_daily (无 iter_daily) 的源: 包装为单帧, 同样接管。""" + provider = _FakeProvider(frames=[_index_df()], use_get_daily=True) + _route(harness, provider) + provider.iter_daily = None # 摘掉 iter_daily, 强制走 get_daily 分支 + repo = _FakeRepo(harness.tmp_path) + + rows = index_sync.sync_and_persist_index_daily(repo, _Caps(ok=True)) + + assert rows == 2 + assert provider.iter_calls[0]["asset_type"] == "index" + assert harness.tickflow_calls == [] + + +def test_custom_provider_stock_only_falls_back_to_tickflow(harness): + """源对 asset_type=index 空手而归 (fuyao/tdxquant/amber 均如此) → 回退 TickFlow。""" + provider = _FakeProvider(frames=[]) + _route(harness, provider) + repo = _FakeRepo(harness.tmp_path) + + rows = index_sync.sync_and_persist_index_daily(repo, _Caps(ok=True)) + + assert rows == 0 + assert repo.index_rows == [] + assert len(harness.tickflow_calls) >= 1 # 回退到 TickFlow 路径 + assert repo.refreshed == 1 + + +def test_provider_without_daily_dataset_falls_back(harness): + from app.data_providers import custom as custom_sources + + harness.monkeypatch.setattr(custom_sources, "provider_has_dataset", lambda name, ds: False) + repo = _FakeRepo(harness.tmp_path) + + rows = index_sync.sync_and_persist_index_daily(repo, _Caps(ok=True)) + + assert rows == 0 + assert len(harness.tickflow_calls) >= 1 + + +def test_tickflow_provider_keeps_original_path(harness): + harness.monkeypatch.setattr(preferences, "get_daily_data_provider", lambda: "tickflow") + repo = _FakeRepo(harness.tmp_path) + + rows = index_sync.sync_and_persist_index_daily(repo, _Caps(ok=True)) + + assert rows == 0 + assert len(harness.tickflow_calls) >= 1 + assert repo.index_rows == [] + + +# ---------------------------------------------------------------- ETF 日K +def test_custom_provider_takes_over_etf_daily(harness): + provider = _FakeProvider(frames=[_etf_df()]) + _route(harness, provider) + repo = _FakeRepo(harness.tmp_path) + + rows = index_sync.sync_and_persist_etf_daily(repo, _Caps(ok=False)) + + assert rows == 1 + assert provider.iter_calls[0]["asset_type"] == "etf" + assert repo.etf_rows and repo.etf_rows[0].height == 1 + assert repo.etf_enriched and repo.etf_enriched[0].height == 1 + assert harness.tickflow_calls == [] + + +def test_custom_etf_stock_only_falls_back_to_tickflow(harness): + provider = _FakeProvider(frames=[]) + _route(harness, provider) + repo = _FakeRepo(harness.tmp_path) + + rows = index_sync.sync_and_persist_etf_daily(repo, _Caps(ok=True)) + + assert rows == 0 + assert len(harness.tickflow_calls) >= 1