all-project/用户信息同步统一身份认证/db_helper.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""数据库操作辅助脚本，供 shell 同步脚本调用。"""

import argparse
import sys
from pathlib import Path

try:
    import pymysql
except ImportError:
    print("ERROR: 需要安装 pymysql: pip install pymysql", file=sys.stderr)
    sys.exit(1)


def connect_mysql(host, port, user, password, database):
    return pymysql.connect(
        host=host,
        port=int(port),
        user=user,
        password=password,
        database=database,
        charset="utf8mb4",
        autocommit=True,
        cursorclass=pymysql.cursors.Cursor,
    )


def exec_sql_file(conn, sql_file):
    sql = Path(sql_file).read_text(encoding="utf-8")
    statements = [s.strip() for s in sql.split(";") if s.strip()]
    with conn.cursor() as cur:
        for stmt in statements:
            if stmt.upper().startswith("SET "):
                cur.execute(stmt)
                continue
            cur.execute(stmt)


def exec_sql(conn, sql):
    with conn.cursor() as cur:
        cur.execute(sql)
        if cur.description:
            return cur.fetchall()
    return None


def main():
    parser = argparse.ArgumentParser(description="MySQL 数据库辅助工具")
    parser.add_argument("--host", required=True)
    parser.add_argument("--port", default="3306")
    parser.add_argument("--user", required=True)
    parser.add_argument("--password", required=True)
    parser.add_argument("--database", required=True)
    sub = parser.add_subparsers(dest="action", required=True)

    p_file = sub.add_parser("file", help="执行 SQL 文件")
    p_file.add_argument("sql_file")

    p_query = sub.add_parser("query", help="执行单条 SQL 并输出结果")
    p_query.add_argument("sql")

    p_scalar = sub.add_parser("scalar", help="执行 SQL 并输出单个值")
    p_scalar.add_argument("sql")

    args = parser.parse_args()
    conn = connect_mysql(args.host, args.port, args.user, args.password, args.database)
    try:
        if args.action == "file":
            exec_sql_file(conn, args.sql_file)
            print("OK")
        elif args.action == "query":
            rows = exec_sql(conn, args.sql)
            for row in rows or []:
                print("\t".join("" if v is None else str(v) for v in row))
        elif args.action == "scalar":
            rows = exec_sql(conn, args.sql)
            if rows:
                print(rows[0][0])
    finally:
        conn.close()


if __name__ == "__main__":
    main()