diff --git a/apps/assets/tree/asset_tree.py b/apps/assets/tree/asset_tree.py index ee1f6069b..e15f8902b 100644 --- a/apps/assets/tree/asset_tree.py +++ b/apps/assets/tree/asset_tree.py @@ -5,6 +5,7 @@ from orgs.utils import current_org from orgs.models import Organization from assets.models import Asset, Node, Platform from assets.const.category import Category +from assets.const.types import AllTypes from common.utils import get_logger, timeit, lazyproperty from .tree import TreeNode, Tree @@ -69,7 +70,7 @@ class AssetTree(Tree): TreeNode = AssetTreeNode - def __init__(self, assets_q_object: Q = None, category=None, org=None, + def __init__(self, assets_q_object: Q = None, asset_category=None, asset_type=None, org=None, with_assets_all=False, with_assets_node_id=None, with_assets_node_levels=None, with_assets_limit=None, full_tree=True): @@ -88,8 +89,9 @@ class AssetTree(Tree): super().__init__() ## 通过资产构建节点树, 支持 Q, category, org 等过滤条件 ## self._assets_q_object: Q = assets_q_object or Q() - self._category = self._check_category(category) - self._category_platform_ids = set() + self._asset_category = self._check_asset_category(asset_category) + self._asset_type = self._check_asset_type(asset_type) + self._asset_platform_ids = set() self._org: Organization = org or current_org # org 下全量节点属性映射, 构建资产树时根据完整的节点进行构建 @@ -109,18 +111,23 @@ class AssetTree(Tree): # 初始化时构建树 self.build() - def _check_category(self, category): - if category is None: - return None + def _check_asset_category(self, category): if category in Category.values: return category logger.warning(f"Invalid category '{category}' for AssetSearchTree.") return None + def _check_asset_type(self, asset_type): + types = list(dict(AllTypes.choices()).keys()) + if asset_type in types: + return asset_type + logger.warning(f"Invalid asset_type '{asset_type}' for AssetSearchTree.") + return None + @timeit def build(self): self._load_nodes_attr_mapper() - self._load_category_platforms_if_needed() + self._load_asset_platforms_if_needed() self._load_nodes_assets_amount() self._init_tree() self._load_nodes_assets_if_needed() @@ -129,13 +136,17 @@ class AssetTree(Tree): self._compute_children_count_total() @timeit - def _load_category_platforms_if_needed(self): - if self._category is None: - return - ids = Platform.objects.filter(category=self._category).values_list('id', flat=True) + def _load_asset_platforms_if_needed(self): + if self._asset_type: + q = Q(type=self._asset_type) + elif self._asset_category: + q = Q(category=self._asset_category) + else: + q = Q() + ids = ids = Platform.objects.filter(q).values_list('id', flat=True) ids = self._uuids_to_string(ids) - self._category_platform_ids = ids - + self._asset_platform_ids = set(ids) + @timeit def _load_nodes_attr_mapper(self): nodes = Node.objects.filter(org_id=self._org.id).values('id', 'key', 'value') @@ -158,8 +169,8 @@ class AssetTree(Tree): @timeit def _make_assets_q_object(self) -> Q: q = Q(org_id=self._org.id) - if self._category_platform_ids: - q &= Q(platform_id__in=self._category_platform_ids) + if self._asset_platform_ids: + q &= Q(platform_id__in=self._asset_platform_ids) if self._assets_q_object: q &= self._assets_q_object return q @@ -267,5 +278,5 @@ class AssetTree(Tree): def print(self, count=20, simple=True): print('org_name: ', getattr(self._org, 'name', 'No-org')) - print(f'asset_category: {self._category}') + print(f'asset_category: {self._asset_category}') super().print(count=count, simple=simple) diff --git a/apps/assets/tree/tree.py b/apps/assets/tree/tree.py index 3c2fa5ff6..93ea37b5c 100644 --- a/apps/assets/tree/tree.py +++ b/apps/assets/tree/tree.py @@ -1,3 +1,4 @@ +from collections import deque from common.utils import get_logger, lazyproperty, timeit @@ -42,6 +43,35 @@ class TreeNode(object): def level(self): return self.key.count(':') + 1 + def get_ancestor_keys(self): + if self.is_root: + return [] + + ancestor_keys = [] + parts = self.key.split(':') + for i in range(1, len(parts)): + ancestor_key = ':'.join(parts[:i]) + ancestor_keys.append(ancestor_key) + return ancestor_keys + + def get_descendants(self, node: 'TreeNode'): + """ + 返回指定节点的所有子孙节点（不包含自身），非递归实现，按层级从近到远排序。 + 返回列表，空列表表示没有子孙或节点为 None。 + """ + if not node: + return [] + + descendants = [] + dq = deque(node.children) + while dq: + cur = dq.popleft() + descendants.append(cur) + # 复制 children 以防在遍历过程中被修改 + for ch in list(cur.children): + dq.append(ch) + return descendants + @property def children_count(self): return len(self.children) @@ -139,6 +169,57 @@ class Tree(object): parent.remove_child(node) self.nodes.pop(node.key, None) + def search_nodes(self, keyword, only_top_level=False): + if not keyword: + return [] + keyword = keyword.strip().lower() + nodes = {} + for node in self.nodes.values(): + node: TreeNode + node_value = str(node.value).strip().lower() + if keyword in node_value: + nodes[node.key] = node + if not only_top_level: + return list(nodes.values()) + + # TODO: 优化性能 + node_keys = list(nodes.keys()) + children_keys = [] + for node_key in node_keys: + _children_keys = [k for k in node_keys if k.startswith(f"{node_key}:")] + children_keys.extend(_children_keys) + for child_key in children_keys: + nodes.pop(child_key, None) + return list(nodes.values()) + + def remove_nodes_descendants(self, nodes: list[TreeNode]): + descendants = self.get_nodes_descendants(nodes) + for node in reversed(descendants): + self.remove_node(node) + + def get_nodes_descendants(self, nodes: list[TreeNode]): + descendants = [] + for node in nodes: + ds = node.get_descendants(node) + descendants.extend(ds) + return descendants + + def get_nodes_ancestors(self, nodes: list[TreeNode]): + ancestors = set() + for node in nodes: + ancestor_keys = node.get_ancestor_keys() + _ancestors = self.get_nodes_by_keys(ancestor_keys) + ancestors.update(_ancestors) + return list(ancestors) + + def get_nodes_by_keys(self, keys): + nodes = [] + for key in keys: + node = self.get_node(key) + if node: + nodes.append(node) + return nodes + def get_nodes(self, levels=None): nodes = list(self.nodes.values()) if levels: diff --git a/apps/perms/api/user_permission/tree.py b/apps/perms/api/user_permission/tree.py index deb4e3224..42c01348c 100644 --- a/apps/perms/api/user_permission/tree.py +++ b/apps/perms/api/user_permission/tree.py @@ -29,18 +29,58 @@ class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, def list(self, request, *args, **kwargs): with_assets = request.query_params.get('assets', '0') == '1' search = request.query_params.get('search') - key = request.query_params.get('key') - - if not with_assets: - return self.init_user_perm_tree() - + node_key = request.query_params.get('key') + asset_category = request.query_params.get('category') + asset_type = request.query_params.get('type') + # test if search: - return self.search_user_perm_tree_with_assets(search) + search_list = search.split() + for s in search_list: + node = s.split('node:') + if len(node) == 2: + search = node[1] + with_assets = False + break - if key: - return self.expand_tree_node_with_assets(key) + asset = s.split('asset:') + if len(asset) == 2: + search = asset[1] + continue + asset_category = s.split('category:') + if len(asset_category) == 2: + asset_category = asset_category[1] + continue + asset_type = s.split('type:') + if len(asset_type) == 2: + asset_type = asset_type[1] + continue - return self.init_user_perm_tree_with_assets() + if node_key: + with_assets = True + search = None + asset_category = None + asset_type = None + + if with_assets: + if node_key: + return self.expand_tree_node_with_assets( + node_key, asset_category, asset_type + ) + elif search: + # search assets + return self.search_user_perm_tree_with_assets( + search, asset_category, asset_type + ) + else: + return self.init_user_perm_tree_with_assets( + asset_category, asset_type + ) + else: + if search: + # search nodes + return self.search_user_perm_tree(search) + else: + return self.init_user_perm_tree() def init_user_perm_tree(self): ''' 初始化用户权限树 - 不包含资产 @@ -73,8 +113,44 @@ class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, data = nodes data = self.add_favorites_and_ungrouped_node(data, with_assets=False) return Response(data=data) + + def search_user_perm_tree(self, search): + ''' 搜索用户授权树 - 不包含资产 + 全局组织: 返回所有匹配节点以及祖先节点，不返回匹配节点的子孙节点，不返回资产，展开所有祖先节点，不展开匹配节点 + 实体组织: 同上 + ''' + if current_org.is_root(): + orgs = self.user.orgs.all() + else: + orgs = self.user.orgs.filter(id=current_org.id) + + search_nodes = [] + nodes_ancestors = [] + for org in orgs: + tree = UserPermTree(user=self.user, org=org) + _search_nodes = tree.search_nodes(search, only_top_level=True) + # tree.remove_nodes_descendants(_search_nodes) + _nodes_ancestors = tree.get_nodes_ancestors(_search_nodes) + search_nodes.extend(_search_nodes) + nodes_ancestors.extend(_nodes_ancestors) + + # 不展开搜索节点 + expand_level = 0 + # 如果有资产，则允许展开 is_parent=True 的节点 + with_assets = True + serialized_search_nodes = self.serialize_nodes( + search_nodes, with_asset_amount=True, expand_level=expand_level, + with_assets=with_assets + ) + # 展开所有祖先节点 + expand_level = 10000 + serialized_nodes_ancestors = self.serialize_nodes( + nodes_ancestors, with_asset_amount=True, expand_level=expand_level + ) + data = [*serialized_nodes_ancestors, *serialized_search_nodes] + return Response(data=data) - def init_user_perm_tree_with_assets(self): + def init_user_perm_tree_with_assets(self, asset_category=None, asset_type=None): ''' 初始化用户权限资产树 - 包含资产 全局组织: 返回第1级节点，不返回资产，不展开节点 实体组织：返回第1级和第2级节点，返回第1级节点的资产，展开第1级节点 @@ -102,7 +178,9 @@ class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, for org in orgs: tree = UserPermTree( - user=self.user, org=org, with_assets_node_levels=with_assets_node_levels + user=self.user, + asset_category=asset_category, asset_type=asset_type, org=org, + with_assets_node_levels=with_assets_node_levels ) _nodes = tree.get_nodes(levels=nodes_level) nodes.extend(_nodes) @@ -117,19 +195,21 @@ class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, data = self.add_favorites_and_ungrouped_node(data, with_assets=True) return Response(data=data) - def expand_tree_node_with_assets(self, key): + def expand_tree_node_with_assets(self, node_key, asset_category=None, asset_type=None): ''' 展开用户权限资产树节点 - 包含资产 全局组织: 返回展开节点的直接孩子节点，返回展开节点的资产，不展开其他节点 实体组织: 同上 ''' expand_level = 0 - node = get_object_or_404(Node, key=key) + node = get_object_or_404(Node, key=node_key) org = self.user.orgs.filter(id=node.org_id).first() if not org: return Response(data=[]) tree = UserPermTree( - user=self.user, org=node.org, with_assets_node_id=node.id + user=self.user, + asset_category=asset_category, asset_type=asset_type, + org=node.org, with_assets_node_id=node.id ) tree_node = tree.get_node(node.key) if not tree_node: @@ -144,7 +224,7 @@ class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, data = [*nodes, *assets] return Response(data=data) - def search_user_perm_tree_with_assets(self, search): + def search_user_perm_tree_with_assets(self, search, asset_category=None, asset_type=None): ''' 初始化用户权限资产搜索树 - 包含资产 全局组织: 返回所有节点，返回所有资产，展开所有节点，搜索资产 (最大 1000， n 个组织，每个组织分配1000/n个资产) 实体组织: 同上，最大资产数 1000 @@ -160,13 +240,14 @@ class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, if not orgs.exists(): return Response(data=[]) - + assets_q_object = Q(name__icontains=search) | Q(address__icontains=search) nodes = [] assets = [] for org in orgs: tree = UserPermTree( - user=self.user, assets_q_object=assets_q_object, org=org, + user=self.user, assets_q_object=assets_q_object, + org=org, asset_category=asset_category, asset_type=asset_type, with_assets_all=with_assets_all, with_assets_limit=with_assets_limit )