fix(index-sync): 指数/ETF 日K同步支持自定义数据源 (#502)

sync_and_persist_index_daily / sync_and_persist_etf_daily 原经 kline_sync.sync_daily_batch 硬编码走 TickFlow, 无视日K数据源偏好。现与个股日K同一路由语义: 声明 daily 数据集的自定义源接管 (按契约传 asset_type=index/etf), 源空手而归 (fuyao/tdxquant/amber 仅服务股票) 告警回退 TickFlow; 路由判定移至 TickFlow 能力检查之前, 纯自定义源部署也能同步指数/ETF。落盘逻辑复用 TickFlow 路径的帧处理, 不另起一套。
2026-11-02 10:26:13 +00:00 · 2026-10-21 18:05:45 +08:00
parent 78d8d54282
commit 24ce7719e2
2 changed files with 352 additions and 30 deletions
@@ -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
@@ -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