Dev beta2 (#3177)

* [Update] 添加loading

* [Update] stash

* [Update] 修改permission

apps/perms/api/user_group_permission.py

@@ -10,9 +10,10 @@ from . import user_permission as uapi
 __all__ = [
     'UserGroupGrantedAssetsApi', 'UserGroupGrantedNodesApi',
     'UserGroupGrantedNodeAssetsApi', 'UserGroupGrantedNodeChildrenApi',
-    'UserGroupGrantedNodeChildrenAsTreeApi', 'UserGroupGrantedNodesWithAssetsAsTreeApi',
+    'UserGroupGrantedNodeChildrenAsTreeApi',
+    'UserGroupGrantedNodeChildrenWithAssetsAsTreeApi',
     'UserGroupGrantedAssetSystemUsersApi',
-    # 'UserGroupGrantedNodesWithAssetsAsTreeApi',
+    # 'UserGroupGrantedNodeChildrenWithAssetsAsTreeApi',
 ]
 
 
@@ -45,7 +46,7 @@ class UserGroupGrantedNodeChildrenAsTreeApi(UserGroupPermissionMixin, uapi.UserG
     pass
 
 
-class UserGroupGrantedNodesWithAssetsAsTreeApi(UserGroupPermissionMixin, uapi.UserGrantedNodesWithAssetsAsTreeApi):
+class UserGroupGrantedNodeChildrenWithAssetsAsTreeApi(UserGroupPermissionMixin, uapi.UserGrantedNodeChildrenWithAssetsAsTreeApi):
     pass
 
 

apps/perms/api/user_permission.py

@@ -23,12 +23,19 @@ from ..models import Action
 logger = get_logger(__name__)
 
 __all__ = [
-    'UserGrantedAssetsApi', 'UserGrantedNodesApi',
+    'UserGrantedAssetsApi',
+    'UserGrantedAssetsAsTreeApi',
     'UserGrantedNodeAssetsApi',
+    'UserGrantedNodesApi',
+    'UserGrantedNodesAsTreeApi',
+    'UserGrantedNodesWithAssetsAsTreeApi',
+    'UserGrantedNodeChildrenApi',
+    'UserGrantedNodeChildrenAsTreeApi',
+    'UserGrantedNodeChildrenWithAssetsAsTreeApi',
+    'RefreshAssetPermissionCacheApi',
+    'UserGrantedAssetSystemUsersApi',
     'ValidateUserAssetPermissionApi',
-    'UserGrantedNodesWithAssetsAsTreeApi', 'GetUserAssetPermissionActionsApi',
-    'RefreshAssetPermissionCacheApi', 'UserGrantedAssetSystemUsersApi',
-    'UserGrantedNodeChildrenAsTreeApi', 'UserGrantedNodesWithAssetsAsTreeApi',
+    'GetUserAssetPermissionActionsApi',
 ]
 
 
@@ -58,6 +65,66 @@ class UserPermissionMixin:
         return super().get_permissions()
 
 
+class UserNodePermissionMixin(UserPermissionMixin):
+    util = None
+
+    def initial(self, *args, **kwargs):
+        super().initial(*args, *kwargs)
+        self.util = AssetPermissionUtilV2(self.obj)
+
+
+class UserNodeTreeMixin:
+    serializer_class = TreeNodeSerializer
+    nodes_only_fields = ParserNode.nodes_only_fields
+    tree = None
+
+    def parse_nodes_to_queryset(self, nodes):
+        nodes = nodes.only(*self.nodes_only_fields)
+        _queryset = []
+
+        tree = self.util.get_user_tree()
+        for node in nodes:
+            assets_amount = tree.assets_amount(node.key)
+            if assets_amount == 0:
+                continue
+            node._assets_amount = assets_amount
+            data = ParserNode.parse_node_to_tree_node(node)
+            _queryset.append(data)
+        return _queryset
+
+    def get_serializer_queryset(self, queryset):
+        queryset = self.parse_nodes_to_queryset(queryset)
+        return queryset
+
+    def get_serializer(self, queryset, many=True, **kwargs):
+        queryset = self.get_serializer_queryset(queryset)
+        queryset.sort()
+        return super().get_serializer(queryset, many=many, **kwargs)
+
+
+class UserAssetTreeMixin:
+    serializer_class = TreeNodeSerializer
+    nodes_only_fields = ParserNode.assets_only_fields
+
+    @staticmethod
+    def parse_assets_to_queryset(assets, node):
+        _queryset = []
+        for asset in assets:
+            data = ParserNode.parse_asset_to_tree_node(node, asset)
+            _queryset.append(data)
+        return _queryset
+
+    def get_serializer_queryset(self, queryset):
+        queryset = queryset.only(*self.nodes_only_fields)
+        _queryset = self.parse_assets_to_queryset(queryset, None)
+        return _queryset
+
+    def get_serializer(self, queryset, many=True, **kwargs):
+        queryset = self.get_serializer_queryset(queryset)
+        queryset.sort()
+        return super().get_serializer(queryset, many=many, **kwargs)
+
+
 class UserGrantedAssetsApi(UserPermissionMixin, ListAPIView):
     permission_classes = (IsOrgAdminOrAppUser,)
     serializer_class = serializers.AssetGrantedSerializer
@@ -89,6 +156,10 @@ class UserGrantedAssetsApi(UserPermissionMixin, ListAPIView):
         return queryset
 
 
+class UserGrantedAssetsAsTreeApi(UserAssetTreeMixin, UserGrantedAssetsApi):
+    pass
+
+
 class UserGrantedNodeAssetsApi(UserGrantedAssetsApi):
     def get_queryset(self):
         node_id = self.kwargs.get("node_id")
@@ -100,18 +171,13 @@ class UserGrantedNodeAssetsApi(UserGrantedAssetsApi):
         return queryset
 
 
-class UserGrantedNodesApi(UserPermissionMixin, ListAPIView):
+class UserGrantedNodesApi(UserNodePermissionMixin, ListAPIView):
     """
     查询用户授权的所有节点的API
     """
     permission_classes = (IsOrgAdminOrAppUser,)
     serializer_class = serializers.NodeGrantedSerializer
-    only_fields = NodeSerializer.Meta.only_fields
-    util = None
-
-    def get(self, request, *args, **kwargs):
-        self.util = AssetPermissionUtilV2(self.obj)
-        return super().get(request, *args, **kwargs)
+    nodes_only_fields = NodeSerializer.Meta.only_fields
 
     def get_serializer_context(self):
         context = super().get_serializer_context()
@@ -120,21 +186,35 @@ class UserGrantedNodesApi(UserPermissionMixin, ListAPIView):
 
     def get_queryset(self):
         node_keys = self.util.get_nodes()
-        queryset = Node.objects.filter(key__in=node_keys)
+        queryset = Node.objects.filter(key__in=node_keys)\
+            .only(*self.nodes_only_fields)
         return queryset
 
 
+class UserGrantedNodesAsTreeApi(UserNodeTreeMixin, UserGrantedNodesApi):
+    pass
+
+
+class UserGrantedNodesWithAssetsAsTreeApi(UserGrantedNodesAsTreeApi):
+    def get_serializer_queryset(self, queryset):
+        _queryset = super().get_serializer_queryset(queryset)
+        for node in queryset:
+            assets = self.util.get_nodes_assets(node)
+            _queryset.extend(
+                UserAssetTreeMixin.parse_assets_to_queryset(assets, node)
+            )
+        return _queryset
+
+
 class UserGrantedNodeChildrenApi(UserGrantedNodesApi):
     node = None
-    util = None
     tree = None
-    root_keys = None
+    root_keys = None  # 如果是第一次访问，则需要把二级节点添加进去，这个 roots_keys
 
     def get(self, request, *args, **kwargs):
         key = self.request.query_params.get("key")
         pk = self.request.query_params.get("id")
         system_user_id = self.request.query_params.get("system_user")
-        self.util = AssetPermissionUtilV2(self.obj)
         if system_user_id:
             self.util.filter_permissions(system_users=system_user_id)
         self.tree = self.util.get_user_tree()
@@ -161,26 +241,16 @@ class UserGrantedNodeChildrenApi(UserGrantedNodesApi):
         return queryset
 
 
-class UserGrantedNodeChildrenAsTreeApi(UserGrantedNodeChildrenApi):
-    serializer_class = TreeNodeSerializer
-    only_fields = ParserNode.nodes_only_fields
-
-    def get_queryset(self):
-        nodes = super().get_queryset()
-        queryset = []
-        for node in nodes:
-            node._assets_amount = self.tree.assets_amount(node.key)
-            data = ParserNode.parse_node_to_tree_node(node)
-            queryset.append(data)
-        return queryset
+class UserGrantedNodeChildrenAsTreeApi(UserNodeTreeMixin, UserGrantedNodeChildrenApi):
+    pass
 
 
-class UserGrantedNodesWithAssetsAsTreeApi(UserGrantedNodeChildrenAsTreeApi):
+class UserGrantedNodeChildrenWithAssetsAsTreeApi(UserGrantedNodeChildrenAsTreeApi):
     nodes_only_fields = ParserNode.nodes_only_fields
     assets_only_fields = ParserNode.assets_only_fields
 
-    def get_queryset(self):
-        queryset = super().get_queryset()
+    def get_serializer_queryset(self, queryset):
+        _queryset = super().get_serializer_queryset(queryset)
         nodes = []
         if self.node:
             nodes.append(self.node)
@@ -188,12 +258,13 @@ class UserGrantedNodesWithAssetsAsTreeApi(UserGrantedNodeChildrenAsTreeApi):
             nodes = Node.objects.filter(key__in=self.root_keys)
 
         for node in nodes:
-            assets = self.util.get_nodes_assets(node).only(*self.assets_only_fields)
-            for asset in assets:
-                data = ParserNode.parse_asset_to_tree_node(node, asset)
-                queryset.append(data)
-        queryset = sorted(queryset)
-        return queryset
+            assets = self.util.get_nodes_assets(node).only(
+                *self.assets_only_fields
+            )
+            _queryset.extend(
+                UserAssetTreeMixin.parse_assets_to_queryset(assets, node)
+            )
+        return _queryset
 
 
 class GetUserAssetPermissionActionsApi(RetrieveAPIView):

apps/perms/urls/api_urls.py

@@ -17,16 +17,34 @@ asset_permission_urlpatterns = [
     path('users/<uuid:pk>/assets/', api.UserGrantedAssetsApi.as_view(), name='user-assets'),
     path('users/assets/', api.UserGrantedAssetsApi.as_view(), name='my-assets'),
 
+    # Assets as tree
+    path('users/<uuid:pk>/assets/tree/', api.UserGrantedAssetsAsTreeApi.as_view(), name='user-assets-as-tree'),
+    path('users/assets/tree/', api.UserGrantedAssetsAsTreeApi.as_view(), name='my-assets-as-tree'),
+
     # Nodes
     path('users/<uuid:pk>/nodes/', api.UserGrantedNodesApi.as_view(), name='user-nodes'),
     path('users/nodes/', api.UserGrantedNodesApi.as_view(), name='my-nodes'),
+
+    # Node children
     path('users/<uuid:pk>/nodes/children/', api.UserGrantedNodesApi.as_view(), name='user-nodes-children'),
     path('users/nodes/children/', api.UserGrantedNodesApi.as_view(), name='my-nodes-children'),
 
     # Node as tree
+    path('users/<uuid:pk>/nodes/tree/', api.UserGrantedNodesAsTreeApi.as_view(), name='user-nodes-as-tree'),
+    path('users/nodes/tree/', api.UserGrantedNodesAsTreeApi.as_view(), name='my-nodes-as-tree'),
+
+    # Node with assets as tree
+    path('users/<uuid:pk>/nodes-with-assets/tree/', api.UserGrantedNodesWithAssetsAsTreeApi.as_view(), name='user-nodes-with-assets-as-tree'),
+    path('users/nodes-with-assets/tree/', api.UserGrantedNodesWithAssetsAsTreeApi.as_view(), name='my-nodes-with-assets-as-tree'),
+
+    # Node children as tree
     path('users/<uuid:pk>/nodes/children/tree/', api.UserGrantedNodeChildrenAsTreeApi.as_view(), name='user-nodes-children-as-tree'),
     path('users/nodes/children/tree/', api.UserGrantedNodeChildrenAsTreeApi.as_view(), name='my-nodes-children-as-tree'),
 
+    # Node children with assets as tree
+    path('users/<uuid:pk>/nodes/children-with-assets/tree/', api.UserGrantedNodeChildrenWithAssetsAsTreeApi.as_view(), name='user-nodes-children-with-assets-as-tree'),
+    path('users/nodes/children-with-assets/tree/', api.UserGrantedNodeChildrenWithAssetsAsTreeApi.as_view(), name='my-nodes-children-with-assets-as-tree'),
+
     # Node assets
     path('users/<uuid:pk>/nodes/<uuid:node_id>/assets/', api.UserGrantedNodeAssetsApi.as_view(), name='user-node-assets'),
     path('users/nodes/<uuid:node_id>/assets/', api.UserGrantedNodeAssetsApi.as_view(), name='my-node-assets'),
@@ -35,10 +53,6 @@ asset_permission_urlpatterns = [
     path('users/<uuid:pk>/assets/<uuid:asset_id>/system-users/', api.UserGrantedAssetSystemUsersApi.as_view(), name='user-asset-system-users'),
     path('users/assets/<uuid:asset_id>/system-users/', api.UserGrantedAssetSystemUsersApi.as_view(), name='my-asset-system-users'),
 
-    # Node assets as tree
-    path('users/<uuid:pk>/nodes/children-with-assets/tree/', api.UserGrantedNodesWithAssetsAsTreeApi.as_view(), name='user-nodes-children-with-assets-as-tree'),
-    path('users/nodes/children-with-assets/tree/', api.UserGrantedNodesWithAssetsAsTreeApi.as_view(), name='my-nodes-children-with-assets-as-tree'),
-
     # 查询某个用户组授权的资产和资产组
     path('user-groups/<uuid:pk>/assets/', api.UserGroupGrantedAssetsApi.as_view(), name='user-group-assets'),
     path('user-groups/<uuid:pk>/nodes/', api.UserGroupGrantedNodesApi.as_view(), name='user-group-nodes'),

apps/perms/utils/asset_permission.py

@@ -1,8 +1,11 @@
 # coding: utf-8
 
+import pickle
+import threading
 from collections import defaultdict
 from functools import reduce
 
+from django.core.cache import cache
 from django.db.models import Q
 from django.conf import settings
 
@@ -12,7 +15,6 @@ from common.tree import TreeNode
 from assets.utils import TreeService
 from ..models import AssetPermission
 from ..hands import Node, Asset, SystemUser
-from .. import const
 
 logger = get_logger(__file__)
 
@@ -69,9 +71,12 @@ class AssetPermissionUtilV2:
         'id', 'hostname', 'ip', "platform", "domain_id",
         'comment', 'is_active', 'os', 'org_id'
     )
+    user_tree_cache_key = 'USER_PERM_TREE_{}'
+    user_tree_cache_ttl = 3600
 
     def __init__(self, obj, cache_policy='0'):
         self.object = obj
+        self.cache_policy = cache_policy
         self.obj_id = str(obj.id)
         self._permissions = None
         self._permissions_id = None  # 标记_permission的唯一值
@@ -84,6 +89,7 @@ class AssetPermissionUtilV2:
         self._nodes_direct = None
         self._user_tree = None
         self.full_tree = Node.tree()
+        self.mutex = threading.Lock()
 
     @staticmethod
     def change_org_if_need():
@@ -103,6 +109,31 @@ class AssetPermissionUtilV2:
     def filter_permissions(self, **filters):
         self._permissions = self.permissions.filter(**filters)
 
+    @classmethod
+    def get_user_tree_from_cache(cls, obj_id):
+        return None
+        key = cls.user_tree_cache_key.format(obj_id)
+        data = cache.get(key)
+        if not data:
+            return None
+        user_tree = pickle.loads(data)
+        return user_tree
+
+    @classmethod
+    def expire_user_tree_cache(cls, obj_id):
+        if obj_id == 'all':
+            key = cls.user_tree_cache_key.format('*')
+            cache.delete_pattern(key)
+        else:
+            key = cls.user_tree_cache_key.format(obj_id)
+            cache.delete(key)
+
+    @classmethod
+    def set_user_tree_to_cache(cls, obj_id, user_tree):
+        data = pickle.dumps(user_tree)
+        key = cls.user_tree_cache_key.format(obj_id)
+        cache.set(key, data, cls.user_tree_cache_ttl)
+
     @property
     def user_tree(self):
         return self.get_user_tree()
@@ -237,20 +268,26 @@ class AssetPermissionUtilV2:
 
     @timeit
     def get_user_tree(self):
-        if self._user_tree:
-            return self._user_tree
-        user_tree = TreeService()
-        full_tree_root = self.full_tree.root_node()
-        user_tree.create_node(
-            tag=full_tree_root.tag,
-            identifier=full_tree_root.identifier
-        )
-        self.add_direct_nodes_to_user_tree(user_tree)
-        self.add_single_assets_node_to_user_tree(user_tree)
-        self.parse_user_tree_to_full_tree(user_tree)
-        self.add_empty_node_if_need(user_tree)
-        self._user_tree = user_tree
-        return user_tree
+        with self.mutex:
+            if self._user_tree:
+                return self._user_tree
+            print(id(self), self._user_tree)
+            user_tree = self.__class__.get_user_tree_from_cache(self.obj_id)
+            if user_tree:
+                self._user_tree = user_tree
+                return user_tree
+            user_tree = TreeService()
+            full_tree_root = self.full_tree.root_node()
+            user_tree.create_node(
+                tag=full_tree_root.tag,
+                identifier=full_tree_root.identifier
+            )
+            self.add_direct_nodes_to_user_tree(user_tree)
+            self.add_single_assets_node_to_user_tree(user_tree)
+            self.parse_user_tree_to_full_tree(user_tree)
+            self.add_empty_node_if_need(user_tree)
+            self.__class__.set_user_tree_to_cache(self.obj_id, user_tree)
+            return user_tree
 
     # Todo: 是否可以获取多个资产的系统用户
     def get_asset_system_users_with_actions(self, asset):
@@ -401,15 +438,17 @@ class ParserNode:
     @staticmethod
     def parse_asset_to_tree_node(node, asset):
         icon_skin = 'file'
-        if asset.platform.lower() == 'windows':
+        platform = asset.platform.lower()
+        if platform == 'windows':
             icon_skin = 'windows'
-        elif asset.platform.lower() == 'linux':
+        elif platform == 'linux':
             icon_skin = 'linux'
+        parent_id = node.key if node else ''
         data = {
             'id': str(asset.id),
             'name': asset.hostname,
             'title': asset.ip,
-            'pId': node.key,
+            'pId': parent_id,
             'isParent': False,
             'open': False,
             'iconSkin': icon_skin,