fastapi-demo/database_manager.py

"""
数据库管理脚本
用于初始化数据库、运行迁移、插入测试数据等
"""
import os
import sys
import asyncio
from sqlalchemy import create_engine, text
from app.core.config import settings
from app.database.models import Base


def create_database():
    """创建数据库（如果不存在）"""
    # 连接到MySQL服务器（不指定数据库）
    server_url = f"mysql+pymysql://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/"
    engine = create_engine(server_url)
    
    with engine.connect() as conn:
        # 创建数据库
        conn.execute(text(f"CREATE DATABASE IF NOT EXISTS `{settings.DB_NAME}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"))
        print(f"✅ 数据库 '{settings.DB_NAME}' 创建成功")


def create_tables():
    """创建数据库表"""
    engine = create_engine(settings.DATABASE_URL)
    Base.metadata.create_all(bind=engine)
    print("✅ 数据库表创建成功")


def run_sql_file(file_path):
    """运行SQL文件"""
    if not os.path.exists(file_path):
        print(f"❌ SQL文件不存在: {file_path}")
        return
    
    engine = create_engine(settings.DATABASE_URL)
    
    with open(file_path, 'r', encoding='utf-8') as f:
        sql_content = f.read()
    
    with engine.connect() as conn:
        # 分割SQL语句并执行
        statements = [stmt.strip() for stmt in sql_content.split(';') if stmt.strip()]
        for statement in statements:
            if statement:
                conn.execute(text(statement))
        conn.commit()
    
    print(f"✅ SQL文件执行成功: {file_path}")


def init_database():
    """初始化数据库"""
    print("🚀 开始初始化数据库...")
    
    # 1. 创建数据库
    create_database()
    
    # 2. 创建表
    create_tables()
    
    # 3. 插入示例数据
    sample_data_file = "sql/seeds/sample_data.sql"
    if os.path.exists(sample_data_file):
        run_sql_file(sample_data_file)
    
    print("🎉 数据库初始化完成！")


def reset_database():
    """重置数据库"""
    print("⚠️  重置数据库...")
    
    # 删除并重新创建数据库
    server_url = f"mysql+pymysql://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/"
    engine = create_engine(server_url)
    
    with engine.connect() as conn:
        conn.execute(text(f"DROP DATABASE IF EXISTS `{settings.DB_NAME}`"))
        print(f"🗑️  数据库 '{settings.DB_NAME}' 已删除")
    
    # 重新初始化
    init_database()


def show_help():
    """显示帮助信息"""
    print("""
数据库管理脚本使用说明:

python database_manager.py init      - 初始化数据库（创建数据库和表）
python database_manager.py reset     - 重置数据库（删除并重新创建）
python database_manager.py create-db - 仅创建数据库
python database_manager.py create-tables - 仅创建表
python database_manager.py seed      - 仅插入示例数据
python database_manager.py help      - 显示此帮助信息

环境变量配置:
- DB_HOST: 数据库主机 (默认: localhost)
- DB_PORT: 数据库端口 (默认: 3306)
- DB_USER: 数据库用户名 (默认: root)
- DB_PASSWORD: 数据库密码 (默认: 空)
- DB_NAME: 数据库名称 (默认: fastapi_demo)
""")


def main():
    """主函数"""
    if len(sys.argv) < 2:
        show_help()
        return
    
    command = sys.argv[1].lower()
    
    try:
        if command == "init":
            init_database()
        elif command == "reset":
            reset_database()
        elif command == "create-db":
            create_database()
        elif command == "create-tables":
            create_tables()
        elif command == "seed":
            run_sql_file("sql/seeds/sample_data.sql")
        elif command == "help":
            show_help()
        else:
            print(f"❌ 未知命令: {command}")
            show_help()
    except Exception as e:
        print(f"❌ 执行失败: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()