"""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())