kis_bot/kis_trader/scripts/prune_env_config_strategy_columns.py

#!/usr/bin/env python3
"""
env_config fat 컬럼 정리 — 전략별 config_* 로 이전된 컬럼 제거.

1) 백업 (JSON): merged env, env_config 전체 행, 삭제 대상 KV, config_*, env_config_ext
2) env_config 를 ENV_GLOBAL_KEYS 만 가진 슬림 테이블로 교체
3) _migrate_config_table_columns 로 누락 글로벌 컬럼 추가

실행:
  cd ~/kis_bot && python3 -m kis_trader.scripts.prune_env_config_strategy_columns
  cd ~/kis_bot && python3 -m kis_trader.scripts.prune_env_config_strategy_columns --dry-run
"""
from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Set

_ROOT = Path(__file__).resolve().parents[2]
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from config_schema import classify_config_key  # noqa: E402
from database import ENV_CONFIG_KEYS, ENV_GLOBAL_KEYS, TradeDB  # noqa: E402


def _row_to_dict(row: Any) -> Dict[str, Any]:
    if row is None:
        return {}
    if hasattr(row, "keys"):
        return {k: row[k] for k in row.keys()}
    return dict(row)


def _json_safe(val: Any) -> Any:
    if val is None:
        return None
    if isinstance(val, (int, float, bool, str)):
        return val
    return str(val)


def backup_all(db: TradeDB, out_dir: Path) -> Dict[str, Any]:
    out_dir.mkdir(parents=True, exist_ok=True)
    meta: Dict[str, Any] = {"created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

    merged = db.get_merged_env_snapshot() or {}
    (out_dir / "merged_env_snapshot.json").write_text(
        json.dumps(merged, ensure_ascii=False, indent=2), encoding="utf-8",
    )
    meta["merged_keys"] = len(merged)

    phys_cols = set(db.conn.get_columns("env_config"))
    phys_cols -= {"id", "created_at"}
    global_set = set(ENV_GLOBAL_KEYS)
    drop_cols = sorted(c for c in phys_cols if c not in global_set)
    keep_cols = [k for k in ENV_GLOBAL_KEYS if k in phys_cols]
    meta["physical_cols"] = len(phys_cols)
    meta["keep_cols"] = len(keep_cols)
    meta["drop_cols"] = len(drop_cols)
    (out_dir / "drop_column_names.json").write_text(
        json.dumps(drop_cols, ensure_ascii=False, indent=2), encoding="utf-8",
    )

    rows = db.conn.execute("SELECT * FROM env_config ORDER BY id").fetchall()
    full_rows: List[Dict[str, Any]] = []
    dropped_kv: List[Dict[str, Any]] = []
    for raw in rows or []:
        row = _row_to_dict(raw)
        rid = row.get("id")
        full_rows.append({k: _json_safe(v) for k, v in row.items()})
        kv: Dict[str, Any] = {"id": rid, "created_at": row.get("created_at")}
        for col in drop_cols:
            v = row.get(col)
            if v is not None and str(v).strip() != "":
                kv[col] = _json_safe(v)
        if len(kv) > 2:
            dropped_kv.append(kv)
    (out_dir / "env_config_full_rows.json").write_text(
        json.dumps(full_rows, ensure_ascii=False, indent=2), encoding="utf-8",
    )
    (out_dir / "env_config_dropped_columns_kv.json").write_text(
        json.dumps(dropped_kv, ensure_ascii=False, indent=2), encoding="utf-8",
    )
    meta["env_config_rows"] = len(full_rows)

    for tbl in (
        "config_scalp", "config_short", "config_momentum", "config_us_momentum", "config_breakout",
        "config_range_break", "config_updow", "config_dbband", "env_config_ext",
    ):
        try:
            trows = db.conn.execute(f"SELECT * FROM {tbl}").fetchall()
            data = [_row_to_dict(r) for r in (trows or [])]
            (out_dir / f"{tbl}.json").write_text(
                json.dumps(data, ensure_ascii=False, indent=2, default=str),
                encoding="utf-8",
            )
            meta[f"{tbl}_rows"] = len(data)
        except Exception as ex:
            meta[f"{tbl}_error"] = str(ex)

    (out_dir / "backup_meta.json").write_text(
        json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8",
    )
    return meta


def prune_env_config(db: TradeDB, dry_run: bool = False) -> None:
    phys = set(db.conn.get_columns("env_config")) - {"id", "created_at"}
    keep_in_phys = [k for k in ENV_GLOBAL_KEYS if k in phys]
    if not keep_in_phys:
        raise RuntimeError("env_config 에 유지할 글로벌 컬럼 없음")

    cols_sql = ", ".join(f"`{k}`" for k in keep_in_phys)
    gcols_all = ", ".join(f"`{k}` TEXT" for k in ENV_GLOBAL_KEYS)

    print(f"유지 컬럼 {len(keep_in_phys)} / ENV_GLOBAL 전체 {len(ENV_GLOBAL_KEYS)}")
    print(f"삭제 컬럼 {len(phys - set(ENV_GLOBAL_KEYS))} 개")

    if dry_run:
        print("[dry-run] 슬림 교체 스킵")
        return

    db.conn.execute("DROP TABLE IF EXISTS env_config_new")
    db.conn.execute(
        f"CREATE TABLE env_config_new ("
        f"id INT NOT NULL AUTO_INCREMENT PRIMARY KEY, "
        f"created_at VARCHAR(30) NOT NULL, {gcols_all}"
        f") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4",
    )
    db.conn.execute(
        f"INSERT INTO env_config_new (created_at, {cols_sql}) "
        f"SELECT created_at, {cols_sql} FROM env_config",
    )
    cnt = db.conn.execute("SELECT COUNT(*) AS c FROM env_config_new").fetchone()
    n = cnt["c"] if isinstance(cnt, dict) else cnt[0]
    print(f"env_config_new 적재 {n}행")

    db.conn.execute("DROP TABLE env_config")
    db.conn.execute("RENAME TABLE env_config_new TO env_config")
    print("env_config 슬림 교체 완료")

    if hasattr(db, "_env_config_cols_cache"):
        db._env_config_cols_cache = None
    db._migrate_config_table_columns()

    try:
        from kis_trader.utils.env import invalidate_merged_env_cache
        invalidate_merged_env_cache()
    except Exception:
        pass

    after = len(db.conn.get_columns("env_config"))
    print(f"env_config 컬럼 수: {after} (id/created_at 포함)")


def main() -> int:
    ap = argparse.ArgumentParser(description="env_config 전략 컬럼 백업 후 제거")
    ap.add_argument("--dry-run", action="store_true", help="백업만, 테이블 교체 안 함")
    ap.add_argument(
        "--backup-dir",
        default="",
        help="백업 디렉터리 (기본: backups/env_config_prune_YYYYMMDD_HHMMSS)",
    )
    args = ap.parse_args()

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_dir = Path(args.backup_dir) if args.backup_dir else _ROOT / "backups" / f"env_config_prune_{ts}"

    db = TradeDB()
    try:
        print(f"백업 → {out_dir}")
        meta = backup_all(db, out_dir)
        print(json.dumps(meta, ensure_ascii=False, indent=2))
        prune_env_config(db, dry_run=args.dry_run)
        if not args.dry_run:
            snap = db.get_merged_env_snapshot()
            print(f"병합 snapshot 키 수: {len(snap)} (기능 유지 확인)")
    finally:
        db.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())