pref: 类型树支持资产

apps/assets/api/tree.py

@@ -1,6 +1,5 @@
 # ~*~ coding: utf-8 ~*~
 
-from django.core.exceptions import PermissionDenied
 from rest_framework.generics import get_object_or_404
 from rest_framework.response import Response
 
@@ -12,44 +11,20 @@ from orgs.utils import current_org
 from .mixin import SerializeToTreeNodeMixin
 from .. import serializers
 from ..const import AllTypes
-from ..models import Node
+from ..models import Node, Platform, Asset
 
 logger = get_logger(__file__)
 __all__ = [
     'NodeChildrenApi',
-    'NodeListAsTreeApi',
     'NodeChildrenAsTreeApi',
     'CategoryTreeApi',
 ]
 
 
-class NodeListAsTreeApi(generics.ListAPIView):
-    """
-    获取节点列表树
-    [
-      {
-        "id": "",
-        "name": "",
-        "pId": "",
-        "meta": ""
-      }
-    ]
-    """
-    model = Node
-    serializer_class = TreeNodeSerializer
-
-    @staticmethod
-    def to_tree_queryset(queryset):
-        queryset = [node.as_tree_node() for node in queryset]
-        return queryset
-
-    def filter_queryset(self, queryset):
-        queryset = super().filter_queryset(queryset)
-        queryset = self.to_tree_queryset(queryset)
-        return queryset
-
-
 class NodeChildrenApi(generics.ListCreateAPIView):
+    """
+    节点的增删改查
+    """
     serializer_class = serializers.NodeSerializer
     search_fields = ('value',)
 
@@ -141,33 +116,38 @@ class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
     def list(self, request, *args, **kwargs):
         nodes = self.filter_queryset(self.get_queryset()).order_by('value')
         nodes = self.serialize_nodes(nodes, with_asset_amount=True)
-        assets = self.get_assets()
+        assets = self.get_assets_as_node()
         data = [*nodes, *assets]
         return Response(data=data)
 
-    def get_assets(self):
+    def get_assets_as_node(self):
         include_assets = self.request.query_params.get('assets', '0') == '1'
         if not self.instance or not include_assets:
             return []
-        assets = self.instance.get_assets().only(
-            "id", "name", "address", "platform_id",
-            "org_id", "is_active",
-        ).prefetch_related('platform')
+        assets = self.instance.get_assets_for_tree()
         return self.serialize_assets(assets, self.instance.key)
 
 
 class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
     serializer_class = TreeNodeSerializer
+    rbac_perms = {
+        'GET': 'assets.view_asset',
+        'list': 'assets.view_asset',
+    }
 
-    def check_permissions(self, request):
-        if not request.user.has_perm('assets.view_asset'):
-            raise PermissionDenied
-        return True
+    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):
-        if request.query_params.get('key'):
-            nodes = []
+        include_asset = bool(self.request.query_params.get('assets'))
+
+        if include_asset and self.request.query_params.get('key'):
+            nodes = self.get_assets()
         else:
-            nodes = AllTypes.to_tree_nodes()
-        serializer = self.get_serializer(nodes, many=True)
-        return Response(data=serializer.data)
+            nodes = AllTypes.to_tree_nodes(include_asset)
+        return Response(data=nodes)

apps/assets/const/types.py

@@ -1,8 +1,9 @@
 from collections import defaultdict
 from copy import deepcopy
 
+from django.utils.translation import gettext as _
+
 from common.db.models import ChoicesMixin
-from common.tree import TreeNode
 from .category import Category
 from .cloud import CloudTypes
 from .database import DatabaseTypes
@@ -134,34 +135,34 @@ class AllTypes(ChoicesMixin):
 
     @staticmethod
     def choice_to_node(choice, pid, opened=True, is_parent=True, meta=None):
-        node = TreeNode(**{
+        node = {
             'id': pid + '_' + choice.name,
             'name': choice.label,
             'title': choice.label,
             'pId': pid,
             'open': opened,
             'isParent': is_parent,
-        })
+        }
         if meta:
-            node.meta = meta
+            node['meta'] = meta
         return node
 
     @classmethod
-    def platform_to_node(cls, p, pid):
-        node = TreeNode(**{
+    def platform_to_node(cls, p, pid, include_asset):
+        node = {
             'id': '{}'.format(p.id),
             'name': p.name,
             'title': p.name,
             'pId': pid,
-            'isParent': True,
+            'isParent': include_asset,
             'meta': {
                 'type': 'platform'
             }
-        })
+        }
         return node
 
     @classmethod
-    def to_tree_nodes(cls):
+    def to_tree_nodes(cls, include_asset):
         from ..models import Asset, Platform
         asset_platforms = Asset.objects.all().values_list('platform_id', flat=True)
         platform_count = defaultdict(int)
@@ -177,26 +178,29 @@ class AllTypes(ChoicesMixin):
             category_type_mapper[p.category] += platform_count[p.id]
             tp_platforms[p.category + '_' + p.type].append(p)
 
-        root = TreeNode(id='ROOT', name='所有类型', title='所有类型', open=True, isParent=True)
+        root = dict(id='ROOT', name=_('All types'), title='所有类型', open=True, isParent=True)
         nodes = [root]
         for category, type_cls in cls.category_types():
+            # Category 格式化
             meta = {'type': 'category', 'category': category.value}
             category_node = cls.choice_to_node(category, 'ROOT', meta=meta)
             category_count = category_type_mapper.get(category, 0)
-            category_node.name += f'({category_count})'
+            category_node['name'] += f'({category_count})'
             nodes.append(category_node)
 
-            tps = type_cls.get_types()
-            for tp in tps:
+            # Type 格式化
+            types = type_cls.get_types()
+            for tp in types:
                 meta = {'type': 'type', 'category': category.value, '_type': tp.value}
-                tp_node = cls.choice_to_node(tp, category_node.id, opened=False, meta=meta)
+                tp_node = cls.choice_to_node(tp, category_node['id'], opened=False, meta=meta)
                 tp_count = category_type_mapper.get(category + '_' + tp, 0)
-                tp_node.name += f'({tp_count})'
+                tp_node['name'] += f'({tp_count})'
                 nodes.append(tp_node)
 
+                # Platform 格式化
                 for p in tp_platforms.get(category + '_' + tp, []):
-                    platform_node = cls.platform_to_node(p, tp_node.id)
-                    platform_node.name += f'({platform_count.get(p.id, 0)})'
+                    platform_node = cls.platform_to_node(p, tp_node['id'], include_asset)
+                    platform_node['name'] += f'({platform_count.get(p.id, 0)})'
                     nodes.append(platform_node)
         return nodes
 

apps/assets/models/node.py

@@ -430,6 +430,12 @@ class NodeAssetsMixin(NodeAllAssetsMappingMixin):
         assets = Asset.objects.filter(nodes=self)
         return assets.distinct()
 
+    def get_assets_for_tree(self):
+        return self.get_assets().only(
+            "id", "name", "address", "platform_id",
+            "org_id", "is_active"
+        ).prefetch_related('platform')
+
     def get_valid_assets(self):
         return self.get_assets().valid()
 

apps/assets/urls/api_urls.py

@@ -50,7 +50,7 @@ urlpatterns = [
          name='account-secret-history'),
 
     path('nodes/category/tree/', api.CategoryTreeApi.as_view(), name='asset-category-tree'),
-    path('nodes/tree/', api.NodeListAsTreeApi.as_view(), name='node-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'),
     path('nodes/children/', api.NodeChildrenApi.as_view(), name='node-children-2'),