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