perf: UserPermedAssetTreeAPI about get_special_nodes

apps/assets/api/tree/base.py

@@ -5,7 +5,7 @@ from rest_framework.response import Response
 from rest_framework.generics import get_object_or_404
 
 from users.models import User
-from common.utils import lazyproperty
+from common.utils import lazyproperty, timeit
 from common.exceptions import APIException
 from orgs.utils import current_org
 from orgs.models import Organization
@@ -15,15 +15,36 @@ from assets.models import Node
 from .mixin import SerializeToTreeNodeMixin
 from .const import RenderTreeType, RenderTreeTypeChoices
 
+
+__all__ = ['AbstractAssetTreeAPI']
+
+
 class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
 
     permission_classes = (RBACPermission,)
 
+    # TODO: 再确认一下 API 所需的权限位
     perm_model = Node
 
-    render_tree_type: RenderTreeType
-    tree_user: User
+    # query parameters keys
+    query_search_key = 'search'
+    query_search_key_value_sep = ':'
+    query_tree_type_key = 'tree_type'
+    query_asset_category_key = 'category'
+    query_asset_type_key = 'type'
+    query_search_asset_key = 'search_asset'
+    query_search_node_key = 'search_node'
+    # 以上参数均支持通过 search 参数传递，格式如 search=key:value
+    # 也支持直接通过对应的 query parameter 传递，如 ?search_node=demo_node&type=linux
+    query_expand_node_key = 'key'
 
+    # 每个组织树搜索资产时，返回的资产数量限制
+    search_assets_per_org_limit_max = 1000
+    search_assets_per_org_limit_min = 100
+
+    render_tree_type: RenderTreeType
+
+    tree_user: User
 
     def initial(self, request, *args, **kwargs):
         super().initial(request, *args, **kwargs)
@@ -31,70 +52,159 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
         self.tree_user = self.get_tree_user()
 
     def initial_render_tree_type(self):
-        tree_type = self.get_query_value('tree_type')
+        tree_type = self.get_query_value(self.query_tree_type_key)
         if not tree_type:
+            # 兼容 assets=1 参数
             with_assets = self.request.query_params.get('assets', '0') == '1'
-            tree_type = RenderTreeTypeChoices.asset if with_assets else RenderTreeTypeChoices.node
+            if with_assets:
+                tree_type = RenderTreeTypeChoices.asset
+            else:
+                tree_type = RenderTreeTypeChoices.node
         return RenderTreeType(tree_type)
 
     def get_tree_user(self) -> User:
+        # 抽象方法: 获取为哪个用户渲染树 #
         raise NotImplementedError
 
+
+    def get_query_value(self, query_key):
+        query_value = self.request.query_params.get(query_key)
+        if not query_value:
+            query_value = self.get_query_value_from_search(query_key)
+        return query_value
+    
+    def get_query_value_from_search(self, query_key):
+        # 辅助方法：获取查询的参数，支持在 search 参数中以 key:value 形式传递 #
+        search = self.request.query_params.get(self.query_search_key, '')
+        if not search:
+            return None
+        
+        sep = self.query_search_key_value_sep
+
+        search_list = search.split()
+        for _search in search_list:
+            if f'{query_key}{sep}' not in _search:
+                continue
+            query_value = _search.replace(f'{query_key}{sep}', '').strip()
+            return query_value
+    
+    def get_org_asset_tree(self, **kwargs) -> AssetTree:
+        # 抽象方法: 获取组织的资产树 #
+        raise NotImplementedError
+
+    @lazyproperty
+    def org_is_global(self):
+        return current_org.is_root()
+
+    def get_tree_user_orgs(self):
+        # 重要: 获取用户有权限渲染树的组织列表 #
+        user = self.tree_user
+        if self.org_is_global:
+            # 如果是全局组织，返回用户所在的所有实体组织
+            orgs = user.orgs.all()
+        else:
+            # 如果时实体组织，从用户所在的实体组织中返回该实体组织
+            orgs = user.orgs.filter(id=current_org.id)
+        if not orgs.exists():
+            raise APIException(
+                'No organization available for rendering the tree'
+            )
+        return orgs
+
+    @timeit
     def list(self, request, *args, **kwargs):
-        asset_category = self.get_query_value('asset_category')
-        asset_type = self.get_query_value('asset_type')
+        # 渲染资产树 API 接口 #
+        # 支持渲染节点树和资产树两种类型
+        # 节点树: 只返回节点
+        # 资产树: 返回节点和节点下的资产
+        asset_category = self.get_query_value(self.query_asset_category_key)
+        asset_type = self.get_query_value(self.query_asset_type_key)
         with_asset_amount = True
 
         if self.render_tree_type.is_node_tree:
-            data = self.render_node_tree(asset_category, asset_type, with_asset_amount)
+            data = self.render_node_tree(
+                asset_category=asset_category, asset_type=asset_type, 
+                with_asset_amount=with_asset_amount
+            )
         elif self.render_tree_type.is_asset_tree:
-            data = self.render_asset_tree(asset_category, asset_type, with_asset_amount)
+            data = self.render_asset_tree(
+                asset_category=asset_category, asset_type=asset_type, 
+                with_asset_amount=with_asset_amount
+            )
         else:
-            raise APIException(f'Invalid tree type: {self.query_tree_type}')
+            raise APIException(
+                f'Invalid tree type: {self.render_tree_type}'
+            )
         return Response(data=data)
 
-    def render_node_tree(self, asset_category, asset_type, with_asset_amount):
-        data = self.init_node_tree(asset_category, asset_type, with_asset_amount)
+    @timeit
+    def render_node_tree(self, asset_category=None, asset_type=None, with_asset_amount=True):
+        # 渲染节点树 #
+        # 节点树返回所有节点的数据，不返回资产
+        # 节点树只有初始化
+        # 搜索节点和展开节点的动作，在前端页面执行
+        data = self.init_node_tree(
+            asset_category=asset_category, 
+            asset_type=asset_type, 
+            with_asset_amount=with_asset_amount
+        )
         return data
     
-    def render_asset_tree(self, asset_category, asset_type, with_asset_amount):
-        expand_node_key = self.get_query_value('key')
-        search_node = self.get_query_value('search_node')
-        search_asset = self.get_query_value('search_asset')
+    @timeit
+    def render_asset_tree(self, asset_category=None, asset_type=None, with_asset_amount=True):
+        # 渲染资产树 #
+        expand_node_key = self.get_query_value(self.query_expand_node_key)
+        search_node = self.get_query_value(self.query_search_node_key)
+        search_asset = self.get_query_value(self.query_search_asset_key)
         data = self._render_asset_tree(
-            expand_node_key=expand_node_key, search_node=search_node, search_asset=search_asset, 
+            expand_node_key=expand_node_key, 
+            search_node=search_node, search_asset=search_asset, 
             asset_category=asset_category, asset_type=asset_type, 
             with_asset_amount=with_asset_amount
         )
         return data
 
+    @timeit
     def _render_asset_tree(self, expand_node_key=None, search_node=None, search_asset=None, 
-                           asset_category=None, asset_type=None, with_asset_amount=None):
+                           asset_category=None, asset_type=None, with_asset_amount=True):
+        # 渲染资产树内部方法，支持子类重载 #
+        # 此方法包含渲染资产树的所有参数
+        # 资产树支持初始化、搜索节点、搜索资产和展开节点等动作
 
-        if not search_asset and self.render_tree_type.is_asset_tree:
-            search = self.get_query_value('search') or ''
-            if ':' not in search:
+        if not search_asset:
+            # 兼容 search 为搜索资产
+            search = self.get_query_value(self.query_search_key) or ''
+            sep = self.query_search_key_value_sep
+            if sep not in search:
                 search_asset = search
 
         if expand_node_key:
             data = self.expand_asset_tree_node(
-                expand_node_key, asset_category, asset_type, with_asset_amount
+                node_key=expand_node_key, 
+                asset_category=asset_category, asset_type=asset_type, 
+                with_asset_amount=with_asset_amount
             )
         elif search_node:
             data = self.search_asset_tree_nodes(
-                search_node, asset_category, asset_type, with_asset_amount
+                search_node=search_node, 
+                asset_category=asset_category, asset_type=asset_type, 
+                with_asset_amount=with_asset_amount
             )
         elif search_asset:
             data = self.search_asset_tree_assets(
-                search_asset, asset_category, asset_type, with_asset_amount
+                search_asset=search_asset, 
+                asset_category=asset_category, asset_type=asset_type, 
+                with_asset_amount=with_asset_amount
             )
         else:
             data = self.init_asset_tree(
-                asset_category, asset_type, with_asset_amount
+                asset_category=asset_category, asset_type=asset_type, 
+                with_asset_amount=with_asset_amount
             )
         return data
     
-    def init_node_tree(self, asset_category, asset_type, with_asset_amount):
+    @timeit
+    def init_node_tree(self, asset_category=None, asset_type=None, with_asset_amount=True):
         orgs = self.get_tree_user_orgs()
 
         if self.org_is_global:
@@ -111,12 +221,13 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
             nodes.extend(_nodes)
 
         data = self.serialize_nodes(
-            nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
-            expand_level=expand_level
+            nodes=nodes, tree_type=self.render_tree_type, 
+            with_asset_amount=with_asset_amount, expand_level=expand_level
         )
         return data
     
-    def init_asset_tree(self, asset_category, asset_type, with_asset_amount):
+    @timeit
+    def init_asset_tree(self, asset_category=None, asset_type=None, with_asset_amount=True):
         orgs = self.get_tree_user_orgs()
         if self.org_is_global:
             node_levels = [1]
@@ -130,7 +241,9 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
         assets = []
         for org in orgs:
             tree: AssetTree = self.get_org_asset_tree(
-                asset_category=asset_category, asset_type=asset_type, org=org, 
+                asset_category=asset_category, 
+                asset_type=asset_type, 
+                org=org, 
                 with_assets_node_levels=with_assets_node_levels
             )
             _nodes = tree.get_nodes(levels=node_levels)
@@ -139,26 +252,34 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
             assets.extend(_assets)
         
         serialized_nodes = self.serialize_nodes(
-            nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
+            nodes=nodes, tree_type=self.render_tree_type, 
+            with_asset_amount=with_asset_amount, 
             expand_level=expand_level
         )
         serialized_assets = self.serialize_assets(assets)
         data = serialized_nodes + serialized_assets
         return data
 
-    def expand_asset_tree_node(self, node_key,  asset_category, asset_type, with_asset_amount):
+    @timeit
+    def expand_asset_tree_node(self, node_key,  asset_category=None, asset_type=None, 
+                               with_asset_amount=True):
+        # 展开资产树节点 #
+        # 展开节点时，返回该节点的直接子节点和直接资产
+
         node = get_object_or_404(Node, key=node_key)
         orgs = self.get_tree_user_orgs()
-        # 确保用户有权限展开该节点所在组织的树
         org = orgs.filter(id=node.org_id).first()
         if not org:
+            # 确保用户有权限展开该节点所在组织的树
             raise APIException(
                 f'No permission to expand the node in this organization: {node.org_name}'
             )
         
         with_assets_node_id = str(node.id)
         tree: AssetTree = self.get_org_asset_tree(
-            asset_category=asset_category, asset_type=asset_type, org=org,
+            asset_category=asset_category, 
+            asset_type=asset_type, 
+            org=org,
             with_assets_node_id=with_assets_node_id
         )
         node_children = tree.get_node_children(node.key)
@@ -167,21 +288,33 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
         # (展开节点)的孩子节点不展开
         expand_level = 0
         serialized_nodes = self.serialize_nodes(
-            node_children, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
+            nodes=node_children, 
+            tree_type=self.render_tree_type, 
+            with_asset_amount=with_asset_amount, 
             expand_level=expand_level
         )
         serialized_assets = self.serialize_assets(node_assets)
         data = serialized_nodes + serialized_assets
         return data
 
-    def search_asset_tree_nodes(self, search_node, asset_category, asset_type, with_asset_amount):
+    @timeit
+    def search_asset_tree_nodes(self, search_node, asset_category=None, asset_type=None, 
+                                with_asset_amount=True):
+        # 搜索资产树节点 #
+        # 搜索节点时，返回匹配的节点和匹配节点的所有祖先节点
+        # 匹配节点中如果包含有父子关系的节点，只返回最上一级的父节点
+        # 祖先节点全部展开
+        # 匹配节点不展开
+
         orgs = self.get_tree_user_orgs()
 
         matched_nodes = []
         matched_nodes_ancestors = []
         for org in orgs:
             tree: AssetTree = self.get_org_asset_tree(
-                asset_category=asset_category, asset_type=asset_type, org=org
+                asset_category=asset_category, 
+                asset_type=asset_type, 
+                org=org
             )
             # 如果匹配的节点中包含有父子关系的节点，只返回最上一级的父节点
             _matched_nodes = tree.search_nodes(search_node, only_top_level=True)
@@ -197,32 +330,45 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
         # 匹配的节点不展开
         expand_level = 0
         serialized_matched_nodes = self.serialize_nodes(
-            matched_nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
+            nodes=matched_nodes, tree_type=self.render_tree_type, 
+            with_asset_amount=with_asset_amount, 
             expand_level=expand_level
         )
 
         # 匹配节点的祖先节点全部展开
-        expand_level = 10000
+        expand_all = True
         serialized_ancestors = self.serialize_nodes(
-            matched_nodes_ancestors, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
-            expand_level=expand_level
+            nodes=matched_nodes_ancestors, 
+            tree_type=self.render_tree_type, 
+            with_asset_amount=with_asset_amount, 
+            expand_level=expand_level,
+            expand_all=expand_all
         )
         data = serialized_matched_nodes + serialized_ancestors
         return data
 
-    def search_asset_tree_assets(self, search_asset, asset_category, asset_type, with_asset_amount):
+    @timeit
+    def search_asset_tree_assets(self, search_asset, asset_category=None, asset_type=None, 
+                                 with_asset_amount=True):
+        # 搜索资产树资产 #
+        # 搜索资产时，返回包含匹配资产所在的节点及其祖先节点，以及匹配的资产
+        # 所有节点全部展开
+
         orgs = self.get_tree_user_orgs()
 
         # 搜索时，展开所有节点
-        expand_level = 10000
+        expand_all = True
 
         # 资产树搜索范围
         assets_q_object = Q(name__icontains=search_asset) | Q(address__icontains=search_asset)
 
         # 限制每个组织搜索返回的资产数量
-        with_assets_limit_max = 1000
-        with_assets_limit_min = 100
-        with_assets_limit = max(with_assets_limit_min, with_assets_limit_max // max(1, orgs.count()))
+        with_assets_limit_max = self.search_assets_per_org_limit_max
+        with_assets_limit_min = self.search_assets_per_org_limit_min
+        with_assets_limit = max(
+            with_assets_limit_min, 
+            with_assets_limit_max // max(1, orgs.count())
+        )
 
         # 搜索树只返回包含资产的节点
         full_tree = False
@@ -232,8 +378,10 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
         for org in orgs:
             tree: AssetTree = self.get_org_asset_tree(
                 assets_q_object=assets_q_object,
-                asset_category=asset_category, asset_type=asset_type, org=org,
-                with_assets_all=True, with_assets_limit=with_assets_limit, 
+                asset_category=asset_category, 
+                asset_type=asset_type, org=org,
+                with_assets_all=True, 
+                with_assets_limit=with_assets_limit, 
                 full_tree=full_tree
             )
             _nodes = tree.get_nodes()
@@ -241,54 +389,10 @@ class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
             _assets = tree.get_assets()
             assets.extend(_assets)
         serialized_nodes = self.serialize_nodes(
-            nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
-            expand_level=expand_level
+            nodes=nodes, tree_type=self.render_tree_type, 
+            with_asset_amount=with_asset_amount, 
+            expand_all=expand_all
         )
         serialized_assets = self.serialize_assets(assets)
         data = serialized_nodes + serialized_assets
         return data
-
-    # tree 抽象方法 #
-    
-    def get_org_asset_tree(self, **kwargs) -> AssetTree:
-        raise NotImplementedError
-
-    # tree 辅助方法 #
-
-    @lazyproperty
-    def org_is_global(self):
-        return current_org.is_root()
-
-    def get_tree_user_orgs(self):
-        ''' 重要: 获取用户有权限渲染树的组织列表 '''
-        user = self.tree_user
-        if self.org_is_global:
-            # 如果是全局组织，返回用户所在的所有实体组织
-            orgs = user.orgs.all()
-        else:
-            # 如果时实体组织，从用户所在的实体组织中返回该实体组织
-            orgs = user.orgs.filter(id=current_org.id)
-
-        if not orgs:
-            raise APIException('No organization available for rendering the tree')
-        return orgs
-
-    # view 辅助方法 #
-
-    def get_query_value(self, query_key):
-        query_value = self.request.query_params.get(query_key)
-        if not query_value:
-            query_value = self.get_query_value_from_search(query_key)
-        return query_value
-    
-    def get_query_value_from_search(self, query_key):
-        search = self.request.query_params.get('search', '')
-        if not search:
-            return None
-
-        search_list = search.split()
-        for _search in search_list:
-            if f'{query_key}:' not in _search:
-                continue
-            query_value = _search.replace(f'{query_key}:', '').strip()
-            return query_value

apps/assets/api/tree/mixin.py

@@ -17,7 +17,7 @@ class SerializeToTreeNodeMixin:
 
     @timeit
     def serialize_nodes(self, nodes: List[AssetTreeNode], tree_type: RenderTreeType, 
-                        with_asset_amount=False, expand_level=1):
+                        with_asset_amount=False, expand_level=1, expand_all=False):
         if not nodes:
             return []
 
@@ -43,7 +43,7 @@ class SerializeToTreeNodeMixin:
                 'title': _name(node),
                 'pId': node.parent_key,
                 'isParent': is_parent(node),
-                'open': node.level <= expand_level,
+                'open': expand_all or node.level <= expand_level,
                 'meta': {
                     'type': 'node',
                     'data': {

apps/perms/api/user_permission/assets.py

@@ -12,7 +12,7 @@ from orgs.utils import tmp_to_root_org
 from perms import serializers
 from perms.pagination import NodePermedAssetPagination, AllPermedAssetPagination
 from perms.utils import UserPermAssetUtil, PermAssetDetailUtil
-from perms.utils.utils import UserPermUtil
+from perms.utils.utils import UserPermedAssetUtil
 from perms.tree import UserPermAssetTreeNode
 from .mixin import (
     SelfOrPKUserMixin
@@ -80,26 +80,26 @@ class UserAllPermedAssetsApi(BaseUserPermedAssetsApi):
 
         if node_id == UserPermAssetTreeNode.SpecialKey.FAVORITE:
             # TODO: Support asset_category, asset_type
-            return UserPermUtil.get_favorite_assets(user=self.user)
+            return UserPermedAssetUtil.get_favorite_assets(user=self.user)
 
         if node_id == UserPermAssetTreeNode.SpecialKey.UNGROUPED:
-            _util = UserPermUtil(user=self.user)
+            _util = UserPermedAssetUtil(user=self.user)
             return _util.get_ungrouped_assets()
 
         node = get_object_or_none(Node, id=node_id)
         if node:
-            _util= UserPermUtil(user=self.user, org=node.org)
+            _util= UserPermedAssetUtil(user=self.user, org=node.org)
             assets = _util.get_node_all_assets(node)
             return assets
 
-        assets = UserPermUtil.get_all_assets(user=self.user)
+        assets = UserPermedAssetUtil.get_all_assets(user=self.user)
         return assets
 
 
 class UserDirectPermedAssetsApi(BaseUserPermedAssetsApi):
 
     def get_assets(self):
-        _util = UserPermUtil(user=self.user)
+        _util = UserPermedAssetUtil(user=self.user)
         assets = _util.get_ungrouped_assets()
         return assets
 
@@ -107,7 +107,7 @@ class UserDirectPermedAssetsApi(BaseUserPermedAssetsApi):
 class UserFavoriteAssetsApi(BaseUserPermedAssetsApi):
 
     def get_assets(self):
-        assets = UserPermUtil.get_favorite_assets(user=self.user)
+        assets = UserPermedAssetUtil.get_favorite_assets(user=self.user)
         return assets
 
 

apps/perms/api/user_permission/tree.py

@@ -7,7 +7,7 @@ from common.utils import get_logger, timeit
 from assets.api.tree import AbstractAssetTreeAPI
 from assets.tree.asset_tree import AssetTreeNodeAsset
 from perms.tree import UserPermAssetTree, UserPermAssetTreeNode
-from perms.utils.utils import UserPermUtil
+from perms.utils.utils import UserPermedAssetUtil
 
 from .mixin import SelfOrPKUserMixin
 
@@ -24,118 +24,141 @@ class UserPermedAssetTreeAPI(SelfOrPKUserMixin, AbstractAssetTreeAPI):
         return self.user
 
     def get_org_asset_tree(self, **kwargs) -> UserPermAssetTree:
+        # 重写父类方法，返回用户授权的组织资产树
         return self.get_user_org_asset_tree(**kwargs)
 
     def get_user_org_asset_tree(self, **kwargs) -> UserPermAssetTree:
+        # 获取用户授权的组织资产树
         tree = UserPermAssetTree(user=self.user, **kwargs)
         return tree
     
     def _render_asset_tree(self, **kwargs):
+        # 重写父类方法，返回特殊节点
+        # 特殊节点如：收藏夹、未分组节点
         data = super()._render_asset_tree(**kwargs)
         expand_node_key = kwargs.pop('expand_node_key', None)
         if expand_node_key:
-            # 特殊节点不需要展开，资产树中特殊节点下的资产已经随着节点一起返回
+            # 展开其他节点时，不返回特殊节点
+            # 特殊节点不允许异步展开
             return data
-        
-        special_nodes = self.get_special_nodes(with_assets=True, expand_level=0, **kwargs)
+        special_nodes = self.get_special_nodes(**kwargs)
         data = special_nodes + data
         return data
     
-    def render_node_tree(self, asset_category, asset_type, with_asset_amount):
-        data = super().render_node_tree(asset_category, asset_type, with_asset_amount)
-        special_nodes = self.get_special_nodes(
-            with_assets=False, expand_level=0, asset_category=asset_category, asset_type=asset_type, 
-            with_asset_amount=with_asset_amount
-        )
+    def render_node_tree(self, **kwargs):
+        # 重写父类方法，返回特殊节点
+        data = super().render_node_tree(**kwargs)
+        special_nodes = self.get_special_nodes(**kwargs)
         data = special_nodes + data
         return data
     
+    @timeit
     def get_special_nodes(self, search_asset=None, search_node=None, 
-                          asset_category=None, asset_type=None,
-                          with_assets=False, expand_level=0, with_asset_amount=False):
+                          asset_category=None, asset_type=None, with_asset_amount=True):
+        # 获取特殊节点数据
+        # 特殊节点如：收藏夹、未分组节点
+        
+        # 默认不包含资产
+        with_assets = False
+        # 默认不展开特殊节点
+        expand_level = 0
+        if self.render_tree_type.is_asset_tree:
+            with_assets = True
+            if search_asset:
+                # 资产树中，搜索资产时，展开一级
+                expand_level = 1
+
         f_node, f_assets = self.get_favorite_node(
             search_asset=search_asset, search_node=search_node, 
-            asset_category=asset_category, asset_type=asset_type, with_assets=with_assets
+            asset_category=asset_category, asset_type=asset_type, 
+            with_assets=with_assets
         )
         u_node, u_assets = self.get_ungrouped_node_if_need(
             search_asset=search_asset, search_node=search_node, 
-            asset_category=asset_category, asset_type=asset_type, with_assets=with_assets
+            asset_category=asset_category, asset_type=asset_type, 
+            with_assets=with_assets
         )
 
-        nodes = []
-        if f_node:
-            nodes.append(f_node)
-        if u_node:
-            nodes.append(u_node)
-
+        nodes = [n for n in [f_node, u_node] if n]
         if not nodes:
             return []
         
-        if search_asset:
-            expand_level = 1
-        
         serialized_nodes = self.serialize_nodes(
             nodes, tree_type=self.render_tree_type,
             with_asset_amount=with_asset_amount, expand_level=expand_level,
         )
+        
         if with_assets:
-            assets = f_assets + u_assets
+            assets = [*f_assets, *u_assets]
             serialized_assets = self.serialize_assets(assets)
         else:
             serialized_assets = []
-        
+
         data = serialized_nodes + serialized_assets
         return data
 
+    @timeit
     def get_favorite_node(self, search_asset=None, search_node=None, asset_category=None, 
                           asset_type=None, with_assets=False):
-        assets = UserPermUtil.get_favorite_assets(
-            user=self.user, search_asset=search_asset,
-            asset_category=asset_category, asset_type=asset_type
+
+        assets = UserPermedAssetUtil.get_favorite_assets(
+            user=self.tree_user, 
+            search_asset=search_asset,
+            asset_category=asset_category, 
+            asset_type=asset_type
         )
         assets_amount = assets.count()
-        f_node = UserPermAssetTreeNode.favorite(
-            assets_amount=assets_amount, assets_amount_total=assets_amount
+        node = UserPermAssetTreeNode.favorite(
+            assets_amount=assets_amount, 
+            assets_amount_total=assets_amount
         )
-        if not f_node.match(search_node):
+
+        if search_node and not node.match(search_node):
             return None, []
+
+        if not with_assets:
+            return node, []
         
         if assets_amount == 0:
-            return f_node, []
+            return node, []
 
-        if with_assets:
-            assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values))
-            assets = f_node.init_assets(assets_attrs)
-        else:
-            assets = []
-        return f_node, assets
+        assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values))
+        assets = node.init_assets(assets_attrs)
+        return node, assets
+
+    @timeit
+    def get_ungrouped_node_if_need(self, search_asset=None, search_node=None, 
+                                   asset_category=None, asset_type=None, with_assets=False):
 
-    def get_ungrouped_node_if_need(self, search_asset=None, search_node=None, asset_category=None, 
-                                   asset_type=None, with_assets=False):
         if not settings.PERM_SINGLE_ASSET_TO_UNGROUP_NODE:
+            # 未分组节点功能关闭
             return None, []
+
         if self.org_is_global:
             # 全局组织不返回未分组节点
             return None, []
+
         org = self.get_tree_user_orgs().first()
-        util = UserPermUtil(user=self.user, org=org)
+        util = UserPermedAssetUtil(user=self.user, org=org)
         assets = util.get_ungrouped_assets(
             search_asset=search_asset, 
-            asset_category=asset_category, asset_type=asset_type
+            asset_category=asset_category, 
+            asset_type=asset_type
         )
         assets_amount = assets.count()
-        u_node = UserPermAssetTreeNode.ungrouped(
-            assets_amount=assets_amount, assets_amount_total=assets_amount
+        node = UserPermAssetTreeNode.ungrouped(
+            assets_amount=assets_amount, 
+            assets_amount_total=assets_amount
         )
-        if not u_node.match(search_node):
+        if search_node and not node.match(search_node):
             return None, []
         
         if assets_amount == 0:
-            return u_node, []
+            return node, []
 
-        if with_assets:
-            assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values))
-            assets = u_node.init_assets(assets_attrs)
-        else:
-            assets = []
-        return u_node, assets
+        if not with_assets:
+            return node, []
+
+        assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values))
+        assets = node.init_assets(assets_attrs)
+        return node, assets

apps/perms/tree.py

@@ -7,7 +7,7 @@ from common.utils import get_logger
 from users.models import User
 from assets.models import FavoriteAsset, Asset
 from assets.tree.asset_tree import AssetTree, AssetTreeNode, AssetTreeNodeAsset
-from perms.utils.utils import UserPermUtil
+from perms.utils.utils import UserPermedAssetUtil
 
 
 __all__ = ['UserPermAssetTree']
@@ -69,7 +69,7 @@ class UserPermAssetTree(AssetTree):
     def __init__(self, user, **kwargs):
 
         self._user: User = user
-        self._util = UserPermUtil(user, org=kwargs.get('org'))
+        self._util = UserPermedAssetUtil(user, org=kwargs.get('org'))
         kwargs.update({
             # 用户授权树只返回有资产的节点
             'full_tree': False,

apps/perms/utils/utils.py

@@ -11,10 +11,10 @@ from perms.models import AssetPermission
 logger = get_logger(__name__)
 
 
-__all__ = ['UserPermUtil']
+__all__ = ['UserPermedAssetUtil']
 
 
-class UserPermUtil(object):
+class UserPermedAssetUtil(object):
 
     UserGroupThrough = User.groups.through
     PermUserThrough = AssetPermission.users.through
@@ -157,28 +157,50 @@ class UserPermUtil(object):
         return assets
 
     @classmethod
-    def get_favorite_assets(cls, user: User, search_asset=None, asset_category=None, asset_type=None):
+    def get_favorite_assets(cls, user, search_asset=None, asset_category=None, asset_type=None):
         asset_ids = FavoriteAsset.get_user_favorite_asset_ids(user)
-        q = Q(id__in=asset_ids)
-        if search_asset:
-            q &= Q(name__icontains=search_asset) | Q(address__icontains=search_asset)
-        if asset_category:
-            q &= Q(platform__category=asset_category)
-        if asset_type:
-            q &= Q(platform__type=asset_type)
-        assets = Asset.objects.filter(q).valid()
+        assets = cls.filter_assets(
+            asset_ids=asset_ids,
+            search_asset=search_asset, 
+            asset_category=asset_category, 
+            asset_type=asset_type
+        )
         return assets
     
     def get_ungrouped_assets(self, search_asset=None, asset_category=None, asset_type=None):
-        q = Q(id__in=self._user_direct_asset_ids)
+        asset_ids = self._user_direct_asset_ids
+        assets = self.filter_assets(
+            asset_ids=asset_ids, 
+            search_asset=search_asset, 
+            asset_category=asset_category, 
+            asset_type=asset_type
+        )
+        return assets
+    
+    @classmethod
+    def filter_assets(cls, asset_ids, search_asset=None, asset_category=None, asset_type=None):
+        q = cls._make_assets_q_object(
+            asset_ids=asset_ids,
+            search_asset=search_asset, 
+            asset_category=asset_category, 
+            asset_type=asset_type
+        )
+        assets = Asset.objects.filter(q).valid()
+        return assets
+    
+    @staticmethod
+    def _make_assets_q_object(asset_ids=None, search_asset=None, asset_category=None, 
+                              asset_type=None):
+        q = Q()
+        if asset_ids:
+            q &= Q(id__in=asset_ids)
         if search_asset:
             q &= Q(name__icontains=search_asset) | Q(address__icontains=search_asset)
         if asset_category:
             q &= Q(platform__category=asset_category)
         if asset_type:
             q &= Q(platform__type=asset_type)
-        assets = Asset.objects.filter(q).valid()
-        return assets
+        return q
     
     def _uuids_to_string(self, uuids):
         return [ str(u) for u in uuids ]