Files
auto_cursor/migrate_db.py
2025-04-01 15:43:27 +08:00

177 lines
7.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
SQLite 到 MySQL 数据迁移脚本
用于将旧的 SQLite 数据库迁移到新的 MySQL 数据库
"""
import asyncio
import sqlite3
import sys
from typing import List, Dict, Any
import aiomysql
import yaml
from loguru import logger
async def migrate_data():
"""迁移数据主函数"""
# 读取配置文件
logger.info("读取配置文件")
with open("config.yaml", "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
# 获取数据库配置
db_config = config.get("database", {})
sqlite_path = db_config.get("path", "cursor.db")
mysql_config = {
"host": db_config.get("host", "localhost"),
"port": db_config.get("port", 3306),
"user": db_config.get("username", "auto_cursor_reg"), # 使用正确的默认用户名
"password": db_config.get("password", "this_password_jiaqiao"), # 使用正确的默认密码
"db": db_config.get("database", "auto_cursor_reg"), # 使用正确的默认数据库名
"charset": "utf8mb4"
}
logger.info(f"MySQL配置: 主机={mysql_config['host']}:{mysql_config['port']}, 用户={mysql_config['user']}, 数据库={mysql_config['db']}")
# 连接SQLite数据库
logger.info(f"连接SQLite数据库: {sqlite_path}")
sqlite_conn = None
mysql_conn = None
try:
sqlite_conn = sqlite3.connect(sqlite_path)
sqlite_conn.row_factory = sqlite3.Row # 启用字典行工厂
except Exception as e:
logger.error(f"无法连接SQLite数据库: {e}")
return
# 连接MySQL数据库
logger.info(f"尝试连接MySQL数据库...")
try:
mysql_conn = await aiomysql.connect(**mysql_config)
logger.info("MySQL数据库连接成功")
except Exception as e:
logger.error(f"无法连接MySQL数据库: {e}")
logger.error(f"请确认MySQL服务已启动且用户名密码正确")
logger.info(f"您可能需要创建MySQL用户和数据库")
logger.info(f" CREATE USER '{mysql_config['user']}'@'localhost' IDENTIFIED BY '{mysql_config['password']}';")
logger.info(f" CREATE DATABASE {mysql_config['db']} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;")
logger.info(f" GRANT ALL PRIVILEGES ON {mysql_config['db']}.* TO '{mysql_config['user']}'@'localhost';")
logger.info(f" FLUSH PRIVILEGES;")
if sqlite_conn:
sqlite_conn.close()
return
try:
# 检查email_accounts表是否存在
logger.info("检查并创建MySQL表结构")
async with mysql_conn.cursor() as cursor:
await cursor.execute('''
CREATE TABLE IF NOT EXISTS email_accounts (
id INT AUTO_INCREMENT PRIMARY KEY,
email VARCHAR(255) UNIQUE NOT NULL,
password VARCHAR(255) NOT NULL,
client_id VARCHAR(255) NOT NULL,
refresh_token TEXT NOT NULL,
in_use BOOLEAN DEFAULT 0,
cursor_password VARCHAR(255),
cursor_cookie TEXT,
cursor_token TEXT,
sold BOOLEAN DEFAULT 0,
status VARCHAR(20) DEFAULT 'pending',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
INDEX idx_status_inuse_sold (status, in_use, sold)
)
''')
await mysql_conn.commit()
# 从SQLite读取数据
logger.info("从SQLite读取数据")
try:
sqlite_cursor = sqlite_conn.cursor()
sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='email_accounts'")
if not sqlite_cursor.fetchone():
logger.warning("SQLite数据库中不存在email_accounts表无需迁移")
return
# 获取表结构信息
sqlite_cursor.execute("PRAGMA table_info(email_accounts)")
columns_info = sqlite_cursor.fetchall()
column_names = [col["name"] for col in columns_info]
# 读取所有数据
sqlite_cursor.execute("SELECT * FROM email_accounts")
rows = sqlite_cursor.fetchall()
logger.info(f"从SQLite读取到 {len(rows)} 条数据")
except Exception as e:
logger.error(f"读取SQLite数据失败: {e}")
return
# 迁移数据到MySQL
if rows:
logger.info("开始迁移数据到MySQL")
try:
async with mysql_conn.cursor() as cursor:
# 检查MySQL表中是否已有数据
await cursor.execute("SELECT COUNT(*) FROM email_accounts")
result = await cursor.fetchone()
if result and result[0] > 0:
logger.warning("MySQL表中已存在数据是否继续(y/n)")
response = input().strip().lower()
if response != 'y':
logger.info("用户取消迁移")
return
# 构建插入查询
placeholders = ", ".join(["%s"] * len(column_names))
columns = ", ".join(column_names)
query = f"INSERT INTO email_accounts ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE id=id"
# 批量插入数据
batch_size = 100
for i in range(0, len(rows), batch_size):
batch = rows[i:i+batch_size]
batch_values = []
for row in batch:
# 将行数据转换为列表
row_values = [row[name] for name in column_names]
batch_values.append(row_values)
await cursor.executemany(query, batch_values)
await mysql_conn.commit()
logger.info(f"已迁移 {min(i+batch_size, len(rows))}/{len(rows)} 条数据")
logger.success(f"数据迁移完成,共迁移 {len(rows)} 条数据")
except Exception as e:
logger.error(f"迁移数据到MySQL失败: {e}")
else:
logger.warning("SQLite数据库中没有数据需要迁移")
finally:
# 安全关闭连接
if sqlite_conn:
sqlite_conn.close()
logger.debug("SQLite连接已关闭")
if mysql_conn:
await mysql_conn.close()
logger.debug("MySQL连接已关闭")
if __name__ == "__main__":
# Windows平台特殊处理强制使用SelectorEventLoop
if sys.platform.startswith('win'):
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
# 设置日志
logger.remove()
logger.add(sys.stderr, level="INFO")
logger.add("migrate.log", rotation="1 MB", level="DEBUG")
# 执行迁移
logger.info("开始执行数据迁移")
asyncio.run(migrate_data())
logger.info("数据迁移脚本执行完毕")