perf: 授权类型树 (#10390)

Co-authored-by: feng <1304903146@qq.com>

apps/assets/api/mixin.py

@@ -69,7 +69,7 @@ class SerializeToTreeNodeMixin:
             return 'file'
 
     @timeit
-    def serialize_assets(self, assets, node_key=None):
+    def serialize_assets(self, assets, node_key=None, pid=None):
         sftp_enabled_platform = PlatformProtocol.objects \
             .filter(name='ssh', setting__sftp_enabled=True) \
             .values_list('platform', flat=True) \
@@ -83,8 +83,10 @@ class SerializeToTreeNodeMixin:
             {
                 'id': str(asset.id),
                 'name': asset.name,
-                'title': f'{asset.address}\n{asset.comment}',
-                'pId': get_pid(asset),
+                'title':
+                    f'{asset.address}\n{asset.comment}'
+                    if asset.comment else asset.address,
+                'pId': pid or get_pid(asset),
                 'isParent': False,
                 'open': False,
                 'iconSkin': self.get_icon(asset),

apps/assets/const/types.py

@@ -193,15 +193,38 @@ class AllTypes(ChoicesMixin):
         }
         return node
 
-    @classmethod
-    def to_tree_nodes(cls, include_asset, count_resource='asset'):
-        from accounts.models import Account
-        from ..models import Asset, Platform
-        if count_resource == 'account':
-            resource_platforms = Account.objects.all().values_list('asset__platform_id', flat=True)
-        else:
-            resource_platforms = Asset.objects.all().values_list('platform_id', flat=True)
 
+    @classmethod
+    def asset_to_node(cls, asset, pid):
+        node = {
+            'id': '{}'.format(asset.id),
+            'name': asset.name,
+            'title': f'{asset.address}\n{asset.comment}',
+            'pId': pid,
+            'isParent': False,
+            'open': False,
+            'iconSkin': asset.type,
+            'chkDisabled': not asset.is_active,
+            'meta': {
+                'type': 'platform',
+                'data': {
+                    'platform_type': asset.platform.type,
+                    'org_name': asset.org_name,
+                    # 'sftp': asset.platform_id in sftp_enabled_platform,
+                    'name': asset.name,
+                    'address': asset.address
+                },
+            }
+        }
+        return node
+
+    @classmethod
+    def get_root_nodes(cls):
+        return dict(id='ROOT', name=_('All types'), title=_('All types'), open=True, isParent=True)
+
+    @classmethod
+    def get_tree_nodes(cls, resource_platforms, include_asset=False):
+        from ..models import Platform
         platform_count = defaultdict(int)
         for platform_id in resource_platforms:
             platform_count[platform_id] += 1
@@ -215,8 +238,7 @@ class AllTypes(ChoicesMixin):
             category_type_mapper[p.category] += platform_count[p.id]
             tp_platforms[p.category + '_' + p.type].append(p)
 
-        root = dict(id='ROOT', name=_('All types'), title=_('All types'), open=True, isParent=True)
-        nodes = [root]
+        nodes = [cls.get_root_nodes()]
         for category, type_cls in cls.category_types():
             # Category 格式化
             meta = {'type': 'category', 'category': category.value}
@@ -244,6 +266,16 @@ class AllTypes(ChoicesMixin):
                     nodes.append(platform_node)
         return nodes
 
+    @classmethod
+    def to_tree_nodes(cls, include_asset, count_resource='asset'):
+        from accounts.models import Account
+        from ..models import Asset
+        if count_resource == 'account':
+            resource_platforms = Account.objects.all().values_list('asset__platform_id', flat=True)
+        else:
+            resource_platforms = Asset.objects.all().values_list('platform_id', flat=True)
+        return cls.get_tree_nodes(resource_platforms, include_asset)
+
     @classmethod
     def get_type_default_platform(cls, category, tp):
         constraints = cls.get_constraints(category, tp)

apps/perms/api/user_permission/tree/node_with_asset.py

@@ -1,4 +1,6 @@
 import abc
+import re
+from collections import defaultdict
 from urllib.parse import parse_qsl
 
 from django.conf import settings
@@ -11,6 +13,7 @@ from rest_framework.response import Response
 
 from accounts.const import AliasAccount
 from assets.api import SerializeToTreeNodeMixin
+from assets.const import AllTypes
 from assets.models import Asset
 from assets.utils import KubernetesTree
 from authentication.models import ConnectionToken
@@ -26,7 +29,8 @@ from ..mixin import SelfOrPKUserMixin
 __all__ = [
     'UserGrantedK8sAsTreeApi',
     'UserPermedNodesWithAssetsAsTreeApi',
-    'UserPermedNodeChildrenWithAssetsAsTreeApi'
+    'UserPermedNodeChildrenWithAssetsAsTreeApi',
+    'UserPermedNodeChildrenWithAssetsAsCategoryTreeApi',
 ]
 
 
@@ -137,6 +141,75 @@ class UserPermedNodeChildrenWithAssetsAsTreeApi(BaseUserNodeWithAssetAsTreeApi):
         return self.query_node_key or self.default_unfolded_node_key
 
 
+class UserPermedNodeChildrenWithAssetsAsCategoryTreeApi(
+    SelfOrPKUserMixin, SerializeToTreeNodeMixin, ListAPIView
+):
+    @property
+    def is_sync(self):
+        sync = self.request.query_params.get('sync', 0)
+        return int(sync) == 1
+
+    @property
+    def tp(self):
+        return self.request.query_params.get('type')
+
+    def get_assets(self):
+        query_asset_util = UserPermAssetUtil(self.user)
+        node = PermNode.objects.filter(
+            granted_node_rels__user=self.user, parent_key='').first()
+        if node:
+            __, assets = query_asset_util.get_node_all_assets(node.id)
+        else:
+            assets = Asset.objects.none()
+        return assets
+
+    def to_tree_nodes(self, assets):
+        if not assets:
+            return []
+        assets = assets.annotate(tp=F('platform__type'))
+        asset_type_map = defaultdict(list)
+        for asset in assets:
+            asset_type_map[asset.tp].append(asset)
+        tp = self.tp
+        if tp:
+            assets = asset_type_map.get(tp, [])
+            if not assets:
+                return []
+            pid = f'ROOT_{str(assets[0].category).upper()}_{tp}'
+            return self.serialize_assets(assets, pid=pid)
+
+        resource_platforms = assets.values_list('platform_id', flat=True)
+        node_all = AllTypes.get_tree_nodes(resource_platforms)
+        pattern = re.compile(r'\(0\)?')
+        nodes = []
+        for node in node_all:
+            meta = node.get('meta', {})
+            if pattern.search(node['name']) or meta.get('type') == 'platform':
+                continue
+            _type = meta.get('_type')
+            if _type:
+                node['type'] = _type
+            nodes.append(node)
+
+        if not self.is_sync:
+            return nodes
+
+        asset_nodes = []
+        for node in nodes:
+            node['open'] = True
+            tp = node.get('meta', {}).get('_type')
+            if not tp:
+                continue
+            assets = asset_type_map.get(tp, [])
+            asset_nodes += self.serialize_assets(assets, pid=node['id'])
+        return nodes + asset_nodes
+
+    def list(self, request, *args, **kwargs):
+        assets = self.get_assets()
+        nodes = self.to_tree_nodes(assets)
+        return Response(data=nodes)
+
+
 class UserGrantedK8sAsTreeApi(SelfOrPKUserMixin, ListAPIView):
     """ 用户授权的K8s树 """
 

apps/perms/urls/user_permission.py

@@ -37,6 +37,9 @@ user_permission_urlpatterns = [
     path('<str:user>/nodes/children-with-assets/tree/',
          api.UserPermedNodeChildrenWithAssetsAsTreeApi.as_view(),
          name='user-node-children-with-assets-as-tree'),
+    path('<str:user>/nodes/children-with-assets/category/tree/',
+         api.UserPermedNodeChildrenWithAssetsAsCategoryTreeApi.as_view(),
+         name='user-node-children-with-assets-as-category-tree'),
     # 同步树
     path('<str:user>/nodes/all-with-assets/tree/',
          api.UserPermedNodesWithAssetsAsTreeApi.as_view(),

apps/perms/utils/user_perm.py

@@ -1,15 +1,11 @@
-from assets.models import FavoriteAsset, Asset
-
 from django.conf import settings
 from django.db.models import Q
 
+from assets.models import FavoriteAsset, Asset
 from common.utils.common import timeit
-
 from perms.models import AssetPermission, PermNode, UserAssetGrantedTreeNodeRelation
-
 from .permission import AssetPermissionUtil
 
-
 __all__ = ['AssetPermissionPermAssetUtil', 'UserPermAssetUtil', 'UserPermNodeUtil']