update:1.更新根据分析建议重新生成章节内容
This commit is contained in:
xiamuceer
@@ -0,0 +1,300 @@
|
||||
"""
|
||||
用户数据迁移脚本 - 从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))
|
||||
Reference in New Issue
Block a user