charging-cabinet/tools/db.py

#!/usr/bin/env python3
"""
PMS 数据库工具
用法：
  需先运行 `pwsh tools/dev-tunnel.ps1` 建立 SSH 隧道（默认连 127.0.0.1:13306）；
  可用环境变量 PMS_MYSQL_HOST / PMS_MYSQL_PORT 指向其它可达地址。
  ⚠️ 默认目标是**测试库 `pms_dev`**（可用 PMS_MYSQL_DB 覆盖；取 `pms` = 生产库，
     脚本会额外打印一行醒目警告）。

  python tools/db.py import <sql文件>     # 导入SQL文件
  python tools/db.py query <SQL语句>      # 执行查询
  python tools/db.py exec <SQL语句>       # 执行写入
  python tools/db.py tables               # 列出所有表
  python tools/db.py count                # 统计各表数据量
"""

import os
import sys
import pymysql

# 数据库配置：默认指向 SSH 隧道（tools/dev-tunnel.ps1），不再使用已移除的 ZeroTier/nginx 转发。
# 可用环境变量 PMS_MYSQL_HOST / PMS_MYSQL_PORT / PMS_MYSQL_DB 覆盖。
# ⚠️ 默认库是**测试库 `pms_dev`** —— 本工具能执行写操作（`exec`/`import`），
#    默认绝不能落在生产库上（历史上文档写 pms_dev、代码却是 pms：以为在测测试库、实际打生产）。
PROD_DB = 'pms'
DB_CONFIG = {
    'host': os.environ.get('PMS_MYSQL_HOST', '127.0.0.1'),
    'port': int(os.environ.get('PMS_MYSQL_PORT', '13306')),
    'user': 'root',
    'password': 'Hbhyg731024@',
    'database': os.environ.get('PMS_MYSQL_DB', 'pms_dev'),
    'charset': 'utf8mb4',
}

# 连不上时的统一提示（隧道没起是最常见原因，绝不能让它变成一段栈回溯）
TUNNEL_HINT = ('连接被拒：请先运行 `pwsh tools/dev-tunnel.ps1` 建立 SSH 隧道'
               '（或设置 PMS_MYSQL_HOST/PMS_MYSQL_PORT 指向其它可达地址）')


def announce_target():
    """每次连接前醒目提示当前目标库（防止"以为在测试库、实际打生产"）"""
    db = DB_CONFIG['database']
    print(f"当前目标库：{db}（可用 PMS_MYSQL_DB 覆盖）")
    if db == PROD_DB:
        print(f"⚠️  你正在操作正式库 {PROD_DB} —— 写操作（exec/import）前请再次确认；"
              f"测试库请用 PMS_MYSQL_DB=pms_dev")


def get_conn():
    """建立连接；连不上时打印可读提示并退出（不抛栈回溯）"""
    announce_target()
    try:
        return pymysql.connect(**DB_CONFIG)
    except pymysql.MySQLError as e:
        print(f"✗ MySQL 连接失败 {DB_CONFIG['host']}:{DB_CONFIG['port']}：{e}", file=sys.stderr)
        print(TUNNEL_HINT, file=sys.stderr)
        sys.exit(2)

def import_sql(filepath):
    """导入SQL文件"""
    with open(filepath, 'r', encoding='utf-8') as f:
        sql = f.read()
    
    conn = get_conn()
    cur = conn.cursor()
    
    # 按分号分割SQL语句
    statements = [s.strip() for s in sql.split(';') if s.strip()]
    
    success = 0
    fail = 0
    for stmt in statements:
        try:
            cur.execute(stmt)
            success += 1
        except Exception as e:
            print(f"  ✗ 失败: {str(e)[:80]}")
            fail += 1
    
    conn.commit()
    conn.close()
    print(f"导入完成: 成功 {success}, 失败 {fail}")

def query(sql):
    """执行查询"""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute(sql)
    columns = [desc[0] for desc in cur.description]
    rows = cur.fetchall()
    conn.close()
    
    if not rows:
        print("无数据")
        return
    
    # 打印表头
    print('\t'.join(columns))
    print('-' * 50)
    for row in rows:
        print('\t'.join(str(v) for v in row))

def execute(sql):
    """执行写入"""
    conn = get_conn()
    cur = conn.cursor()
    try:
        cur.execute(sql)
        conn.commit()
        print(f"执行成功, 影响行数: {cur.rowcount}")
    except Exception as e:
        print(f"执行失败: {e}")
        conn.rollback()
    finally:
        conn.close()

def tables():
    """列出所有表"""
    query("SHOW TABLES")

def count():
    """统计各表数据量"""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SHOW TABLES")
    tables = [row[0] for row in cur.fetchall()]
    
    print("表名\t\t数据量")
    print("-" * 30)
    for t in tables:
        cur.execute(f"SELECT COUNT(*) FROM `{t}`")
        cnt = cur.fetchone()[0]
        print(f"{t}\t{cnt}")
    conn.close()

def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return
    
    cmd = sys.argv[1]
    
    if cmd == 'import' and len(sys.argv) >= 3:
        import_sql(sys.argv[2])
    elif cmd == 'query' and len(sys.argv) >= 3:
        query(' '.join(sys.argv[2:]))
    elif cmd == 'exec' and len(sys.argv) >= 3:
        execute(' '.join(sys.argv[2:]))
    elif cmd == 'tables':
        tables()
    elif cmd == 'count':
        count()
    else:
        print(__doc__)

if __name__ == '__main__':
    main()