from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, text from models import Role, User, Permission from utils.password import get_password_hash async def ensure_permission_table_structure(db: AsyncSession): """确保 permission 表结构符合最新要求""" print("检查并更新 permission 表结构...") # 检查表是否存在 result = await db.execute( text("SELECT table_name FROM information_schema.tables WHERE table_name = 'permission'") ) table_exists = result.scalar_one_or_none() if not table_exists: print("permission 表不存在，将由 SQLAlchemy 创建") return # 检查是否缺少 parentid 字段 result = await db.execute( text("SELECT column_name FROM information_schema.columns WHERE table_name = 'permission'") ) existing_columns = {row[0] for row in result.fetchall()} # 添加 parentid 字段 if 'parentid' not in existing_columns: try: await db.execute(text("ALTER TABLE permission ADD COLUMN IF NOT EXISTS parentid INTEGER DEFAULT 0")) print("添加 parentid 字段成功") except Exception as e: print(f"添加 parentid 字段失败: {e}") # 删除旧字段（如果存在） for old_column in ['module', 'submodule', 'action']: if old_column in existing_columns: try: await db.execute(text(f"ALTER TABLE permission DROP COLUMN IF EXISTS {old_column}")) print(f"删除旧字段 {old_column} 成功") except Exception as e: print(f"删除旧字段 {old_column} 失败: {e}") await db.commit() print("permission 表结构更新完成") async def init_permissions(db: AsyncSession): """初始化权限数据（树形结构）""" # 定义需要初始化的权限（树形结构） # parentid: 0 = 顶级模块, >0 = 子权限 permission_tree = [ # 顶级模块 {"name": "权限管理", "code": "perm", "parentid": 0, "description": "权限管理模块"}, # 二级子模块 {"name": "用户管理", "code": "perm:user", "parentid": 1, "description": "用户管理子模块"}, {"name": "角色管理", "code": "perm:role", "parentid": 1, "description": "角色管理子模块"}, {"name": "权限配置", "code": "perm:config", "parentid": 1, "description": "权限配置子模块"}, {"name": "系统设置", "code": "perm:system", "parentid": 1, "description": "系统设置子模块"}, # 用户管理操作 {"name": "查看列表", "code": "perm:user:list", "parentid": 2, "description": "查看用户列表"}, {"name": "查看详情", "code": "perm:user:detail", "parentid": 2, "description": "查看用户详情"}, {"name": "搜索", "code": "perm:user:search", "parentid": 2, "description": "搜索用户"}, {"name": "创建", "code": "perm:user:create", "parentid": 2, "description": "创建新用户"}, {"name": "编辑", "code": "perm:user:edit", "parentid": 2, "description": "编辑用户信息"}, {"name": "删除", "code": "perm:user:delete", "parentid": 2, "description": "删除用户"}, {"name": "导出", "code": "perm:user:export", "parentid": 2, "description": "导出用户数据"}, # 角色管理操作 {"name": "查看列表", "code": "perm:role:list", "parentid": 3, "description": "查看角色列表"}, {"name": "查看详情", "code": "perm:role:detail", "parentid": 3, "description": "查看角色详情"}, {"name": "搜索", "code": "perm:role:search", "parentid": 3, "description": "搜索角色"}, {"name": "创建", "code": "perm:role:create", "parentid": 3, "description": "创建新角色"}, {"name": "编辑", "code": "perm:role:edit", "parentid": 3, "description": "编辑角色信息"}, {"name": "删除", "code": "perm:role:delete", "parentid": 3, "description": "删除角色"}, {"name": "导出", "code": "perm:role:export", "parentid": 3, "description": "导出角色数据"}, {"name": "分配权限", "code": "perm:role:assign", "parentid": 3, "description": "为角色分配权限"}, # 权限配置操作 {"name": "查看列表", "code": "perm:config:list", "parentid": 4, "description": "查看权限列表"}, {"name": "查看详情", "code": "perm:config:detail", "parentid": 4, "description": "查看权限详情"}, {"name": "搜索", "code": "perm:config:search", "parentid": 4, "description": "搜索权限"}, {"name": "创建", "code": "perm:config:create", "parentid": 4, "description": "创建新权限"}, {"name": "编辑", "code": "perm:config:edit", "parentid": 4, "description": "编辑权限信息"}, {"name": "删除", "code": "perm:config:delete", "parentid": 4, "description": "删除权限"}, {"name": "导出", "code": "perm:config:export", "parentid": 4, "description": "导岀权限数据"}, # 系统设置操作（parentid 为 5，对应系统设置子模块） {"name": "查看列表", "code": "perm:system:list", "parentid": 5, "description": "查看系统设置列表"}, {"name": "查看详情", "code": "perm:system:detail", "parentid": 5, "description": "查看系统设置详情"}, {"name": "创建", "code": "perm:system:create", "parentid": 5, "description": "创建设置项"}, {"name": "编辑", "code": "perm:system:edit", "parentid": 5, "description": "编辑设置项"}, {"name": "删除", "code": "perm:system:delete", "parentid": 5, "description": "删除设置项"}, ] # 检查现有权限 result = await db.execute(select(Permission)) existing_permissions = result.scalars().all() existing_permission_codes = {perm.code: perm for perm in existing_permissions} # 创建或更新权限 created_count = 0 updated_count = 0 for perm_data in permission_tree: perm_code = perm_data["code"] if perm_code not in existing_permission_codes: # 权限不存在，需要创建 db.add(Permission(**perm_data)) created_count += 1 else: # 权限存在，检查是否需要更新 existing_perm = existing_permission_codes[perm_code] update_needed = False if existing_perm.name != perm_data["name"]: existing_perm.name = perm_data["name"] update_needed = True if existing_perm.parentid != perm_data["parentid"]: existing_perm.parentid = perm_data["parentid"] update_needed = True if existing_perm.description != perm_data["description"]: existing_perm.description = perm_data["description"] update_needed = True if update_needed: updated_count += 1 if created_count > 0 or updated_count > 0: await db.commit() print(f"创建了 {created_count} 个权限，更新了 {updated_count} 个权限") else: 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 ensure_permission_table_structure(db) await init_roles(db) await init_permissions(db) await init_role_permissions(db) await init_admin_user(db) await init_test_user(db) print("数据库数据初始化完成")