BookSystem/backend/auth/init_data.py

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from models import Role, User, Permission
from utils.password import get_password_hash


async def init_permissions(db: AsyncSession):
    """初始化权限数据"""
    # 定义需要初始化的权限
    required_permissions = [
        {"name": "用户管理", "code": "user:manage", "description": "用户管理权限"},
        {"name": "用户查询", "code": "user:read", "description": "用户查询权限"},
        {"name": "用户创建", "code": "user:create", "description": "用户创建权限"},
        {"name": "用户更新", "code": "user:update", "description": "用户更新权限"},
        {"name": "用户删除", "code": "user:delete", "description": "用户删除权限"},
        {"name": "角色管理", "code": "role:manage", "description": "角色管理权限"},
        {"name": "角色查询", "code": "role:read", "description": "角色查询权限"},
        {"name": "角色创建", "code": "role:create", "description": "角色创建权限"},
        {"name": "角色更新", "code": "role:update", "description": "角色更新权限"},
        {"name": "角色删除", "code": "role:delete", "description": "角色删除权限"},
        {"name": "权限管理", "code": "permission:manage", "description": "权限管理权限"},
        {"name": "权限查询", "code": "permission:read", "description": "权限查询权限"},
        {"name": "权限创建", "code": "permission:create", "description": "权限创建权限"},
        {"name": "权限更新", "code": "permission:update", "description": "权限更新权限"},
        {"name": "权限删除", "code": "permission:delete", "description": "权限删除权限"},
    ]

    # 检查现有权限
    result = await db.execute(select(Permission))
    existing_permissions = result.scalars().all()
    existing_permission_codes = {perm.code for perm in existing_permissions}

    # 检查需要创建或更新的权限
    permissions_to_add = []
    permissions_to_update = []

    for perm_data in required_permissions:
        perm_code = perm_data["code"]
        if perm_code not in existing_permission_codes:
            # 权限不存在，需要创建
            permissions_to_add.append(Permission(**perm_data))
        else:
            # 权限存在，检查是否需要更新
            existing_perm = next(
                perm for perm in existing_permissions if perm.code == perm_code
            )
            if existing_perm.name != perm_data["name"] or existing_perm.description != perm_data["description"]:
                existing_perm.name = perm_data["name"]
                existing_perm.description = perm_data["description"]
                permissions_to_update.append(existing_perm)

    # 执行创建和更新操作
    if permissions_to_add:
        for perm in permissions_to_add:
            db.add(perm)
        await db.commit()
        print(f"创建了 {len(permissions_to_add)} 个权限")

    if permissions_to_update:
        await db.commit()
        print(f"更新了 {len(permissions_to_update)} 个权限")

    if not permissions_to_add and not permissions_to_update:
        print("权限数据已存在且无需更新，跳过初始化")

    print("权限数据初始化/更新完成")


async def init_role_permissions(db: AsyncSession):
    """初始化角色权限分配"""
    from sqlalchemy.orm import selectinload
    
    # 获取管理员角色（使用预加载避免懒加载问题）
    result = await db.execute(
        select(Role)
        .filter(Role.name == "管理员")
        .options(selectinload(Role.permissions))
    )
    admin_role = result.scalar_one_or_none()
    if not admin_role:
        print("管理员角色不存在，请先初始化角色数据")
        return
    
    # 获取所有权限
    result = await db.execute(select(Permission))
    all_permissions = result.scalars().all()
    
    if not all_permissions:
        print("没有找到任何权限，请先初始化权限数据")
        return
    
    # 获取管理员角色当前的权限（使用预加载的数据）
    admin_permission_codes = set()
    if admin_role.permissions:
        admin_permission_codes = {perm.code for perm in admin_role.permissions}
    
    # 检查是否需要添加权限
    permissions_to_add = []
    for perm in all_permissions:
        if perm.code not in admin_permission_codes:
            permissions_to_add.append(perm)

    if permissions_to_add:
        # 直接设置权限（替换而非追加，确保权限完整）
        admin_role.permissions = all_permissions
        await db.commit()
        print(f"为管理员角色分配了 {len(permissions_to_add)} 个权限")
    else:
        print("管理员角色权限已完整，无需更新")

    print("角色权限初始化/更新完成")


async def init_roles(db: AsyncSession):
    """初始化角色数据"""
    # 定义需要初始化的角色
    required_roles = [
        {"name": "管理员", "creator": "system"},
        {"name": "普通用户", "creator": "system"},
        {"name": "访客", "creator": "system"},
        {"name": "会员", "creator": "system"},
    ]

    # 检查现有角色
    result = await db.execute(select(Role))
    existing_roles = result.scalars().all()
    existing_role_names = {role.name for role in existing_roles}

    # 检查需要创建或更新的角色
    roles_to_add = []
    roles_to_update = []

    for role_data in required_roles:
        role_name = role_data["name"]
        if role_name not in existing_role_names:
            # 角色不存在，需要创建
            roles_to_add.append(Role(**role_data))
        else:
            # 角色存在，检查是否需要更新
            existing_role = next(
                role for role in existing_roles if role.name == role_name
            )
            if existing_role.creator != role_data["creator"]:
                # 需要更新
                existing_role.creator = role_data["creator"]
                roles_to_update.append(existing_role)

    # 执行创建和更新操作
    if roles_to_add:
        for role in roles_to_add:
            db.add(role)
        await db.commit()
        print(f"创建了 {len(roles_to_add)} 个角色")

    if roles_to_update:
        await db.commit()
        print(f"更新了 {len(roles_to_update)} 个角色")

    if not roles_to_add and not roles_to_update:
        print("角色数据已存在且无需更新，跳过初始化")

    print("角色数据初始化/更新完成")


async def init_admin_user(db: AsyncSession):
    """初始化管理员用户"""
    # 获取管理员角色
    result = await db.execute(select(Role).filter(Role.name == "管理员"))
    admin_role = result.scalar_one_or_none()
    if not admin_role:
        print("管理员角色不存在，请先初始化角色数据")
        return

    # 定义需要的管理员用户数据
    admin_data = {
        "username": "admin",
        "nickname": "系统管理员",
        "email": "admin@example.com",
        "phone": "13800138000",
        "password_hash": get_password_hash("admin"),
        "role_id": admin_role.id,
        "avatar": "",
        "wx_openid": "",
    }

    # 检查管理员用户是否存在
    result = await db.execute(select(User).filter(User.username == "admin"))
    existing_admin = result.scalar_one_or_none()

    if not existing_admin:
        # 管理员不存在，创建
        admin_user = User(**admin_data)
        db.add(admin_user)
        await db.commit()
        print("创建了管理员用户")
    else:
        # 管理员存在，检查是否需要更新
        update_needed = False

        # 检查字段是否需要更新（不包括密码，因为密码只在首次创建时设置）
        if existing_admin.nickname != admin_data["nickname"]:
            existing_admin.nickname = admin_data["nickname"]
            update_needed = True
        if existing_admin.email != admin_data["email"]:
            existing_admin.email = admin_data["email"]
            update_needed = True
        if existing_admin.phone != admin_data["phone"]:
            existing_admin.phone = admin_data["phone"]
            update_needed = True
        if existing_admin.role_id != admin_data["role_id"]:
            existing_admin.role_id = admin_data["role_id"]
            update_needed = True
        if existing_admin.avatar != admin_data["avatar"]:
            existing_admin.avatar = admin_data["avatar"]
            update_needed = True
        if existing_admin.wx_openid != admin_data["wx_openid"]:
            existing_admin.wx_openid = admin_data["wx_openid"]
            update_needed = True

        if update_needed:
            await db.commit()
            print("更新了管理员用户")
        else:
            print("管理员用户已存在且无需更新，跳过初始化")

    print("管理员用户初始化/更新完成")


async def init_test_user(db: AsyncSession):
    """初始化测试用户"""
    # 获取普通用户角色
    result = await db.execute(select(Role).filter(Role.name == "普通用户"))
    user_role = result.scalar_one_or_none()
    if not user_role:
        print("普通用户角色不存在，请先初始化角色数据")
        return

    # 定义需要的测试用户数据
    test_data = {
        "username": "test",
        "nickname": "测试用户",
        "email": "test@example.com",
        "phone": "13800138001",
        "password_hash": get_password_hash("test"),
        "role_id": user_role.id,
        "avatar": "",
        "wx_openid": "",
    }

    # 检查测试用户是否存在
    result = await db.execute(select(User).filter(User.username == "test"))
    existing_test_user = result.scalar_one_or_none()

    if not existing_test_user:
        # 测试用户不存在，创建
        test_user = User(**test_data)
        db.add(test_user)
        await db.commit()
        print("创建了测试用户")
    else:
        # 测试用户存在，检查是否需要更新
        update_needed = False

        # 检查字段是否需要更新（不包括密码，因为密码只在首次创建时设置）
        if existing_test_user.nickname != test_data["nickname"]:
            existing_test_user.nickname = test_data["nickname"]
            update_needed = True
        if existing_test_user.email != test_data["email"]:
            existing_test_user.email = test_data["email"]
            update_needed = True
        if existing_test_user.phone != test_data["phone"]:
            existing_test_user.phone = test_data["phone"]
            update_needed = True
        if existing_test_user.role_id != test_data["role_id"]:
            existing_test_user.role_id = test_data["role_id"]
            update_needed = True
        if existing_test_user.avatar != test_data["avatar"]:
            existing_test_user.avatar = test_data["avatar"]
            update_needed = True
        if existing_test_user.wx_openid != test_data["wx_openid"]:
            existing_test_user.wx_openid = test_data["wx_openid"]
            update_needed = True

        if update_needed:
            await db.commit()
            print("更新了测试用户")
        else:
            print("测试用户已存在且无需更新，跳过初始化")

    print("测试用户初始化/更新完成")


async def init_all_data(db: AsyncSession):
    """初始化所有数据"""
    print("开始初始化数据库数据...")
    await init_roles(db)
    await init_permissions(db)
    await init_role_permissions(db)
    await init_admin_user(db)
    await init_test_user(db)
    print("数据库数据初始化完成")