FinDataHub/Scripts/bin/create_missing_tables.py

#!/usr/bin/env python3
"""create_missing_tables.py — 为「排期已引用但库中不存在」的表推导 DDL 并建表

背景
----
排期前飞行检查发现：`tushare_sync.sh` 的 monthly 段引用了 **21 张库中不存在的表**
（fut_* / fx_* / hk_* / us_* / opt_daily / sge_daily / ths_index / rate_* …）。
由于缺表只记 `写入失败` 日志、退出码仍为 0，脚本带 pipefail 也不会 abort ——
表现为**这些行长期静默空转**（与用户报的 stk_holdertrade 停摆是同一类缺陷）。

为什么需要这个工具
------------------
`tables_registry.json` 只登记了库表名/唯一键/日期列，**不含字段定义**，
仓库里也没有通用的 DDL 生成器（`Scripts/bin/init_db.py` 是手写 DDL，只覆盖少数表）。
因此这里从 **Tushare 接口实际返回的 DataFrame** 推导列与类型，保证 DDL 与写入时的
`df.columns` 一致 —— 否则 upsert 会因缺列而整批写入失败。

类型映射（与既有表风格一致，参考 tushare_data.stock_rewards）
------------------------------------------------------------
    object / string / datetime → text        （唯一键用前缀长度索引，如 ts_code(200)）
    float / int                → double DEFAULT NULL
    列注释取列名，便于与 add_comments.py 后续补中文注释

用法
----
    set -a; . ~/.config/findatahub/env; set +a
    unset HTTP_PROXY HTTPS_PROXY http_proxy https_proxy ALL_PROXY all_proxy   # 境内直连
    python3 Scripts/bin/create_missing_tables.py --list            # 列出缺表及其接口
    python3 Scripts/bin/create_missing_tables.py --dry-run         # 只打印将执行的 DDL
    python3 Scripts/bin/create_missing_tables.py --create          # 实际建表（幂等）
    python3 Scripts/bin/create_missing_tables.py --create --only fut_daily

安全性
------
- **幂等**：表已存在则跳过，不 DROP、不改动既有表结构。
- 只建「排期引用但不存在」的表；不触碰任何已有数据。
- 唯一键取自 `tables_registry.json`，与 `db.upsert` 的写法保持一致（否则幂等写入失效）。
"""
from __future__ import annotations

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

import pymysql

FD_ROOT = Path("/Applications/FinDataHub.app/Contents/Resources/tushare_fetcher")
SCHEMA = "tushare_data"
TEXT_COLS = ("ts_code", "trade_date", "ann_date", "end_date", "cal_date", "date", "symbol")


# ── 解析 TABLE_MAP：key → (库表名, api_name) ──
def load_table_map(fd: Path) -> dict[str, tuple[str, str | None]]:
    src = (fd / "src/tushare_fetcher/core/fetcher.py").read_text(encoding="utf-8")
    blk = src[src.index("TABLE_MAP"):]
    out: dict[str, tuple[str, str | None]] = {}
    for m in re.finditer(r'^\s{8}"([a-z0-9_]+)":\s*\(([^)]*)\)', blk, re.M):
        key, rest = m.group(1), m.group(2)
        quoted = re.findall(r'"([A-Za-z0-9_]+)"', rest)
        db_table = quoted[0] if quoted else key
        api = quoted[1] if len(quoted) > 1 else None
        out[key] = (db_table, api)
    return out


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


def scheduled_keys(fd: Path) -> set[str]:
    cli = (fd / "src/tushare_fetcher/cli.py").read_text(encoding="utf-8")
    keys: set[str] = set()
    for _g, items in dict(re.findall(r'"(\w+)":\s*\[([^\]]*)\]', cli, re.S)).items():
        keys |= {x.strip().strip('"') for x in items.split(",") if x.strip()}
    sh = (fd / "scripts/tushare_sync.sh").read_text(encoding="utf-8")
    for m in re.findall(r"--tables\s+([A-Za-z0-9_,]+)", sh):
        keys |= {x.strip() for x in m.split(",") if x.strip()}
    keys.add("dividend")
    return keys


def sample_param_ladder(key: str, reg: dict) -> list[dict]:
    """取样入参阶梯：按顺序尝试，取第一个能返回非空 DataFrame 的组合。

    为什么需要阶梯：不同接口要求的入参不同（有的要 ts_code、有的要 trade_date、
    有的必须给 start_date/end_date），而且部分接口有 1 次/分钟的频率限制 ——
    单次失败不代表接口不可用。实测：hibor / shibor_lpr 首次调用只是撞上
    1 次/分钟的限流，加间隔重试即可成功；fx_daily 则需要 start/end 区间。
    """
    ladder: list[dict] = []
    mode = reg.get("param_mode")
    if mode in ("ts_code", "fund_ts_code"):
        ladder.append({"ts_code": "000001.SZ"})
    if reg.get("date_column") in ("trade_date", "ann_date", "cal_date"):
        ladder.append({"trade_date": "20260918"})
    ladder += [
        {},                                             # 无参（分页默认）
        {"start_date": "20260901", "end_date": "20260918"},
        {"ts_code": "000001.SZ"},
        {"trade_date": "20260918"},
        {"date": "20260918"},
    ]
    # 去重且保序
    seen, out = set(), []
    for d in ladder:
        t = tuple(sorted(d.items()))
        if t not in seen:
            seen.add(t)
            out.append(d)
    return out


def derive_ddl(db_table: str, df, unique_keys: list[str]) -> str:
    """从接口样本推导 CREATE TABLE。

    列类型判断**必须用 pandas 自己的类型谓词**，不能对 `str(dtype)` 做字符串匹配：
    pandas 3.x 的字符串列 dtype 显示为 `str`（而非 `object`），
    早期版本用 `"object" in str(dtype)` 判断会把 ts_code/name 这类字符串列
    误判成数值列并生成 `double` —— 实测踩到过。
    """
    import pandas as pd

    cols = []
    names = []
    for c in df.columns:
        name = str(c)
        if name == "index":
            continue
        names.append(name)
        s = df[c]
        if pd.api.types.is_bool_dtype(s):
            sql_type = "tinyint(1) DEFAULT NULL"
        elif pd.api.types.is_numeric_dtype(s):
            sql_type = "double DEFAULT NULL"
        else:
            sql_type = "text"
        cols.append(f"  `{name}` {sql_type} COMMENT '{name}'")

    kept_uk = [k for k in (unique_keys or []) if k in names]
    if kept_uk:
        ucols = ", ".join(f"`{k}`(200)" for k in kept_uk)
        cols.append(f"  UNIQUE KEY `uk_{db_table[:40]}` ({ucols})")

    body = ",\n".join(cols)
    return (f"CREATE TABLE `{SCHEMA}`.`{db_table}` (\n{body}\n) "
            f"ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci")


def main() -> int:
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--list", action="store_true", help="列出排期引用但不存在的表")
    g.add_argument("--dry-run", action="store_true", help="打印将执行的 DDL，不建表")
    g.add_argument("--create", action="store_true", help="实际建表（幂等）")
    ap.add_argument("--only", default=None, help="只处理指定 key")
    args = ap.parse_args()

    pwd = os.environ.get("MYSQL_PASSWORD")
    if not pwd:
        print("✗ 需要 MYSQL_PASSWORD")
        return 2
    # 境内接口走直连，避免会话代理端口变动导致失败
    for k in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"):
        os.environ.pop(k, None)

    fd = FD_ROOT
    tmap = load_table_map(fd)
    reg = load_registry(fd)
    keys = scheduled_keys(fd)

    conn = pymysql.connect(host="127.0.0.1", port=3306, user="root", password=pwd,
                           charset="utf8mb4", autocommit=True)
    cur = conn.cursor()
    cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema=%s", (SCHEMA,))
    existing = {r[0] for r in cur.fetchall()}

    missing = sorted(k for k in keys
                     if k in tmap and tmap[k][0] not in existing
                     and (args.only is None or k == args.only))
    if not missing:
        print("✓ 排期引用的表在库中均存在")
        return 0

    print(f"排期引用但库中不存在的表：{len(missing)} 张\n")
    if args.list:
        for k in missing:
            db_table, api = tmap[k]
            print(f"  {k:<16} 表={db_table:<20} 接口={api}")
        return 0

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

    created, skipped = [], []
    for k in missing:
        db_table, api = tmap[k]
        if not api:
            print(f"  {k:<16} ✗ TABLE_MAP 未接 api_name，跳过"); skipped.append(k); continue
        # 按入参阶梯取样，取第一个非空结果
        df, err, used = None, None, None
        for params in sample_param_ladder(k, reg.get(k, {})):
            try:
                cand = getattr(pro, api)(**params)
                if cand is not None and not cand.empty:
                    df, used = cand, params
                    break
                err = err or "返回空"
            except Exception as e:                   # noqa: BLE001
                err = str(e)[:90]
            time.sleep(1.2)                          # 礼让 1 次/分钟的限流
        if df is None:
            print(f"  {k:<16} ✗ 取样失败（{api}）：{err}"); skipped.append(k); continue

        uk = (reg.get(k) or {}).get("unique_keys") or []
        ddl = derive_ddl(db_table, df, uk)
        if args.dry_run:
            print(f"  ── {k} → {db_table}（{api}，{len(df.columns)} 列）──")
            print("    " + ddl.replace("\n", "\n    "))
            continue
        try:
            cur.execute(ddl)
            ncols = len([c for c in df.columns if str(c) != "index"])
            print(f"  {k:<16} ✓ 已建 {db_table}（{ncols} 列，唯一键 {uk}，入参 {used}）")
            created.append(k)
        except Exception as e:                       # noqa: BLE001
            print(f"  {k:<16} ✗ 建表失败：{str(e)[:90]}")
            skipped.append(k)

    if args.create:
        print(f"\n建表完成：成功 {len(created)} 张，跳过 {len(skipped)} 张")
        if skipped:
            print(f"跳过（多为无接口权限 / 需特定入参）：{', '.join(skipped)}")
    conn.close()
    return 0


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