Tolan/backend/init_database.py

"""
数据库表初始化脚本
运行: python init_database.py
"""
import asyncio
import aiomysql
from config import settings
import sys


# 用户表
CREATE_USERS_TABLE = """
CREATE TABLE IF NOT EXISTS users (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    user_id VARCHAR(64) UNIQUE NOT NULL COMMENT '用户唯一标识(UUID)',
    
    -- 基本信息
    username VARCHAR(100) NOT NULL COMMENT '用户名/昵称',
    user_type ENUM('guest', 'registered') DEFAULT 'guest' COMMENT '用户类型',
    bot_name VARCHAR(100) COMMENT '用户给机器人取的名字',
    
    -- 注册信息（可选）
    phone VARCHAR(20) UNIQUE COMMENT '手机号',
    wechat_openid VARCHAR(100) UNIQUE COMMENT '微信OpenID',
    google_email VARCHAR(200) UNIQUE COMMENT 'Google邮箱',
    
    -- 认证信息
    is_verified BOOLEAN DEFAULT FALSE COMMENT '是否已认证',
    registered_at TIMESTAMP NULL COMMENT '注册时间',
    
    -- 统计信息
    total_messages INT DEFAULT 0 COMMENT '总消息数',
    total_rounds INT DEFAULT 0 COMMENT '总对话轮次',
    total_speech_duration_ms BIGINT DEFAULT 0 COMMENT '累积说话时长（毫秒）',
    last_active_at TIMESTAMP NULL COMMENT '最后活跃时间',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    
    INDEX idx_user_id (user_id),
    INDEX idx_phone (phone),
    INDEX idx_wechat_openid (wechat_openid),
    INDEX idx_google_email (google_email),
    INDEX idx_user_type (user_type)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='用户表'
"""

# 消息表
CREATE_MESSAGES_TABLE = """
CREATE TABLE IF NOT EXISTS messages (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    message_id VARCHAR(64) UNIQUE NOT NULL COMMENT '消息ID',
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    
    -- 消息内容
    role ENUM('user', 'assistant', 'system') NOT NULL COMMENT '角色',
    content TEXT NOT NULL COMMENT '消息内容',
    
    -- 轮次信息
    round_number INT NOT NULL COMMENT '对话轮次(从1开始)',
    
    -- 扩展信息
    emotion VARCHAR(50) COMMENT '情感标签',
    metadata JSON COMMENT '扩展信息(音频路径、识别耗时等)',
    token_count INT COMMENT 'Token数量',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    
    INDEX idx_user_id (user_id),
    INDEX idx_user_round (user_id, round_number),
    INDEX idx_created_at (created_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='消息表'
"""

# 对话总结表
CREATE_SUMMARIES_TABLE = """
CREATE TABLE IF NOT EXISTS conversation_summaries (
    id BIGINT PRIMARY KEY AUTO_INCREMENT,
    summary_id VARCHAR(64) UNIQUE NOT NULL COMMENT '总结ID',
    user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
    
    -- 总结范围
    start_round INT NOT NULL COMMENT '起始轮次',
    end_round INT NOT NULL COMMENT '结束轮次',
    
    -- 总结内容
    summary_text TEXT NOT NULL COMMENT '总结内容',
    summary_type ENUM('auto', 'manual') DEFAULT 'auto' COMMENT '总结类型',
    
    -- 统计信息
    message_count INT COMMENT '总结的消息数量',
    token_count INT COMMENT '总结文本Token数',
    
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    
    INDEX idx_user_id (user_id),
    INDEX idx_user_end_round (user_id, end_round),
    UNIQUE KEY uk_user_round_range (user_id, start_round, end_round)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
COMMENT='对话总结表'
"""


async def init_tables():
    """初始化数据库表"""
    
    try:
        # 创建数据库连接
        conn = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset='utf8mb4'
        )
        
        async with conn.cursor() as cursor:
            # 创建用户表
            await cursor.execute(CREATE_USERS_TABLE)
            
            # 创建消息表
            await cursor.execute(CREATE_MESSAGES_TABLE)
            
            # 创建总结表
            await cursor.execute(CREATE_SUMMARIES_TABLE)
            
            await conn.commit()
        
        conn.close()
        
        return True
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return False


async def check_tables():
    """检查表是否存在"""
    try:
        conn = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset='utf8mb4'
        )
        
        async with conn.cursor() as cursor:
            await cursor.execute("SHOW TABLES")
            tables = await cursor.fetchall()
        
        conn.close()
        
    except Exception as e:
        pass


if __name__ == "__main__":
    # 确认是否继续
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续初始化数据库表? (y/n): ")
    
    if confirm.lower() != 'y':
        sys.exit(0)
    
    # 执行初始化
    success = asyncio.run(init_tables())
    
    if success:
        # 检查创建的表
        asyncio.run(check_tables())
    else:
        sys.exit(1)