market_sync/bin/archive/migrate_mysql_to_pg.py

"""[DEPRECATED 2026-06-16] MySQL → PostgreSQL 数据迁移工具。

⚠️ 项目从 2026-06-16 起放弃 MySQL，只支持 PostgreSQL。
   本脚本已归档到 ``bin/archive/``，仅作历史参考 / 应急回滚。
   如果你想从 MySQL 备份恢复数据到 PG，可以临时把它移回 ``bin/`` 用一次。

同步目标：把 MySQL `market_data_sync_db` 库 15 张表的数据搬到 PG `market_data.schema` 下。

策略（不重新建表，建表由 pg_bootstrap 负责）：
  - 用 pymysql 读 MySQL 流式分页（避免 1 亿行一次占满内存）
  - 用 psycopg2.copy_from 或 executemany 写 PG（按表选择最高效方式）
  - 顺序：先小表（配置/字典）→ 后大表（kline_5min 等）

数据量预估（按当时 MySQL 实测，2026-06-16）：
  config:                  20 行
  dataset_registry:        15 行
  industry:              5,207 行
  sectors:                  83 行
  stock_sector_map:     5,207 行
  indices:                  6 行
  market_regime_daily:  2,046 行
  share:                6,141 行
  stocks:               5,210 行
  sector_indices:     103,084 行
  sector_features_daily: 416,521 行
  kline_index:          37,480 行
  moneyflow:        3,188,301 行
  kline_stock:     11,605,541 行
  kline_5min:     121,945,849 行   ← 最大，约 1.2 亿

用法（应急时）：
    mv bin/archive/migrate_mysql_to_pg.py bin/migrate_mysql_to_pg.py
    .venv/bin/python -m bin.migrate_mysql_to_pg
"""
from __future__ import annotations

import argparse
import sys
import time
from io import StringIO
from typing import Iterable

import pymysql
import psycopg2
from dotenv import load_dotenv

load_dotenv(".env", override=False)

# 15 张表的列定义（MySQL → PG 列名映射）
# 大部分列名一致；少数 PG 用 JSONB（dataset_registry.dependency_ids）需要 JSON 解析
TABLES = [
    # (table_name, mysql_order_by, list_of_column_names)
    ("config",                   "key",      ["key", "value", "category", "description", "updated_at"]),
    ("dataset_registry",         "dataset_id", ["dataset_id", "name", "description", "storage_uri", "storage_layer",
                                             "management_role", "source", "sync_script", "dependency_ids",
                                             "enabled", "sort_order", "status", "trigger_source",
                                             "started_at", "finished_at", "last_success_at", "last_failure_at",
                                             "message", "last_error", "needs_resync", "progress_current",
                                             "progress_total", "current_step", "updated_at"]),
    ("indices",                  "index_code", ["index_code", "index_name", "market", "category", "source", "enabled", "updated_at"]),
    ("industry",                 "code",      ["code", "industry_name", "industry_classification", "update_date", "updated_at"]),
    ("sectors",                  "sector_key", ["sector_key", "sector_name", "taxonomy", "level", "source", "enabled", "updated_at"]),
    ("stock_sector_map",         "stock_code", ["stock_code", "sector_key", "updated_at"]),
    ("stocks",                   "code",      ["code", "name", "exchange", "list_date", "listing_status", "industry",
                                             "total_share", "float_share", "share_updated_at", "kline_synced_at",
                                             "updated_at"]),
    ("market_regime_daily",      "trade_date", ["trade_date", "advancers", "decliners", "advance_ratio", "turnover",
                                             "turnover_avg_5d", "turnover_ratio_5d", "source", "is_extreme_panic",
                                             "updated_at"]),
    ("share",                    "stock_code,trade_date", ["stock_code", "trade_date", "total_share", "float_share"]),
    ("sector_indices",           "trade_date,sector_name", ["trade_date", "sector_name", "close", "sector_amplitude"]),
    ("sector_features_daily",    "trade_date,sector_name", ["trade_date", "sector_name", "sector_ret", "sector_amplitude",
                                                       "close", "ema10", "ema20", "ema200", "score"]),
    ("kline_index",              "index_code,trade_date", ["index_code", "trade_date", "open", "high", "low", "close", "volume"]),
    ("moneyflow",                "stock_code,trade_date", ["stock_code", "trade_date", "main_net_inflow",
                                                       "large_net_inflow", "medium_net_inflow", "small_net_inflow"]),
    ("kline_stock",              "stock_code,trade_date", ["stock_code", "trade_date", "open", "high", "low", "close", "volume"]),
    ("kline_5min",               "stock_code,bar_time",   ["stock_code", "bar_time", "open", "high", "low", "close",
                                                       "volume", "amount", "turnover_rate"]),
]


def mysql_conn():
    return pymysql.connect(
        host=__import__("os").environ.get("MYSQL_HOST", "127.0.0.1"),
        port=int(__import__("os").environ.get("MYSQL_PORT", "3306")),
        user=__import__("os").environ.get("MYSQL_USER", "root"),
        password=__import__("os").environ.get("MYSQL_PASSWORD", ""),
        database=__import__("os").environ.get("MYSQL_DATABASE", "market_data_sync_db"),
        charset="utf8mb4",
        cursorclass=pymysql.cursors.SSDictCursor,  # 服务端 cursor，流式不撑内存
    )


def pg_conn():
    import os as _os
    pg_url = _os.environ.get("PG_URL", "").strip()
    # SQLAlchemy 格式 → psycopg2 格式（剥掉 +psycopg2）
    if pg_url.startswith("postgresql+psycopg2://"):
        pg_url = pg_url.replace("postgresql+psycopg2://", "postgresql://", 1)
    if not pg_url:
        host = _os.environ.get("PG_HOST", "127.0.0.1")
        port = _os.environ.get("PG_PORT", "5432")
        user = _os.environ.get("PG_USER", "market_sync")
        pw = _os.environ.get("PG_PASSWORD", "market_sync")
        db = _os.environ.get("PG_DB_NAME", "market_data")
        pg_url = f"postgresql://{user}:{pw}@{host}:{port}/{db}"
    return psycopg2.connect(pg_url)


def _convert_value(col: str, v):
    """MySQL → PG 类型兼容转换。"""
    if v is None:
        return None
    if col == "dependency_ids" and isinstance(v, str):
        # MySQL 存的是 Python list repr（如 ['stock_info']），不是合法 JSON。
        # 先后备方案：先 json.loads（合法 JSON），失败再 ast.literal_eval（Python repr）。
        import ast
        import json
        if not v.strip():
            return []
        # 试 1: 当 JSON 解析
        try:
            return json.loads(v)
        except json.JSONDecodeError:
            pass
        # 试 2: 当 Python list/tuple repr 解析
        try:
            parsed = ast.literal_eval(v)
            if isinstance(parsed, (list, tuple)):
                return list(parsed)
        except (ValueError, SyntaxError):
            pass
        # 试 3: 逗号分隔的纯字符串列表（兜底）
        return [s.strip().strip("'\"") for s in v.strip("[]").split(",") if s.strip()]
    if isinstance(v, str) and v == "":
        return None
    return v


def _copy_from_path(conn_pg, table_name: str, columns: list[str], tmp_path: str) -> int:
    """从已写好的 temp file COPY 到 PG。

    temp file 由调用方在游标活跃时填充好。
    """
    cur = conn_pg.cursor()
    with open(tmp_path, "rb") as fb:
        full_table = f"market_data.{table_name}"
        cur.copy_expert(
            f"COPY {full_table} ({','.join(columns)}) FROM STDIN WITH (FORMAT text, NULL '')",
            fb,
        )
    conn_pg.commit()
    return 0  # 实际行数已经在外层统计了


def _truncate_pg_table(conn_pg, table_name: str) -> None:
    """清空 PG 表。用 DELETE 而不是 TRUNCATE（TRUNCATE 需要额外权限）。

    DELETE 在大表上慢，但仅迁移时用一次可以接受。
    """
    cur = conn_pg.cursor()
    cur.execute(f"DELETE FROM market_data.{table_name}")
    conn_pg.commit()


def migrate_table(mysql, pg, table_name: str, order_by: str, columns: list[str],
                 chunk_size: int = 100_000) -> int:
    """单表迁移：先 TRUNCATE，再流式读 MySQL + COPY 写 PG。"""
    t0 = time.time()
    col_list = ", ".join(f"`{c}`" for c in columns)

    with mysql.cursor() as cur:
        cur.execute(f"SELECT COUNT(*) FROM {table_name}")
        n_total = cur.fetchone()["COUNT(*)"]

    if n_total == 0:
        print(f"  {table_name:<25} 0 行（跳过）")
        return 0

    # 先清空 PG 表（确保是干净迁移，不重复）
    _truncate_pg_table(pg, table_name)

    # 读 MySQL（流式：SSCursor 不缓存全表）→ 写 PG 临时文件 → COPY
    # 关键：迭代 + 写 temp file 必须在 with 块内完成（cur 关闭后再迭代会失败）
    print(f"  {table_name:<25} {n_total:>12,} 行 ...", end=" ", flush=True)
    import json as _json
    import tempfile
    JSONB_COLS = {"dependency_ids"}
    tmp = tempfile.NamedTemporaryFile(
        mode="w", encoding="utf-8", prefix=f"migrate_{table_name}_", suffix=".tsv",
        delete=False,
    )
    tmp_path = tmp.name
    written = 0
    try:
        with mysql.cursor() as cur:
            order_by_quoted = ", ".join(f"`{c.strip()}`" for c in order_by.split(","))
            cur.execute(f"SELECT {col_list} FROM `{table_name}` ORDER BY {order_by_quoted}")
            for r in cur:  # SSCursor：一次一行，不缓存全表
                cells = []
                for col, v in zip(columns, (r.get(col) for col in columns)):
                    if v is None:
                        cells.append("")
                    elif col in JSONB_COLS and not isinstance(v, str):
                        cells.append(_json.dumps(v, ensure_ascii=False))
                    else:
                        s = str(v)
                        cells.append(s.replace("\\", "\\\\").replace("\t", "\\t").replace("\n", "\\n"))
                tmp.write("\t".join(cells))
                tmp.write("\n")
                written += 1
        # with 块结束，cur 已关闭。temp file 写完。关 file 准备 COPY。
        tmp.flush()
        tmp.close()
        # COPY 从磁盘 temp file 读
        with open(tmp_path, "rb") as fb:
            cur_pg = pg.cursor()
            full_table = f"market_data.{table_name}"
            cur_pg.copy_expert(
                f"COPY {full_table} ({','.join(columns)}) FROM STDIN WITH (FORMAT text, NULL '')",
                fb,
            )
            pg.commit()
    finally:
        try:
            import os as _os
            _os.unlink(tmp_path)
        except OSError:
            pass
    elapsed = time.time() - t0
    print(f"→ {written:>12,} 行  ({elapsed:.1f}s)")
    return written


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--only", help="逗号分隔的表名列表，只迁指定的")
    p.add_argument("--chunk-size", type=int, default=100_000)
    p.add_argument("--skip", help="逗号分隔要跳过的表")
    args = p.parse_args()

    only = set((args.only or "").split(",")) if args.only else None
    skip = set((args.skip or "").split(",")) if args.skip else set()

    targets = [
        t for t in TABLES
        if (only is None or t[0] in only) and t[0] not in skip
    ]

    print(f"=== MySQL → PG 数据迁移 ===")
    print(f"目标表: {len(targets)} 张")
    print(f"源: MySQL market_data_sync_db (127.0.0.1:3306)")
    print(f"目标: PG market_data 库 market_data schema")
    print()

    mysql = mysql_conn()
    pg = pg_conn()

    grand_total = 0
    grand_t0 = time.time()
    for table_name, order_by, columns in targets:
        try:
            n = migrate_table(mysql, pg, table_name, order_by, columns, args.chunk_size)
            grand_total += n
        except Exception as e:
            print(f"  ✗ {table_name} 失败: {e}")
    grand_elapsed = time.time() - grand_t0
    print()
    print(f"=== 全部完成：{grand_total:,} 行 / 耗时 {grand_elapsed:.1f}s ===")

    # 简单校验：每张表行数对比
    print("\n=== 校验：MySQL vs PG 行数对比 ===")
    with mysql.cursor() as cur, pg.cursor() as cur_pg:
        for table_name, _, _ in targets:
            cur.execute(f"SELECT COUNT(*) AS n FROM {table_name}")
            n_mysql = cur.fetchone()["n"]
            cur_pg.execute(f"SELECT COUNT(*) FROM market_data.{table_name}")
            n_pg = cur_pg.fetchone()[0]
            ok = "✓" if n_mysql == n_pg else "✗"
            print(f"  {ok} {table_name:<25} MySQL={n_mysql:>12,}  PG={n_pg:>12,}")

    mysql.close()
    pg.close()


if __name__ == "__main__":
    main()