aitsc/apply_database_optimization.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库性能优化脚本
应用索引优化，提升查询性能
"""
import pymysql
import sys
import os

# 设置环境变量
os.environ['SECRET_KEY'] = 'dev-key'

# 直接使用数据库连接字符串
DATABASE_URI = 'mysql+pymysql://root:123456@localhost:3306/pro_db?charset=utf8mb4'

def apply_database_optimization():
    """应用数据库优化"""
    try:
        # 解析数据库连接信息
        db_uri = DATABASE_URI
        # 从 mysql+pymysql://root:123456@localhost:3306/pro_db?charset=utf8mb4 解析
        if 'mysql+pymysql://' in db_uri:
            uri_part = db_uri.replace('mysql+pymysql://', '')
            auth_host, db_part = uri_part.split('@')
            username, password = auth_host.split(':')
            host_port, db_name = db_part.split('/')
            if ':' in host_port:
                host, port = host_port.split(':')
                port = int(port)
            else:
                host = host_port
                port = 3306
            db_name = db_name.split('?')[0]
        else:
            print("❌ 无法解析数据库连接URI")
            return False
        
        print(f"🔗 连接数据库: {host}:{port}/{db_name}")
        
        # 连接数据库
        conn = pymysql.connect(
            host=host,
            port=port,
            user=username,
            password=password,
            database=db_name,
            charset='utf8mb4'
        )
        cursor = conn.cursor()
        
        print("✅ 数据库连接成功")
        
        # 执行索引优化
        optimization_queries = [
            # 用户表索引
            "CREATE INDEX IF NOT EXISTS idx_user_created_time ON user(created_time)",
            "CREATE INDEX IF NOT EXISTS idx_user_status ON user(status)",
            "CREATE INDEX IF NOT EXISTS idx_user_login_name ON user(login_name)",
            
            # 提示词表索引
            "CREATE INDEX IF NOT EXISTS idx_prompt_created_at ON prompt(created_at)",
            "CREATE INDEX IF NOT EXISTS idx_prompt_user_id ON prompt(user_id)",
            "CREATE INDEX IF NOT EXISTS idx_prompt_wx_user_id ON prompt(wx_user_id)",
            "CREATE INDEX IF NOT EXISTS idx_prompt_created_at_user_id ON prompt(created_at, user_id)",
            
            # 模板表索引
            "CREATE INDEX IF NOT EXISTS idx_prompt_template_is_default ON prompt_template(is_default)",
            "CREATE INDEX IF NOT EXISTS idx_prompt_template_category ON prompt_template(category)",
            
            # 反馈表索引
            "CREATE INDEX IF NOT EXISTS idx_feedback_created_at ON feedback(created_at)",
            "CREATE INDEX IF NOT EXISTS idx_feedback_user_id ON feedback(user_id)",
            
            # 收藏表索引
            "CREATE INDEX IF NOT EXISTS idx_favorites_created_time ON favorites(created_time)",
            "CREATE INDEX IF NOT EXISTS idx_favorites_user_id ON favorites(user_id)",
            "CREATE INDEX IF NOT EXISTS idx_favorites_template_id ON favorites(template_id)",
            
            # 复合索引优化
            "CREATE INDEX IF NOT EXISTS idx_prompt_date_user ON prompt(DATE(created_at), user_id)",
            "CREATE INDEX IF NOT EXISTS idx_user_date_status ON user(DATE(created_time), status)"
        ]
        
        print("📊 开始创建数据库索引...")
        
        for i, query in enumerate(optimization_queries, 1):
            try:
                cursor.execute(query)
                print(f"✅ [{i:2d}/{len(optimization_queries)}] 索引创建成功")
            except Exception as e:
                if "Duplicate key name" in str(e) or "already exists" in str(e):
                    print(f"ℹ️  [{i:2d}/{len(optimization_queries)}] 索引已存在，跳过")
                else:
                    print(f"⚠️  [{i:2d}/{len(optimization_queries)}] 索引创建失败: {str(e)}")
        
        # 提交更改
        conn.commit()
        
        # 显示索引信息
        print("\n📋 当前数据库索引状态:")
        cursor.execute("SHOW INDEX FROM user")
        user_indexes = cursor.fetchall()
        print(f"   user表索引数: {len(user_indexes)}")
        
        cursor.execute("SHOW INDEX FROM prompt")
        prompt_indexes = cursor.fetchall()
        print(f"   prompt表索引数: {len(prompt_indexes)}")
        
        cursor.execute("SHOW INDEX FROM prompt_template")
        template_indexes = cursor.fetchall()
        print(f"   prompt_template表索引数: {len(template_indexes)}")
        
        cursor.close()
        conn.close()
        
        print("\n🎉 数据库优化完成！")
        return True
        
    except Exception as e:
        print(f"❌ 数据库优化失败: {str(e)}")
        return False

if __name__ == "__main__":
    print("🚀 开始数据库性能优化...")
    success = apply_database_optimization()
    if success:
        print("✅ 优化成功，现在可以重启应用服务")
        sys.exit(0)
    else:
        print("❌ 优化失败，请检查数据库连接")
        sys.exit(1)