128 lines
4.9 KiB
Python
128 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
更新数据库结构，删除api_keys表，添加roles表，修改users表
|
||
"""
|
||
|
||
from sqlalchemy.orm import Session
|
||
from sqlalchemy import text
|
||
from app.models.database import engine, Base, SessionLocal
|
||
from app.models.models import User, Role
|
||
from app.services.user import UserService
|
||
|
||
|
||
def update_db():
|
||
"""更新数据库结构"""
|
||
db = SessionLocal()
|
||
|
||
try:
|
||
# 1. 删除api_keys表
|
||
print("删除api_keys表...")
|
||
try:
|
||
db.execute(text("DROP TABLE IF EXISTS api_keys CASCADE"))
|
||
db.commit()
|
||
print("✅ api_keys表删除成功")
|
||
except Exception as e:
|
||
print(f"⚠️ 删除api_keys表时出错: {e}")
|
||
db.rollback()
|
||
|
||
# 2. 创建roles表
|
||
print("\n创建roles表...")
|
||
try:
|
||
# 直接执行SQL创建表，避免依赖模型的顺序
|
||
db.execute(text("""
|
||
CREATE TABLE IF NOT EXISTS roles (
|
||
id VARCHAR PRIMARY KEY,
|
||
name VARCHAR UNIQUE NOT NULL,
|
||
description TEXT DEFAULT '',
|
||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||
)
|
||
"""))
|
||
db.commit()
|
||
print("✅ roles表创建成功")
|
||
except Exception as e:
|
||
print(f"⚠️ 创建roles表时出错: {e}")
|
||
db.rollback()
|
||
|
||
# 3. 修改users表，添加role_id字段，删除role字段
|
||
print("\n修改users表...")
|
||
try:
|
||
# 检查是否存在role_id字段
|
||
result = db.execute(text("SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'role_id'")).fetchone()
|
||
|
||
if not result:
|
||
# 添加role_id字段
|
||
db.execute(text("ALTER TABLE users ADD COLUMN role_id VARCHAR"))
|
||
print("✅ 添加role_id字段成功")
|
||
|
||
# 初始化默认角色
|
||
UserService.init_default_roles(db)
|
||
print("✅ 默认角色初始化成功")
|
||
|
||
# 获取默认角色
|
||
admin_role = UserService.get_role_by_name(db, "admin")
|
||
user_role = UserService.get_role_by_name(db, "user")
|
||
|
||
if admin_role and user_role:
|
||
# 更新现有用户的role_id字段
|
||
db.execute(text(f"UPDATE users SET role_id = CASE WHEN role = 'admin' THEN '{admin_role.id}' ELSE '{user_role.id}' END"))
|
||
print("✅ 更新用户role_id字段成功")
|
||
|
||
# 删除role字段
|
||
result = db.execute(text("SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'role'")).fetchone()
|
||
if result:
|
||
db.execute(text("ALTER TABLE users DROP COLUMN role"))
|
||
print("✅ 删除role字段成功")
|
||
|
||
# 添加外键约束
|
||
db.execute(text("ALTER TABLE users ADD CONSTRAINT fk_users_role FOREIGN KEY (role_id) REFERENCES roles(id)"))
|
||
print("✅ 添加外键约束成功")
|
||
|
||
db.commit()
|
||
except Exception as e:
|
||
print(f"⚠️ 修改users表时出错: {e}")
|
||
db.rollback()
|
||
|
||
# 4. 检查并创建默认管理员账号
|
||
print("\n检查默认管理员账号...")
|
||
try:
|
||
# 检查是否已存在管理员账号
|
||
admin_user = db.query(User).filter(User.username == "admin").first()
|
||
|
||
if not admin_user:
|
||
# 获取默认管理员角色
|
||
admin_role = UserService.get_role_by_name(db, "admin")
|
||
if admin_role:
|
||
# 创建默认管理员账号
|
||
admin_user = User(
|
||
id="user-admin",
|
||
username="admin",
|
||
email="admin@example.com",
|
||
password_hash=UserService.get_password_hash("admin123"),
|
||
role_id=admin_role.id,
|
||
status="active"
|
||
)
|
||
db.add(admin_user)
|
||
db.commit()
|
||
db.refresh(admin_user)
|
||
print("✅ 默认管理员账号创建成功")
|
||
print(f"用户名: admin")
|
||
print(f"密码: admin123")
|
||
else:
|
||
print("❌ 无法创建管理员账号，因为admin角色不存在")
|
||
else:
|
||
print("⚠️ 管理员账号已存在")
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ 检查管理员账号时出错: {e}")
|
||
db.rollback()
|
||
|
||
print("\n✅ 数据库结构更新完成")
|
||
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
update_db()
|