MobileModels/tools/load_mysql_seed.py

#!/usr/bin/env python3
"""Load MobileModels schema and seed data into MySQL."""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import time
from pathlib import Path

from project_layout import PROJECT_ROOT


def mysql_env(password: str) -> dict[str, str]:
    env = os.environ.copy()
    env["MYSQL_PWD"] = password
    return env


def mysql_command(user: str, host: str, port: int, database: str | None = None) -> list[str]:
    command = [
        "mysql",
        f"--host={host}",
        f"--port={port}",
        f"--user={user}",
        "--protocol=TCP",
        "--default-character-set=utf8mb4",
    ]
    if database:
        command.append(database)
    return command


def mysqladmin_ping(user: str, password: str, host: str, port: int) -> bool:
    proc = subprocess.run(
        [
            "mysqladmin",
            f"--host={host}",
            f"--port={port}",
            f"--user={user}",
            "--protocol=TCP",
            "ping",
            "--silent",
        ],
        env=mysql_env(password),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        check=False,
    )
    return proc.returncode == 0


def wait_for_mysql(user: str, password: str, host: str, port: int, timeout: int) -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        if mysqladmin_ping(user, password, host, port):
            return
        time.sleep(2)
    raise RuntimeError(f"MySQL 未在 {timeout}s 内就绪: {host}:{port}")


def run_sql_file(user: str, password: str, host: str, port: int, path: Path, database: str | None = None) -> None:
    sql_text = path.read_text(encoding="utf-8")
    proc = subprocess.run(
        mysql_command(user, host, port, database=database),
        env=mysql_env(password),
        input=sql_text,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if proc.returncode != 0:
        message = proc.stderr.strip() or proc.stdout.strip() or f"mysql exited with {proc.returncode}"
        raise RuntimeError(f"执行 SQL 文件失败 {path}: {message}")


def sql_string(value: str) -> str:
    return value.replace("\\", "\\\\").replace("'", "''")


def ensure_reader_user(
    user: str,
    password: str,
    host: str,
    port: int,
    database: str,
    reader_user: str,
    reader_password: str,
) -> None:
    sql = f"""
CREATE USER IF NOT EXISTS '{sql_string(reader_user)}'@'%' IDENTIFIED BY '{sql_string(reader_password)}';
ALTER USER '{sql_string(reader_user)}'@'%' IDENTIFIED BY '{sql_string(reader_password)}';
GRANT SELECT ON `{database}`.* TO '{sql_string(reader_user)}'@'%';
GRANT SELECT ON `python_services_test`.* TO '{sql_string(reader_user)}'@'%';
FLUSH PRIVILEGES;
"""
    proc = subprocess.run(
        mysql_command(user, host, port),
        env=mysql_env(password),
        input=sql,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if proc.returncode != 0:
        message = proc.stderr.strip() or proc.stdout.strip() or f"mysql exited with {proc.returncode}"
        raise RuntimeError(f"创建只读账号失败: {message}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Load MobileModels schema and seed data into MySQL.")
    parser.add_argument("--schema", type=Path, default=Path("sql/mobilemodels_mysql_schema.sql"))
    parser.add_argument("--seed", type=Path, default=Path("dist/mobilemodels_mysql_seed.sql"))
    parser.add_argument("--host", default=os.environ.get("MYSQL_HOST", "mysql"))
    parser.add_argument("--port", type=int, default=int(os.environ.get("MYSQL_PORT", "3306")))
    parser.add_argument("--user", default=os.environ.get("MYSQL_ROOT_USER", "root"))
    parser.add_argument("--password", default=os.environ.get("MYSQL_ROOT_PASSWORD", "mobilemodels_root"))
    parser.add_argument("--database", default=os.environ.get("MYSQL_DATABASE", "mobilemodels"))
    parser.add_argument("--reader-user", default=os.environ.get("MYSQL_READER_USER", ""))
    parser.add_argument("--reader-password", default=os.environ.get("MYSQL_READER_PASSWORD", ""))
    parser.add_argument("--wait-timeout", type=int, default=120)
    parser.add_argument("--check-only", action="store_true", help="Only check MySQL readiness")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    schema_path = args.schema if args.schema.is_absolute() else PROJECT_ROOT / args.schema
    seed_path = args.seed if args.seed.is_absolute() else PROJECT_ROOT / args.seed

    wait_for_mysql(args.user, args.password, args.host, args.port, args.wait_timeout)

    if args.check_only:
        print(f"MySQL ready: {args.host}:{args.port}")
        return 0

    run_sql_file(args.user, args.password, args.host, args.port, schema_path)
    run_sql_file(args.user, args.password, args.host, args.port, seed_path)

    if args.reader_user and args.reader_password:
        ensure_reader_user(
            args.user,
            args.password,
            args.host,
            args.port,
            args.database,
            args.reader_user,
            args.reader_password,
        )

    print(f"Loaded schema: {schema_path}")
    print(f"Loaded seed: {seed_path}")
    if args.reader_user:
        print(f"Ensured reader user: {args.reader_user}")
    return 0


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