harness_zbt/backend/app/commands/init_databases.py

import argparse
import getpass
import re

import pymysql

from app.core.config import Settings

SAFE_ACCOUNT = re.compile(r"^[A-Za-z0-9_]+$")


def initialize_databases(root_user: str, root_password: str, settings: Settings) -> None:
    if not SAFE_ACCOUNT.fullmatch(settings.mysql_user):
        raise ValueError("MYSQL_USER 只能包含字母、数字和下划线")
    if not settings.mysql_password:
        raise ValueError("请先在 .env 中设置 MYSQL_PASSWORD")

    connection = pymysql.connect(
        host=settings.mysql_host,
        port=settings.mysql_port,
        user=root_user,
        password=root_password,
        autocommit=True,
    )
    databases = (
        settings.mysql_core_database,
        settings.mysql_agent_database,
        settings.mysql_analytics_database,
    )
    account = f"'{settings.mysql_user}'@'localhost'"
    escaped_password = connection.escape(settings.mysql_password)
    with connection, connection.cursor() as cursor:
        for database in databases:
            cursor.execute(
                f"CREATE DATABASE IF NOT EXISTS `{database}` "
                "CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci"
            )
        cursor.execute(f"CREATE USER IF NOT EXISTS {account} IDENTIFIED BY {escaped_password}")
        for database in databases:
            cursor.execute(f"GRANT ALL PRIVILEGES ON `{database}`.* TO {account}")
        cursor.execute("FLUSH PRIVILEGES")


def main() -> None:
    parser = argparse.ArgumentParser(description="初始化智保通第一阶段 MySQL 三库")
    parser.add_argument("--root-user", default="root")
    args = parser.parse_args()
    password = getpass.getpass(f"MySQL {args.root_user} 密码：")
    settings = Settings()
    initialize_databases(args.root_user, password, settings)
    print(
        "已初始化数据库："
        f"{settings.mysql_core_database}、{settings.mysql_agent_database}、"
        f"{settings.mysql_analytics_database}"
    )


if __name__ == "__main__":
    main()