118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
#!/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("数据库升级脚本执行完毕") |