perf: add AbstractAssetTreeAPI for AssetTreeAPI and UserPermedAssetTreeAPI, but need support UserGroupPermedAssetTreeAPI in the future

apps/assets/api/tree.py

@@ -1,267 +0,0 @@
-# ~*~ coding: utf-8 ~*~
-
-from django.db.models import Q, Count
-from django.utils.translation import gettext_lazy as _
-from rest_framework.generics import get_object_or_404
-from rest_framework.response import Response
-
-from assets.locks import NodeAddChildrenLock
-from common.exceptions import JMSException
-from common.tree import TreeNodeSerializer
-from common.utils import get_logger
-from orgs.mixins import generics
-from orgs.utils import current_org
-from orgs.models import Organization
-from .mixin import SerializeToTreeNodeMixin
-from .. import serializers
-from ..const import AllTypes
-from ..models import Node, Platform, Asset
-from assets.tree.asset_tree import AssetTree
-
-
-logger = get_logger(__file__)
-
-__all__ = [
-    'NodeChildrenApi',
-    'NodeChildrenAsTreeApi',
-    'CategoryTreeApi',
-]
-
-
-class NodeChildrenApi(generics.ListCreateAPIView):
-    ''' 节点的增删改查 '''
-    serializer_class = serializers.NodeSerializer
-    search_fields = ('value',)
-
-    instance = None
-    is_initial = False
-    perm_model = Node
-
-    def initial(self, request, *args, **kwargs):
-        super().initial(request, *args, **kwargs)
-        self.initial_org_root_node_if_need()
-        self.instance = self.get_object()
-    
-    def initial_org_root_node_if_need(self):
-        if current_org.is_root():
-            return
-        Node.org_root()
-
-    def perform_create(self, serializer):
-        data = serializer.validated_data
-        _id = data.get("id")
-        value = data.get("value")
-        if value:
-            children = self.instance.get_children()
-            if children.filter(value=value).exists():
-                raise JMSException(_('The same level node name cannot be the same'))
-        else:
-            value = self.instance.get_next_child_preset_name()
-        with NodeAddChildrenLock(self.instance):
-            node = self.instance.create_child(value=value, _id=_id)
-            # 避免查询 full value
-            node._full_value = node.value
-            serializer.instance = node
-
-    def get_object(self):
-        pk = self.kwargs.get('pk') or self.request.query_params.get('id')
-        key = self.request.query_params.get("key")
-
-        if not pk and not key:
-            self.is_initial = True
-            if current_org.is_root():
-                node = None
-            else:
-                node = Node.org_root()
-            return node
-
-        if pk:
-            node = get_object_or_404(Node, pk=pk)
-        else:
-            node = get_object_or_404(Node, key=key)
-        return node
-
-
-class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
-    ''' 节点子节点作为树返回，
-    [
-      {
-        "id": "",
-        "name": "",
-        "pId": "",
-        "meta": ""
-      }
-    ]
-    '''
-
-    model = Node
-
-    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_asset_tree()
-        
-        if search:
-            # 初始化资产搜索树 - 包含资产
-            return self.search_asset_tree_with_assets(search)
-        
-        if key:
-            # 展开资产树节点 - 包含资产
-            return self.expand_asset_tree_node_with_assets(key)
-
-        # 初始化资产树 - 包含资产
-        return self.init_asset_tree_with_assets()
-    
-    def init_asset_tree(self):
-        ''' 初始化资产树 - 不包含资产
-        返回所有节点，前端本地展开和搜索
-        全局组织: 不展开节点
-        实体组织: 展开第1级节点
-        '''
-        if current_org.is_root():
-            orgs = Organization.objects.all()
-            expand_level = 0
-        else:
-            orgs = [current_org]
-            expand_level = 1
-        nodes = []
-        for org in orgs:
-            tree = AssetTree(org=org)
-            _nodes = tree.get_nodes()
-            nodes.extend(_nodes)
-        nodes = self.serialize_nodes(
-            nodes, with_asset_amount=True, expand_level=expand_level
-        )
-        return Response(data=nodes)
-    
-    def init_asset_tree_with_assets(self):
-        ''' 初始化资产树 - 包含资产
-        全局组织: 返回第1级节点，不返回资产，不展开
-        实体组织: 返回第1级节点和第2级节点，返回第1级节点的资产，展开第1级节点
-        '''
-        if current_org.is_root():
-            orgs = Organization.objects.all()
-            node_levels = [1]
-            with_assets_node_levels = None
-            expand_level = 0
-        else:
-            orgs = [current_org]
-            node_levels = [1, 2]
-            with_assets_node_levels = [1]
-            expand_level = 1
-
-        nodes = []
-        assets = []
-        for org in orgs:
-            tree = AssetTree(
-                org=org, with_assets_node_levels=with_assets_node_levels
-            )
-            _nodes = tree.get_nodes(levels=node_levels)
-            nodes.extend(_nodes)
-            _assets = tree.get_assets()
-            assets.extend(_assets)
-        
-        nodes = self.serialize_nodes(
-            nodes, with_asset_amount=True, expand_level=expand_level, 
-            with_assets=True
-        )
-        assets = self.serialize_assets(assets)
-        data = [*nodes, *assets]
-        return Response(data=data)
-    
-    def expand_asset_tree_node_with_assets(self, key):
-        ''' 展开资产树节点 - 包含资产
-        全局组织: 返回展开节点的直接孩子节点，返回展开节点的资产，不展开节点
-        实体组织: 同上
-        '''
-        node = get_object_or_404(Node, key=key)
-        org = node.org
-        with_assets_node_id = node.id
-        expand_level = 0
-
-        tree = AssetTree(with_assets_node_id=with_assets_node_id, org=org)
-        tree_node = tree.get_node(key=node.key)
-        if not tree_node:
-            return Response(data=[])
-        _nodes = tree_node.children
-        nodes = self.serialize_nodes(
-            _nodes, with_asset_amount=True, expand_level=expand_level, 
-            with_assets=True
-        )
-        _assets = tree.get_assets()
-        assets = self.serialize_assets(_assets)
-        data = [*nodes, *assets]
-        return Response(data=data)
-    
-    def search_asset_tree_with_assets(self, search):
-        ''' 初始化资产搜索树 - 包含资产
-        不反回完整树，资产数量为0的节点不返回
-        全局组织: 返回所有节点，返回所有资产，展开所有节点，限制资产总数量 1000（n 个组织，每个组织分配1000/n个资产）
-        实体组织：同上，限制资产总数量 1000
-        '''
-        # 展开所有节点
-        expand_level = 10000
-        with_assets_all = True
-        with_assets_limit = 1000
-        full_tree = False
-        assets_q_object = Q(name__icontains=search) | Q(address__icontains=search)
-        if current_org.is_root():
-            orgs = list(Organization.objects.all())
-            with_assets_limit = max(100, with_assets_limit // max(1, orgs.count()))
-        else:
-            orgs = [current_org]
-
-        nodes = []
-        assets = []
-        for org in orgs:
-            tree = AssetTree(
-                assets_q_object=assets_q_object, org=org, 
-                with_assets_all=with_assets_all, 
-                with_assets_limit=with_assets_limit,
-                full_tree=full_tree
-            )
-            _nodes = tree.get_nodes()
-            nodes.extend(_nodes)
-            _assets = tree.get_assets()
-            assets.extend(_assets)
-
-        nodes = self.serialize_nodes(
-            nodes, with_asset_amount=True, expand_level=expand_level, 
-            with_assets=True
-        )
-        assets = self.serialize_assets(assets)
-        data = [*nodes, *assets]
-        return Response(data=data)
-
-class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
-    serializer_class = TreeNodeSerializer
-    rbac_perms = {
-        'GET': 'assets.view_asset',
-        'list': 'assets.view_asset',
-    }
-    queryset = Node.objects.none()
-
-    def get_assets(self):
-        key = self.request.query_params.get('key')
-        platform = Platform.objects.filter(id=key).first()
-        if not platform:
-            return []
-        assets = Asset.objects.filter(platform=platform).prefetch_related('platform')
-        return self.serialize_assets(assets, key)
-
-    def list(self, request, *args, **kwargs):
-        include_asset = self.request.query_params.get('assets', '0') == '1'
-        # 资源数量统计可选项 (asset, account)
-        count_resource = self.request.query_params.get('count_resource', 'asset')
-
-        if not self.request.query_params.get('key'):
-            nodes = AllTypes.to_tree_nodes(include_asset, count_resource=count_resource)
-        elif include_asset:
-            nodes = self.get_assets()
-        else:
-            nodes = []
-        return Response(data=nodes)

apps/assets/api/tree/__init__.py

@@ -0,0 +1,3 @@
+from .tree import *
+from .base import *
+from .const import *

apps/assets/api/tree/base.py

@@ -0,0 +1,294 @@
+
+from django.db.models import Q
+from rest_framework import generics
+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.exceptions import APIException
+from orgs.utils import current_org
+from orgs.models import Organization
+from rbac.permissions import RBACPermission
+from assets.tree.asset_tree import AssetTree
+from assets.models import Node
+from .mixin import SerializeToTreeNodeMixin
+from .const import RenderTreeType, RenderTreeTypeChoices
+
+class AbstractAssetTreeAPI(SerializeToTreeNodeMixin, generics.ListAPIView):
+
+    permission_classes = (RBACPermission,)
+
+    perm_model = Node
+
+    render_tree_type: RenderTreeType
+    tree_user: User
+
+
+    def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+        self.render_tree_type = self.initial_render_tree_type()
+        self.tree_user = self.get_tree_user()
+
+    def initial_render_tree_type(self):
+        tree_type = self.get_query_value('tree_type')
+        if not tree_type:
+            with_assets = self.request.query_params.get('assets', '0') == '1'
+            tree_type = RenderTreeTypeChoices.asset if with_assets else RenderTreeTypeChoices.node
+        return RenderTreeType(tree_type)
+
+    def get_tree_user(self) -> User:
+        raise NotImplementedError
+
+    def list(self, request, *args, **kwargs):
+        asset_category = self.get_query_value('asset_category')
+        asset_type = self.get_query_value('asset_type')
+        with_asset_amount = True
+
+        if self.render_tree_type.is_node_tree:
+            data = self.render_node_tree(asset_category, asset_type, with_asset_amount)
+        elif self.render_tree_type.is_asset_tree:
+            data = self.render_asset_tree(asset_category, asset_type, with_asset_amount)
+        else:
+            raise APIException(f'Invalid tree type: {self.query_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)
+        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')
+        data = self._render_asset_tree(
+            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
+
+    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):
+
+        if not search_asset and self.render_tree_type.is_asset_tree:
+            search = self.get_query_value('search') or ''
+            if ':' 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
+            )
+        elif search_node:
+            data = self.search_asset_tree_nodes(
+                search_node, asset_category, asset_type, with_asset_amount
+            )
+        elif search_asset:
+            data = self.search_asset_tree_assets(
+                search_asset, asset_category, asset_type, with_asset_amount
+            )
+        else:
+            data = self.init_asset_tree(
+                asset_category, asset_type, with_asset_amount
+            )
+        return data
+    
+    def init_node_tree(self, asset_category, asset_type, with_asset_amount):
+        orgs = self.get_tree_user_orgs()
+
+        if self.org_is_global:
+            expand_level = 0
+        else:
+            expand_level = 1
+
+        nodes = []
+        for org in orgs:
+            tree = self.get_org_asset_tree(
+                asset_category=asset_category, asset_type=asset_type, org=org
+            )
+            _nodes = tree.get_nodes()
+            nodes.extend(_nodes)
+
+        data = self.serialize_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):
+        orgs = self.get_tree_user_orgs()
+        if self.org_is_global:
+            node_levels = [1]
+            with_assets_node_levels = None
+            expand_level = 0
+        else:
+            node_levels = [1, 2]
+            with_assets_node_levels = [1]
+            expand_level = 1
+        nodes = []
+        assets = []
+        for org in orgs:
+            tree: AssetTree = self.get_org_asset_tree(
+                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)
+            nodes.extend(_nodes)
+            _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
+        )
+        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):
+        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,
+            with_assets_node_id=with_assets_node_id
+        )
+        node_children = tree.get_node_children(node.key)
+        node_assets = tree.get_assets()
+
+        # (展开节点)的孩子节点不展开
+        expand_level = 0
+        serialized_nodes = self.serialize_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):
+        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
+            )
+            # 如果匹配的节点中包含有父子关系的节点，只返回最上一级的父节点
+            _matched_nodes = tree.search_nodes(search_node, only_top_level=True)
+            if not _matched_nodes:
+                continue
+            matched_nodes.extend(_matched_nodes)
+            _ancestors = tree.get_nodes_ancestors(_matched_nodes)
+            matched_nodes_ancestors.extend(_ancestors)
+        
+        if not matched_nodes:
+            return []
+        
+        # 匹配的节点不展开
+        expand_level = 0
+        serialized_matched_nodes = self.serialize_nodes(
+            matched_nodes, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
+            expand_level=expand_level
+        )
+
+        # 匹配节点的祖先节点全部展开
+        expand_level = 10000
+        serialized_ancestors = self.serialize_nodes(
+            matched_nodes_ancestors, tree_type=self.render_tree_type, with_asset_amount=with_asset_amount, 
+            expand_level=expand_level
+        )
+        data = serialized_matched_nodes + serialized_ancestors
+        return data
+
+    def search_asset_tree_assets(self, search_asset, asset_category, asset_type, with_asset_amount):
+        orgs = self.get_tree_user_orgs()
+
+        # 搜索时，展开所有节点
+        expand_level = 10000
+
+        # 资产树搜索范围
+        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()))
+
+        # 搜索树只返回包含资产的节点
+        full_tree = False
+
+        nodes = []
+        assets = []
+        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, 
+                full_tree=full_tree
+            )
+            _nodes = tree.get_nodes()
+            nodes.extend(_nodes)
+            _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
+        )
+        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/const.py

@@ -0,0 +1,28 @@
+from django.db.models import TextChoices
+
+
+__all__ = ['RenderTreeType', 'RenderTreeTypeChoices']
+
+
+class RenderTreeTypeChoices(TextChoices):
+    node = 'node', 'Node'
+    asset = 'asset', 'Asset'
+
+
+class RenderTreeType:
+
+    def __init__(self, _type):
+        if _type not in RenderTreeTypeChoices.values:
+            raise ValueError(f'Invalid tree type: {_type}')
+        self._type: RenderTreeTypeChoices = _type
+
+    @property
+    def is_asset_tree(self):
+        return self._type == RenderTreeTypeChoices.asset
+
+    @property
+    def is_node_tree(self):
+        return self._type == RenderTreeTypeChoices.node
+    
+    def __str__(self):
+        return self._type.value

apps/assets/api/tree/mixin.py

@@ -0,0 +1,125 @@
+from typing import List
+
+from rest_framework.request import Request
+
+from assets.models import Platform, Protocol, MyAsset
+from common.utils import lazyproperty, timeit
+from assets.tree.asset_tree import AssetTreeNode, AssetTreeNodeAsset
+
+from .const import RenderTreeType
+
+
+__all__ = ['SerializeToTreeNodeMixin']
+
+
+class SerializeToTreeNodeMixin:
+    request: Request
+
+    @timeit
+    def serialize_nodes(self, nodes: List[AssetTreeNode], tree_type: RenderTreeType, 
+                        with_asset_amount=False, expand_level=1):
+        if not nodes:
+            return []
+
+        def _name(node: AssetTreeNode):
+            v = node.value
+            if not with_asset_amount:
+                return v
+            v = f'{v} ({node.assets_amount_total})'
+            return v
+        
+        def is_parent(node: AssetTreeNode):
+            if tree_type.is_asset_tree:
+                return node.assets_amount > 0 or not node.is_leaf
+            elif tree_type.is_node_tree:
+                return not node.is_leaf
+            else:
+                return True
+
+        data = [
+            {
+                'id': node.key,
+                'name': _name(node),
+                'title': _name(node),
+                'pId': node.parent_key,
+                'isParent': is_parent(node),
+                'open': node.level <= expand_level,
+                'meta': {
+                    'type': 'node',
+                    'data': {
+                        "id": node.id,
+                        "key": node.key,
+                        "value": node.value,
+                        "assets_amount": node.assets_amount,
+                        "assets_amount_total": node.assets_amount_total,
+                        "children_count_total": node.children_count_total,
+                    },
+                }
+            }
+            for node in nodes
+        ]
+        return data
+
+    @lazyproperty
+    def support_types(self):
+        from assets.const import AllTypes
+        return AllTypes.get_types_values(exclude_custom=True)
+
+    def get_icon(self, asset):
+        if asset.category == 'device':
+            return 'switch'
+        if asset.type in self.support_types:
+            return asset.type
+        else:
+            return 'file'
+
+    @timeit
+    def serialize_assets(self, assets, node_key=None, get_pid=None):
+        if not assets:
+            return []
+
+        if not get_pid and not node_key:
+            get_pid = lambda asset, platform: getattr(asset, 'parent_key', '')
+
+        sftp_asset_ids = Protocol.objects.filter(name='sftp') \
+            .values_list('asset_id', flat=True)
+        sftp_asset_ids = set(sftp_asset_ids)
+        platform_map = {p.id: p for p in Platform.objects.all()}
+
+        data = []
+        root_assets_count = 0
+        MyAsset.set_asset_custom_value(assets, self.request.user)
+        for asset in assets:
+            asset: AssetTreeNodeAsset
+            platform = platform_map.get(asset.platform_id)
+            if not platform:
+                continue
+            pid = node_key or get_pid(asset, platform)
+            if not pid:
+                continue
+            # 根节点最多显示 1000 个资产
+            if pid.isdigit():
+                if root_assets_count > 1000:
+                    continue
+                root_assets_count += 1
+            data.append({
+                'id': str(asset.id),
+                'name': asset.name,
+                'title': f'{asset.address}\n{asset.comment}'.strip(),
+                'pId': pid,
+                'isParent': False,
+                'open': False,
+                'iconSkin': self.get_icon(platform),
+                'chkDisabled': not asset.is_active,
+                'meta': {
+                    'type': 'asset',
+                    'data': {
+                        'platform_type': platform.type,
+                        'org_name': asset.org_name,
+                        'sftp': asset.id in sftp_asset_ids,
+                        'name': asset.name,
+                        'address': asset.address
+                    },
+                }
+            })
+        return data

apps/assets/api/tree/tree.py

@@ -0,0 +1,135 @@
+# ~*~ coding: utf-8 ~*~
+
+from django.utils.translation import gettext_lazy as _
+from rest_framework.generics import get_object_or_404
+from rest_framework.response import Response
+
+from assets.locks import NodeAddChildrenLock
+from common.exceptions import JMSException
+from common.tree import TreeNodeSerializer
+from common.utils import get_logger
+from orgs.mixins import generics
+from orgs.utils import current_org, tmp_to_org
+from ..mixin import SerializeToTreeNodeMixin
+from ... import serializers
+from ...const import AllTypes
+from ...models import Node, Platform, Asset
+from assets.tree.asset_tree import AssetTree
+from .base import AbstractAssetTreeAPI
+
+
+logger = get_logger(__file__)
+
+__all__ = [
+    'NodeChildrenApi',
+    'AssetTreeAPI',
+    'CategoryTreeApi',
+]
+
+
+class NodeChildrenApi(generics.ListCreateAPIView):
+    ''' 节点的增删改查 '''
+    serializer_class = serializers.NodeSerializer
+    search_fields = ('value',)
+
+    instance = None
+    is_initial = False
+    perm_model = Node
+
+    def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+        self.initial_org_root_node_if_need()
+        self.instance = self.get_object()
+    
+    def initial_org_root_node_if_need(self):
+        if current_org.is_root():
+            return
+        Node.org_root()
+
+    def perform_create(self, serializer):
+        data = serializer.validated_data
+        _id = data.get("id")
+        value = data.get("value")
+        if value:
+            children = self.instance.get_children()
+            if children.filter(value=value).exists():
+                raise JMSException(_('The same level node name cannot be the same'))
+        else:
+            value = self.instance.get_next_child_preset_name()
+        with NodeAddChildrenLock(self.instance):
+            node = self.instance.create_child(value=value, _id=_id)
+            # 避免查询 full value
+            node._full_value = node.value
+            serializer.instance = node
+
+    def get_object(self):
+        pk = self.kwargs.get('pk') or self.request.query_params.get('id')
+        key = self.request.query_params.get("key")
+
+        if not pk and not key:
+            self.is_initial = True
+            if current_org.is_root():
+                node = None
+            else:
+                node = Node.org_root()
+            return node
+
+        if pk:
+            node = get_object_or_404(Node, pk=pk)
+        else:
+            node = get_object_or_404(Node, key=key)
+        return node
+
+
+class AssetTreeAPI(AbstractAssetTreeAPI):
+
+    def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+        self.initial_org_root_node_if_need()
+
+    def initial_org_root_node_if_need(self):
+        if current_org.is_root():
+            orgs = self.request.user.orgs.all()
+        else:
+            orgs = [current_org]
+
+        for org in orgs:
+            with tmp_to_org(org):
+                Node.org_root()
+
+    def get_tree_user(self):
+        return self.request.user
+
+    def get_org_asset_tree(self, **kwargs) -> AssetTree:
+        tree = AssetTree(**kwargs)
+        return tree
+
+
+class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
+    serializer_class = TreeNodeSerializer
+    rbac_perms = {
+        'GET': 'assets.view_asset',
+        'list': 'assets.view_asset',
+    }
+    queryset = Node.objects.none()
+
+    def get_assets(self):
+        key = self.request.query_params.get('key')
+        platform = Platform.objects.filter(id=key).first()
+        if not platform:
+            return []
+        assets = Asset.objects.filter(platform=platform).prefetch_related('platform')
+        return self.serialize_assets(assets, key)
+
+    def list(self, request, *args, **kwargs):
+        include_asset = self.request.query_params.get('assets', '0') == '1'
+        # 资源数量统计可选项 (asset, account)
+        count_resource = self.request.query_params.get('count_resource', 'asset')
+
+        if not self.request.query_params.get('key'):
+            nodes = AllTypes.to_tree_nodes(include_asset, count_resource=count_resource)
+        elif include_asset:
+            nodes = self.get_assets()
+        else:
+            nodes = []
+        return Response(data=nodes)

apps/assets/tree/asset_tree.py

@@ -19,7 +19,8 @@ __all__ = ['AssetTree', 'AssetTreeNode']
 class AssetTreeNodeAsset:
 
     model_values = [
-        'id', 'name', 'node_id', 'platform_id', 'address', 'is_active', 'comment', 'org_id'
+        'id', 'name', 'node_id', 'platform_id', 'address', 
+        'is_active', 'comment', 'org_id'
     ]
 
     def __init__(self, parent_key, **kwargs):
@@ -39,18 +40,19 @@ class AssetTreeNodeAsset:
 
 class AssetTreeNode(TreeNode):
 
-    def __init__(self, _id, key, value, assets_amount=0):
+    def __init__(self, _id, key, value, assets_amount=0, assets_amount_total=0):
         super().__init__(_id, key, value)
         self.assets_amount = assets_amount
-        self.assets_amount_total = 0
+        self.assets_amount_total = assets_amount_total
         self.assets: list[AssetTreeNodeAsset] = []
     
-    def init_assets(self, assets):
-        if not assets:
+    def init_assets(self, assets_attrs):
+        if not assets_attrs:
             return
-        for asset in assets:
-            asset['parent_key'] = self.key
-            self.assets.append(AssetTreeNodeAsset(**asset))
+        for asset_attrs in assets_attrs:
+            asset_attrs['parent_key'] = self.key
+            asset = AssetTreeNodeAsset(**asset_attrs)
+            self.assets.append(asset)
         return self.assets
     
     def get_assets(self):

apps/assets/tree/tree.py

@@ -17,6 +17,13 @@ class TreeNode(object):
         self.parent = None
         self.children_count_total = 0
     
+    def match(self, keyword):
+        if not keyword:
+            return True
+        keyword = str(keyword).strip().lower()
+        node_value = str(self.value).strip().lower()
+        return keyword in node_value
+    
     @lazyproperty
     def parent_key(self):
         if self.is_root:
@@ -176,12 +183,14 @@ class Tree(object):
         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 node.match(keyword):
+                continue
+            nodes[node.key] = node
+
         if not only_top_level:
             return list(nodes.values())
 
+        # 如果匹配的节点中包含有父子关系的节点，只返回最上一级的父节点
         # TODO: 优化性能
         node_keys = list(nodes.keys())
         children_keys = []

apps/assets/urls/api_urls.py

@@ -41,7 +41,7 @@ urlpatterns = [
          api.AssetPermUserGroupPermissionsListApi.as_view(), name='asset-perm-user-group-permission-list'),
 
     path('nodes/category/tree/', api.CategoryTreeApi.as_view(), name='asset-category-tree'),
-    path('nodes/children/tree/', api.NodeChildrenAsTreeApi.as_view(), name='node-children-tree'),
+    path('nodes/children/tree/', api.AssetTreeAPI.as_view(), name='node-children-tree'),
     path('nodes/<uuid:pk>/children/', api.NodeChildrenApi.as_view(), name='node-children'),
     path('nodes/children/', api.NodeChildrenApi.as_view(), name='node-children-2'),
     path('nodes/<uuid:pk>/children/add/', api.NodeAddChildrenApi.as_view(), name='node-add-children'),

apps/perms/api/user_permission/assets.py

@@ -70,6 +70,7 @@ class BaseUserPermedAssetsApi(SelfOrPKUserMixin, ExtraFilterFieldsMixin, ListAPI
 
 
 class UserAllPermedAssetsApi(BaseUserPermedAssetsApi):
+    # TODO: support filter by asset_category and asset_type
 
     def get_assets(self):
         if self.user.is_superuser and self.request.query_params.get('id'):
@@ -78,6 +79,7 @@ class UserAllPermedAssetsApi(BaseUserPermedAssetsApi):
         node_id = self.request.query_params.get('node_id')
 
         if node_id == UserPermAssetTreeNode.SpecialKey.FAVORITE:
+            # TODO: Support asset_category, asset_type
             return UserPermUtil.get_favorite_assets(user=self.user)
 
         if node_id == UserPermAssetTreeNode.SpecialKey.UNGROUPED:

apps/perms/api/user_permission/tree.py

@@ -2,14 +2,9 @@ from django.db.models import Q
 from django.utils.translation import gettext_lazy as _
 from django.conf import settings
 
-from rest_framework.generics import ListAPIView
-from rest_framework.response import Response
-from rest_framework.generics import get_object_or_404
 
 from common.utils import get_logger, timeit
-from orgs.utils import current_org
-from assets.api import SerializeToTreeNodeMixin
-from assets.models import Node, FavoriteAsset, Asset
+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
@@ -19,361 +14,128 @@ from .mixin import SelfOrPKUserMixin
 logger = get_logger(__name__)
 
 __all__ = [
-    'UserPermNodeChildrenAsTreeApi'
+    'UserPermedAssetTreeAPI'
 ]
 
 
-class RenderTreeType:
-    NODE = 'node'
-    ASSET = 'asset'
+class UserPermedAssetTreeAPI(SelfOrPKUserMixin, AbstractAssetTreeAPI):
 
-    def __init__(self, tp):
-        if tp not in [self.NODE, self.ASSET]:
-            raise ValueError(f'Invalid tree type: {tp}')
-        self._type = tp
+    def get_tree_user(self):
+        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 is_asset_tree(self):
-        return self._type == self.ASSET
-
-    def is_node_tree(self):
-        return self._type == self.NODE
-    
-    def __str__(self):
-        return self._type
-
-
-class UserPermNodeChildrenAsTreeApi(SelfOrPKUserMixin, SerializeToTreeNodeMixin, ListAPIView):
-
-    def get_query_value(self, query_key):
-        query_value = self.request.query_params.get(query_key)
-        if query_value:
-            return query_value
-
-        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()
-            break
-        return query_value
-    
-    @property
-    def render_tree_type(self):
-        tree_type = self.get_query_value('tree_type')
-
-        if tree_type:
-            return RenderTreeType(tree_type)
-
-        if self.request.query_params.get('assets', '0') == '1':
-            tree_type = RenderTreeType.ASSET
-        else:
-            tree_type = RenderTreeType.NODE
-
-        return RenderTreeType(tree_type)
-
-    @timeit
-    def list(self, request, *args, **kwargs):
-        # node / asset
-        expand_node_key = self.get_query_value('key')
-        search_node = self.get_query_value('search_node')
-        search_asset = self.get_query_value('search_asset')
-        search = self.get_query_value('search')
-        asset_category = self.get_query_value('asset_category')
-        asset_type = self.get_query_value('asset_type')
-
-        if self.render_tree_type.is_asset_tree():
-            if expand_node_key:
-                return self.expand_user_perm_asset_tree_node(expand_node_key, asset_category, asset_type)
-            elif search_node:
-                # search nodes
-                return self.search_user_perm_node_tree(search_node, asset_category, asset_type)
-            elif search_asset or search:
-                search_asset = search_asset or search
-                # search assets
-                return self.search_user_perm_asset_tree(search_asset, asset_category, asset_type)
-            else:
-                return self.init_user_perm_asset_tree(asset_category, asset_type)
-        else:  # if self.render_tree_type.is_node_tree():
-            if expand_node_key:
-                raise NotImplementedError(_('Expanding node in node tree is not supported yet'))
-            if search_node:
-                raise NotImplementedError(_('Searching node in node tree is not supported yet'))
-            if search_asset:
-                raise NotImplementedError(_('Searching asset in node tree is not supported yet'))
-            else:
-                return self.init_user_perm_node_tree(asset_category, asset_type)
+    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
         
-    
-    def init_user_perm_node_tree(self, asset_category=None, asset_type=None):
-        ''' 初始化用户权限树 - 不包含资产
-        前端搜索
-
-        全局组织: 返回所有节点，不返回资产，不展开节点
-        实体组织：返回所有节点，不返回资产，展开第1级节点
-
-        返回收藏节点
-        返回未分组节点 (如果需要)
-        '''
-        if current_org.is_root():
-            orgs = self.user.orgs.all()
-            expand_level = 0
-        else:
-            orgs = self.user.orgs.filter(id=current_org.id)
-            expand_level = 1
-        
-        if not orgs.exists():
-            return Response(data=[])
-        
-        nodes = []
-        for org in orgs:
-            tree = UserPermAssetTree(
-                user=self.user, asset_category=asset_category, 
-                asset_type=asset_type, org=org
-            )
-            _nodes = tree.get_nodes()
-            nodes.extend(_nodes)
-        nodes = self.serialize_nodes(
-            nodes, with_asset_amount=True, expand_level=expand_level,
-            tree_type=str(self.render_tree_type)
-        )
-        data = nodes
-        data = self.add_favorites_and_ungrouped_node(data, with_assets=False)
-        return Response(data=data)
-    
-    def search_user_perm_node_tree(self, search, asset_category=None, asset_type=None):
-        ''' 搜索用户授权树 - 不包含资产
-        全局组织: 返回所有匹配节点以及祖先节点，不返回匹配节点的子孙节点，不返回资产，展开所有祖先节点，不展开匹配节点
-        实体组织: 同上
-        '''
-        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 = UserPermAssetTree(
-                user=self.user, asset_category=asset_category, 
-                asset_type=asset_type, 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, 
-            tree_type=str(self.render_tree_type)
-        )
-        # 展开所有祖先节点
-        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_asset_tree(self, asset_category=None, asset_type=None):
-        ''' 初始化用户权限资产树 - 包含资产
-        全局组织: 返回第1级节点，不返回资产，不展开节点
-        实体组织：返回第1级和第2级节点，返回第1级节点的资产，展开第1级节点
-
-        返回收藏节点和资产
-        返回未分组节点和资产 (如果需要)
-        '''
-
-        if current_org.is_root():
-            orgs = self.user.orgs.all()
-            nodes_level = [1]
-            with_assets_node_levels = None
-            expand_level = 0
-        else:
-            orgs = self.user.orgs.filter(id=current_org.id)
-            nodes_level = [1, 2]
-            with_assets_node_levels = [1]
-            expand_level = 1
-        
-        if not orgs.exists():
-            return Response(data=[])
-        
-        nodes = []
-        assets = []
-
-        for org in orgs:
-            tree = UserPermAssetTree(
-                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)
-            _assets = tree.get_assets()
-            assets.extend(_assets)
-        
-        nodes = self.serialize_nodes(
-            nodes, with_asset_amount=True, expand_level=expand_level, 
-            tree_type=str(self.render_tree_type)
-        )
-        assets = self.serialize_assets(assets)
-        data = [*nodes, *assets]
-        data = self.add_favorites_and_ungrouped_node(data, with_assets=True)
-        return Response(data=data)
-    
-    def expand_user_perm_asset_tree_node(self, node_key, asset_category=None, asset_type=None):
-        ''' 展开用户权限资产树节点 - 包含资产
-        全局组织: 返回展开节点的直接孩子节点，返回展开节点的资产，不展开其他节点
-        实体组织: 同上
-        '''
-        expand_level = 0
-        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 = UserPermAssetTree(
-            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:
-            return Response(data=[])
-        
-        _nodes = tree_node.children
-        nodes = self.serialize_nodes(
-            _nodes, with_asset_amount=True, expand_level=expand_level,
-            tree_type=str(self.render_tree_type)
-        )
-        _assets = tree.get_assets()
-        assets = self.serialize_assets(_assets)
-        data = [*nodes, *assets]
-        return Response(data=data)
-    
-    def search_user_perm_asset_tree(self, search, asset_category=None, asset_type=None):
-        ''' 初始化用户权限资产搜索树 - 包含资产 
-        全局组织: 返回所有节点，返回所有资产，展开所有节点，搜索资产 (最大 1000， n 个组织，每个组织分配1000/n个资产)
-        实体组织: 同上，最大资产数 1000
-        '''
-        expand_level = 10000
-        with_assets_all = True
-        with_assets_limit = 1000
-        if current_org.is_root():
-            orgs = self.user.orgs.all()
-            with_assets_limit = max(100, with_assets_limit // max(1, orgs.count()))
-        else:
-            orgs = self.user.orgs.filter(id=current_org.id)
-
-        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 = UserPermAssetTree(
-                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
-            )
-            _nodes = tree.get_nodes()
-            nodes.extend(_nodes)
-            _assets = tree.get_assets()
-            assets.extend(_assets)
-
-        nodes = self.serialize_nodes(
-            nodes, with_asset_amount=True, expand_level=expand_level,
-            tree_type=str(self.render_tree_type)
-        )
-        assets = self.serialize_assets(assets)
-        data = [*nodes, *assets]
-        return Response(data=data)
-
-    def add_favorites_and_ungrouped_node(self, data: list, with_assets=False):
-        # 未分组节点和资产
-        u_node, u_assets = self.get_ungrouped_node_if_need()
-        if u_node:
-            data.insert(0, u_node)
-            data.extend(u_assets)
-
-        # 收藏节点和资产
-        f_node, f_assets = self.get_favorite_node()
-        if f_node:
-            data.insert(0, f_node)
-            data.extend(f_assets)
+        special_nodes = self.get_special_nodes(with_assets=True, expand_level=0, **kwargs)
+        data = special_nodes + data
         return data
     
-    def get_favorite_node(self, with_assets=False):
-        assets = UserPermUtil.get_favorite_assets(self.user)
-        assets_amount = assets.count()
-        node = UserPermAssetTreeNode(
-            _id=UserPermAssetTreeNode.SpecialKey.FAVORITE.value,
-            key=UserPermAssetTreeNode.SpecialKey.FAVORITE.value,
-            value=UserPermAssetTreeNode.SpecialKey.FAVORITE.label,
-            assets_amount=assets_amount
+    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
         )
-        node.assets_amount_total = assets_amount
-        nodes = self.serialize_nodes(
-            [node], with_asset_amount=True, expand_level=0,
-            tree_type=str(self.render_tree_type)
+        data = special_nodes + data
+        return data
+    
+    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):
+        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
+        )
+        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
         )
-        serialized_node = nodes[0] if nodes else None
 
-        if not with_assets:
-            return serialized_node, []
+        nodes = []
+        if f_node:
+            nodes.append(f_node)
+        if u_node:
+            nodes.append(u_node)
 
-        if assets_amount == 0:
-            return serialized_node, []
+        if not nodes:
+            return []
         
-        assets = list(assets.values(*AssetTreeNodeAsset.model_values))
-        assets = node.init_assets(assets)
-        serialized_assets = self.serialize_assets(assets)
-        return serialized_node, serialized_assets
+        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
+            serialized_assets = self.serialize_assets(assets)
+        else:
+            serialized_assets = []
+        
+        data = serialized_nodes + serialized_assets
+        return data
 
-    def get_ungrouped_node_if_need(self, with_assets=False):
+    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_amount = assets.count()
+        f_node = UserPermAssetTreeNode.favorite(
+            assets_amount=assets_amount, assets_amount_total=assets_amount
+        )
+        if not f_node.match(search_node):
+            return None, []
+        
+        if assets_amount == 0:
+            return f_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
+
+    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 current_org.is_root():
+        if self.org_is_global:
+            # 全局组织不返回未分组节点
             return None, []
-
-        org = self.user.orgs.filter(id=current_org.id).first()
-        if not org:
-            return None, []
-
-        _util = UserPermUtil(user=self.user, org=org)
-        assets = _util.get_ungrouped_assets()
+        org = self.get_tree_user_orgs().first()
+        util = UserPermUtil(user=self.user, org=org)
+        assets = util.get_ungrouped_assets(
+            search_asset=search_asset, 
+            asset_category=asset_category, asset_type=asset_type
+        )
         assets_amount = assets.count()
-        node = UserPermAssetTreeNode(
-            _id=UserPermAssetTreeNode.SpecialKey.UNGROUPED.value,
-            key=UserPermAssetTreeNode.SpecialKey.UNGROUPED.value,
-            value=UserPermAssetTreeNode.SpecialKey.UNGROUPED.label,
-            assets_amount=assets_amount
+        u_node = UserPermAssetTreeNode.ungrouped(
+            assets_amount=assets_amount, assets_amount_total=assets_amount
         )
-        node.assets_amount_total = assets_amount
-        nodes = self.serialize_nodes(
-            [node], with_asset_amount=True, expand_level=0,
-            tree_type=str(self.render_tree_type)
-        )
-        serialized_node = nodes[0] if nodes else None
-
-        if not with_assets:
-            return serialized_node, []
-
+        if not u_node.match(search_node):
+            return None, []
+        
         if assets_amount == 0:
-            return serialized_node, []
+            return u_node, []
 
-        assets = assets.values(*AssetTreeNodeAsset.model_values)
-        assets = node.init_assets(list(assets))
-        serialized_assets = self.serialize_assets(assets)
-        return serialized_node, serialized_assets
+        if with_assets:
+            assets_attrs = list(assets.values(*AssetTreeNodeAsset.model_values))
+            assets = u_node.init_assets(assets_attrs)
+        else:
+            assets = []
+        return u_node, assets

apps/perms/tree.py

@@ -18,28 +18,48 @@ logger = get_logger(__name__)
 
 class UserPermAssetTreeNode(AssetTreeNode):
 
-    class Type:
-        # Neither a permission node nor a node with direct permission assets
-        BRIDGE = 'bridge'
-        # Node with direct permission
-        DN = 'dn'
-        # Node with only direct permission assets
-        DA = 'da'
-    
+    class Type(TextChoices):
+        DN = 'direct_node', 'Direct Node'
+        DA = 'direct_asset', 'Direct Asset'
+        BRIDGE = 'bridge', 'Bridge'
+
     class SpecialKey(TextChoices):
         FAVORITE = 'favorite', _('Favorite')
         UNGROUPED = 'ungrouped', _('Ungrouped')
-
-    def __init__(self, tp=None, **kwargs):
-        self.type = tp or self.Type.BRIDGE
+    
+    def __init__(self, tp, **kwargs):
+        self.tp = tp
         super().__init__(**kwargs)
     
+    @classmethod
+    def favorite(cls, **kwargs):
+        node = cls(
+            tp=cls.Type.BRIDGE,
+            _id=cls.SpecialKey.FAVORITE.value,
+            key=cls.SpecialKey.FAVORITE.value,
+            value=cls.SpecialKey.FAVORITE.label,
+            **kwargs
+        )
+        return node
+    
+    @classmethod
+    def ungrouped(cls, **kwargs):
+        node = cls(
+            tp=cls.Type.BRIDGE,
+            _id=cls.SpecialKey.UNGROUPED.value,
+            key=cls.SpecialKey.UNGROUPED.value,
+            value=cls.SpecialKey.UNGROUPED.label,
+            **kwargs
+        )
+        return node
+    
     def as_dict(self, simple=True):
-        data = super().as_dict(simple=simple)
+        data = super().as_dict(simple)
         data.update({
-            'type': self.type,
+            'type': self.tp,
         })
         return data
+
     
 
 class UserPermAssetTree(AssetTree):

apps/perms/urls/user_permission.py

@@ -5,7 +5,7 @@ from .. import api
 user_permission_urlpatterns = [
     path('<str:user>/assets/<uuid:pk>/', api.UserPermedAssetRetrieveApi.as_view(), name='user-permed-asset'),
     path('<str:user>/assets/', api.UserAllPermedAssetsApi.as_view(), name='user-all-assets'),
-    path('<str:user>/nodes/children/tree/', api.UserPermNodeChildrenAsTreeApi.as_view(), name='user-perm-node-children-tree'),
+    path('<str:user>/nodes/children/tree/', api.UserPermedAssetTreeAPI.as_view(), name='user-perm-node-children-tree'),
 ]
 
 user_group_permission_urlpatterns = [

apps/perms/utils/utils.py

@@ -157,13 +157,27 @@ class UserPermUtil(object):
         return assets
 
     @classmethod
-    def get_favorite_assets(cls, user: User):
+    def get_favorite_assets(cls, user: User, search_asset=None, asset_category=None, asset_type=None):
         asset_ids = FavoriteAsset.get_user_favorite_asset_ids(user)
-        assets = Asset.objects.filter(id__in=asset_ids).valid()
+        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
     
-    def get_ungrouped_assets(self):
-        assets = Asset.objects.filter(id__in=self._user_direct_asset_ids).valid()
+    def get_ungrouped_assets(self, search_asset=None, asset_category=None, asset_type=None):
+        q = Q(id__in=self._user_direct_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
     
     def _uuids_to_string(self, uuids):