fix_database.py

# -*- coding: utf-8 -*-
# fix_database.py
"""
修复数据库：为 chat_history 表添加 username 字段
"""
from sqlalchemy import create_engine, text

# 【请修改】你的数据库配置（与 server.py 保持一致）
SQLALCHEMY_DATABASE_URL = "mysql+pymysql://root:123456@localhost:3306/rag_db"

engine = create_engine(SQLALCHEMY_DATABASE_URL)


def fix_chat_history_table():
    """为 chat_history 表添加 username 字段"""
    print("=" * 50)
    print("🔧 正在修复数据库表结构...")
    print("=" * 50)
    
    try:
        with engine.connect() as conn:
            # 检查字段是否已存在
            result = conn.execute(text("""
                SELECT COUNT(*) as count 
                FROM information_schema.COLUMNS 
                WHERE TABLE_SCHEMA = 'rag_db' 
                AND TABLE_NAME = 'chat_history' 
                AND COLUMN_NAME = 'username'
            """))
            
            exists = result.fetchone()[0] > 0
            
            if exists:
                print("✅ username 字段已存在，无需修复")
            else:
                print("⚠️  检测到缺少 username 字段，正在添加...")
                
                # 添加 username 字段
                conn.execute(text("""
                    ALTER TABLE chat_history 
                    ADD COLUMN username VARCHAR(50) DEFAULT 'default' AFTER id
                """))
                conn.commit()
                print("✅ 已添加 username 字段")
                
                # 添加索引以提升查询性能
                try:
                    conn.execute(text("""
                        ALTER TABLE chat_history 
                        ADD INDEX idx_username (username)
                    """))
                    conn.commit()
                    print("✅ 已添加 username 索引")
                except Exception as e:
                    print(f"⚠️  索引可能已存在: {e}")
                
                # 更新现有记录的 username 为 'default'
                conn.execute(text("""
                    UPDATE chat_history 
                    SET username = 'default' 
                    WHERE username IS NULL OR username = ''
                """))
                conn.commit()
                print("✅ 已更新现有记录的 username")
        
        print("\n" + "=" * 50)
        print("✅ 数据库修复完成！")
        print("=" * 50)
        return True
        
    except Exception as e:
        print(f"\n❌ 修复失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    fix_chat_history_table()