skills/bms-mysql-prd/scripts/bms-mysql-prd.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
BMS MySQL Production Database Query Tool

Read-only access to the BMS production MySQL database.
- SELECT/SHOW/DESC/EXPLAIN: execute directly
- INSERT/UPDATE/DELETE/DDL: require confirmation popup (will fail with permission error)
"""

import argparse
import re
import sys

import pymysql


# ─── Connection Config ───────────────────────────────────────────────────────

DB_CONFIG = {
    "host": "47.106.125.251",
    "port": 9696,
    "user": "i7bSEtwU",
    "password": "",
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}

# ─── Query Classification ────────────────────────────────────────────────────

READ_ONLY_PATTERNS = re.compile(
    r"^\s*(SELECT|SHOW|DESC|DESCRIBE|EXPLAIN|USE|SET)\b", re.IGNORECASE
)

WRITE_DDL_PATTERNS = re.compile(
    r"^\s*(INSERT|UPDATE|DELETE|REPLACE|CREATE|ALTER|DROP|TRUNCATE|RENAME|GRANT|REVOKE|LOAD\s+DATA)\b",
    re.IGNORECASE,
)


def is_read_only(sql: str) -> bool:
    """Return True if the SQL is a read-only query."""
    return bool(READ_ONLY_PATTERNS.match(sql))


def is_write_or_ddl(sql: str) -> bool:
    """Return True if the SQL modifies data or schema."""
    return bool(WRITE_DDL_PATTERNS.match(sql))


# ─── Confirmation Dialog ─────────────────────────────────────────────────────

def ask_confirmation(sql: str) -> bool:
    """Show a Windows popup dialog asking the user to confirm the SQL execution.

    Returns True if the user clicks 'Yes', False otherwise.
    """
    import tkinter as tk
    from tkinter import messagebox

    root = tk.Tk()
    root.withdraw()  # Hide the main window

    title = "BMS MySQL 生产 - 写操作确认"
    message = (
        "即将执行以下写操作 / DDL 语句（生产库只读，该操作将报错），是否继续？\n\n"
        f"SQL:\n{sql}\n"
    )

    result = messagebox.askyesno(title, message, icon="warning")
    root.destroy()
    return result


def show_error_popup(sql: str, error: str) -> None:
    """Show a Windows popup dialog when a write operation is rejected."""
    import tkinter as tk
    from tkinter import messagebox

    root = tk.Tk()
    root.withdraw()

    title = "BMS MySQL 生产 - 操作被拒绝"
    message = (
        "生产库为只读权限，以下写操作已被数据库拒绝：\n\n"
        f"SQL:\n{sql}\n\n"
        f"错误信息:\n{error}"
    )

    messagebox.showwarning(title, message)
    root.destroy()


# ─── Table Formatting ────────────────────────────────────────────────────────

def format_table(rows: list, columns: list) -> str:
    """Format query results as an aligned text table."""
    if not rows:
        return "0 rows returned."

    # Calculate column widths
    widths = {col: len(str(col)) for col in columns}
    for row in rows:
        for col in columns:
            val = str(row.get(col, ""))
            widths[col] = max(widths[col], len(val))

    # Build header
    header = " | ".join(str(col).ljust(widths[col]) for col in columns)
    separator = "-+-".join("-" * widths[col] for col in columns)

    # Build rows
    lines = [header, separator]
    for row in rows:
        line = " | ".join(
            str(row.get(col, "")).ljust(widths[col]) for col in columns
        )
        lines.append(line)

    lines.append(f"\n{len(rows)} row(s) returned.")
    return "\n".join(lines)


# ─── Main Execution ──────────────────────────────────────────────────────────

def execute_query(sql: str) -> None:
    """Connect to the database and execute the given SQL query."""

    print(f"\n{'='*60}")
    print(f"BMS MySQL Production (47.106.125.251:9696)")
    print(f"{'='*60}")
    print(f"SQL: {sql}\n")

    # Classify the query
    if is_read_only(sql):
        query_type = "READ-ONLY"
    elif is_write_or_ddl(sql):
        query_type = "WRITE / DDL (只读库，将报错)"
    else:
        query_type = "UNKNOWN"

    print(f"类型: {query_type}\n")

    # For write/DDL, require confirmation
    if is_write_or_ddl(sql):
        if not ask_confirmation(sql):
            print("用户已取消操作。")
            sys.exit(0)

    # Connect and execute
    try:
        conn = pymysql.connect(**DB_CONFIG)
        cursor = conn.cursor()
        cursor.execute(sql)

        if is_read_only(sql):
            rows = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            print(format_table(rows, columns))
        else:
            conn.commit()
            affected = cursor.rowcount
            print(f"操作成功，影响 {affected} 行。")

        cursor.close()
        conn.close()

    except pymysql.MySQLError as e:
        error_msg = str(e)
        print(f"数据库错误: {error_msg}", file=sys.stderr)
        if is_write_or_ddl(sql):
            show_error_popup(sql, error_msg)
        sys.exit(1)
    except Exception as e:
        print(f"执行失败: {e}", file=sys.stderr)
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="BMS MySQL 生产数据库查询工具（只读）"
    )
    parser.add_argument(
        "--query", "-q",
        required=True,
        help="要执行的 SQL 语句",
    )
    args = parser.parse_args()
    execute_query(args.query)


if __name__ == "__main__":
    main()