From 77876c7ee4338a46b1e213dd442bc7a8fd75283a Mon Sep 17 00:00:00 2001
From: shy3130 <415333856@qq.com>
Date: Wed, 15 Jul 2026 21:34:44 +0800
Subject: [PATCH] fix(data): avoid destructive enriched rebuild cleanup

---
 backend/app/indicators/pipeline.py          | 17 +++--
 backend/tests/test_enriched_full_rebuild.py | 80 +++++++++++++++++++++
 2 files changed, 92 insertions(+), 5 deletions(-)
 create mode 100644 backend/tests/test_enriched_full_rebuild.py

diff --git a/backend/app/indicators/pipeline.py b/backend/app/indicators/pipeline.py
index 7ba70f1..c1606d0 100644
--- a/backend/app/indicators/pipeline.py
+++ b/backend/app/indicators/pipeline.py
@@ -1097,8 +1097,7 @@ def run_pipeline(data_dir: Path | None = None,
     SYM_BATCH = prefs_mod.get_enriched_batch_size()  # 每批 N 只 × ~244 天, 可在设置中调整
     total_batches = (total_syms + SYM_BATCH - 1) // SYM_BATCH
 
-    # 全量模式: 先清理旧 enriched 目录, 最后一次性按日期写入
-    # 收集所有批次结果, 按日期分区写入
+    # 全量模式: 收集所有批次结果, 最后按日期分区覆盖写入
     from collections import defaultdict
     date_buffers: dict[str, list[pl.DataFrame]] = defaultdict(list)
 
@@ -1165,9 +1164,17 @@ def run_pipeline(data_dir: Path | None = None,
 
     # 全量模式: 按日期分区写入
     if not symbols and date_buffers:
-        if base.exists():
-            import shutil
-            shutil.rmtree(base)
+        existing_dates = {
+            p.name.removeprefix("date=")
+            for p in base.glob("date=*")
+            if p.is_dir()
+        }
+        rebuilt_dates = set(date_buffers)
+        missing_dates = existing_dates - rebuilt_dates
+        if missing_dates:
+            sample = ", ".join(sorted(missing_dates)[:5])
+            raise RuntimeError(f"全量重建结果缺少已有日期分区,拒绝覆盖: {sample}")
+
         base.mkdir(parents=True, exist_ok=True)
 
         for ds, dfs in date_buffers.items():
diff --git a/backend/tests/test_enriched_full_rebuild.py b/backend/tests/test_enriched_full_rebuild.py
new file mode 100644
index 0000000..913a15d
--- /dev/null
+++ b/backend/tests/test_enriched_full_rebuild.py
@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from datetime import date
+
+import polars as pl
+import pytest
+
+from app.indicators import pipeline
+
+
+def _write_daily(data_dir, ds: str, close: float) -> None:
+    out = data_dir / "kline_daily" / f"date={ds}" / "part.parquet"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    pl.DataFrame({
+        "symbol": ["600000.SH"],
+        "date": [date.fromisoformat(ds)],
+        "open": [close],
+        "high": [close],
+        "low": [close],
+        "close": [close],
+        "volume": [100.0],
+        "amount": [1000.0],
+        "quote_ts": [0],
+    }).write_parquet(out)
+
+
+def _write_existing(data_dir, ds: str, close: float) -> None:
+    out = data_dir / "kline_daily_enriched" / f"date={ds}" / "part.parquet"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    pl.DataFrame({
+        "symbol": ["600000.SH"],
+        "date": [date.fromisoformat(ds)],
+        "close": [close],
+    }).write_parquet(out)
+
+
+def _fake_compute_enriched(raw: pl.DataFrame, **_kwargs) -> pl.DataFrame:
+    return raw.with_columns(
+        pl.col("close").alias("raw_close"),
+        pl.col("high").alias("raw_high"),
+        pl.col("low").alias("raw_low"),
+        pl.lit(None, dtype=pl.Float64).alias("turnover_rate"),
+        pl.lit(0, dtype=pl.UInt32).alias("consecutive_limit_ups"),
+        pl.lit(0, dtype=pl.UInt32).alias("consecutive_limit_downs"),
+    )
+
+
+def test_full_rebuild_overwrites_existing_partitions_without_deleting_base(tmp_path, monkeypatch):
+    _write_daily(tmp_path, "2026-07-14", 14.0)
+    _write_daily(tmp_path, "2026-07-15", 15.0)
+    _write_existing(tmp_path, "2026-07-15", 1.0)
+    marker = tmp_path / "kline_daily_enriched" / "keep.txt"
+    marker.write_text("keep", encoding="utf-8")
+    monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
+
+    written = pipeline.run_pipeline(data_dir=tmp_path)
+
+    assert written == 2
+    assert marker.read_text(encoding="utf-8") == "keep"
+    assert pl.read_parquet(
+        tmp_path / "kline_daily_enriched" / "date=2026-07-14" / "part.parquet"
+    )["close"].to_list() == [14.0]
+    assert pl.read_parquet(
+        tmp_path / "kline_daily_enriched" / "date=2026-07-15" / "part.parquet"
+    )["close"].to_list() == [15.0]
+
+
+def test_full_rebuild_rejects_missing_existing_dates_before_writing(tmp_path, monkeypatch):
+    _write_daily(tmp_path, "2026-07-15", 15.0)
+    _write_existing(tmp_path, "2026-07-14", 14.0)
+    _write_existing(tmp_path, "2026-07-15", 1.0)
+    monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
+
+    with pytest.raises(RuntimeError, match="缺少已有日期分区"):
+        pipeline.run_pipeline(data_dir=tmp_path)
+
+    existing = pl.read_parquet(
+        tmp_path / "kline_daily_enriched" / "date=2026-07-15" / "part.parquet"
+    )
+    assert existing["close"].to_list() == [1.0]