fix(data): 成份股 merge/migrate pipeline 可重跑(idempotent, 读 bs_index_constituent_old)

This commit is contained in:
2026-07-23 09:04:00 +08:00
parent ae3a768091
commit 48c67d05d2
3 changed files with 442 additions and 106 deletions
+60 -22
View File
@@ -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)
+138 -84
View File
@@ -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)
@@ -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"]