From c0cb58c0019f2a744c0ebb907ebfa083e98db3b2 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Mon, 8 Aug 2022 10:41:37 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E6=B7=BB=E5=8A=A0=20tree=20api?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/assets/api/node.py      | 20 +++++++++++++-------
 apps/assets/const.py         | 27 +++++++++++++++++++++++++++
 apps/assets/urls/api_urls.py |  1 +
 apps/jumpserver/urls.py      | 12 ------------
 4 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/apps/assets/api/node.py b/apps/assets/api/node.py
index 8ddb79e71..59a1da793 100644
--- a/apps/assets/api/node.py
+++ b/apps/assets/api/node.py
@@ -1,6 +1,7 @@
 # ~*~ coding: utf-8 ~*~
 from functools import partial
 from collections import namedtuple, defaultdict
+from django.core.exceptions import PermissionDenied
 
 from rest_framework import status
 from rest_framework.serializers import ValidationError
@@ -35,9 +36,8 @@ logger = get_logger(__file__)
 __all__ = [
     'NodeViewSet', 'NodeChildrenApi', 'NodeAssetsApi',
     'NodeAddAssetsApi', 'NodeRemoveAssetsApi', 'MoveAssetsToNodeApi',
-    'NodeAddChildrenApi', 'NodeListAsTreeApi',
-    'NodeChildrenAsTreeApi',
-    'NodeTaskCreateApi',
+    'NodeAddChildrenApi', 'NodeListAsTreeApi', 'NodeChildrenAsTreeApi',
+    'NodeTaskCreateApi', 'CategoryTreeApi',
 ]
 
 
@@ -199,11 +199,17 @@ class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
 
 
 class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
+    serializer_class = TreeNodeSerializer
+
+    def check_permissions(self, request):
+        if not request.user.has_perm('assets.view_asset'):
+            raise PermissionDenied
+        return True
+
     def list(self, request, *args, **kwargs):
-        category_types = AllTypes.category_types()
-        nodes = self.get_queryset().order_by('value')
-        nodes = self.serialize_nodes(nodes, with_asset_amount=True)
-        return Response(data=nodes)
+        nodes = AllTypes.to_tree_nodes()
+        serializer = self.get_serializer(nodes, many=True)
+        return Response(data=serializer.data)
 
 
 class NodeAssetsApi(generics.ListAPIView):
diff --git a/apps/assets/const.py b/apps/assets/const.py
index 9069a6c57..4da356115 100644
--- a/apps/assets/const.py
+++ b/apps/assets/const.py
@@ -1,6 +1,7 @@
 from django.db import models
 from django.utils.translation import gettext_lazy as _
 from common.db.models import IncludesTextChoicesMeta
+from common.tree import TreeNode
 
 
 __all__ = [
@@ -160,6 +161,32 @@ class AllTypes(metaclass=IncludesTextChoicesMeta):
         title = ['value', 'display_name']
         return [dict(zip(title, choice)) for choice in choices]
 
+    @staticmethod
+    def choice_to_node(choice, pid, opened=True, is_parent=True, meta=None):
+        node = TreeNode(**{
+            'id': choice.name,
+            'name': choice.label,
+            'title': choice.label,
+            'pId': pid,
+            'open': opened,
+            'isParent': is_parent,
+        })
+        if meta:
+            node.meta = meta
+        return node
+
+    @classmethod
+    def to_tree_nodes(cls):
+        root = TreeNode(id='ROOT', name='类型节点', title='类型节点')
+        nodes = [root]
+        for category, types in cls.category_types():
+            category_node = cls.choice_to_node(category, 'ROOT', meta={'type': 'category'})
+            nodes.append(category_node)
+            for tp in types:
+                tp_node = cls.choice_to_node(tp, category_node.id, meta={'type': 'type'})
+                nodes.append(tp_node)
+        return nodes
+
 
 class Protocol(models.TextChoices):
     ssh = 'ssh', 'SSH'
diff --git a/apps/assets/urls/api_urls.py b/apps/assets/urls/api_urls.py
index 5070cd44d..3e6b1309b 100644
--- a/apps/assets/urls/api_urls.py
+++ b/apps/assets/urls/api_urls.py
@@ -39,6 +39,7 @@ urlpatterns = [
 
     path('accounts/tasks/', api.AccountTaskCreateAPI.as_view(), name='account-task-create'),
 
+    path('nodes/category/tree/', api.CategoryTreeApi.as_view(), name='asset-category-tree'),
     path('nodes/tree/', api.NodeListAsTreeApi.as_view(), name='node-tree'),
     path('nodes/children/tree/', api.NodeChildrenAsTreeApi.as_view(), name='node-children-tree'),
     path('nodes/<uuid:pk>/children/', api.NodeChildrenApi.as_view(), name='node-children'),
diff --git a/apps/jumpserver/urls.py b/apps/jumpserver/urls.py
index ef654fddf..cd4d2311d 100644
--- a/apps/jumpserver/urls.py
+++ b/apps/jumpserver/urls.py
@@ -45,12 +45,6 @@ if settings.XPACK_ENABLED:
     )
 
 
-apps = [
-    'users', 'assets', 'perms', 'terminal', 'ops', 'audits',
-    'orgs', 'auth', 'applications', 'tickets', 'settings', 'xpack',
-    'flower', 'luna', 'koko', 'ws', 'docs', 'redocs',
-]
-
 urlpatterns = [
     path('', views.IndexView.as_view(), name='index'),
     path('api/v1/', include(api_v1)),
@@ -86,11 +80,5 @@ if os.environ.get('DEBUG_TOOLBAR', False):
     ]
 
 
-# 兼容之前的
-old_app_pattern = '|'.join(apps)
-old_app_pattern = r'^{}'.format(old_app_pattern)
-urlpatterns += [re_path(old_app_pattern, views.redirect_old_apps_view)]
-
-
 handler404 = 'jumpserver.views.handler404'
 handler500 = 'jumpserver.views.handler500'