Kamixitong/setup_mysql.py

#!/usr/bin/env python3
"""
MySQL数据库快速配置脚本

此脚本将帮助您：
1. 检查MySQL配置
2. 创建数据库（如果不存在）
3. 测试数据库连接
4. 初始化表结构
5. 创建默认超级管理员（如果需要）
"""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from app import create_app, db
from app.models import Admin
from werkzeug.security import generate_password_hash

def parse_database_url(url):
    """解析数据库URL"""
    # 格式: mysql+pymysql://user:password@host:port/database
    try:
        # 移除mysql+pymysql://前缀
        url = url.replace('mysql+pymysql://', '')

        # 找到@符号前的用户名密码
        at_index = url.index('@')
        user_pass = url[:at_index]
        url = url[at_index + 1:]

        # 分割用户名和密码
        user, password = user_pass.split(':', 1)

        # 找到数据库名
        if '/' in url:
            host_port, database = url.split('/', 1)
            # 提取端口
            if ':' in host_port:
                host, port = host_port.split(':', 1)
            else:
                host = host_port
                port = '3306'
        else:
            raise ValueError("URL格式不正确")

        return {
            'user': user,
            'password': password,
            'host': host,
            'port': port,
            'database': database
        }
    except Exception as e:
        print(f"❌ 无法解析数据库URL: {e}")
        return None

def check_mysql_dependencies():
    """检查MySQL依赖"""
    print("\n" + "=" * 60)
    print("1. 检查MySQL依赖")
    print("=" * 60)

    try:
        import pymysql
        print(f"✅ PyMySQL已安装 (版本: {pymysql.__version__})")
        return True
    except ImportError:
        print("❌ PyMySQL未安装")
        print("请运行: pip install PyMySQL")
        return False

def check_dotenv():
    """检查.env文件加载"""
    print("\n" + "=" * 60)
    print("2. 检查.env文件")
    print("=" * 60)

    env_file = '.env'
    if os.path.exists(env_file):
        print(f"✅ 找到.env文件: {os.path.abspath(env_file)}")

        # 读取DATABASE_URL
        with open(env_file, 'r', encoding='utf-8') as f:
            for line in f:
                if line.startswith('DATABASE_URL='):
                    url = line.split('=', 1)[1].strip()
                    print(f"✅ DATABASE_URL已配置: {url}")
                    return url
        print("⚠️  .env文件中未找到DATABASE_URL")
        return None
    else:
        print(f"❌ 未找到.env文件: {os.path.abspath(env_file)}")
        return None

def test_database_connection():
    """测试数据库连接"""
    print("\n" + "=" * 60)
    print("3. 测试数据库连接")
    print("=" * 60)

    app = create_app()

    with app.app_context():
        try:
            # 测试连接
            db.create_all()
            print("✅ 数据库连接成功！")
            return True
        except Exception as e:
            print(f"❌ 数据库连接失败: {e}")
            return False

def create_database():
    """创建数据库（如果不存在）"""
    print("\n" + "=" * 60)
    print("4. 创建数据库")
    print("=" * 60)

    # 从.env文件获取数据库配置
    env_file = '.env'
    if not os.path.exists(env_file):
        print("❌ .env文件不存在")
        return False

    with open(env_file, 'r', encoding='utf-8') as f:
        for line in f:
            if line.startswith('DATABASE_URL='):
                db_url = line.split('=', 1)[1].strip()
                break

    if not db_url:
        print("❌ 未找到DATABASE_URL配置")
        return False

    # 解析URL
    db_config = parse_database_url(db_url)
    if not db_config:
        return False

    print(f"数据库信息:")
    print(f"   主机: {db_config['host']}")
    print(f"   端口: {db_config['port']}")
    print(f"   数据库: {db_config['database']}")
    print(f"   用户: {db_config['user']}")

    try:
        import pymysql

        # 连接到MySQL服务器（不指定数据库）
        connection = pymysql.connect(
            host=db_config['host'],
            port=int(db_config['port']),
            user=db_config['user'],
            password=db_config['password'],
            charset='utf8mb4'
        )

        with connection.cursor() as cursor:
            # 检查数据库是否存在
            cursor.execute(f"SHOW DATABASES LIKE '{db_config['database']}'")
            result = cursor.fetchone()

            if not result:
                print(f"📦 正在创建数据库 '{db_config['database']}'...")
                cursor.execute(
                    f"CREATE DATABASE {db_config['database']} "
                    "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
                )
                print(f"✅ 数据库 '{db_config['database']}' 创建成功")
            else:
                print(f"✅ 数据库 '{db_config['database']}' 已存在")

        connection.close()
        return True

    except ImportError:
        print("❌ PyMySQL未安装，请运行: pip install PyMySQL")
        return False
    except Exception as e:
        print(f"❌ 创建数据库失败: {e}")
        print("请手动创建数据库，或检查MySQL服务是否运行")
        return False

def init_database():
    """初始化数据库表"""
    print("\n" + "=" * 60)
    print("5. 初始化数据库表")
    print("=" * 60)

    app = create_app()

    with app.app_context():
        try:
            # 检查admin表
            admin_columns = [c.name for c in Admin.__table__.columns]
            print(f"📋 Admin表字段: {', '.join(admin_columns)}")

            # 检查新字段
            if 'is_deleted' not in admin_columns:
                print("⚠️  缺少软删除字段，正在添加...")
                db.session.execute("ALTER TABLE admin ADD COLUMN is_deleted INT NOT NULL DEFAULT 0")
                db.session.execute("ALTER TABLE admin ADD COLUMN delete_time DATETIME NULL")
                db.session.execute("CREATE INDEX ix_admin_is_deleted ON admin(is_deleted)")
                db.session.commit()
                print("✅ 软删除字段添加成功")

            # 检查audit_log表
            try:
                from app.models import AuditLog
                AuditLog.query.count()
                print("✅ audit_log表存在")
            except Exception as e:
                print(f"⚠️  audit_log表不存在，正在创建...")
                db.session.execute("""
                    CREATE TABLE audit_log (
                        log_id INT AUTO_INCREMENT PRIMARY KEY,
                        admin_id INT NOT NULL,
                        action VARCHAR(32) NOT NULL,
                        target_type VARCHAR(32) NOT NULL,
                        target_id INT NULL,
                        details TEXT NULL,
                        ip_address VARCHAR(32) NULL,
                        user_agent VARCHAR(256) NULL,
                        create_time DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
                        FOREIGN KEY (admin_id) REFERENCES admin(admin_id)
                    )
                """)
                db.session.execute("CREATE INDEX ix_audit_log_admin_id ON audit_log(admin_id)")
                db.session.execute("CREATE INDEX ix_audit_log_action ON audit_log(action)")
                db.session.execute("CREATE INDEX ix_audit_log_create_time ON audit_log(create_time)")
                db.session.commit()
                print("✅ audit_log表创建成功")

            print("✅ 数据库表初始化成功")
            return True

        except Exception as e:
            db.session.rollback()
            print(f"❌ 数据库表初始化失败: {e}")
            import traceback
            traceback.print_exc()
            return False

def create_default_admin():
    """创建默认超级管理员（如果需要）"""
    print("\n" + "=" * 60)
    print("6. 创建默认超级管理员")
    print("=" * 60)

    app = create_app()

    with app.app_context():
        try:
            # 检查是否已有超级管理员
            super_admin = Admin.query.filter_by(role=1, status=1, is_deleted=0).first()

            if super_admin:
                print(f"✅ 已存在超级管理员: {super_admin.username}")
                return True

            print("📝 没有超级管理员，创建默认账号...")
            print("   用户名: admin")
            print("   密码: admin123456")
            print("   ⚠️  登录后请立即修改密码！")

            admin = Admin(
                username='admin',
                email='admin@example.com',
                role=1,
                status=1
            )
            admin.set_password('admin123456')

            db.session.add(admin)
            db.session.commit()

            print("✅ 默认超级管理员创建成功")
            return True

        except Exception as e:
            db.session.rollback()
            print(f"❌ 创建超级管理员失败: {e}")
            import traceback
            traceback.print_exc()
            return False

def main():
    """主函数"""
    print("\n" + "=" * 60)
    print("🚀 MySQL数据库快速配置工具")
    print("=" * 60)

    success = True

    # 1. 检查依赖
    if not check_mysql_dependencies():
        print("\n⚠️  请先安装PyMySQL: pip install PyMySQL")
        return 1

    # 2. 检查.env文件
    db_url = check_dotenv()
    if not db_url:
        print("\n⚠️  请配置.env文件中的DATABASE_URL")
        return 1

    # 3. 创建数据库
    if not create_database():
        success = False

    # 4. 测试连接
    if success and not test_database_connection():
        success = False

    # 5. 初始化表
    if success and not init_database():
        success = False

    # 6. 创建默认管理员
    if success and not create_default_admin():
        success = False

    # 总结
    print("\n" + "=" * 60)
    if success:
        print("🎉 MySQL配置完成！")
        print("=" * 60)
        print("\n✅ 现在可以启动应用:")
        print("   python run.py")
        print("\n📖 默认管理员登录信息:")
        print("   用户名: admin")
        print("   密码: admin123456")
        print("   ⚠️  登录后请立即修改密码！")
        return 0
    else:
        print("❌ 配置未完成，请查看错误信息")
        print("=" * 60)
        return 1

if __name__ == '__main__':
    sys.exit(main())