a337cba3f7
- 更新迁移脚本以确保在数据库中添加缺失的字段（user_id、phone_number、avatar） - 优化异常处理，确保在迁移过程中发生错误时能够正确回滚 - 统一字符串引号格式，提升代码可读性 - 添加必要的导入语句，确保脚本正常运行
145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
用户表字段迁移脚本
|
||
为现有用户添加新字段：user_id, phone_number, avatar
|
||
将现有的 username 作为 user_id 的默认值
|
||
"""
|
||
|
||
# ruff: noqa: E402
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
from sqlalchemy import text
|
||
|
||
# 添加项目根目录到Python路径
|
||
PROJECT_ROOT = Path(__file__).parent
|
||
sys.path.insert(0, str(PROJECT_ROOT))
|
||
|
||
from server.db_manager import db_manager
|
||
from server.models.user_model import User as User
|
||
|
||
|
||
def migrate_user_fields():
|
||
"""执行用户字段迁移"""
|
||
print("开始用户字段迁移...")
|
||
|
||
try:
|
||
# 获取数据库会话
|
||
db = db_manager.get_session()
|
||
|
||
# 1. 添加新字段（如果不存在）
|
||
print("检查并添加新字段...")
|
||
|
||
# 检查字段是否存在的SQL
|
||
check_columns_sql = """
|
||
SELECT column_name FROM information_schema.columns
|
||
WHERE table_name = 'users' AND table_schema = DATABASE()
|
||
"""
|
||
|
||
try:
|
||
result = db.execute(text(check_columns_sql))
|
||
existing_columns = [row[0] for row in result.fetchall()]
|
||
print(f"现有字段: {existing_columns}")
|
||
|
||
# 添加缺失的字段
|
||
if "user_id" not in existing_columns:
|
||
print("添加 user_id 字段...")
|
||
db.execute(text("ALTER TABLE users ADD COLUMN user_id VARCHAR(255)"))
|
||
|
||
if "phone_number" not in existing_columns:
|
||
print("添加 phone_number 字段...")
|
||
db.execute(text("ALTER TABLE users ADD COLUMN phone_number VARCHAR(255)"))
|
||
|
||
if "avatar" not in existing_columns:
|
||
print("添加 avatar 字段...")
|
||
db.execute(text("ALTER TABLE users ADD COLUMN avatar VARCHAR(500)"))
|
||
|
||
db.commit()
|
||
print("字段添加完成")
|
||
|
||
except Exception as e:
|
||
print(f"字段检查/添加失败: {e}")
|
||
# 对于SQLite，尝试直接添加字段
|
||
try:
|
||
db.execute(text("ALTER TABLE users ADD COLUMN user_id TEXT"))
|
||
db.execute(text("ALTER TABLE users ADD COLUMN phone_number TEXT"))
|
||
db.execute(text("ALTER TABLE users ADD COLUMN avatar TEXT"))
|
||
db.commit()
|
||
print("字段添加完成（SQLite模式）")
|
||
except Exception as sqlite_e:
|
||
print(f"SQLite字段添加也失败: {sqlite_e}")
|
||
print("字段可能已存在，继续执行...")
|
||
|
||
# 2. 为现有用户设置默认 user_id
|
||
print("为现有用户设置默认 user_id...")
|
||
|
||
# 查询所有没有 user_id 的用户
|
||
users_without_user_id = db.execute(
|
||
text("SELECT id, username FROM users WHERE user_id IS NULL OR user_id = ''")
|
||
).fetchall()
|
||
|
||
print(f"找到 {len(users_without_user_id)} 个需要设置 user_id 的用户")
|
||
|
||
for user_id, username in users_without_user_id:
|
||
# 将 username 作为默认的 user_id
|
||
print(f"为用户 {username} (ID: {user_id}) 设置 user_id: {username}")
|
||
db.execute(text("UPDATE users SET user_id = :user_id WHERE id = :id"), {"user_id": username, "id": user_id})
|
||
|
||
db.commit()
|
||
|
||
# 3. 添加唯一索引
|
||
print("添加唯一索引...")
|
||
try:
|
||
# 先检查索引是否存在
|
||
try:
|
||
db.execute(text("CREATE UNIQUE INDEX idx_users_user_id ON users(user_id)"))
|
||
print("创建 user_id 唯一索引")
|
||
except Exception:
|
||
print("user_id 索引可能已存在")
|
||
|
||
try:
|
||
db.execute(
|
||
text(
|
||
"CREATE UNIQUE INDEX idx_users_phone_number ON users(phone_number) "
|
||
"WHERE phone_number IS NOT NULL"
|
||
)
|
||
)
|
||
print("创建 phone_number 唯一索引")
|
||
except Exception:
|
||
print("phone_number 索引可能已存在")
|
||
|
||
db.commit()
|
||
|
||
except Exception as e:
|
||
print(f"索引创建失败: {e}")
|
||
print("继续执行...")
|
||
|
||
# 4. 验证迁移结果
|
||
print("验证迁移结果...")
|
||
total_users = db.execute(text("SELECT COUNT(*) FROM users")).scalar()
|
||
users_with_user_id = db.execute(
|
||
text("SELECT COUNT(*) FROM users WHERE user_id IS NOT NULL AND user_id != ''")
|
||
).scalar()
|
||
|
||
print(f"总用户数: {total_users}")
|
||
print(f"已设置 user_id 的用户数: {users_with_user_id}")
|
||
|
||
if total_users == users_with_user_id:
|
||
print("✅ 迁移成功完成！")
|
||
else:
|
||
print("❌ 迁移可能有问题，请检查数据库")
|
||
|
||
except Exception as e:
|
||
print(f"迁移过程中发生错误: {e}")
|
||
if "db" in locals():
|
||
db.rollback()
|
||
raise
|
||
finally:
|
||
if "db" in locals():
|
||
db.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
migrate_user_fields()
|