FinDataHub/Scripts/bin/verify_unique_keys.py

#!/usr/bin/env python3
"""verify_unique_keys.py — 校验 tables_registry.json 声明的 unique_keys 是否真的唯一

背景
----
`create_missing_tables.py` 依据 registry 的 `unique_keys` 生成 UNIQUE KEY，
`db.upsert` 也依赖同一份声明做幂等写入。但 registry 是**人工维护的声明**，
未必与接口实际返回的数据一致。

实测踩到：`us_basic` 声明唯一键为 `["ts_code"]`，而该接口每只美股返回**多行**
（不同交易所/币种），建表后首次写入即报
    (1062, "Duplicate entry ... for key 'uk_us_basic'")
—— 整个 upsert 批次失败。

本脚本对每张表取样一次，检查声明唯一键在真实数据上是否真的唯一，
并给出「实际可用的唯一键」建议（在原声明基础上追加列直到唯一）。

用法
----
    set -a; . ~/.config/findatahub/env; set +a
    unset HTTP_PROXY HTTPS_PROXY http_proxy https_proxy ALL_PROXY all_proxy
    python3 Scripts/bin/verify_unique_keys.py [--keys fut_basic,us_basic ...]
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
from pathlib import Path

FD_ROOT = Path("/Applications/FinDataHub.app/Contents/Resources/tushare_fetcher")

# 取样入参阶梯：不同接口要求的入参不同，取第一个非空结果
PARAM_LADDER = [
    {},
    {"trade_date": "20260918"},
    {"start_date": "20260901", "end_date": "20260918"},
    {"ts_code": "000001.SZ"},
    {"date": "20260918"},
]


def load_table_map(fd: Path) -> dict[str, str | None]:
    """TABLE_MAP: key → api_name（元组里第 2 个引号串；第 1 个是库表名）。"""
    src = (fd / "src/tushare_fetcher/core/fetcher.py").read_text(encoding="utf-8")
    blk = src[src.index("TABLE_MAP"):]
    out: dict[str, str | None] = {}
    for m in re.finditer(r'^\s{8}"([a-z0-9_]+)":\s*\(([^)]*)\)', blk, re.M):
        quoted = re.findall(r'"([A-Za-z0-9_]+)"', m.group(2))
        out[m.group(1)] = quoted[1] if len(quoted) > 1 else (quoted[0] if quoted else None)
    return out


def load_registry(fd: Path) -> dict[str, dict]:
    entries = json.loads((fd / "tables_registry.json").read_text(encoding="utf-8"))["tables"]
    return {e["key"]: e for e in entries if e.get("key")}


def sample(pro, api: str, ladder=PARAM_LADDER):
    for params in ladder:
        try:
            df = getattr(pro, api)(**params)
            if df is not None and not df.empty:
                return df, params
        except Exception:                            # noqa: BLE001
            pass
        time.sleep(0.5)
    return None, None


def check_one(pro, key: str, api: str, declared: list[str]) -> tuple[str, str]:
    """返回 (结论标记, 说明)."""
    if not api:
        return "SKIP", "TABLE_MAP 未接 api_name"
    df, used = sample(pro, api)
    if df is None:
        return "SKIP", "取不到样本"
    if not declared:
        return "SKIP", "registry 未声明唯一键"
    missing = [c for c in declared if c not in df.columns]
    if missing:
        return "BAD", f"声明的唯一键列不在返回里: {missing}"

    dup = int(df.duplicated(subset=declared).sum())
    if dup == 0:
        return "OK", f"唯一键={declared} 样本 {len(df)} 行"

    # 找出与重复相关的额外列，给出建议
    extra = [c for c in df.columns if c not in declared]
    suggestion = None
    for c in extra:
        cand = declared + [c]
        if int(df.duplicated(subset=cand).sum()) == 0:
            suggestion = cand
            break
    if suggestion is None:
        # 单列追加不够时，在声明键之外做**组合搜索**（2~3 列），取最小的能唯一化的组合
        import itertools
        best = None
        for size in (2, 3):
            for combo in itertools.combinations(extra, size):
                cand = declared + list(combo)
                if int(df.duplicated(subset=cand).sum()) == 0:
                    best = cand
                    break
            if best:
                break
        suggestion = best
    hint = f"建议改为 {suggestion}" if suggestion else "组合搜索（≤3 列）仍未唯一，需人工确认"
    return "DUP", f"唯一键={declared} 样本 {len(df)} 行，重复 {dup} 行 → {hint}"


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--keys", default=None, help="逗号分隔的表 key；缺省则检查全部")
    ap.add_argument("--root", default=str(FD_ROOT))
    args = ap.parse_args()

    for k in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"):
        os.environ.pop(k, None)

    fd = Path(args.root)
    key2api = load_table_map(fd)
    declarations = load_registry(fd)

    keys = ([x.strip() for x in args.keys.split(",") if x.strip()]
            if args.keys else sorted(key2api))

    import tushare as ts
    ts.set_token(os.environ["TUSHARE_TOKEN"])
    pro = ts.pro_api()

    counts = {"OK": 0, "DUP": 0, "BAD": 0, "SKIP": 0}
    problems: list[tuple[str, str, str]] = []
    for key in keys:
        if key not in key2api:
            continue
        declared = (declarations.get(key) or {}).get("unique_keys") or []
        flag, note = check_one(pro, key, key2api[key], declared)
        counts[flag] += 1
        if flag in ("DUP", "BAD"):
            problems.append((key, flag, note))
        if flag != "OK":
            print(f"  {key:<16} [{flag}] {note}")

    print(f"\n总计 {len(keys)} 张：唯一键正确 {counts['OK']}，"
          f"重复冲突 {counts['DUP']}，列缺失 {counts['BAD']}，跳过 {counts['SKIP']}")
    if problems:
        print("\n★ 需修正的表：")
        for key, flag, note in problems:
            print(f"    {key:<16} {note}")
    return 1 if problems else 0


if __name__ == "__main__":
    sys.exit(main())