github/jumpserver
1
0
Fork 0
mirror of https://github.com/jumpserver/jumpserver.git synced 2026-03-18 19:12:07 +00:00
Code Issues Projects Releases Wiki Activity

perf: UserPermTreeAPI support search nodes, search assetes by category or type

Browse Source
This commit is contained in:
Bai
2025-12-27 21:45:50 +08:00
parent 62ec308d1e
commit bc9c9c30cd
3 changed files with 206 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
apps/assets/tree/asset_tree.py
View File

@@ -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)

81
apps/assets/tree/tree.py
View File

@@ -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:

115
apps/perms/api/user_permission/tree.py
View File

@@ -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
)