chuanqi/cq2_sdk/scripts/init_db.py

"""Initialize the 996cq MySQL database schema.

Reads connection from .env (CQ2_DB_HOST/PORT/USER/PASSWORD/NAME),
creates 5 tables (users, servers, zones, runs, task_logs), then verifies.

Idempotent: uses ``CREATE TABLE IF NOT EXISTS`` so re-running is safe.

Usage::

    python scripts/init_db.py             # create all tables
    python scripts/init_db.py --drop      # DROP all first, then create (DESTRUCTIVE)
    python scripts/init_db.py --check     # only check existing tables
"""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path

# allow running without `uv run` from anywhere
_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE.parent / "src"))


def _load_env() -> dict[str, str]:
    """Minimal .env loader (avoid dependency on cq2_sdk.env_loader
    so this script can run before the SDK is fully wired up)."""
    env = {}
    env_path = _HERE.parent / ".env"
    if env_path.exists():
        for line in env_path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            env[k.strip()] = v.strip().strip('"').strip("'")
    # env vars override file
    for k in ("CQ2_DB_HOST", "CQ2_DB_PORT", "CQ2_DB_USER", "CQ2_DB_PASSWORD", "CQ2_DB_NAME"):
        if k in os.environ:
            env[k] = os.environ[k]
    required = ["CQ2_DB_HOST", "CQ2_DB_PORT", "CQ2_DB_USER", "CQ2_DB_PASSWORD", "CQ2_DB_NAME"]
    missing = [k for k in required if k not in env]
    if missing:
        print(f"ERROR: missing env vars: {missing}")
        print(f"  set them in .env or as shell env vars")
        sys.exit(1)
    return env


DDL_STATEMENTS = [
    # users (no FK)
    """
    CREATE TABLE IF NOT EXISTS users (
        id                INT PRIMARY KEY AUTO_INCREMENT,
        username          VARCHAR(64)  UNIQUE,
        phone             VARCHAR(32)  UNIQUE,
        password_hash     VARCHAR(255),
        token             TEXT,
        token_expires_at  TIMESTAMP NULL,
        last_login_at     TIMESTAMP NULL,
        created_at        TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
    # servers (no FK)
    """
    CREATE TABLE IF NOT EXISTS servers (
        id              INT PRIMARY KEY AUTO_INCREMENT,
        server_name     VARCHAR(128) UNIQUE,
        server_id       INT,
        game_id         INT DEFAULT 206923,
        first_seen_at   TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
    # zones (FK -> servers)
    """
    CREATE TABLE IF NOT EXISTS zones (
        id           INT PRIMARY KEY AUTO_INCREMENT,
        server_id    INT NOT NULL,
        zone_id      VARCHAR(64),
        zone_name    VARCHAR(128),
        host         VARCHAR(128),
        port         INT,
        open_time    TIMESTAMP,
        status       VARCHAR(32),
        fetched_at   TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        UNIQUE KEY uniq_zone (server_id, zone_id, open_time),
        FOREIGN KEY (server_id) REFERENCES servers(id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
    # runs (FK -> users, servers, zones)
    """
    CREATE TABLE IF NOT EXISTS runs (
        id              INT PRIMARY KEY AUTO_INCREMENT,
        user_id         INT NOT NULL,
        server_id       INT NOT NULL,
        zone_id         INT NOT NULL,
        started_at      TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        completed_at    TIMESTAMP NULL,
        current_task    VARCHAR(8),
        status          VARCHAR(32),
        snapshot_path   VARCHAR(255),
        FOREIGN KEY (user_id)   REFERENCES users(id),
        FOREIGN KEY (server_id) REFERENCES servers(id),
        FOREIGN KEY (zone_id)   REFERENCES zones(id),
        INDEX idx_user_status (user_id, status)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
    # task_logs (FK -> runs)
    # NOTE: MySQL 5.7 严格要求每个非首列 TIMESTAMP 必须显式 NULL/default，
    # 否则报 1067 "Invalid default value"。
    """
    CREATE TABLE IF NOT EXISTS task_logs (
        id              INT PRIMARY KEY AUTO_INCREMENT,
        run_id          INT NOT NULL,
        task_id         VARCHAR(8),
        success         TINYINT(1),
        error           TEXT,
        elapsed_seconds DOUBLE,
        started_at      TIMESTAMP NULL,
        finished_at     TIMESTAMP NULL,
        FOREIGN KEY (run_id) REFERENCES runs(id),
        INDEX idx_run_task (run_id, task_id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
]

EXPECTED_TABLES = {"users", "servers", "zones", "runs", "task_logs"}

DROP_ORDER = ["task_logs", "runs", "zones", "servers", "users"]


def main() -> int:
    parser = argparse.ArgumentParser(description="Initialize 996cq schema")
    parser.add_argument("--drop", action="store_true",
                        help="DROP all tables first (DESTRUCTIVE)")
    parser.add_argument("--check", action="store_true",
                        help="Only check existing tables, don't create")
    args = parser.parse_args()

    import pymysql  # imported lazily so --help works without pymysql

    env = _load_env()
    print(f"Connecting to {env['CQ2_DB_HOST']}:{env['CQ2_DB_PORT']} "
          f"as {env['CQ2_DB_USER']}@{env['CQ2_DB_NAME']}")

    conn = pymysql.connect(
        host=env["CQ2_DB_HOST"],
        port=int(env["CQ2_DB_PORT"]),
        user=env["CQ2_DB_USER"],
        password=env["CQ2_DB_PASSWORD"],
        database=env["CQ2_DB_NAME"],
        charset="utf8mb4",
        connect_timeout=10,
        autocommit=True,
    )

    try:
        with conn.cursor() as cur:
            if args.drop:
                print("--drop: dropping existing tables (DESTRUCTIVE)")
                for tbl in DROP_ORDER:
                    cur.execute(f"DROP TABLE IF EXISTS {tbl}")
                    print(f"  dropped: {tbl}")

            if args.check:
                cur.execute("SHOW TABLES")
                existing = {r[0] for r in cur.fetchall()}
                print(f"\n[check] existing tables: {sorted(existing)}")
                missing = EXPECTED_TABLES - existing
                if missing:
                    print(f"[check] MISSING: {sorted(missing)}")
                    return 1
                print("[check] all 5 expected tables present")
                return 0

            # Create
            print("\nCreating tables (idempotent via IF NOT EXISTS):")
            for ddl in DDL_STATEMENTS:
                # Get table name from DDL for nicer logging
                first_line = ddl.strip().splitlines()[0]
                tbl_name = first_line.split("EXISTS ")[1].split(" ")[0].strip("(`")
                cur.execute(ddl)
                print(f"  OK: {tbl_name}")

            # Verify
            cur.execute("SHOW TABLES")
            existing = {r[0] for r in cur.fetchall()}
            print(f"\n[verify] existing tables: {sorted(existing)}")
            missing = EXPECTED_TABLES - existing
            if missing:
                print(f"[verify] FAILED: missing {sorted(missing)}")
                return 1
            print(f"[verify] OK: all 5 tables present")

            # Show schema for each
            print("\nSchema details:")
            for tbl in sorted(EXPECTED_TABLES):
                cur.execute(
                    "SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_KEY "
                    "FROM information_schema.COLUMNS "
                    "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s "
                    "ORDER BY ORDINAL_POSITION",
                    (env["CQ2_DB_NAME"], tbl),
                )
                cols = cur.fetchall()
                print(f"\n  [{tbl}] ({len(cols)} columns)")
                for c in cols:
                    marker = " PK" if "PRI" in c[3] else (" UNI" if "UNI" in c[3] else "")
                    print(f"    {c[0]:<22} {c[1]:<28} NULL={c[2]}{marker}")
        print("\n[done] schema initialized.")
        return 0
    finally:
        conn.close()


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