优化自动化

2025-04-01 16:25:08 +08:00
parent 7dad9f6b2f
commit b20094adfc
10 changed files with 1284 additions and 5 deletions
--- a/upgrade_database.py
+++ b/upgrade_database.py
@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""
+数据库升级脚本
+在已有的数据库中添加新字段和索引
+"""
+import asyncio
+import sys
+import yaml
+import traceback
+from loguru import logger
+
+# Windows平台特殊处理，强制使用SelectorEventLoop
+if sys.platform.startswith('win'):
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+
+from core.config import Config
+from core.database import DatabaseManager
+
+
+async def add_extracted_field(db_manager):
+    """向email_accounts表添加extracted字段和索引"""
+    logger.info("正在检查extracted字段...")
+    
+    # 检查extracted字段是否存在
+    check_query = """
+        SELECT COUNT(*) as field_exists
+        FROM information_schema.COLUMNS 
+        WHERE TABLE_SCHEMA = DATABASE()
+        AND TABLE_NAME = 'email_accounts' 
+        AND COLUMN_NAME = 'extracted'
+    """
+    
+    result = await db_manager.fetch_one(check_query)
+    field_exists = result['field_exists'] > 0 if result else False
+    
+    if field_exists:
+        logger.info("extracted字段已存在，无需添加")
+    else:
+        logger.info("添加extracted字段...")
+        add_field_query = """
+            ALTER TABLE email_accounts
+            ADD COLUMN extracted BOOLEAN DEFAULT 0
+        """
+        await db_manager.execute(add_field_query)
+        logger.success("成功添加extracted字段")
+    
+    # 检查索引是否存在
+    check_index_query = """
+        SELECT COUNT(*) as index_exists
+        FROM information_schema.STATISTICS
+        WHERE TABLE_SCHEMA = DATABASE()
+        AND TABLE_NAME = 'email_accounts'
+        AND INDEX_NAME = 'idx_extracted'
+    """
+    
+    result = await db_manager.fetch_one(check_index_query)
+    index_exists = result['index_exists'] > 0 if result else False
+    
+    if index_exists:
+        logger.info("idx_extracted索引已存在，无需添加")
+    else:
+        logger.info("添加idx_extracted索引...")
+        add_index_query = """
+            ALTER TABLE email_accounts
+            ADD INDEX idx_extracted (extracted, status, sold)
+        """
+        await db_manager.execute(add_index_query)
+        logger.success("成功添加idx_extracted索引")
+    
+    logger.info("数据库升级完成")
+    
+    # 设置所有注册成功且已售出的账号的extracted状态
+    update_query = """
+        UPDATE email_accounts
+        SET extracted = 0
+        WHERE status = 'success' AND sold = 1 AND extracted IS NULL
+    """
+    affected = await db_manager.execute(update_query)
+    logger.info(f"已更新 {affected} 条记录的extracted状态")
+
+
+async def upgrade_database():
+    """升级数据库结构"""
+    logger.info("开始数据库升级...")
+    
+    try:
+        # 加载配置
+        config = Config.from_yaml()
+        
+        # 创建数据库管理器并初始化
+        db_manager = DatabaseManager(config)
+        await db_manager.initialize()
+        
+        try:
+            # 添加extracted字段和索引
+            await add_extracted_field(db_manager)
+            
+            logger.success("数据库升级成功完成")
+            
+        finally:
+            # 清理资源
+            await db_manager.cleanup()
+            
+    except Exception as e:
+        logger.error(f"数据库升级失败: {str(e)}")
+        logger.error(traceback.format_exc())
+
+
+if __name__ == "__main__":
+    # 设置日志
+    logger.remove()
+    logger.add(sys.stderr, level="INFO")
+    logger.add("upgrade_database.log", rotation="1 MB", level="DEBUG")
+    
+    # 执行升级
+    logger.info("启动数据库升级")
+    asyncio.run(upgrade_database())
+    logger.info("数据库升级脚本执行完毕")