diff --git a/apps/common/db/models.py b/apps/common/db/models.py index 7be495538..ae726b85d 100644 --- a/apps/common/db/models.py +++ b/apps/common/db/models.py @@ -9,9 +9,8 @@ - 此文件中添加代码的时候，注意不要跟 `django.db.models` 中的命名冲突 """ -import inspect import uuid -from functools import reduce, partial +from functools import reduce from django.db import models from django.db import transaction @@ -96,87 +95,6 @@ def output_as_string(field_name): return ExpressionWrapper(F(field_name), output_field=models.CharField()) -class UnionQuerySet(QuerySet): - after_union = ['order_by'] - not_return_qs = [ - 'query', 'get', 'create', 'get_or_create', - 'update_or_create', 'bulk_create', 'count', - 'latest', 'earliest', 'first', 'last', 'aggregate', - 'exists', 'update', 'delete', 'as_manager', 'explain', - ] - - def __init__(self, *queryset_list): - self.queryset_list = queryset_list - self.after_union_items = [] - self.before_union_items = [] - - def __execute(self): - queryset_list = [] - for qs in self.queryset_list: - for attr, args, kwargs in self.before_union_items: - qs = getattr(qs, attr)(*args, **kwargs) - queryset_list.append(qs) - union_qs = reduce(lambda x, y: x.union(y), queryset_list) - for attr, args, kwargs in self.after_union_items: - union_qs = getattr(union_qs, attr)(*args, **kwargs) - return union_qs - - def __before_union_perform(self, item, *args, **kwargs): - self.before_union_items.append((item, args, kwargs)) - return self.__clone(*self.queryset_list) - - def __after_union_perform(self, item, *args, **kwargs): - self.after_union_items.append((item, args, kwargs)) - return self.__clone(*self.queryset_list) - - def __clone(self, *queryset_list): - uqs = UnionQuerySet(*queryset_list) - uqs.after_union_items = self.after_union_items - uqs.before_union_items = self.before_union_items - return uqs - - def __getattribute__(self, item): - if item.startswith('__') or item in UnionQuerySet.__dict__ or item in [ - 'queryset_list', 'after_union_items', 'before_union_items' - ]: - return object.__getattribute__(self, item) - - if item in UnionQuerySet.not_return_qs: - return getattr(self.__execute(), item) - - origin_item = object.__getattribute__(self, 'queryset_list')[0] - origin_attr = getattr(origin_item, item, None) - if not inspect.ismethod(origin_attr): - return getattr(self.__execute(), item) - - if item in UnionQuerySet.after_union: - attr = partial(self.__after_union_perform, item) - else: - attr = partial(self.__before_union_perform, item) - return attr - - def __getitem__(self, item): - return self.__execute()[item] - - def __iter__(self): - return iter(self.__execute()) - - def __str__(self): - return str(self.__execute()) - - def __repr__(self): - return repr(self.__execute()) - - @classmethod - def test_it(cls): - from assets.models import Asset - assets1 = Asset.objects.filter(hostname__startswith='a') - assets2 = Asset.objects.filter(hostname__startswith='b') - - qs = cls(assets1, assets2) - return qs - - class MultiTableChildQueryset(QuerySet): def bulk_create(self, objs, batch_size=None): diff --git a/apps/perms/api/asset_permission_relation.py b/apps/perms/api/asset_permission_relation.py index c2c116248..5b85cb971 100644 --- a/apps/perms/api/asset_permission_relation.py +++ b/apps/perms/api/asset_permission_relation.py @@ -9,7 +9,7 @@ from orgs.mixins.api import OrgBulkModelViewSet from orgs.utils import current_org from perms import serializers from perms import models -from perms.utils.user_permission import UserGrantedAssetsQueryUtils +from perms.utils import AssetPermissionPermAssetUtil from assets.serializers import AccountSerializer __all__ = [ @@ -95,8 +95,7 @@ class AssetPermissionAllAssetListApi(generics.ListAPIView): def get_queryset(self): pk = self.kwargs.get("pk") - query_utils = UserGrantedAssetsQueryUtils(None, asset_perm_ids=[pk]) - assets = query_utils.get_all_granted_assets() + assets = AssetPermissionPermAssetUtil(perm_ids=[pk]).get_all_assets() return assets diff --git a/apps/perms/api/user_permission/assets.py b/apps/perms/api/user_permission/assets.py index e494b7a4a..e499a3127 100644 --- a/apps/perms/api/user_permission/assets.py +++ b/apps/perms/api/user_permission/assets.py @@ -6,7 +6,7 @@ from assets.api.asset.asset import AssetFilterSet from perms import serializers from perms.pagination import AllPermedAssetPagination from perms.pagination import NodePermedAssetPagination -from perms.utils.user_permission import UserGrantedAssetsQueryUtils +from perms.utils import UserPermAssetUtil from common.utils import get_logger, lazyproperty from .mixin import ( @@ -43,21 +43,23 @@ class BaseUserPermedAssetsApi(SelfOrPKUserMixin, ListAPIView): def get_assets(self): return Asset.objects.none() + query_asset_util: UserPermAssetUtil + @lazyproperty def query_asset_util(self): - return UserGrantedAssetsQueryUtils(self.user) + return UserPermAssetUtil(self.user) class UserAllPermedAssetsApi(BaseUserPermedAssetsApi): pagination_class = AllPermedAssetPagination def get_assets(self): - return self.query_asset_util.get_all_granted_assets() + return self.query_asset_util.get_all_assets() class UserDirectPermedAssetsApi(BaseUserPermedAssetsApi): def get_assets(self): - return self.query_asset_util.get_direct_granted_assets() + return self.query_asset_util.get_direct_assets() class UserFavoriteAssetsApi(BaseUserPermedAssetsApi): diff --git a/apps/perms/api/user_permission/nodes.py b/apps/perms/api/user_permission/nodes.py index 1972909e9..793ad3858 100644 --- a/apps/perms/api/user_permission/nodes.py +++ b/apps/perms/api/user_permission/nodes.py @@ -7,7 +7,7 @@ from rest_framework.generics import ListAPIView from assets.models import Node from common.utils import get_logger, lazyproperty from perms import serializers -from perms.utils.user_permission import UserGrantedNodesQueryUtils +from perms.utils import UserPermNodeUtil from .mixin import SelfOrPKUserMixin logger = get_logger(__name__) @@ -32,7 +32,7 @@ class BaseUserPermedNodesApi(SelfOrPKUserMixin, ListAPIView): @lazyproperty def query_node_util(self): - return UserGrantedNodesQueryUtils(self.user) + return UserPermNodeUtil(self.user) class UserAllPermedNodesApi(BaseUserPermedNodesApi): diff --git a/apps/perms/api/user_permission/tree/node_with_asset.py b/apps/perms/api/user_permission/tree/node_with_asset.py index f65cfc0b1..204859f36 100644 --- a/apps/perms/api/user_permission/tree/node_with_asset.py +++ b/apps/perms/api/user_permission/tree/node_with_asset.py @@ -17,11 +17,8 @@ from common.utils import get_object_or_none, lazyproperty from common.utils.common import timeit from perms.hands import Node from perms.models import PermNode -from perms.utils import PermAccountUtil -from perms.utils.permission import AssetPermissionUtil -from perms.utils.user_permission import ( - UserGrantedNodesQueryUtils, UserGrantedAssetsQueryUtils, -) +from perms.utils import PermAccountUtil, UserPermNodeUtil, AssetPermissionUtil +from perms.utils import UserPermAssetUtil from .mixin import RebuildTreeMixin from ..mixin import SelfOrPKUserMixin @@ -54,13 +51,12 @@ class BaseUserNodeWithAssetAsTreeApi( class UserPermedNodesWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi): - query_node_util: UserGrantedNodesQueryUtils - query_asset_util: UserGrantedAssetsQueryUtils + query_node_util: UserPermNodeUtil + query_asset_util: UserPermAssetUtil def get_nodes_assets(self): - perm_ids = AssetPermissionUtil().get_permissions_for_user(self.request.user, flat=True) - self.query_node_util = UserGrantedNodesQueryUtils(self.request.user, perm_ids) - self.query_asset_util = UserGrantedAssetsQueryUtils(self.request.user, perm_ids) + self.query_node_util = UserPermNodeUtil(self.request.user) + self.query_asset_util = UserPermAssetUtil(self.request.user) ung_nodes, ung_assets = self._get_nodes_assets_for_ungrouped() fav_nodes, fav_assets = self._get_nodes_assets_for_favorite() all_nodes, all_assets = self._get_nodes_assets_for_all() @@ -89,9 +85,9 @@ class UserPermedNodesWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi): def _get_nodes_assets_for_all(self): nodes = self.query_node_util.get_whole_tree_nodes(with_special=False) if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: - assets = self.query_asset_util.get_direct_granted_nodes_assets() + assets = self.query_asset_util.get_perm_nodes_assets() else: - assets = self.query_asset_util.get_all_granted_assets() + assets = self.query_asset_util.get_all_assets() assets = assets.annotate(parent_key=F('nodes__key')).prefetch_related('platform') return nodes, assets @@ -102,8 +98,8 @@ class UserPermedNodeChildrenWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi): def get_nodes_assets(self): nodes = PermNode.objects.none() assets = Asset.objects.none() - query_node_util = UserGrantedNodesQueryUtils(self.user) - query_asset_util = UserGrantedAssetsQueryUtils(self.user) + query_node_util = UserPermNodeUtil(self.user) + query_asset_util = UserPermAssetUtil(self.user) node_key = self.query_node_key if not node_key: nodes = query_node_util.get_top_level_nodes() @@ -113,7 +109,7 @@ class UserPermedNodeChildrenWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi): assets = query_asset_util.get_favorite_assets() else: nodes = query_node_util.get_node_children(node_key) - assets = query_asset_util.get_node_assets(node_key) + assets = query_asset_util.get_node_assets(key=node_key) assets = assets.prefetch_related('platform') return nodes, assets diff --git a/apps/perms/filters.py b/apps/perms/filters.py index 05f202655..8743c38be 100644 --- a/apps/perms/filters.py +++ b/apps/perms/filters.py @@ -1,7 +1,6 @@ from django_filters import rest_framework as filters from django.db.models import QuerySet, Q -from common.db.models import UnionQuerySet from common.drf.filters import BaseFilterSet from common.utils import get_object_or_none from users.models import User, UserGroup @@ -169,10 +168,10 @@ class AssetPermissionFilter(PermissionBaseFilter): inherit_all_node_ids = Node.objects.filter(key__in=inherit_all_node_keys).values_list('id', flat=True) inherit_all_node_ids = list(inherit_all_node_ids) - qs1 = queryset.filter(assets__in=asset_ids).distinct() - qs2 = queryset.filter(nodes__in=inherit_all_node_ids).distinct() - - qs = UnionQuerySet(qs1, qs2) + qs1_ids = queryset.filter(assets__in=asset_ids).distinct().values_list('id', flat=True) + qs2_ids = queryset.filter(nodes__in=inherit_all_node_ids).distinct().values_list('id', flat=True) + qs_ids = list(qs1_ids) + list(qs2_ids) + qs = queryset.filter(id__in=qs_ids) return qs def filter_effective(self, queryset): diff --git a/apps/perms/models/asset_permission.py b/apps/perms/models/asset_permission.py index ac6fecdb4..989f10cbc 100644 --- a/apps/perms/models/asset_permission.py +++ b/apps/perms/models/asset_permission.py @@ -5,14 +5,14 @@ from django.db.models import Q from django.utils import timezone from django.utils.translation import ugettext_lazy as _ +from users.models import User from assets.models import Asset, Account -from common.db.models import UnionQuerySet -from common.utils import date_expired_default -from common.utils.timezone import local_now from orgs.mixins.models import JMSOrgBaseModel from orgs.mixins.models import OrgManager +from common.utils import date_expired_default +from common.utils.timezone import local_now + from perms.const import ActionChoices -from users.models import User __all__ = ['AssetPermission', 'ActionChoices'] @@ -104,9 +104,10 @@ class AssetPermission(JMSOrgBaseModel): group_ids = self.user_groups.all().values_list('id', flat=True) user_ids = list(user_ids) group_ids = list(group_ids) - qs1 = User.objects.filter(id__in=user_ids).distinct() - qs2 = User.objects.filter(groups__id__in=group_ids).distinct() - qs = UnionQuerySet(qs1, qs2) + qs1_ids = User.objects.filter(id__in=user_ids).distinct().values_list('id', flat=True) + qs2_ids = User.objects.filter(groups__id__in=group_ids).distinct().values_list('id', flat=True) + qs_ids = list(qs1_ids) + list(qs2_ids) + qs = User.objects.filter(id__in=qs_ids) return qs def get_all_assets(self, flat=False): diff --git a/apps/perms/models/perm_node.py b/apps/perms/models/perm_node.py index 89354a06b..ed4b4b213 100644 --- a/apps/perms/models/perm_node.py +++ b/apps/perms/models/perm_node.py @@ -36,15 +36,14 @@ class UserAssetGrantedTreeNodeRelation(FamilyMixin, JMSOrgBaseModel): return self.node_parent_key @classmethod - def get_node_granted_status(cls, user, key): + def get_node_from_with_node(cls, user, key): ancestor_keys = set(cls.get_node_ancestor_keys(key, with_self=True)) - ancestor_rel_nodes = cls.objects.filter(user=user, node_key__in=ancestor_keys) - - for rel_node in ancestor_rel_nodes: - if rel_node.key == key: - return rel_node.node_from, rel_node - if rel_node.node_from == cls.NodeFrom.granted: - return cls.NodeFrom.granted, None + ancestor_nodes = cls.objects.filter(user=user, node_key__in=ancestor_keys) + for node in ancestor_nodes: + if node.key == key: + return node.node_from, node + if node.node_from == cls.NodeFrom.granted: + return node.node_from, None return '', None @@ -91,15 +90,16 @@ class PermNode(Node): node.assets_amount = assets_amount return node - def get_granted_status(self, user): - status, rel_node = UserAssetGrantedTreeNodeRelation.get_node_granted_status(user, self.key) - self.node_from = status - if rel_node: - self.granted_assets_amount = rel_node.node_assets_amount - return status + def compute_node_from_and_assets_amount(self, user): + node_from, node = UserAssetGrantedTreeNodeRelation.get_node_from_with_node( + user, self.key + ) + self.node_from = node_from + if node: + self.granted_assets_amount = node.node_assets_amount def save(self): - # 这是个只读 Model + """ 这是个只读 Model """ raise NotImplementedError diff --git a/apps/perms/utils/__init__.py b/apps/perms/utils/__init__.py index 8563b21e9..808bdd1b8 100644 --- a/apps/perms/utils/__init__.py +++ b/apps/perms/utils/__init__.py @@ -1,4 +1,4 @@ from .permission import * -from .user_permission import * from .account import * from .user_perm_tree import * +from .user_perm import * diff --git a/apps/perms/utils/user_perm.py b/apps/perms/utils/user_perm.py new file mode 100644 index 000000000..a5f956de8 --- /dev/null +++ b/apps/perms/utils/user_perm.py @@ -0,0 +1,219 @@ +from assets.models import FavoriteAsset, Asset + +from django.conf import settings +from django.db.models import Q + +from common.utils.common import timeit + +from perms.models import AssetPermission, PermNode, UserAssetGrantedTreeNodeRelation + +from .permission import AssetPermissionUtil + + +__all__ = ['AssetPermissionPermAssetUtil', 'UserPermAssetUtil', 'UserPermNodeUtil'] + + +class AssetPermissionPermAssetUtil: + + def __init__(self, perm_ids): + self.perm_ids = perm_ids + + def get_all_assets(self): + """ 获取所有授权的资产 """ + node_asset_ids = self.get_perm_nodes_assets(flat=True) + direct_asset_ids = self.get_direct_assets(flat=True) + asset_ids = list(node_asset_ids) + list(direct_asset_ids) + assets = Asset.objects.filter(id__in=asset_ids) + return assets + + def get_perm_nodes_assets(self, flat=False): + """ 获取所有授权节点下的资产 """ + node_ids = AssetPermission.nodes.through.objects \ + .filter(assetpermission_id__in=self.perm_ids) \ + .values_list('node_id', flat=True) \ + .distinct() + node_ids = list(node_ids) + nodes = PermNode.objects.filter(id__in=node_ids).only('id', 'key') + assets = PermNode.get_nodes_all_assets(*nodes) + if flat: + return assets.values_list('id', flat=True) + return assets + + def get_direct_assets(self, flat=False): + """ 获取直接授权的资产 """ + assets = Asset.objects.order_by() \ + .filter(granted_by_permissions__id__in=self.perm_ids) \ + .distinct() + if flat: + return assets.values_list('id', flat=True) + return assets + + +class UserPermAssetUtil(AssetPermissionPermAssetUtil): + + def __init__(self, user): + self.user = user + perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True) + super().__init__(perm_ids) + + def get_ungroup_assets(self): + return self.get_direct_assets() + + def get_favorite_assets(self): + assets = self.get_all_assets() + asset_ids = FavoriteAsset.objects.filter(user=self.user).values_list('asset_id', flat=True) + assets = assets.filter(id__in=list(asset_ids)) + return assets + + def get_node_assets(self, key): + node = PermNode.objects.get(key=key) + node.compute_node_from_and_assets_amount(self.user) + if node.node_from == node.NodeFrom.granted: + assets = Asset.objects.filter(nodes__id=node.id).order_by() + elif node.node_from == node.NodeFrom.asset: + assets = self._get_indirect_perm_node_assets(node) + else: + assets = Asset.objects.none() + assets = assets.order_by('name') + return assets + + def get_node_all_assets(self, node_id): + """ 获取节点下的所有资产 """ + node = PermNode.objects.get(id=node_id) + node.compute_node_from_and_assets_amount(self.user) + if node.node_from == node.NodeFrom.granted: + assets = PermNode.get_nodes_all_assets() + elif node.node_from in (node.NodeFrom.asset, node.NodeFrom.child): + node.assets_amount = node.granted_assets_amount + assets = self._get_indirect_perm_node_all_assets(node) + else: + node.assets_amount = 0 + assets = Asset.objects.none() + return node, assets + + def _get_indirect_perm_node_assets(self, node): + """ 获取间接授权节点下的直接资产 """ + assets = self.get_direct_assets() + assets = assets.filter(nodes__id=node.id).order_by().distinct() + return assets + + def _get_indirect_perm_node_all_assets(self, node): + """ 获取间接授权节点下的所有资产 + 此算法依据 `UserAssetGrantedTreeNodeRelation` 的数据查询 + 1. 查询该节点下的直接授权节点 + 2. 查询该节点下授权资产关联的节点 + """ + # 查询节点下直接授权的子节点 + asset_ids = set() + children_from_granted = UserAssetGrantedTreeNodeRelation.objects \ + .filter(user=self.user) \ + .filter(node_key__startwith=f'{node.key}:', node_from=node.NodeFrom.granted) \ + .only('node_id', 'node_key') + for n in children_from_granted: + n.id = n.node_id + _assets = PermNode.get_nodes_all_assets(*children_from_granted) + _asset_ids = _assets.values_list('id', flat=True) + asset_ids.update(list(_asset_ids)) + + # 查询节点下资产授权的节点 + children_from_assets = UserAssetGrantedTreeNodeRelation.objects \ + .filter(user=self.user) \ + .filter(node_key__startwith=f'{node.key}:', node_from=node.NodeFrom.asset) \ + .values_list('node_id', flat=True) + children_from_assets = set(children_from_assets) + if node.node_from == node.NodeFrom.asset: + children_from_assets.add(node.id) + _asset_ids = Asset.objects \ + .filter(node__id__in=children_from_assets) \ + .filter(granted_by_permissions__id__in=self.perm_ids) \ + .distinct() \ + .order_by() \ + .values_list('id', flat=True) + asset_ids.update(list(_asset_ids)) + + return Asset.objects.filter(id__in=asset_ids) + + +class UserPermNodeUtil: + + def __init__(self, user): + self.user = user + self.perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True) + + def get_favorite_node(self): + assets_amount = UserPermAssetUtil(self.user).get_favorite_assets().count() + return PermNode.get_favorite_node(assets_amount) + + def get_ungrouped_node(self): + assets_amount = UserPermAssetUtil(self.user).get_direct_assets().count() + return PermNode.get_favorite_node(assets_amount) + + def get_top_level_nodes(self): + nodes = self.get_special_nodes() + real_nodes = self._get_indirect_perm_node_children(key='') + if len(real_nodes) == 1: + children = self.get_node_children(real_nodes[0].key) + nodes.extend(children) + return nodes + + def get_special_nodes(self): + nodes = [] + if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: + ung_node = self.get_ungrouped_node() + nodes.append(ung_node) + fav_node = self.get_favorite_node() + nodes.append(fav_node) + return nodes + + def get_node_children(self, key): + if not key: + return self.get_top_level_nodes() + + if key in [PermNode.FAVORITE_NODE_KEY, PermNode.UNGROUPED_NODE_KEY]: + return PermNode.objects.none() + + node = PermNode.objects.get(key=key) + node.compute_node_from_and_assets_amount(self.user) + if node.node_from == node.NodeFrom.granted: + children = PermNode.objects.filter(parent_key=key) + elif node.node_from in (node.NodeFrom.asset, node.NodeFrom.child): + children = self._get_indirect_perm_node_children(key) + else: + children = PermNode.objects.none() + children = sorted(children, key=lambda x: x.value) + return children + + def _get_indirect_perm_node_children(self, key): + """ 获取未直接授权节点的子节点 """ + children = PermNode.objects.filter(granted_node_rels__user=self.user, parent_key=key) + children = children.annotate(**PermNode.annotate_granted_node_rel_fields).distinct() + for node in children: + node.assets_amount = node.granted_assets_amount + return children + + @timeit + def get_whole_tree_nodes(self, with_special=True): + user_nodes = PermNode.objects.filter(granted_node_rels__user=self.user) + user_nodes = user_nodes.annotate(**PermNode.annotate_granted_node_rel_fields).distinct() + + key_node_mapper = {} + q_nodes_descendant = Q() + for node in user_nodes: + node.assets_amount = node.granted_assets_amount + key_node_mapper[node.key] = node + if node.node_from == node.NodeFrom.granted: + """ 直接授权的节点, 增加后代节点的过滤条件 """ + q_nodes_descendant |= Q(key__startswith=f'{node.key}:') + if q_nodes_descendant: + descendant_nodes = PermNode.objects.filter(q_nodes_descendant) + for node in descendant_nodes: + key_node_mapper[node.key] = node + + nodes = [] + if with_special: + special_nodes = self.get_special_nodes() + nodes.extend(special_nodes) + nodes.extend(list(key_node_mapper.values())) + + return nodes + diff --git a/apps/perms/utils/user_permission.py b/apps/perms/utils/user_permission.py deleted file mode 100644 index 8ee5992dd..000000000 --- a/apps/perms/utils/user_permission.py +++ /dev/null @@ -1,280 +0,0 @@ -from collections import defaultdict -from typing import List, Tuple - -from django.conf import settings -from django.db.models import Q, QuerySet -from django.utils.translation import gettext as _ - -from users.models import User -from assets.utils import NodeAssetsUtil -from assets.models import ( - Asset, - FavoriteAsset, - AssetQuerySet, - NodeQuerySet -) -from orgs.utils import ( - tmp_to_org, - current_org, - ensure_in_real_or_default_org, -) -from common.db.models import output_as_string, UnionQuerySet -from common.utils import get_logger -from common.utils.common import lazyproperty, timeit - -from perms.models import ( - AssetPermission, - PermNode, - UserAssetGrantedTreeNodeRelation -) -from .permission import AssetPermissionUtil - -NodeFrom = UserAssetGrantedTreeNodeRelation.NodeFrom -NODE_ONLY_FIELDS = ('id', 'key', 'parent_key', 'org_id') - -logger = get_logger(__name__) - - -class UserGrantedUtilsBase: - user: User - - def __init__(self, user, asset_perm_ids=None): - self.user = user - self._asset_perm_ids = asset_perm_ids and set(asset_perm_ids) - - @lazyproperty - def asset_perm_ids(self) -> set: - if self._asset_perm_ids: - return self._asset_perm_ids - - asset_perm_ids = AssetPermissionUtil().get_permissions_for_user(self.user, flat=True) - return asset_perm_ids - - -class UserGrantedAssetsQueryUtils(UserGrantedUtilsBase): - - def get_favorite_assets(self) -> QuerySet: - assets = self.get_all_granted_assets() - asset_ids = FavoriteAsset.objects.filter(user=self.user).values_list('asset_id', flat=True) - assets = assets.filter(id__in=list(asset_ids)) - return assets - - def get_ungroup_assets(self) -> AssetQuerySet: - return self.get_direct_granted_assets() - - def get_direct_granted_assets(self) -> AssetQuerySet: - queryset = Asset.objects.order_by().filter( - granted_by_permissions__id__in=self.asset_perm_ids - ).distinct() - return queryset - - def get_direct_granted_nodes_assets(self) -> AssetQuerySet: - granted_node_ids = AssetPermission.nodes.through.objects.filter( - assetpermission_id__in=self.asset_perm_ids - ).values_list('node_id', flat=True).distinct() - granted_node_ids = list(granted_node_ids) - granted_nodes = PermNode.objects.filter(id__in=granted_node_ids).only('id', 'key') - queryset = PermNode.get_nodes_all_assets(*granted_nodes) - return queryset - - def get_all_granted_assets(self) -> QuerySet: - nodes_assets = self.get_direct_granted_nodes_assets() - assets = self.get_direct_granted_assets() - # queryset = UnionQuerySet(nodes_assets, assets) - # return queryset - node_asset_ids = nodes_assets.values_list('id', flat=True) - direct_asset_ids = assets.values_list('id', flat=True) - asset_ids = list(node_asset_ids) + list(direct_asset_ids) - asset = Asset.objects.filter(id__in=asset_ids) - return asset - - def get_node_all_assets(self, id) -> Tuple[PermNode, QuerySet]: - node = PermNode.objects.get(id=id) - granted_status = node.get_granted_status(self.user) - if granted_status == NodeFrom.granted: - assets = PermNode.get_nodes_all_assets(node) - return node, assets - elif granted_status in (NodeFrom.asset, NodeFrom.child): - node.use_granted_assets_amount() - assets = self._get_indirect_granted_node_all_assets(node) - return node, assets - else: - node.assets_amount = 0 - return node, Asset.objects.none() - - def get_node_assets(self, key) -> AssetQuerySet: - node = PermNode.objects.get(key=key) - granted_status = node.get_granted_status(self.user) - - if granted_status == NodeFrom.granted: - assets = Asset.objects.order_by().filter(nodes__id=node.id) - elif granted_status == NodeFrom.asset: - assets = self._get_indirect_granted_node_assets(node.id) - else: - assets = Asset.objects.none() - assets = assets.order_by('name') - return assets - - def _get_indirect_granted_node_assets(self, id) -> AssetQuerySet: - assets = Asset.objects.order_by().filter(nodes__id=id).distinct() & self.get_direct_granted_assets() - return assets - - def _get_indirect_granted_node_all_assets(self, node) -> QuerySet: - """ - 此算法依据 `UserAssetGrantedTreeNodeRelation` 的数据查询 - 1. 查询该节点下的直接授权节点 - 2. 查询该节点下授权资产关联的节点 - """ - user = self.user - - # 查询该节点下的授权节点 - granted_nodes = UserAssetGrantedTreeNodeRelation.objects.filter( - user=user, node_from=NodeFrom.granted - ).filter( - Q(node_key__startswith=f'{node.key}:') - ).only('node_id', 'node_key') - - for n in granted_nodes: - n.id = n.node_id - - node_assets = PermNode.get_nodes_all_assets(*granted_nodes) - - # 查询该节点下的资产授权节点 - only_asset_granted_node_ids = UserAssetGrantedTreeNodeRelation.objects.filter( - user=user, node_from=NodeFrom.asset - ).filter( - Q(node_key__startswith=f'{node.key}:') - ).values_list('node_id', flat=True) - - only_asset_granted_node_ids = list(only_asset_granted_node_ids) - if node.node_from == NodeFrom.asset: - only_asset_granted_node_ids.append(node.id) - - assets = Asset.objects.filter( - nodes__id__in=only_asset_granted_node_ids, - granted_by_permissions__id__in=self.asset_perm_ids - ).distinct().order_by() - granted_assets = UnionQuerySet(node_assets, assets) - return granted_assets - - -class UserGrantedNodesQueryUtils(UserGrantedUtilsBase): - def sort(self, nodes): - nodes = sorted(nodes, key=lambda x: x.value) - return nodes - - def get_node_children(self, key): - if not key: - return self.get_top_level_nodes() - - nodes = PermNode.objects.none() - if key in [PermNode.FAVORITE_NODE_KEY, PermNode.UNGROUPED_NODE_KEY]: - return nodes - - node = PermNode.objects.get(key=key) - granted_status = node.get_granted_status(self.user) - if granted_status == NodeFrom.granted: - nodes = PermNode.objects.filter(parent_key=key) - elif granted_status in (NodeFrom.asset, NodeFrom.child): - nodes = self.get_indirect_granted_node_children(key) - nodes = self.sort(nodes) - return nodes - - def get_indirect_granted_node_children(self, key): - """ - 获取用户授权树中未授权节点的子节点 - 只匹配在 `UserAssetGrantedTreeNodeRelation` 中存在的节点 - """ - user = self.user - nodes = PermNode.objects.filter( - granted_node_rels__user=user, - parent_key=key - ).annotate( - **PermNode.annotate_granted_node_rel_fields - ).distinct() - - # 设置节点授权资产数量 - for node in nodes: - node.use_granted_assets_amount() - return nodes - - def get_top_level_nodes(self): - nodes = self.get_special_nodes() - real_nodes = self.get_indirect_granted_node_children('') - nodes.extend(real_nodes) - if len(real_nodes) == 1: - children = self.get_node_children(real_nodes[0].key) - nodes.extend(children) - return nodes - - def get_ungrouped_node(self): - assets_util = UserGrantedAssetsQueryUtils(self.user, self.asset_perm_ids) - assets_amount = assets_util.get_direct_granted_assets().count() - return PermNode.get_ungrouped_node(assets_amount) - - def get_favorite_node(self): - assets_query_utils = UserGrantedAssetsQueryUtils(self.user, self.asset_perm_ids) - assets_amount = assets_query_utils.get_favorite_assets().values_list('id').count() - return PermNode.get_favorite_node(assets_amount) - - @staticmethod - def get_root_node(): - name = _('My assets') - node = { - 'id': '', - 'name': name, - 'title': name, - 'pId': '', - 'open': True, - 'isParent': True, - 'meta': { - 'type': 'root' - } - } - return node - - def get_special_nodes(self): - nodes = [] - if settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE: - ungrouped_node = self.get_ungrouped_node() - nodes.append(ungrouped_node) - favorite_node = self.get_favorite_node() - nodes.append(favorite_node) - return nodes - - @timeit - def get_whole_tree_nodes(self, with_special=True): - """ - 这里的 granted nodes, 是整棵树需要的node，推算出来的也算 - :param with_special: - :return: - """ - nodes = PermNode.objects.filter(granted_node_rels__user=self.user) \ - .annotate(**PermNode.annotate_granted_node_rel_fields) \ - .distinct() - - key_to_node_mapper = {} - nodes_descendant_q = Q() - - for node in nodes: - node.use_granted_assets_amount() - key_to_node_mapper[node.key] = node - - if node.node_from == NodeFrom.granted: - # 直接授权的节点 - # 增加查询后代节点的过滤条件 - nodes_descendant_q |= Q(key__startswith=f'{node.key}:') - - if nodes_descendant_q: - descendant_nodes = PermNode.objects.filter( - nodes_descendant_q - ) - for node in descendant_nodes: - key_to_node_mapper[node.key] = node - - all_nodes = [] - if with_special: - special_nodes = self.get_special_nodes() - all_nodes.extend(special_nodes) - all_nodes.extend(key_to_node_mapper.values()) - return all_nodes