MediaCrawler/tools/clone_database.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Clone one schema to a new database on the same MySQL server, and provision a
scoped account for it.

Written for the production cut-over: the monitor data lives in `mediacrawler` and
the server deployment reads `mediacrawler_prod`. Both sit on the same host, so the
copy is a cross-schema ``INSERT ... SELECT`` rather than a dump-and-reload -- and
since neither this workstation nor the server has a mysql client, that is also
the only option available.

Two things this deliberately does NOT do, both because the server hosts ~22
unrelated production databases and the provisioning account is a full admin:

* it never writes outside the two schemas named on the command line;
* it hands the application its own account, scoped to the new schema, rather than
  reusing that admin account for the app.

Foreign keys are disabled only for the duration of the copy. Source and target
definitions are identical, so ordering is the sole thing at stake, and turning
the checks off is what makes an arbitrary table order safe.
"""

import argparse
import asyncio
import secrets
import string
import sys
from pathlib import Path

import aiomysql

PASSWORD_ALPHABET = string.ascii_letters + string.digits


def generate_password(length: int = 32) -> str:
    """Alphanumeric only: it ends up in a .env value and a SQL literal, and a
    symbol that needs escaping in either place is a support ticket waiting."""
    return "".join(secrets.choice(PASSWORD_ALPHABET) for _ in range(length))


async def _admin_conn(args, db=None):
    return await aiomysql.connect(
        host=args.host,
        port=args.port,
        user=args.admin_user,
        password=args.admin_password,
        db=db,
        charset="utf8mb4",
        autocommit=True,
    )


async def create_schema(args) -> None:
    conn = await _admin_conn(args)
    try:
        cur = await conn.cursor()
        await cur.execute(
            f"CREATE DATABASE IF NOT EXISTS `{args.target_db}` "
            "DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
        )
        await cur.execute(
            "SELECT DEFAULT_CHARACTER_SET_NAME, DEFAULT_COLLATION_NAME "
            "FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = %s",
            (args.target_db,),
        )
        charset, collation = await cur.fetchone()
        print(f"[库] {args.target_db}  {charset} / {collation}")
        await cur.close()
    finally:
        conn.close()


async def provision_app_user(args, password: str) -> None:
    conn = await _admin_conn(args)
    try:
        cur = await conn.cursor()
        # The host is bound as a parameter instead of written as a literal '%':
        # aiomysql applies %-substitution whenever arguments are supplied, so a
        # bare '%' in the SQL is read as a format specifier and raises
        # "unsupported format character".
        await cur.execute(
            "CREATE USER IF NOT EXISTS %s@%s IDENTIFIED BY %s",
            (args.app_user, "%", password),
        )
        # Idempotent re-runs: if the account already existed, make the password
        # the one we are about to write into .env rather than a stale one.
        await cur.execute(
            "ALTER USER %s@%s IDENTIFIED BY %s", (args.app_user, "%", password)
        )
        await cur.execute(
            f"GRANT ALL PRIVILEGES ON `{args.target_db}`.* TO %s@%s",
            (args.app_user, "%"),
        )
        await cur.execute("FLUSH PRIVILEGES")
        print(f"[账号] {args.app_user}@% 授权范围仅 `{args.target_db}`.*")
        await cur.close()
    finally:
        conn.close()


async def copy_tables(args) -> None:
    src = await _admin_conn(args, db=args.source_db)
    dst = await _admin_conn(args, db=args.target_db)
    try:
        scur = await src.cursor()
        await scur.execute("SHOW TABLES")
        tables = [row[0] for row in await scur.fetchall()]
        print(f"[表] 源库共 {len(tables)} 张")

        dcur = await dst.cursor()
        # Order is arbitrary on purpose -- checks are off for the whole copy, so
        # a table may safely be created before the table it references.
        await dcur.execute("SET FOREIGN_KEY_CHECKS=0")

        for table in tables:
            await scur.execute(f"SHOW CREATE TABLE `{table}`")
            ddl = (await scur.fetchone())[1]
            # `IF NOT EXISTS` makes the script re-runnable after a partial run.
            await dcur.execute(ddl.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1))

            await dcur.execute(f"DELETE FROM `{table}`")
            await dcur.execute(
                f"INSERT INTO `{table}` SELECT * FROM `{args.source_db}`.`{table}`"
            )
            copied = dcur.rowcount

            await scur.execute(f"SELECT COUNT(*) FROM `{table}`")
            expected = (await scur.fetchone())[0]
            flag = "OK" if copied == expected else "!! 行数不符"
            print(f"     {table:<26} {copied:>6} / {expected:<6} {flag}")

        await dcur.execute("SET FOREIGN_KEY_CHECKS=1")
        await dcur.close()
        await scur.close()
    finally:
        src.close()
        dst.close()


def write_env(path: Path, args, password: str) -> None:
    path.write_text(
        f"""# 服务器部署环境配置（已被 .gitignore 忽略，不会提交）
#
# 由 tools/clone_database.py 生成。库名改了之后必须同时确认账号对该库有权限：
# 应用启动时会校验「实际连到的库」是否等于下面的 MYSQL_DB_NAME，不符会拒绝启动。
MC_HOST=0.0.0.0
MC_PORT=18051

# --- 数据库（监控层） ---
MYSQL_DB_HOST={args.host}
MYSQL_DB_PORT={args.port}
MYSQL_DB_USER={args.app_user}
MYSQL_DB_PWD={password}
MYSQL_DB_NAME={args.target_db}

# 登录鉴权：留空则首次启动自动生成随机密码并打印在启动日志里
# MC_PASSWORD=

# 面板走 HTTPS 时才打开；局域网明文 HTTP 下必须保持注释，
# 否则浏览器丢弃 Cookie，表现为登录页反复刷新且无任何报错
# MC_COOKIE_SECURE=1
""",
        encoding="utf-8",
    )
    path.chmod(0o600)


async def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--host", default="192.168.2.27")
    ap.add_argument("--port", type=int, default=3306)
    ap.add_argument("--admin-user", required=True)
    ap.add_argument("--admin-password", required=True)
    ap.add_argument("--source-db", required=True)
    ap.add_argument("--target-db", required=True)
    ap.add_argument("--app-user", required=True)
    ap.add_argument("--env-file", type=Path)
    args = ap.parse_args()

    if args.source_db == args.target_db:
        print("源库与目标库相同，拒绝执行", file=sys.stderr)
        return 2

    password = generate_password()

    await create_schema(args)
    await provision_app_user(args, password)
    await copy_tables(args)

    if args.env_file:
        write_env(args.env_file, args, password)
        print(f"[.env] 已写入 {args.env_file}（权限 600，密码未回显）")

    print("\n完成。")
    return 0


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