ai_game/server/sql/init_db.py

"""
数据库初始化脚本
在IDEA中直接运行此文件即可建库
"""
import os
import pymysql
from pathlib import Path

# 获取当前脚本所在目录
SQL_DIR = Path(__file__).parent

# 从.env文件读取配置
def load_env():
    env_file = SQL_DIR.parent / '.env'
    config = {}
    if env_file.exists():
        for line in env_file.read_text(encoding='utf-8').splitlines():
            if '=' in line and not line.startswith('#'):
                k, v = line.split('=', 1)
                config[k.strip()] = v.strip()
    return config

env_config = load_env()

# 数据库配置（优先从.env读取）
DB_CONFIG = {
    'host': env_config.get('DB_HOST', os.getenv('DB_HOST', 'localhost')),
    'port': int(env_config.get('DB_PORT', os.getenv('DB_PORT', 3306))),
    'user': env_config.get('DB_USER', os.getenv('DB_USER', 'root')),
    'password': env_config.get('DB_PASSWORD', os.getenv('DB_PASSWORD', '')),
    'charset': 'utf8mb4'
}


def read_sql_file(filename: str) -> str:
    """读取SQL文件"""
    filepath = SQL_DIR / filename
    with open(filepath, 'r', encoding='utf-8') as f:
        return f.read()


def execute_sql(cursor, sql: str, description: str):
    """执行SQL语句（支持多语句）"""
    print(f'{description}...')
    # 按分号分割SQL语句
    statements = [s.strip() for s in sql.split(';') if s.strip()]
    for stmt in statements:
        if stmt:
            try:
                cursor.execute(stmt)
            except pymysql.Error as e:
                # 忽略某些错误（如数据库已存在）
                if e.args[0] not in [1007, 1050]:  # 数据库已存在、表已存在
                    print(f'  警告: {e.args[1]}')
    print(f'  {description}完成！')


def init_database():
    """初始化数据库"""
    print('=' * 50)
    print('星域故事汇 - 数据库初始化')
    print('=' * 50)
    print(f'连接信息: {DB_CONFIG["user"]}@{DB_CONFIG["host"]}:{DB_CONFIG["port"]}')
    print()

    # 连接MySQL
    connection = pymysql.connect(**DB_CONFIG)
    cursor = connection.cursor()

    try:
        # 1. 执行schema.sql（创建数据库和表）
        schema_sql = read_sql_file('schema.sql')
        execute_sql(cursor, schema_sql, '创建数据库表结构')
        connection.commit()

        # 2. 执行种子数据第1部分
        seed1_sql = read_sql_file('seed_stories_part1.sql')
        execute_sql(cursor, seed1_sql, '导入种子数据（第1部分）')
        connection.commit()

        # 3. 执行种子数据第2部分
        seed2_sql = read_sql_file('seed_stories_part2.sql')
        execute_sql(cursor, seed2_sql, '导入种子数据（第2部分）')
        connection.commit()

        print()
        print('=' * 50)
        print('数据库初始化完成！')
        print('共创建10个种子故事，包含66个剧情节点和多个结局分支。')
        print('=' * 50)
        print()
        print('现在可以启动后端服务了：')
        print('  cd "D:\\idea project\\ai_game\\server"')
        print('  python -m app.main')

    except Exception as e:
        print(f'初始化失败: {e}')
        connection.rollback()
        raise
    finally:
        cursor.close()
        connection.close()


if __name__ == '__main__':
    init_database()