From 48c67d05d257f51c9b4530701e03c22664a2c1b8 Mon Sep 17 00:00:00 2001 From: claude_dev Date: Thu, 23 Jul 2026 09:04:00 +0800 Subject: [PATCH] =?UTF-8?q?fix(data):=20=E6=88=90=E4=BB=BD=E8=82=A1=20merg?= =?UTF-8?q?e/migrate=20pipeline=20=E5=8F=AF=E9=87=8D=E8=B7=91(idempotent,?= =?UTF-8?q?=20=E8=AF=BB=20bs=5Findex=5Fconstituent=5Fold)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/data_platform/merge_constituent.py | 82 ++++-- scripts/data_platform/migrate_constituent.py | 222 ++++++++++------ .../portfolio/test_merge_constituent_rerun.py | 244 ++++++++++++++++++ 3 files changed, 442 insertions(+), 106 deletions(-) create mode 100644 tests/portfolio/test_merge_constituent_rerun.py diff --git a/scripts/data_platform/merge_constituent.py b/scripts/data_platform/merge_constituent.py index 467b6bc..bd55152 100644 --- a/scripts/data_platform/merge_constituent.py +++ b/scripts/data_platform/merge_constituent.py @@ -1,32 +1,70 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -"""merge_constituent.py — 单元3 合并: staging -> constituent_unified, bs_index_constituent -> _old。 +"""merge_constituent.py — 单元3 合并: staging -> constituent_unified (幂等版)。 -前提: migrate_constituent.py 已建 constituent_unified_staging (验证通过)。 -- constituent_unified_staging RENAME constituent_unified (正式表) -- bs_index_constituent RENAME bs_index_constituent_old (保留 988 时点精度, 不删) -- 建索引 (index_code, code) 加速 get_index_stocks 查询 -回滚: rename 反向 (constituent_unified->staging, _old->bs_index_constituent) +设计: +- 前置: ``migrate_constituent.py`` 已建好 ``constituent_unified_staging``(全集型, 验证通过) +- 本脚本把 staging 内容复制成正式表 ``constituent_unified``(治幸存者偏差的选股池) + +幂等(idempotent, 可重跑, 月度 schtask 安全): +- ``DROP TABLE IF EXISTS constituent_unified`` (旧正式表, 有则删) +- ``CREATE TABLE constituent_unified AS SELECT ... FROM constituent_unified_staging`` +- ``CREATE INDEX IF NOT EXISTS idx_constituent_unified`` (index_code, code 加速 get_index_stocks) +- **绝不用 RENAME**(RENAME 只能跑一次, 再跑必崩 —— 早期方案A首次落地用过, 现已废弃) +- **绝不碰 ``bs_index_constituent_old``**(那是 migrate 的事, 这里只管 staging→unified) +- staging 表保留(migrate 下次跑会 DROP+rebuild, 这里不动) + +环境变量: +- ``SANGUO_DB``: quant_trading.db 路径, 默认 ``C:\\sanguo_vnpy_v2\\data\\quant_trading.db``(VPS) """ +import os import sqlite3 -DB = r"C:\sanguo_vnpy_v2\data\quant_trading.db" -c = sqlite3.connect(DB, timeout=60) -c.execute("PRAGMA busy_timeout = 60000") -c.execute("ALTER TABLE constituent_unified_staging RENAME TO constituent_unified") -print("renamed constituent_unified_staging -> constituent_unified") +def merge(db_path: str) -> None: + """从 constituent_unified_staging (幂等) 重建 constituent_unified 正式表。 -c.execute("ALTER TABLE bs_index_constituent RENAME TO bs_index_constituent_old") -print("renamed bs_index_constituent -> bs_index_constituent_old (时点精度保留)") + Args: + db_path: quant_trading.db 路径(测试可传 tmp sqlite; 生产读 ``SANGUO_DB``)。 -c.execute("CREATE INDEX IF NOT EXISTS idx_constituent_unified " - "ON constituent_unified(index_code, code)") -c.commit() + Raises: + sqlite3.OperationalError: 若 ``constituent_unified_staging`` 不存在 + (前置 migrate 未跑; 提示先跑 migrate_constituent.py)。 + """ + c = sqlite3.connect(db_path, timeout=60) + try: + c.execute("PRAGMA busy_timeout = 60000") -print("constituent_unified rows:", c.execute( - "SELECT COUNT(*) FROM constituent_unified").fetchone()[0]) -print("bs_index_constituent_old rows:", c.execute( - "SELECT COUNT(*) FROM bs_index_constituent_old").fetchone()[0]) -c.close() -print("MERGE DONE") + # 幂等: 先删旧正式表(若存在), 再从 staging CREATE 一张全新的 + # (DROP+CREATE 语义, 不是 RENAME —— 可无限次重跑) + c.execute("DROP TABLE IF EXISTS constituent_unified") + c.execute( + "CREATE TABLE constituent_unified AS " + "SELECT index_code, code, code_name, source, in_current, was_removed " + "FROM constituent_unified_staging" + ) + print("rebuilt constituent_unified from constituent_unified_staging") + + c.execute( + "CREATE INDEX IF NOT EXISTS idx_constituent_unified " + "ON constituent_unified(index_code, code)" + ) + c.commit() + + n_unified = c.execute( + "SELECT COUNT(*) FROM constituent_unified" + ).fetchone()[0] + n_staging = c.execute( + "SELECT COUNT(*) FROM constituent_unified_staging" + ).fetchone()[0] + print(f"constituent_unified rows: {n_unified} (staging: {n_staging})") + finally: + c.close() + print("MERGE DONE") + + +if __name__ == "__main__": + _db = os.environ.get( + "SANGUO_DB", r"C:\sanguo_vnpy_v2\data\quant_trading.db" + ) + merge(_db) diff --git a/scripts/data_platform/migrate_constituent.py b/scripts/data_platform/migrate_constituent.py index 87693b0..da8dc29 100644 --- a/scripts/data_platform/migrate_constituent.py +++ b/scripts/data_platform/migrate_constituent.py @@ -1,17 +1,25 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -"""migrate_constituent.py — 单元3: 合并成份股 -> constituent_unified_staging (全集型)。 +"""migrate_constituent.py — 单元3: 合并成份股 -> constituent_unified_staging (全集型, 幂等)。 设计 (spec §14, 治幸存者偏差选股池): -- baostock (bs_index_constituent 988 时点) -> 聚合成全集 (hs300/zz500/sz50 -> 000300/000905/000016), - in_current=最后时点成份, was_removed=历史入选过但已踢出 +- baostock (``bs_index_constituent_old`` 988 时点全集, 方案A 后的权威历史源) + -> 聚合成全集 (hs300/zz500/sz50 -> 000300/000905/000016), + in_current=最后时点成份, was_removed=历史入选过但已踢出 - 深证/国证 _union.parquet (399001/399006/399005/399330) -> 直入 (akshare cni) - 中证1000/2000 _snapshot.parquet (000852/932000) -> 当前 (akshare csindex) - 新浪 _sina.parquet -> 丢弃 (baostock 300/500/50 已权威) - code 统一 6 位无前缀 (sh.600000 / 600519.SH -> 600519) -输出 constituent_unified_staging(index_code, code, code_name, source, in_current, was_removed)。 -验证 OK 后手动 rename: staging->constituent_unified, bs_index_constituent->_old。 +幂等(可重跑, 月度 schtask 安全): +- ``DROP TABLE IF EXISTS constituent_unified_staging`` + ``CREATE TABLE ...`` 每次重建 +- baostock 源表读取**自动适配**: + * 方案A 后正常只有 ``bs_index_constituent_old``(权威全集) + * 若将来重建了 live ``bs_index_constituent``(新时点增量), UNION ALL 两者去重, 兼容两种状态 + +环境变量: +- ``SANGUO_DB``: quant_trading.db 路径, 默认 ``C:\\sanguo_vnpy_v2\\data\\quant_trading.db``(VPS) +- ``HIST``(模块常量, 测试可 monkeypatch): 深证/中证 parquet 目录 """ import glob import os @@ -20,7 +28,6 @@ import sqlite3 import pandas as pd -DB = r"C:\sanguo_vnpy_v2\data\quant_trading.db" HIST = r"C:\sanguo_vnpy_v2\data\index_const_hist" STAGING = "constituent_unified_staging" BS_MAP = {"hs300": "000300", "zz500": "000905", "sz50": "000016"} @@ -32,91 +39,138 @@ def norm_code(code): return s.zfill(6) if s.isdigit() and len(s) <= 6 else s -def main(): - c = sqlite3.connect(DB, timeout=60) - c.execute("PRAGMA busy_timeout = 60000") +def _read_baostock_constituent(c: sqlite3.Connection) -> pd.DataFrame: + """读 baostock 成份股历史, 自动适配方案A后(_old) / 未来(live)两种状态。 - # 1. baostock -> 全集 - df_bs = pd.read_sql( - "SELECT updateDate, index_code, code, code_name FROM bs_index_constituent", c) - df_bs["index_code"] = df_bs["index_code"].map(BS_MAP) - df_bs["code"] = df_bs["code"].apply(norm_code) - last_sets = {} - for idx, grp in df_bs.groupby("index_code"): - last_d = grp["updateDate"].max() - last_sets[idx] = set(grp[grp["updateDate"] == last_d]["code"]) - pool = (df_bs.groupby(["index_code", "code"])["code_name"] - .first().reset_index()) - pool["in_current"] = pool.apply( - lambda r: r["code"] in last_sets.get(r["index_code"], set()), axis=1) - pool["was_removed"] = ~pool["in_current"] - pool["source"] = "baostock" - print(f"[baostock] pool rows={len(pool)} (300/500/50 全集)") + - 只有 ``bs_index_constituent_old``: 读它(方案A 后常态) + - 只有 ``bs_index_constituent``: 读它(未来重建 live 表) + - 两者都在: UNION ALL 后 drop_duplicates(兼容过渡期) + - 都没有: 返回空 DataFrame(不崩, 由上层决定是否报错) - # 2. 深证 union - deep = [] - for f in sorted(glob.glob(os.path.join(HIST, "*_union.parquet"))): - d = pd.read_parquet(f)[["index_code", "code", "code_name", - "in_current", "was_removed"]] - d["source"] = "akshare_cni" - deep.append(d) - df_deep = pd.concat(deep, ignore_index=True) if deep else pd.DataFrame( - columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"]) - df_deep["code"] = df_deep["code"].apply(norm_code) - print(f"[深证 union] rows={len(df_deep)}") + 返回字段: ``updateDate, index_code, code, code_name``。 + """ + cur = c.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND " + "name IN ('bs_index_constituent_old', 'bs_index_constituent')" + ) + tables = {row[0] for row in cur.fetchall()} - # 3. 中证 snapshot - snap = [] - for f in [os.path.join(HIST, "000852_snapshot.parquet"), - os.path.join(HIST, "932000_snapshot.parquet")]: - if os.path.exists(f): - d = pd.read_parquet(f)[["index_code", "code", "code_name"]] - d["in_current"] = True - d["was_removed"] = False - d["source"] = "akshare_csindex" - snap.append(d) - df_snap = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame( - columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"]) - df_snap["code"] = df_snap["code"].apply(norm_code) - print(f"[中证 snapshot] rows={len(df_snap)}") + parts = [] + if "bs_index_constituent_old" in tables: + parts.append(pd.read_sql( + "SELECT updateDate, index_code, code, code_name " + "FROM bs_index_constituent_old", c, + )) + print("[baostock] read from bs_index_constituent_old (方案A 权威历史全集)") + if "bs_index_constituent" in tables: + parts.append(pd.read_sql( + "SELECT updateDate, index_code, code, code_name " + "FROM bs_index_constituent", c, + )) + print("[baostock] read from bs_index_constituent (live 增量)") - # 合并 + 去重 (同 index+code+source) - all_df = pd.concat([pool, df_deep, df_snap], ignore_index=True) - all_df = all_df.drop_duplicates(["index_code", "code", "source"]) - print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, " - f"{all_df['index_code'].nunique()} indices") + if not parts: + print("[baostock] WARN: 既无 _old 也无 live 表, baostock 段产出 0 行") + return pd.DataFrame(columns=["updateDate", "index_code", "code", "code_name"]) - print("\n各指数分布:") - print(all_df.groupby("index_code").agg( - n=("code", "count"), src=("source", "first"), - in_cur=("in_current", "sum"), removed=("was_removed", "sum"))) + df = pd.concat(parts, ignore_index=True).drop_duplicates() + return df - # 写 staging - c.execute(f"DROP TABLE IF EXISTS {STAGING}") - c.execute(f"""CREATE TABLE {STAGING} ( - index_code TEXT, code TEXT, code_name TEXT, source TEXT, - in_current INTEGER, was_removed INTEGER)""") - work = all_df[["index_code", "code", "code_name", "source", - "in_current", "was_removed"]].copy() - work["in_current"] = work["in_current"].astype(int) - work["was_removed"] = work["was_removed"].astype(int) - c.executemany(f"INSERT INTO {STAGING} VALUES (?,?,?,?,?,?)", - work.itertuples(index=False, name=None)) - c.commit() - n = c.execute(f"SELECT COUNT(*) FROM {STAGING}").fetchone()[0] - # 抽样验证 - print(f"\n[staging] {STAGING}: {n} rows") - print("sample 300:", c.execute( - "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM " - f"{STAGING} WHERE index_code='000300'").fetchone()) - print("sample 399001:", c.execute( - "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM " - f"{STAGING} WHERE index_code='399001'").fetchone()) - print("sample 000852:", c.execute( - "SELECT COUNT(*) FROM " f"{STAGING} WHERE index_code='000852'").fetchone()) - c.close() + +def migrate(db_path: str) -> None: + """(幂等)重建 constituent_unified_staging。 + + Args: + db_path: quant_trading.db 路径(测试可传 tmp sqlite; 生产读 ``SANGUO_DB``)。 + """ + c = sqlite3.connect(db_path, timeout=60) + try: + c.execute("PRAGMA busy_timeout = 60000") + + # 1. baostock -> 全集(自动适配 _old / live / 两者皆在) + df_bs = _read_baostock_constituent(c) + df_bs["index_code"] = df_bs["index_code"].map(BS_MAP) + df_bs["code"] = df_bs["code"].apply(norm_code) + last_sets = {} + for idx, grp in df_bs.groupby("index_code"): + last_d = grp["updateDate"].max() + last_sets[idx] = set(grp[grp["updateDate"] == last_d]["code"]) + pool = (df_bs.groupby(["index_code", "code"])["code_name"] + .first().reset_index()) + pool["in_current"] = pool.apply( + lambda r: r["code"] in last_sets.get(r["index_code"], set()), axis=1) + pool["was_removed"] = ~pool["in_current"] + pool["source"] = "baostock" + print(f"[baostock] pool rows={len(pool)} (300/500/50 全集)") + + # 2. 深证 union + deep = [] + for f in sorted(glob.glob(os.path.join(HIST, "*_union.parquet"))): + d = pd.read_parquet(f)[["index_code", "code", "code_name", + "in_current", "was_removed"]] + d["source"] = "akshare_cni" + deep.append(d) + df_deep = pd.concat(deep, ignore_index=True) if deep else pd.DataFrame( + columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"]) + df_deep["code"] = df_deep["code"].apply(norm_code) + print(f"[深证 union] rows={len(df_deep)}") + + # 3. 中证 snapshot + snap = [] + for f in [os.path.join(HIST, "000852_snapshot.parquet"), + os.path.join(HIST, "932000_snapshot.parquet")]: + if os.path.exists(f): + d = pd.read_parquet(f)[["index_code", "code", "code_name"]] + d["in_current"] = True + d["was_removed"] = False + d["source"] = "akshare_csindex" + snap.append(d) + df_snap = pd.concat(snap, ignore_index=True) if snap else pd.DataFrame( + columns=["index_code", "code", "code_name", "in_current", "was_removed", "source"]) + df_snap["code"] = df_snap["code"].apply(norm_code) + print(f"[中证 snapshot] rows={len(df_snap)}") + + # 合并 + 去重 (同 index+code+source) + all_df = pd.concat([pool, df_deep, df_snap], ignore_index=True) + all_df = all_df.drop_duplicates(["index_code", "code", "source"]) + print(f"\n[TOTAL] constituent_unified: {len(all_df)} rows, " + f"{all_df['index_code'].nunique()} indices") + + print("\n各指数分布:") + print(all_df.groupby("index_code").agg( + n=("code", "count"), src=("source", "first"), + in_cur=("in_current", "sum"), removed=("was_removed", "sum"))) + + # 写 staging(幂等: DROP+CREATE) + c.execute(f"DROP TABLE IF EXISTS {STAGING}") + c.execute(f"""CREATE TABLE {STAGING} ( + index_code TEXT, code TEXT, code_name TEXT, source TEXT, + in_current INTEGER, was_removed INTEGER)""") + work = all_df[["index_code", "code", "code_name", "source", + "in_current", "was_removed"]].copy() + work["in_current"] = work["in_current"].astype(int) + work["was_removed"] = work["was_removed"].astype(int) + c.executemany(f"INSERT INTO {STAGING} VALUES (?,?,?,?,?,?)", + work.itertuples(index=False, name=None)) + c.commit() + n = c.execute(f"SELECT COUNT(*) FROM {STAGING}").fetchone()[0] + # 抽样验证 + print(f"\n[staging] {STAGING}: {n} rows") + print("sample 300:", c.execute( + "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM " + f"{STAGING} WHERE index_code='000300'").fetchone()) + print("sample 399001:", c.execute( + "SELECT COUNT(*), SUM(in_current), SUM(was_removed) FROM " + f"{STAGING} WHERE index_code='399001'").fetchone()) + print("sample 000852:", c.execute( + "SELECT COUNT(*) FROM " f"{STAGING} WHERE index_code='000852'").fetchone()) + finally: + c.close() print("\nMIGRATE STAGING DONE (未 rename, 验证 OK 后单独合并)") if __name__ == "__main__": - main() + _db = os.environ.get( + "SANGUO_DB", r"C:\sanguo_vnpy_v2\data\quant_trading.db" + ) + migrate(_db) diff --git a/tests/portfolio/test_merge_constituent_rerun.py b/tests/portfolio/test_merge_constituent_rerun.py new file mode 100644 index 0000000..c5324d8 --- /dev/null +++ b/tests/portfolio/test_merge_constituent_rerun.py @@ -0,0 +1,244 @@ +"""成份股 merge / migrate pipeline 可重跑(idempotent)回归测试。 + +背景: 方案A 落地后 VPS 上的状态是 +- ``constituent_unified`` 已存在(正式表) +- ``bs_index_constituent_old`` 已存在(baostock 300/500/50 全部历史时点, 权威历史源) +- ``bs_index_constituent`` 不存在(被 rename 走了) +- ``constituent_unified_staging`` 仅在 migrate 跑完后存在 + +旧 ``merge_constituent.py`` 用 RENAME + 旧 ``migrate_constituent.py`` 直接 +``FROM bs_index_constituent`` → 月度 schtask 再跑会崩。这组测试断言: + +1. ``merge`` 跑两次都不崩, 行数稳定(DROP+CREATE from staging 幂等, 不用 RENAME), + 且不碰 ``bs_index_constituent_old``。 +2. ``migrate`` 跑两次都不崩, 从 ``bs_index_constituent_old`` 读, staging 行数稳定。 + +全部在 Mac 本地用 tmp sqlite + 空 HIST 目录跑, 不依赖 VPS / baostock / akshare。 +""" +from __future__ import annotations + +import os +import sqlite3 +import sys + +import pytest + +# 让测试能 import scripts/data_platform/ 下的模块 +_HERE = os.path.dirname(os.path.abspath(__file__)) +_SCRIPT_DIR = os.path.abspath(os.path.join(_HERE, "..", "..", "scripts", "data_platform")) +if _SCRIPT_DIR not in sys.path: + sys.path.insert(0, _SCRIPT_DIR) + +import merge_constituent as mc # noqa: E402 +import migrate_constituent as mig # noqa: E402 + + +# ======================== helpers ======================== +def _create_staging(db_path: str, rows: list[tuple[str, str, str, str, int, int]]) -> None: + """在 tmp db 里造一张 constituent_unified_staging。 + + rows: (index_code, code, code_name, source, in_current, was_removed) + """ + with sqlite3.connect(db_path) as c: + c.execute("PRAGMA busy_timeout = 60000") + c.execute( + "CREATE TABLE constituent_unified_staging (" + "index_code TEXT, code TEXT, code_name TEXT, source TEXT, " + "in_current INTEGER, was_removed INTEGER)" + ) + c.executemany( + "INSERT INTO constituent_unified_staging VALUES (?,?,?,?,?,?)", rows + ) + c.commit() + + +def _create_bs_old(db_path: str, rows: list[tuple[str, str, str, str]]) -> None: + """在 tmp db 里造一张 bs_index_constituent_old(updateDate, index_code, code, code_name)。 + + 默认 index_code 用 baostock 原始命名(hs300/zz500/sz50), migrate 内有 BS_MAP 映射。 + """ + with sqlite3.connect(db_path) as c: + c.execute("PRAGMA busy_timeout = 60000") + c.execute( + "CREATE TABLE bs_index_constituent_old (" + "updateDate TEXT, index_code TEXT, code TEXT, code_name TEXT)" + ) + c.executemany( + "INSERT INTO bs_index_constituent_old VALUES (?,?,?,?)", rows + ) + c.commit() + + +def _table_exists(db_path: str, table: str) -> bool: + with sqlite3.connect(db_path) as c: + row = c.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,) + ).fetchone() + return row is not None + + +def _count_rows(db_path: str, table: str) -> int: + with sqlite3.connect(db_path) as c: + return c.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + + +# ======================== test 1: merge 幂等 ======================== +def test_merge_idempotent_rerun(tmp_path): + """merge 跑两次不崩, constituent_unified 行数 == staging 行数; _old 不被碰。 + + 断言的是 DROP+CREATE from staging 语义(幂等), 而不是 RENAME(只能跑一次)。 + """ + db = tmp_path / "test.db" + db_path = str(db) + + # staging 5 行 (300 + 500 各几只) + staging_rows = [ + ("000300", "600519", "贵州茅台", "baostock", 1, 0), + ("000300", "601318", "中国平安", "baostock", 1, 0), + ("000300", "000001", "平安银行", "baostock", 0, 1), + ("000905", "002415", "海康威视", "baostock", 1, 0), + ("000905", "300750", "宁德时代", "baostock", 1, 0), + ] + _create_staging(db_path, staging_rows) + + # 模拟方案A后状态: bs_index_constituent_old 已存在(不该被 merge 碰) + _create_bs_old(db_path, [ + ("2024-06-17", "hs300", "sh.600519", "贵州茅台"), + ("2024-06-17", "hs300", "sh.601318", "中国平安"), + ]) + bs_old_rows_before = _count_rows(db_path, "bs_index_constituent_old") + + # Act 1: 第一次 merge + mc.merge(db_path) + + # Assert 1: constituent_unified 存在, 行数 == staging 行数 + assert _table_exists(db_path, "constituent_unified") is True + assert _count_rows(db_path, "constituent_unified") == len(staging_rows) + # staging 仍可读(没被 rename 走) + assert _count_rows(db_path, "constituent_unified_staging") == len(staging_rows) + # bs_old 没被碰 + assert _count_rows(db_path, "bs_index_constituent_old") == bs_old_rows_before + + # Act 2: 第二次 merge —— 关键回归(RENAME 方案这里会崩) + mc.merge(db_path) + + # Assert 2: 行数稳定, 没异常 + assert _count_rows(db_path, "constituent_unified") == len(staging_rows) + assert _count_rows(db_path, "constituent_unified_staging") == len(staging_rows) + assert _count_rows(db_path, "bs_index_constituent_old") == bs_old_rows_before + + +# ======================== test 2: migrate 幂等 ======================== +def test_migrate_idempotent_rerun_from_old(tmp_path, monkeypatch): + """migrate 跑两次不崩, 从 bs_index_constituent_old 读, staging 行数稳定。 + + 构造 baostock 原始 index_code (hs300/zz500/sz50), migrate 内 BS_MAP 映射到 + 000300/000905/000016。HIST 目录用空目录(深证/中证 parquet 缺失, migrate + 应优雅产出 0 行不崩)。 + """ + db = tmp_path / "test.db" + db_path = str(db) + + # baostock 历史时点(2 个日期, 模拟一只股被踢出) + bs_rows = [ + # 2024-06-17 时点: hs300 含 600519 / 601318 + ("2024-06-17", "hs300", "sh.600519", "贵州茅台"), + ("2024-06-17", "hs300", "sh.601318", "中国平安"), + # 2024-12-16 时点: hs300 600519 留下, 601318 被踢, 新进 000001 + ("2024-12-16", "hs300", "sh.600519", "贵州茅台"), + ("2024-12-16", "hs300", "sz.000001", "平安银行"), + # zz500 一只 + ("2024-06-17", "zz500", "sh.600036", "招商银行"), + # sz50 一只 + ("2024-06-17", "sz50", "sh.600000", "浦发银行"), + ] + _create_bs_old(db_path, bs_rows) + + # HIST 指向空目录(深证 union / 中证 snapshot 都缺失, migrate 应优雅跳过) + empty_hist = tmp_path / "index_const_hist" + empty_hist.mkdir() + monkeypatch.setattr(mig, "HIST", str(empty_hist)) + + # Act 1 + mig.migrate(db_path) + + # Assert 1: staging 建好, 行数 == 去重后的 _old 行数 + assert _table_exists(db_path, "constituent_unified_staging") is True + # 去重后 6 个唯一 (index_code, code) 对: + # hs300: 600519 / 601318 / 000001 (3) + # zz500: 600036 (1) + # sz50: 600000 (1) + # 注意 600519 在两个时点出现, groupby first() 去重为 1 行 → 共 5 行 + expected_rows = 5 + assert _count_rows(db_path, "constituent_unified_staging") == expected_rows + + # 内容抽查: hs300 映射对了, was_removed 标记对了 + with sqlite3.connect(db_path) as c: + # hs300 应有 3 行 + n_hs300 = c.execute( + "SELECT COUNT(*) FROM constituent_unified_staging WHERE index_code='000300'" + ).fetchone()[0] + assert n_hs300 == 3 + # 601318 在最后时点已不在 → was_removed=1 + row_601318 = c.execute( + "SELECT in_current, was_removed FROM constituent_unified_staging " + "WHERE index_code='000300' AND code='601318'" + ).fetchone() + assert row_601318 == (0, 1) + # 600519 仍在 → in_current=1 + row_600519 = c.execute( + "SELECT in_current, was_removed FROM constituent_unified_staging " + "WHERE index_code='000300' AND code='600519'" + ).fetchone() + assert row_600519 == (1, 0) + + # Act 2: 第二次 migrate —— 关键回归 + mig.migrate(db_path) + + # Assert 2: staging 行数稳定(DROP+rebuild 幂等) + assert _count_rows(db_path, "constituent_unified_staging") == expected_rows + # _old 没被碰 + assert _count_rows(db_path, "bs_index_constituent_old") == len(bs_rows) + + +def test_migrate_reads_bs_old_not_live_when_both_exist(tmp_path, monkeypatch): + """migrate 兼容性: _old 和 live 表都在时, 两边 UNION 不丢数据(robustness)。 + + 方案A 后正常只有 _old; 但若将来有人重建了 live bs_index_constituent, + migrate 应把两者都读(UNION ALL 去重), 兼容两种状态。 + """ + db = tmp_path / "test.db" + db_path = str(db) + + # _old 里有 hs300 600519 + _create_bs_old(db_path, [ + ("2024-06-17", "hs300", "sh.600519", "贵州茅台"), + ]) + # live 表里补一只新的(模拟未来重建 live 后采到新时点) + with sqlite3.connect(db_path) as c: + c.execute( + "CREATE TABLE bs_index_constituent (" + "updateDate TEXT, index_code TEXT, code TEXT, code_name TEXT)" + ) + c.execute( + "INSERT INTO bs_index_constituent VALUES (?,?,?,?)", + ("2025-06-16", "hs300", "sh.688981", "中芯国际"), + ) + c.commit() + + empty_hist = tmp_path / "index_const_hist" + empty_hist.mkdir() + monkeypatch.setattr(mig, "HIST", str(empty_hist)) + + # Act + mig.migrate(db_path) + + # Assert: staging 应同时含 _old 的 600519 + live 的 688981 (UNION 后去重 2 行) + with sqlite3.connect(db_path) as c: + codes = sorted( + r[0] for r in c.execute( + "SELECT code FROM constituent_unified_staging " + "WHERE index_code='000300'" + ).fetchall() + ) + assert codes == ["600519", "688981"]