"""Run a read-only MySQL query and archive the CSV result in ana-data. The script only allows SELECT-like statements and reads the password from TIANXIA_MYSQL_PASSWORD. It never accepts a password argument and never writes to the database. """ from __future__ import annotations import argparse import csv import hashlib import os import re import subprocess import sys from datetime import datetime, timezone from pathlib import Path SCRIPT_PATH = Path(__file__).resolve() PROJECT_ROOT = SCRIPT_PATH.parents[2] DEFAULT_MYSQL_EXE = Path(os.environ.get("TIANXIA_MYSQL_EXE", "M:/mysql/server/bin/mysql.exe")) ALLOWED_PREFIXES = ("select", "with", "show", "describe", "desc", "explain") FORBIDDEN_SQL = re.compile( r"\b(insert|update|delete|drop|alter|create|truncate|replace|load|grant|" r"revoke|set|call|do|handler|lock|unlock|analyze|optimize|repair)\b", re.I, ) def now_iso() -> str: return datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds") def safe_name(value: str, fallback: str = "artifact", max_len: int = 96) -> str: value = re.sub(r"[^A-Za-z0-9._-]+", "_", value.strip()) value = value.strip("._-") return (value or fallback)[:max_len] def sha256_text(value: str) -> str: return hashlib.sha256(value.encode("utf-8")).hexdigest() def project_relative(path: Path) -> str: try: return path.resolve().relative_to(PROJECT_ROOT.resolve()).as_posix() except ValueError: return path.resolve().as_posix() def read_sql(args: argparse.Namespace) -> str: if args.sql_file: sql = Path(args.sql_file).read_text(encoding="utf-8") elif args.sql: sql = args.sql else: raise SystemExit("pass --sql or --sql-file") sql = sql.strip() if sql.endswith(";"): sql = sql[:-1].strip() return sql def validate_read_only(sql: str) -> None: compact = re.sub(r"/\*.*?\*/", " ", sql, flags=re.S) compact = re.sub(r"--.*?$", " ", compact, flags=re.M).strip() if ";" in compact: raise SystemExit("only one SQL statement is allowed") prefix = compact.split(None, 1)[0].lower() if compact else "" if prefix not in ALLOWED_PREFIXES: raise SystemExit(f"SQL must start with one of: {', '.join(ALLOWED_PREFIXES)}") if FORBIDDEN_SQL.search(compact): raise SystemExit("SQL contains a forbidden write/admin keyword") if re.search(r"\binto\s+outfile\b", compact, re.I): raise SystemExit("INTO OUTFILE is not allowed") def run_mysql(args: argparse.Namespace, sql: str) -> str: password = os.environ.get("TIANXIA_MYSQL_PASSWORD") if not password: raise SystemExit("TIANXIA_MYSQL_PASSWORD is required in the environment") mysql_exe = Path(args.mysql_exe) if not mysql_exe.exists(): raise SystemExit(f"mysql executable not found: {mysql_exe}") cmd = [ str(mysql_exe), "--protocol=TCP", "--batch", "--raw", "--quick", "--default-character-set=utf8mb4", f"--host={args.host}", f"--port={args.port}", f"--user={args.user}", args.database, "-e", sql, ] env = os.environ.copy() env["MYSQL_PWD"] = password proc = subprocess.run(cmd, text=True, capture_output=True, env=env, check=False) if proc.returncode != 0: raise SystemExit(proc.stderr.strip() or proc.stdout.strip()) return proc.stdout def write_csv_from_tsv(tsv: str, out_path: Path) -> int: out_path.parent.mkdir(parents=True, exist_ok=True) row_count = 0 with out_path.open("w", encoding="utf-8-sig", newline="") as fh: writer = csv.writer(fh) for line_index, line in enumerate(tsv.splitlines()): cells = line.split("\t") writer.writerow(cells) if line_index > 0: row_count += 1 return row_count def append_manifest(path: Path, row: dict[str, str]) -> None: path.parent.mkdir(parents=True, exist_ok=True) exists = path.exists() with path.open("a", encoding="utf-8-sig", newline="") as fh: writer = csv.DictWriter(fh, fieldnames=list(row.keys())) if not exists: writer.writeheader() writer.writerow(row) def export(args: argparse.Namespace) -> int: sql = read_sql(args) validate_read_only(sql) industry_root = PROJECT_ROOT / "ana-data" / "cases" / args.industry_case output_dir = industry_root / args.output_kind manifest_dir = industry_root / "manifest" output_dir.mkdir(parents=True, exist_ok=True) manifest_dir.mkdir(parents=True, exist_ok=True) artifact = safe_name(args.artifact_name) run_safe = safe_name(args.run_id, "run") query_id = f"{run_safe}_{artifact}" sql_hash = sha256_text(sql) out_path = output_dir / f"{artifact}_{run_safe}_{sql_hash[:12]}.csv" sql_path = manifest_dir / f"mysql_query_{query_id}_{sql_hash[:12]}.sql" manifest_path = manifest_dir / f"mysql_query_manifest_{run_safe}.csv" tsv = run_mysql(args, sql) row_count = write_csv_from_tsv(tsv, out_path) sql_path.write_text(sql + "\n", encoding="utf-8") append_manifest( manifest_path, { "query_id": query_id, "case_id": args.case_id, "batch_id": args.batch_id, "run_id": args.run_id, "artifact_name": artifact, "sql_sha256": sql_hash, "sql_relative_path": project_relative(sql_path), "output_relative_path": project_relative(out_path), "row_count": str(row_count), "executed_at": now_iso(), "db_host": args.host, "db_port": str(args.port), "db_name": args.database, "read_only_guard": "PASS", }, ) print(f"output={project_relative(out_path)} rows={row_count}") print(f"manifest={project_relative(manifest_path)}") return 0 def parse_args(argv: list[str]) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--industry-case", required=True) parser.add_argument("--case-id", required=True) parser.add_argument("--batch-id", required=True) parser.add_argument("--run-id", required=True) parser.add_argument("--artifact-name", required=True) parser.add_argument("--output-kind", choices=["supplement", "evidence", "manifest"], default="supplement") parser.add_argument("--sql") parser.add_argument("--sql-file") parser.add_argument("--mysql-exe", default=str(DEFAULT_MYSQL_EXE)) parser.add_argument("--host", default=os.environ.get("TIANXIA_MYSQL_HOST", "127.0.0.1")) parser.add_argument("--port", default=os.environ.get("TIANXIA_MYSQL_PORT", "3306")) parser.add_argument("--user", default=os.environ.get("TIANXIA_MYSQL_USER", "root")) parser.add_argument("--database", default=os.environ.get("TIANXIA_MYSQL_DB", "tianxia")) return parser.parse_args(argv) def main(argv: list[str] | None = None) -> int: args = parse_args(argv or sys.argv[1:]) return export(args) if __name__ == "__main__": raise SystemExit(main())