149 lines
5.0 KiB
Python
149 lines
5.0 KiB
Python
from typing import List
|
||
|
||
from django.core.cache import cache
|
||
from loguru import logger
|
||
from ninja.errors import HttpError
|
||
|
||
from users.models import User, Menu # 替换为你的实际 app 名
|
||
from functools import wraps
|
||
import asyncio
|
||
from asgiref.sync import sync_to_async
|
||
|
||
|
||
# 你的权限检查函数（不变）
|
||
def has_permission(user_id: int, permission: str) -> bool:
|
||
# 你原来的逻辑
|
||
return True
|
||
|
||
|
||
# ====================== 多权限通用装饰器 ======================
|
||
def require_permissions(*perms: str, match_all: bool = True):
|
||
"""
|
||
多权限检查装饰器
|
||
:param perms: 权限字符串，可传多个
|
||
:param match_all: True = 必须拥有所有权限（AND），False = 拥有任意一个即可（OR）
|
||
用法：
|
||
@require_permissions("role:update", "role:delete") # 两个都要有
|
||
@require_permissions("role:view", "role:list", match_all=False) # 有一个就行
|
||
"""
|
||
|
||
def decorator(view_func):
|
||
@wraps(view_func)
|
||
async def async_wrapper(request, *args, **kwargs):
|
||
user_id = request.auth.id
|
||
check = sync_to_async(has_permission)
|
||
|
||
# 检查所有权限
|
||
results = [await check(user_id, perm) for perm in perms]
|
||
allowed = all(results) if match_all else any(results)
|
||
|
||
if not allowed:
|
||
raise HttpError(403, "权限不足")
|
||
|
||
return await view_func(request, *args, **kwargs)
|
||
|
||
@wraps(view_func)
|
||
def sync_wrapper(request, *args, **kwargs):
|
||
user_id = request.auth.id
|
||
results = [has_permission(user_id, perm) for perm in perms]
|
||
allowed = all(results) if match_all else any(results)
|
||
|
||
if not allowed:
|
||
raise HttpError(403, "权限不足")
|
||
|
||
return view_func(request, *args, **kwargs)
|
||
|
||
# 自动识别同步/异步
|
||
if asyncio.iscoroutinefunction(view_func):
|
||
return async_wrapper
|
||
return sync_wrapper
|
||
|
||
return decorator
|
||
|
||
|
||
def get_user_permissions(user_id: int) -> List[str]:
|
||
"""
|
||
获取用户所有按钮级权限标识（permission_code）
|
||
利用 Django ORM 的多对多关系链式查询
|
||
"""
|
||
cache_key = f"user_permissions_{user_id}"
|
||
permissions = cache.get(cache_key)
|
||
logger.info(f"当前缓存的权限列表：{permissions}")
|
||
if permissions is None:
|
||
try:
|
||
# 链式查询：用户 → 角色 → 菜单（仅按钮类型）
|
||
permissions = list(
|
||
Menu.objects.filter(
|
||
roles__users__id=user_id, # 通过 roles__users 反向关联到用户
|
||
menu_type=2, # 2 = 按钮
|
||
status=True
|
||
)
|
||
.values_list('permission_code', flat=True)
|
||
.distinct()
|
||
)
|
||
logger.info(f"当前获取到的权限列表：{permissions}")
|
||
cache.set(cache_key, permissions, 3600) # 缓存1小时
|
||
except User.DoesNotExist:
|
||
return []
|
||
logger.info(f"当前用户的权限列表：{permissions}")
|
||
return permissions
|
||
|
||
|
||
def has_permission(user_id: int, permission_code: str) -> bool:
|
||
"""检查用户是否拥有指定权限"""
|
||
if not permission_code:
|
||
return True
|
||
return permission_code in get_user_permissions(user_id)
|
||
|
||
|
||
def clear_user_permissions_cache(user_id: int):
|
||
"""清除用户权限缓存（当角色或菜单变更时调用）"""
|
||
cache.delete(f"user_permissions_{user_id}")
|
||
|
||
|
||
# 批量清除缓存的辅助函数（可选）
|
||
def clear_all_users_permissions_cache():
|
||
"""清除所有用户权限缓存（谨慎使用）"""
|
||
# 实际项目中可通过信号或更精细的方式处理
|
||
cache.clear() # 或遍历已知用户ID
|
||
|
||
|
||
def build_menu_tree(menus):
|
||
"""高效构建菜单树（O(n) 时间复杂度）"""
|
||
# 1. 将菜单列表转为字典，key=id
|
||
menu_dict = {}
|
||
for menu in menus:
|
||
menu_dict[menu.id] = {
|
||
'id': menu.id,
|
||
'name': menu.name,
|
||
'code': menu.code,
|
||
'path': menu.path,
|
||
'component': menu.component,
|
||
'icon': menu.icon,
|
||
'order_num': menu.order_num,
|
||
'parent_id': menu.parent_id, # ✅ 只存 ID，避免对象引用
|
||
'permission_code': menu.permission_code,
|
||
'menu_type': menu.menu_type,
|
||
'children': [] # 初始化为空列表
|
||
}
|
||
|
||
# 2. 构建树结构
|
||
tree = []
|
||
for menu_id, node in menu_dict.items():
|
||
parent_id = node['parent_id']
|
||
if parent_id is None or parent_id not in menu_dict:
|
||
# 根节点（parent_id 为 None 或指向不存在的父节点）
|
||
tree.append(node)
|
||
else:
|
||
# 添加到父节点的 children
|
||
menu_dict[parent_id]['children'].append(node)
|
||
|
||
# 3. （可选）按 order_num 排序
|
||
def sort_children(nodes):
|
||
nodes.sort(key=lambda x: x['order_num'] or 0)
|
||
for node in nodes:
|
||
sort_children(node['children'])
|
||
|
||
sort_children(tree)
|
||
return tree
|