xiamuceer
300 lines
9.7 KiB
Python
300 lines
9.7 KiB
Python
"""
|
||
用户数据迁移脚本 - 从JSON文件迁移到PostgreSQL数据库
|
||
|
||
使用方法:
|
||
python migrate_users_to_postgres.py
|
||
python migrate_users_to_postgres.py --db-url postgresql+asyncpg://user:pass@localhost/dbname
|
||
"""
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import sys
|
||
import argparse
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
|
||
# 添加项目根目录到 Python 路径
|
||
project_root = Path(__file__).parent.parent
|
||
sys.path.insert(0, str(project_root))
|
||
|
||
from sqlalchemy import select, text
|
||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||
from app.config import settings, DATA_DIR
|
||
|
||
|
||
async def create_tables(engine):
|
||
"""创建用户相关表"""
|
||
from app.database import Base
|
||
from app.models.user import User, UserPassword
|
||
|
||
print("📋 创建数据库表...")
|
||
async with engine.begin() as conn:
|
||
# 只创建用户相关的表
|
||
await conn.run_sync(User.metadata.create_all)
|
||
await conn.run_sync(UserPassword.metadata.create_all)
|
||
print("✅ 表创建成功")
|
||
|
||
|
||
async def migrate_users(session):
|
||
"""迁移用户数据"""
|
||
from app.models.user import User as UserModel
|
||
|
||
users_file = DATA_DIR / "users.json"
|
||
|
||
if not users_file.exists():
|
||
print("ℹ️ 用户数据文件不存在，跳过迁移")
|
||
return 0
|
||
|
||
try:
|
||
with open(users_file, "r", encoding="utf-8") as f:
|
||
users_data = json.load(f)
|
||
|
||
if not users_data:
|
||
print("ℹ️ 用户数据为空，跳过迁移")
|
||
return 0
|
||
|
||
migrated_count = 0
|
||
for user_id, user_info in users_data.items():
|
||
try:
|
||
# 检查用户是否已存在
|
||
result = await session.execute(
|
||
select(UserModel).where(UserModel.user_id == user_id)
|
||
)
|
||
existing = result.scalar_one_or_none()
|
||
|
||
if existing:
|
||
print(f"ℹ️ 用户已存在，跳过: {user_info['username']} ({user_id})")
|
||
continue
|
||
|
||
# 创建用户记录
|
||
user = UserModel(
|
||
user_id=user_id,
|
||
username=user_info["username"],
|
||
display_name=user_info["display_name"],
|
||
avatar_url=user_info.get("avatar_url"),
|
||
trust_level=user_info.get("trust_level", 0),
|
||
is_admin=user_info.get("is_admin", False),
|
||
linuxdo_id=user_info["linuxdo_id"],
|
||
created_at=datetime.fromisoformat(user_info.get("created_at", datetime.now().isoformat())),
|
||
last_login=datetime.fromisoformat(user_info.get("last_login", datetime.now().isoformat()))
|
||
)
|
||
session.add(user)
|
||
|
||
migrated_count += 1
|
||
print(f"✅ 迁移用户: {user_info['username']} ({user_id})")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 迁移用户 {user_id} 失败: {e}")
|
||
|
||
await session.commit()
|
||
print(f"\n✅ 用户数据迁移完成: {migrated_count}/{len(users_data)} 个用户")
|
||
|
||
return migrated_count
|
||
|
||
except Exception as e:
|
||
print(f"❌ 迁移用户数据失败: {e}")
|
||
await session.rollback()
|
||
return 0
|
||
|
||
|
||
async def migrate_passwords(session):
|
||
"""迁移密码数据"""
|
||
from app.models.user import UserPassword
|
||
|
||
passwords_file = DATA_DIR / "user_passwords.json"
|
||
|
||
if not passwords_file.exists():
|
||
print("ℹ️ 密码数据文件不存在，跳过迁移")
|
||
return 0
|
||
|
||
try:
|
||
with open(passwords_file, "r", encoding="utf-8") as f:
|
||
passwords_data = json.load(f)
|
||
|
||
if not passwords_data:
|
||
print("ℹ️ 密码数据为空，跳过迁移")
|
||
return 0
|
||
|
||
migrated_count = 0
|
||
for user_id, pwd_info in passwords_data.items():
|
||
try:
|
||
# 检查密码是否已存在
|
||
result = await session.execute(
|
||
select(UserPassword).where(UserPassword.user_id == user_id)
|
||
)
|
||
existing = result.scalar_one_or_none()
|
||
|
||
if existing:
|
||
print(f"ℹ️ 密码已存在，跳过: {pwd_info['username']} ({user_id})")
|
||
continue
|
||
|
||
# 创建密码记录
|
||
pwd_record = UserPassword(
|
||
user_id=user_id,
|
||
username=pwd_info["username"],
|
||
password_hash=pwd_info["password_hash"],
|
||
has_custom_password=pwd_info.get("has_custom_password", False),
|
||
created_at=datetime.now(),
|
||
updated_at=datetime.now()
|
||
)
|
||
session.add(pwd_record)
|
||
|
||
migrated_count += 1
|
||
print(f"✅ 迁移密码: {pwd_info['username']} ({user_id})")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 迁移密码 {user_id} 失败: {e}")
|
||
|
||
await session.commit()
|
||
print(f"\n✅ 密码数据迁移完成: {migrated_count}/{len(passwords_data)} 个密码")
|
||
|
||
return migrated_count
|
||
|
||
except Exception as e:
|
||
print(f"❌ 迁移密码数据失败: {e}")
|
||
await session.rollback()
|
||
return 0
|
||
|
||
|
||
async def backup_json_files():
|
||
"""备份原始JSON文件"""
|
||
files_to_backup = ["users.json", "user_passwords.json", "admins.json"]
|
||
|
||
print("\n📦 备份原始文件...")
|
||
for filename in files_to_backup:
|
||
source = DATA_DIR / filename
|
||
if source.exists():
|
||
backup = DATA_DIR / f"{filename}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||
import shutil
|
||
shutil.copy2(source, backup)
|
||
print(f"✅ 备份: {filename} -> {backup.name}")
|
||
|
||
|
||
async def main(db_url=None):
|
||
"""主函数
|
||
|
||
Args:
|
||
db_url: 可选的数据库URL，如果不提供则使用配置文件中的
|
||
"""
|
||
print("=" * 70)
|
||
print("用户数据迁移工具 - JSON 到 PostgreSQL")
|
||
print("=" * 70)
|
||
print()
|
||
|
||
# 确定使用的数据库URL
|
||
target_db_url = db_url if db_url else settings.database_url
|
||
|
||
# 检查数据库配置
|
||
if "postgresql" not in target_db_url:
|
||
print("❌ 错误: 未指定 PostgreSQL 数据库")
|
||
if not db_url:
|
||
print(f" 当前配置: {settings.database_url}")
|
||
print(" 请使用 --db-url 参数指定PostgreSQL数据库，或在 .env 中配置 DATABASE_URL")
|
||
else:
|
||
print(f" 提供的URL: {target_db_url}")
|
||
print()
|
||
print("示例:")
|
||
print(" python migrate_users_to_postgres.py --db-url postgresql+asyncpg://user:pass@localhost/dbname")
|
||
return
|
||
|
||
# 隐藏密码部分显示
|
||
display_url = target_db_url
|
||
if '@' in display_url:
|
||
parts = display_url.split('@')
|
||
if ':' in parts[0]:
|
||
user_part = parts[0].split(':')[0]
|
||
display_url = f"{user_part}:****@{parts[1]}"
|
||
|
||
print(f"📊 目标数据库: {display_url}")
|
||
print()
|
||
|
||
try:
|
||
# 创建数据库引擎
|
||
engine = create_async_engine(
|
||
target_db_url,
|
||
echo=False,
|
||
future=True,
|
||
pool_pre_ping=True,
|
||
)
|
||
|
||
# 创建表
|
||
await create_tables(engine)
|
||
print()
|
||
|
||
# 创建会话
|
||
async_session = async_sessionmaker(
|
||
engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False
|
||
)
|
||
|
||
# 迁移用户
|
||
print("📋 步骤 1/2: 迁移用户数据")
|
||
print("-" * 70)
|
||
async with async_session() as session:
|
||
user_count = await migrate_users(session)
|
||
print()
|
||
|
||
# 迁移密码
|
||
print("📋 步骤 2/2: 迁移密码数据")
|
||
print("-" * 70)
|
||
async with async_session() as session:
|
||
pwd_count = await migrate_passwords(session)
|
||
print()
|
||
|
||
# 备份原文件
|
||
await backup_json_files()
|
||
print()
|
||
|
||
# 总结
|
||
print("=" * 70)
|
||
print("迁移完成")
|
||
print("=" * 70)
|
||
print(f"✅ 用户: {user_count}")
|
||
print(f"✅ 密码: {pwd_count}")
|
||
print()
|
||
print("💡 提示:")
|
||
print(" - 原文件已备份（带时间戳）")
|
||
print(" - 可以安全删除 users.json 和 user_passwords.json")
|
||
print(" - 如需回滚，请从备份文件恢复")
|
||
print()
|
||
|
||
# 关闭引擎
|
||
await engine.dispose()
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ 迁移过程出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 解析命令行参数
|
||
parser = argparse.ArgumentParser(
|
||
description="迁移用户数据从JSON到PostgreSQL数据库",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例:
|
||
# 使用 .env 配置的数据库
|
||
python migrate_users_to_postgres.py
|
||
|
||
# 指定数据库URL
|
||
python migrate_users_to_postgres.py --db-url postgresql+asyncpg://user:pass@localhost/dbname
|
||
|
||
# 使用环境变量
|
||
DATABASE_URL=postgresql+asyncpg://user:pass@localhost/db python migrate_users_to_postgres.py
|
||
"""
|
||
)
|
||
|
||
parser.add_argument(
|
||
"--db-url",
|
||
type=str,
|
||
help="PostgreSQL数据库连接URL (格式: postgresql+asyncpg://user:password@host:port/database)",
|
||
default=None
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 运行迁移
|
||
asyncio.run(main(db_url=args.db_url)) |