FinDataHub/Scripts/bin/migrate_date_format.py

#!/usr/bin/env python3
"""migrate_date_format.py — 统一 trade_date 类日期列的格式为 `YYYYMMDD`

背景（P0-6）
------------
同一列里混存两种格式，导致按标准 `YYYYMMDD` 写的等值/范围过滤**静默丢弃**另一种格式的行。
实测 `tushare_data.stock_daily`（804 万行）中带横线 766.5 万 / 无横线 37.6 万，
2026 年共 94.6 万行，`BETWEEN '20260101' AND '20260930'` 只命中 37.6 万行 → 丢掉 **60%**。

统一口径选 `YYYYMMDD`（无分隔符）：
  - 分析层 `fin_data_clean` 已全部是无横线（1278 万行）
  - 39 个日期样字符列中已统一为无横线的有 21 个，混存 6 个、仅带横线 3 个
  - Tushare 接口原生格式即 `YYYYMMDD`

目标列（经 `audit_date_format.py` 扫描确定）
--------------------------------------------
  混存：akshare_data._table_meta.latest_date
        akshare_data.index_daily.trade_date
        tushare_data.margin_detail.trade_date
        tushare_data.stock_daily.trade_date
        tushare_data.stock_daily_basic.trade_date
        tushare_data.stock_moneyflow.trade_date
  仅带横线：akshare_data.stock_daily.trade_date
            akshare_data.stock_fundamental.found_date
            akshare_data.stock_fundamental.list_date

用法
----
    set -a; . ~/.config/findatahub/env; set +a
    python3 Scripts/bin/migrate_date_format.py --report          # 只统计影响面
    python3 Scripts/bin/migrate_date_format.py --backup          # 导出待改行的 (主键,原值)
    python3 Scripts/bin/migrate_date_format.py --execute         # 按年分批执行
    python3 Scripts/bin/migrate_date_format.py --verify          # 迁移后校验不变量

安全性设计
----------
- 迁移前先做**格式不变量检查**：非空值只允许 `YYYY-MM-DD` 或 `YYYYMMDD` 两种形态，
  出现任何第三种形态即拒绝执行（避免把未知数据一起改坏）。
- 按 `YYYY` 前缀分批 UPDATE，每批一个事务，避免大事务与长锁。
- 变换 `REPLACE(col,'-','')` 是幂等的：重复执行不会产生额外影响。
- `--verify` 校验：非空值全部匹配 `^[0-9]{8}$`、行数未减少、随机抽样值可逆。
"""

from __future__ import annotations

import argparse
import os
import random
import sys
from datetime import datetime
from pathlib import Path

import pymysql

# (库, 表, 日期列, 主键列)  主键列用于备份时定位行
TARGETS: list[tuple[str, str, str, str]] = [
    # (库, 表, 日期列, 分批键列)
    #
    # 分批键选 ts_code 而非自增 id：这些表**都没有自增主键**
    # （tushare_data.* 只有 uk_(ts_code, trade_date)，akshare_data 是 PRIMARY(ts_code, trade_date)），
    # 而 ts_code 是所有相关索引的最左列 —— 按它分批既走索引，又不会在扫描过程中
    # 修改到被用来分片的列本身（按 trade_date 分批会踩这个坑：
    # 更新后行的索引位置变化，可能重复访问或漏行）。
    ("tushare_data", "stock_daily", "trade_date", "ts_code"),
    ("tushare_data", "stock_daily_basic", "trade_date", "ts_code"),
    ("tushare_data", "stock_moneyflow", "trade_date", "ts_code"),
    ("tushare_data", "margin_detail", "trade_date", "ts_code"),
    ("akshare_data", "stock_daily", "trade_date", "ts_code"),
    ("akshare_data", "index_daily", "trade_date", "ts_code"),
    ("akshare_data", "stock_fundamental", "found_date", "ts_code"),
    ("akshare_data", "stock_fundamental", "list_date", "ts_code"),
    ("akshare_data", "_table_meta", "latest_date", "table_name"),
]

PLAIN_RE = "^[0-9]{8}$"
DASHED_RE = "^[0-9]{4}-[0-9]{2}-[0-9]{2}$"
BACKUP_DIR = Path.home() / "Library/Application Support/FinDataHub/backups"


def connect():
    pwd = os.environ.get("MYSQL_PASSWORD")
    if not pwd:
        print("✗ 需要 MYSQL_PASSWORD（可先 `set -a; . ~/.config/findatahub/env; set +a`）")
        sys.exit(2)
    return pymysql.connect(
        host="127.0.0.1", port=3306, user=os.environ.get("MYSQL_USER", "root"),
        password=pwd, charset="utf8mb4", connect_timeout=5, read_timeout=600,
        autocommit=False,
    )


def table_size_mb(cur, sch, tbl) -> int:
    cur.execute(
        "SELECT ROUND((data_length+index_length)/1024/1024) FROM information_schema.tables "
        "WHERE table_schema=%s AND table_name=%s",
        (sch, tbl),
    )
    r = cur.fetchone()
    return int(r[0]) if r and r[0] is not None else 0


def column_stats(cur, sch, tbl, col) -> dict:
    """返回该列的形态统计。

    性能取舍：`trade_date` 这类列的 `LIKE '%-%'` 与 `REGEXP` 都无法走索引，
    每次统计都是一次全表扫描。16.9M 行的表上多跑几次 REGEXP 要几分钟，
    因此 `unexpected` 用**便宜的 LENGTH 判据**近似（正常值只可能是 8 位或 10 位），
    真正强的不变量交给 `plain + empty == total 且 dashed == 0` 表达。

    注意：SQL 里不能出现裸 `%`（会被 pymysql 当成格式符），故 LIKE '%-%' 走参数绑定。
    """
    q = f"`{sch}`.`{tbl}`"
    cur.execute(
        f"SELECT COUNT(*), "
        f"SUM(`{col}` LIKE %s), "
        f"SUM(`{col}` REGEXP %s), "
        f"SUM(`{col}` IS NULL OR `{col}` = ''), "
        f"SUM(`{col}` IS NOT NULL AND `{col}` <> '' AND LENGTH(`{col}`) NOT IN (8, 10)) "
        f"FROM {q}",
        ("%-%", PLAIN_RE),
    )
    total, dashed, plain, empty, unexpected = cur.fetchone()
    return {
        "total": int(total or 0), "dashed": int(dashed or 0),
        "plain": int(plain or 0), "empty": int(empty or 0),
        "unexpected": int(unexpected or 0),
    }


def stats_ok(st: dict) -> bool:
    """强不变量：带横线 0、异常 0、且「无横线 + 空值」覆盖全部行（无行丢失）."""
    return st["dashed"] == 0 and st["unexpected"] == 0 and st["plain"] + st["empty"] == st["total"]


def cmd_report(cur) -> None:
    print(f"{'库.表.列':<48}{'表MB':>7}{'总行':>12}{'带横线':>11}{'无横线':>11}{'异常形态':>9}")
    print("-" * 100)
    grand_dashed = 0
    for sch, tbl, col, _pk in TARGETS:
        st = column_stats(cur, sch, tbl, col)
        mb = table_size_mb(cur, sch, tbl)
        grand_dashed += st["dashed"]
        print(f"{sch+'.'+tbl+'.'+col:<48}{mb:>7}{st['total']:>12,}"
              f"{st['dashed']:>11,}{st['plain']:>11,}{st['unexpected']:>9,}")
        if st["unexpected"]:
            print(f"    ⚠ 存在 {st['unexpected']} 行非 YYYY-MM-DD / YYYYMMDD 形态 —— --execute 将拒绝执行")
    print("-" * 100)
    print(f"{'合计待改（带横线行数）':<48}{'':>7}{'':>12}{grand_dashed:>11,}")


def cmd_backup(cur) -> None:
    """备份：小表全量导出；大表只导出 (主键, 原值) 的随机抽样 + 逐年直方图。

    为什么不做 44.5M 行全量导出：本迁移只做 `REPLACE(col,'-','')`，
    **不删行、不改行数**，且变换是确定的、幂等的。因此真正需要留证的是
    「迁移前后计数一致」与「可复算的端点样本」，而不是逐行快照。
    """
    BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    summary = []
    for sch, tbl, col, key in TARGETS:
        st = column_stats(cur, sch, tbl, col)
        if not st["dashed"]:
            continue
        # 逐年直方图（迁移前后应完全一致）
        cur.execute(
            f"SELECT LEFT(`{col}`,4) AS y, COUNT(*) FROM `{sch}`.`{tbl}` "
            f"WHERE `{col}` LIKE %s GROUP BY y ORDER BY y",
            ("%-%",),
        )
        hist = cur.fetchall()
        # 随机抽样（大表限 500 行，小表全量）
        limit = "" if st["dashed"] <= 5000 else " LIMIT 500"
        cur.execute(
            f"SELECT `{key}`, `{col}` FROM `{sch}`.`{tbl}` WHERE `{col}` LIKE %s{limit}",
            ("%-%",),
        )
        rows = cur.fetchall()
        out = BACKUP_DIR / f"datefmt-sample-{sch}-{tbl}-{col}-{stamp}.tsv"
        with open(out, "w", encoding="utf-8") as f:
            f.write(f"# {sch}.{tbl}.{col} 分批键={key} 迁移前样本/快照\n")
            f.write(f"# 总行={st['total']} 带横线={st['dashed']} 无横线={st['plain']}\n")
            f.write("# 逐年带横线直方图: " + ", ".join(f"{y}={c}" for y, c in hist) + "\n")
            f.write("# 键值\t原值\n")
            for r in rows:
                f.write(f"{r[0]}\t{r[1]}\n")
        summary.append((f"{sch}.{tbl}.{col}", st["total"], st["dashed"],
                        len(rows), out.name))
        print(f"  {sch}.{tbl}.{col}: 总 {st['total']:,} / 带横线 {st['dashed']:,} "
              f"/ 留样 {len(rows):,} → {out.name}")

    # 迁移前基线（供迁移后比对）
    baseline = BACKUP_DIR / f"datefmt-baseline-{stamp}.tsv"
    with open(baseline, "w", encoding="utf-8") as f:
        f.write("# 迁移前基线：库.表.列\t总行\t带横线\t样本数\t样本文件\n")
        for row in summary:
            f.write("\t".join(str(x) for x in row) + "\n")
    print(f"\n基线清单：{baseline}")
    print(f"备份目录：{BACKUP_DIR}")


KEYS_PER_BATCH = 400         # 每批处理多少个分批键（ts_code）


def cmd_execute(cur) -> None:
    # 注：不做全局预检。每张表的 st0 已经包含同样的检查，
    # 单独再跑一轮等于多 9 次全表扫描（大表上是分钟级开销，实测会导致调用超时）。
    for sch, tbl, col, key in TARGETS:
        st0 = column_stats(cur, sch, tbl, col)
        if st0["unexpected"]:
            print(f"─ {sch}.{tbl}.{col}: ✗ 存在 {st0['unexpected']} 行长度异常值，跳过（需人工确认）")
            continue
        if not st0["dashed"]:
            print(f"─ {sch}.{tbl}.{col}: 无需迁移")
            continue
        # 取出全部分批键，按批切片（ts_code 是相关索引最左列 → 每批走索引范围扫描）
        cur.execute(f"SELECT DISTINCT `{key}` FROM `{sch}`.`{tbl}` ORDER BY `{key}`")
        keys = [r[0] for r in cur.fetchall()]
        batches = [keys[i:i + KEYS_PER_BATCH] for i in range(0, len(keys), KEYS_PER_BATCH)]
        print(f"─ {sch}.{tbl}.{col}: 待改 {st0['dashed']:,} / 总 {st0['total']:,}，"
              f"分批键 {key} 共 {len(keys):,} 个 → {len(batches)} 批")
        done = 0
        for bi, batch in enumerate(batches, 1):
            ph = ", ".join(["%s"] * len(batch))
            cur.execute(
                f"UPDATE `{sch}`.`{tbl}` SET `{col}`=REPLACE(`{col}`,'-','') "
                f"WHERE `{key}` IN ({ph}) AND `{col}` LIKE %s",
                tuple(batch) + ("%-%",),
            )
            n = cur.rowcount
            cur.connection.commit()
            done += n
            if bi % 5 == 0 or bi == len(batches):
                print(f"      {bi}/{len(batches)} 批，累计更新 {done:,} 行")
        st1 = column_stats(cur, sch, tbl, col)
        ok = stats_ok(st1) and st1["total"] == st0["total"]
        print(f"    已更新 {done:,} 行；迁移后：带横线 {st1['dashed']} / "
              f"异常 {st1['unexpected']} / 总行 {st1['total']:,}（原 {st0['total']:,}） "
              f"{'✓' if ok else '✗ 需人工检查'}")
    print("\n迁移结束。请执行 --verify 做独立校验。")


def cmd_verify(cur) -> None:
    ok = True
    print(f"{'库.表.列':<48}{'总行':>12}{'带横线':>10}{'无横线':>12}{'异常':>8}  结论"
          f"   （不变量：dashed=0 且 plain+empty=total）")
    print("-" * 100)
    for sch, tbl, col, key in TARGETS:
        st = column_stats(cur, sch, tbl, col)
        good = stats_ok(st)
        ok = ok and good
        print(f"{sch+'.'+tbl+'.'+col:<48}{st['total']:>12,}{st['dashed']:>10,}"
              f"{st['plain']:>12,}{st['unexpected']:>8,}  {'✓' if good else '✗'}")
        # 抽样可逆性：随机取 3 行，检查 8 位形态可还原为 YYYY-MM-DD
        if good and st["plain"]:
            cur.execute(
                f"SELECT `{key}`, `{col}` FROM `{sch}`.`{tbl}` "
                f"WHERE `{col}` REGEXP %s LIMIT 3",
                (PLAIN_RE,),
            )
            for _kv, v in cur.fetchall():
                if not (len(v) == 8 and v.isdigit()):
                    print(f"    ✗ 抽样发现非 8 位值: {v!r}")
                    ok = False
    print("-" * 100)
    print("校验" + ("通过 ✓" if ok else "未通过 ✗"))
    return 0 if ok else 1


def main() -> int:
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--report", action="store_true", help="只统计影响面")
    g.add_argument("--backup", action="store_true", help="导出待改行的样本与基线")
    g.add_argument("--execute", action="store_true", help="执行分批迁移")
    g.add_argument("--verify", action="store_true", help="迁移后校验")
    ap.add_argument("--only", default=None,
                    help="只处理指定目标，格式 库.表 或 库.表.列（便于分批执行/重跑）")
    args = ap.parse_args()

    if args.only:
        parts = args.only.split(".")
        keep = [t for t in TARGETS
                if ".".join(t[:len(parts)]) == ".".join(parts)]
        if not keep:
            print(f"✗ --only {args.only!r} 未匹配任何目标")
            return 2
        TARGETS[:] = keep
        print(f"--only 过滤后目标数：{len(TARGETS)}")

    random.seed(20260918)
    conn = connect()
    try:
        cur = conn.cursor()
        if args.report:
            cmd_report(cur)
        elif args.backup:
            cmd_backup(cur)
        elif args.execute:
            cmd_execute(cur)
        else:
            return cmd_verify(cur)
    finally:
        conn.close()
    return 0


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