[Update] 优化节点

apps/assets/api/asset.py

@@ -73,19 +73,21 @@ class AssetViewSet(LabelFilter, OrgBulkModelViewSet):
         node = get_object_or_404(Node, id=node_id)
         show_current_asset = self.request.query_params.get("show_current_asset") in ('1', 'true')
 
+        # 当前节点是顶层节点, 并且仅显示直接资产
         if node.is_root() and show_current_asset:
             queryset = queryset.filter(
                 Q(nodes=node_id) | Q(nodes__isnull=True)
-            )
+            ).distinct()
+        # 当前节点是顶层节点，显示所有资产
         elif node.is_root() and not show_current_asset:
-            pass
+            return queryset
+        # 当前节点不是鼎城节点，只显示直接资产
         elif not node.is_root() and show_current_asset:
             queryset = queryset.filter(nodes=node)
         else:
-            queryset = queryset.filter(
-                nodes__key__regex='^{}(:[0-9]+)*$'.format(node.key),
-            )
-        return queryset.distinct()
+            children = node.get_all_children(with_self=True)
+            queryset = queryset.filter(nodes__in=children).distinct()
+        return queryset
 
     def filter_admin_user_id(self, queryset):
         admin_user_id = self.request.query_params.get('admin_user_id')

apps/assets/api/node.py

@@ -13,8 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from rest_framework import generics, mixins, viewsets
+import time
+
+from rest_framework import generics, mixins
 from rest_framework.serializers import ValidationError
+from rest_framework.pagination import LimitOffsetPagination
 from rest_framework.views import APIView
 from rest_framework.response import Response
 from django.utils.translation import ugettext_lazy as _
@@ -22,6 +25,7 @@ from django.shortcuts import get_object_or_404
 
 from common.utils import get_logger, get_object_or_none
 from common.tree import TreeNodeSerializer
+from orgs.mixins import OrgModelViewSet
 from ..hands import IsOrgAdmin
 from ..models import Node
 from ..tasks import update_assets_hardware_info_util, test_asset_connectivity_util
@@ -39,29 +43,26 @@ __all__ = [
 ]
 
 
-class NodeViewSet(viewsets.ModelViewSet):
-    filter_fields = ('value', 'key', )
-    search_fields = filter_fields
+class NodeViewSet(OrgModelViewSet):
+    filter_fields = ('value', 'key', 'id')
+    search_fields = ('value', )
     queryset = Node.objects.all()
     permission_classes = (IsOrgAdmin,)
     serializer_class = serializers.NodeSerializer
+    pagination_class = LimitOffsetPagination
 
+    # 仅支持根节点指直接创建，子节点下的节点需要通过children接口创建
     def perform_create(self, serializer):
         child_key = Node.root().get_next_child_key()
         serializer.validated_data["key"] = child_key
         serializer.save()
 
-    def update(self, request, *args, **kwargs):
+    def perform_update(self, serializer):
         node = self.get_object()
-        if node.is_root():
-            node_value = node.value
-            post_value = request.data.get('value')
-            if node_value != post_value:
-                return Response(
-                    {"msg": _("You can't update the root node name")},
-                    status=400
-                )
-        return super().update(request, *args, **kwargs)
+        if node.is_root() and node.value != serializer.validated_data['value']:
+            msg = _("You can't update the root node name")
+            raise ValidationError({"error": msg})
+        return super().perform_update(serializer)
 
 
 class NodeListAsTreeApi(generics.ListAPIView):
@@ -79,17 +80,19 @@ class NodeListAsTreeApi(generics.ListAPIView):
     permission_classes = (IsOrgAdmin,)
     serializer_class = TreeNodeSerializer
 
-    def get_queryset(self):
-        queryset = Node.objects.all()
+    def to_tree_queryset(self, queryset):
         util = NodeUtil()
         nodes = util.get_nodes_by_queryset(queryset)
         queryset = [node.as_tree_node() for node in nodes]
         return queryset
 
-    @staticmethod
-    def refresh_nodes(queryset):
-        Node.expire_nodes_assets_amount()
-        Node.expire_nodes_full_value()
+    def get_queryset(self):
+        queryset = Node.objects.all()
+        return queryset
+
+    def filter_queryset(self, queryset):
+        queryset = super().filter_queryset(queryset)
+        queryset = self.to_tree_queryset(queryset)
         return queryset
 
 
@@ -112,18 +115,28 @@ class NodeChildrenAsTreeApi(generics.ListAPIView):
     is_root = False
 
     def get_queryset(self):
+        t1 = time.time()
         self.check_need_refresh_nodes()
+        t2 = time.time()
+        print("1: ", t2 - t1)
         node_key = self.request.query_params.get('key')
-        util = NodeUtil()
+        # util = NodeUtil()
         # 是否包含自己
         with_self = False
         if not node_key:
             node_key = Node.root().key
             with_self = True
-        self.node = util.get_node_by_key(node_key)
+        # self.node = util.get_node_by_key(node_key)
+        self.node = get_object_or_404(Node, key=node_key)
+        t3 = time.time()
+        print("2: ", t3 - t2)
         queryset = self.node.get_children(with_self=with_self)
+        t4 = time.time()
         queryset = [node.as_tree_node() for node in queryset]
+        print("3: ", t4 - t3)
+        t5 = time.time()
         queryset = sorted(queryset)
+        print("4: ", t5 - t4)
         return queryset
 
     def filter_assets(self, queryset):
@@ -131,7 +144,8 @@ class NodeChildrenAsTreeApi(generics.ListAPIView):
         if not include_assets:
             return queryset
         assets = self.node.get_assets().only(
-            "id", "hostname", "ip", 'platform', "os", "org_id", "protocols",
+            "id", "hostname", "ip", 'platform', "os",
+            "org_id", "protocols",
         )
         for asset in assets:
             queryset.append(asset.as_tree_node(self.node))

apps/assets/forms/asset.py

@@ -29,9 +29,14 @@ class ProtocolForm(forms.Form):
 class AssetCreateForm(OrgModelForm):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        if self.data:
+            return
         nodes_field = self.fields['nodes']
-        nodes_field.choices = ((n.id, n.full_value) for n in
-                               Node.get_queryset())
+        if self.instance:
+            nodes_field.choices = ((n.id, n.full_value) for n in
+                                   self.instance.nodes.all())
+        else:
+            nodes_field.choices = []
 
     class Meta:
         model = Asset
@@ -42,7 +47,7 @@ class AssetCreateForm(OrgModelForm):
         ]
         widgets = {
             'nodes': forms.SelectMultiple(attrs={
-                'class': 'select2', 'data-placeholder': _('Nodes')
+                'class': 'nodes-select2', 'data-placeholder': _('Nodes')
             }),
             'admin_user': forms.Select(attrs={
                 'class': 'select2', 'data-placeholder': _('Admin user')
@@ -68,6 +73,17 @@ class AssetCreateForm(OrgModelForm):
 
 
 class AssetUpdateForm(OrgModelForm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.data:
+            return
+        nodes_field = self.fields['nodes']
+        if self.instance:
+            nodes_field.choices = ((n.id, n.full_value) for n in
+                                   self.instance.nodes.all())
+        else:
+            nodes_field.choices = []
+
     class Meta:
         model = Asset
         fields = [
@@ -77,7 +93,7 @@ class AssetUpdateForm(OrgModelForm):
         ]
         widgets = {
             'nodes': forms.SelectMultiple(attrs={
-                'class': 'select2', 'data-placeholder': _('Node')
+                'class': 'nodes-select2', 'data-placeholder': _('Node')
             }),
             'admin_user': forms.Select(attrs={
                 'class': 'select2', 'data-placeholder': _('Admin user')

apps/assets/serializers/node.py

@@ -13,17 +13,24 @@ __all__ = [
 
 
 class NodeSerializer(BulkOrgResourceModelSerializer):
-    assets_amount = serializers.IntegerField(read_only=True)
     name = serializers.ReadOnlyField(source='value')
+    full_value = serializers.SerializerMethodField(label=_("Full value"))
 
     class Meta:
         model = Node
         only_fields = ['id', 'key', 'value', 'org_id']
-        fields = only_fields + ['name', 'assets_amount']
+        fields = only_fields + ['name', 'full_value']
         read_only_fields = [
-            'key', 'name', 'assets_amount', 'org_id',
+            'key', 'name', 'full_value', 'org_id',
         ]
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tree = Node.tree()
+
+    def get_full_value(self, obj):
+        return self.tree.get_node_full_tag(obj.key)
+
     def validate_value(self, data):
         instance = self.instance if self.instance else Node.root()
         children = instance.parent.get_children()

apps/assets/signals_handler.py

@@ -112,7 +112,8 @@ def on_node_assets_changed(sender, instance=None, **kwargs):
 @receiver(post_save, sender=Node)
 def on_node_update_or_created(sender, instance=None, created=False, **kwargs):
     if instance and not created:
-        instance.expire_full_value()
+        pass
+        # instance.expire_full_value()
 
 
 @receiver(post_save, sender=AuthBook)

apps/assets/templates/assets/_node_tree.html

@@ -236,7 +236,8 @@ function onBodyMouseDown(event){
 }
 
 function onRename(event, treeId, treeNode, isCancel){
-    var url = "{% url 'api-assets:node-detail' pk=DEFAULT_PK %}".replace("{{ DEFAULT_PK }}", current_node_id);
+    var url = "{% url 'api-assets:node-detail' pk=DEFAULT_PK %}"
+        .replace("{{ DEFAULT_PK }}", current_node_id);
     var data = {"value": treeNode.name};
     if (isCancel){
         return
@@ -250,7 +251,7 @@ function onRename(event, treeId, treeNode, isCancel){
             treeNode.name = treeNode.name + ' (' + treeNode.meta.node.assets_amount + ')';
             zTree.updateNode(treeNode);
             console.log("Success: " + treeNode.name)
-        }
+        },
     })
 }
 

apps/assets/templates/assets/asset_create.html

@@ -110,6 +110,27 @@ $(document).ready(function () {
     $('.select2').select2({
         allowClear: true
     });
+    $(".nodes-select2").select2({
+        closeOnSelect: false,
+        ajax: {
+            url: '{% url "api-assets:node-list" %}',
+            data: function (params) {
+                var page = params.page || 1;
+                var query = {
+                   search: params.term,
+                   offset: (page -1) * 10,
+                   limit: 10
+                };
+                return query
+            },
+            processResults: function (data) {
+                var results = $.map(data.results, function (v, i) {
+                    return {id: v.id, text: v.full_value}
+                });
+                return {results: results, pagination: {"more": true}}
+            }
+        },
+    });
     $(".labels").select2({
         allowClear: true,
         templateSelection: format

apps/assets/utils.py

 # ~*~ coding: utf-8 ~*~
 #
-import time
 from functools import reduce
+from treelib import Tree
 from django.db.models import Prefetch, Q
 
-from common.utils import get_object_or_none, get_logger
+from common.utils import get_object_or_none, get_logger, timeit
 from common.struct import Stack
 from .models import SystemUser, Label, Node, Asset
 
@@ -54,10 +54,12 @@ class LabelFilter(LabelFilterMixin):
 
 
 class NodeUtil:
+    full_value_sep = ' / '
+
     def __init__(self, with_assets_amount=False, debug=False):
         self.stack = Stack()
-        self._nodes = {}
         self.with_assets_amount = with_assets_amount
+        self._nodes = {}
         self._debug = debug
         self.init()
 
@@ -65,62 +67,85 @@ class NodeUtil:
     def sorted_by(node):
         return [int(i) for i in node.key.split(':')]
 
+    @timeit
     def get_queryset(self):
-        all_nodes = Node.objects.all()
+        queryset = Node.objects.all().only('id', 'key', 'value')
         if self.with_assets_amount:
-            all_nodes = all_nodes.prefetch_related(
-                Prefetch('assets', queryset=Asset.objects.all().only('id'))
-            )
-            all_nodes = list(all_nodes)
-            for node in all_nodes:
-                node._assets = set(node.assets.all())
-        return all_nodes
+            queryset = queryset.prefetch_related(
+                Prefetch('assets', queryset=Asset.objects.all().only('id')
+            ))
+        return list(queryset)
 
-    def get_all_nodes(self):
-        all_nodes = sorted(self.get_queryset(), key=self.sorted_by)
+    @staticmethod
+    def set_node_default_attr(node):
+        setattr(node, '_full_value', node.value)
+        setattr(node, '_children', [])
+        setattr(node, '_all_children', [])
+        setattr(node, '_parents', [])
 
-        guarder = Node(key='', value='Guarder')
-        guarder._assets = []
-        all_nodes.append(guarder)
-        return all_nodes
+    @timeit
+    def get_all_nodes(self):
+        queryset = sorted(list(self.get_queryset()), key=self.sorted_by)
+        guarder = Node(key='', value='ROOT')
+        self.set_node_default_attr(guarder)
+        queryset.append(guarder)
+        for node in queryset[:-1]:
+            self.set_node_default_attr(node)
+            if not self.with_assets_amount:
+                continue
+            assets = set([str(a.id) for a in node.assets.all()])
+            node._assets = assets
+            node._all_assets = assets
+        return queryset
 
     def push_to_stack(self, node):
         # 入栈之前检查
         # 如果栈是空的，证明是一颗树的根部
         if self.stack.is_empty():
             node._full_value = node.value
-            node._parents = []
         else:
             # 如果不是根节点,
             # 该节点的祖先应该是父节点的祖先加上父节点
             # 该节点的名字是父节点的名字+自己的名字
-            node._parents = [self.stack.top] + self.stack.top._parents
-            node._full_value = ' / '.join(
-                [self.stack.top._full_value, node.value]
+            node._parents = [self.stack.top] + getattr(self.stack.top, '_parents')
+            node._full_value = self.full_value_sep.join(
+                [getattr(self.stack.top, '_full_value', ''), node.value]
             )
-        node._children = []
-        node._all_children = []
-        self.debug("入栈: {}".format(node.key))
+        # self.debug("入栈: {}".format(node.key))
         self.stack.push(node)
 
     # 出栈
+    # @timeit
     def pop_from_stack(self):
         _node = self.stack.pop()
-        self.debug("出栈: {} 栈顶: {}".format(_node.key, self.stack.top.key if self.stack.top else None))
+        # self.debug("出栈: {} 栈顶: {}".format(
+        #     _node.key, self.stack.top.key if self.stack.top else None)
+        # )
         self._nodes[_node.key] = _node
         if not self.stack.top:
             return
-        if self.with_assets_amount:
-            self.stack.top._assets.update(_node._assets)
-            _node._assets_amount = len(_node._assets)
-            delattr(_node, '_assets')
-        self.stack.top._children.append(_node)
-        self.stack.top._all_children.extend([_node] + _node._all_children)
+        parent = self.stack.top
+        parent_children = getattr(parent, '_children')
+        parent_all_children = getattr(parent, '_all_children')
+        node_all_children = getattr(_node, '_all_children')
+        parent_children.append(_node)
+        parent_all_children.extend([_node] + node_all_children)
+
+        if not self.with_assets_amount:
+            return
 
+        node_all_assets = getattr(_node, '_all_assets')
+        parent_all_assets = getattr(parent, '_all_assets')
+        _node._assets_amount = len(node_all_assets)
+        parent_all_assets.update(node_all_assets)
+
+    @timeit
     def init(self):
         all_nodes = self.get_all_nodes()
         for node in all_nodes:
-            self.debug("准备: {} 栈顶: {}".format(node.key, self.stack.top.key if self.stack.top else None))
+            self.debug("准备: {} 栈顶: {}".format(
+                node.key, self.stack.top.key if self.stack.top else None)
+            )
             # 入栈之前检查，该节点是不是栈顶节点的子节点
             # 如果不是，则栈顶出栈
             while self.stack.top and not self.stack.top.is_children(node):
@@ -144,27 +169,19 @@ class NodeUtil:
     def debug(self, msg):
         self._debug and logger.debug(msg)
 
-    def set_assets_amount(self):
-        for node in self._nodes.values():
-            node.assets_amount = node._assets_amount
-
-    def set_full_value(self):
-        for node in self._nodes.values():
-            node.full_value = node._full_value
-
     @property
     def nodes(self):
         return list(self._nodes.values())
 
     def get_family_by_key(self, key):
-        tree_nodes = set()
+        family = set()
         node = self.get_node_by_key(key)
         if not node:
             return []
-        tree_nodes.update(node._parents)
-        tree_nodes.add(node)
-        tree_nodes.update(node._all_children)
-        return list(tree_nodes)
+        family.update(getattr(node, '_parents'))
+        family.add(node)
+        family.update(getattr(node, '_all_children'))
+        return list(family)
 
     # 使用给定节点生成一颗树
     # 找到他们的祖先节点
@@ -179,8 +196,8 @@ class NodeUtil:
     def get_some_nodes_family_by_keys(self, keys):
         family = set()
         for key in keys:
-            family.update(self.get_family_by_key(key))
-        return family
+            family.update(set(self.get_family_by_key(key)))
+        return list(family)
 
     def get_some_nodes_family_keys_by_keys(self, keys):
         family = self.get_some_nodes_family_by_keys(keys)
@@ -191,7 +208,7 @@ class NodeUtil:
         node = self.get_node_by_key(key)
         if not node:
             return []
-        parents.update(set(node._parents))
+        parents.update(set(getattr(node, '_parents')))
         if with_self:
             parents.add(node)
         return list(parents)
@@ -203,21 +220,20 @@ class NodeUtil:
         nodes = self.get_nodes_parents_by_key(key, with_self=with_self)
         return [n.key for n in nodes]
 
-    def get_all_children_by_key(self, key, with_self=True):
-        children = set()
+    def get_node_all_children_by_key(self, key, with_self=True):
         node = self.get_node_by_key(key)
         if not node:
             return []
-        children.update(set(node._all_children))
+        children = set(getattr(node, '_all_children'))
         if with_self:
             children.add(node)
         return list(children)
 
     def get_all_children(self, node, with_self=True):
-        return self.get_all_children_by_key(node.key, with_self=with_self)
+        return self.get_node_all_children_by_key(node.key, with_self=with_self)
 
     def get_all_children_keys_by_key(self, key, with_self=True):
-        nodes = self.get_all_children_by_key(key, with_self=with_self)
+        nodes = self.get_node_all_children_by_key(key, with_self=with_self)
         return [n.key for n in nodes]
 
 
@@ -250,7 +266,42 @@ def test_node_tree():
             )
 
 
+class TreeService(Tree):
+    tag_sep = ' / '
 
-
-
-
+    @classmethod
+    @timeit
+    def new(cls):
+        from .models import Node
+        all_nodes = Node.objects.all()
+        tree = cls()
+        tree.create_node(tag='', identifier='')
+        for node in all_nodes:
+            tree.create_node(
+                tag=node.value, identifier=node.key,
+                parent=node.parent_key, data=node
+            )
+        return tree
+
+    def all_children(self, nid, with_self=True):
+        children_ids = self.expand_tree(nid)
+        if not with_self:
+            next(children_ids)
+        return [self[i] for i in children_ids]
+
+    def ancestors(self, nid, with_self=True):
+        ancestor_ids = list(self.rsearch(nid))
+        ancestor_ids.pop()
+        if not with_self:
+            ancestor_ids.pop(0)
+        return [self.get_node(i) for i in ancestor_ids]
+
+    def get_node_full_tag(self, nid):
+        ancestors = self.ancestors(nid)
+        ancestors.reverse()
+        return self.tag_sep.join(n.tag for n in ancestors)
+
+    def get_family(self, nid):
+        ancestors = self.ancestors(nid, with_self=False)
+        children = self.all_children(nid, with_self=False)
+        return ancestors + [self[nid]] + children

apps/jumpserver/settings.py

@@ -304,10 +304,10 @@ LOGGING = {
             'handlers': ['gunicorn_console', 'gunicorn_file'],
             'level': 'INFO',
         },
-        # 'django.db': {
-        #     'handlers': ['console', 'file'],
-        #     'level': 'DEBUG'
-        # }
+        'django.db': {
+            'handlers': ['console', 'file'],
+            'level': 'DEBUG'
+        }
     }
 }
 

apps/perms/forms/asset_permission.py

@@ -41,17 +41,17 @@ class AssetPermissionForm(OrgModelForm):
         users_field = self.fields.get('users')
         users_field.queryset = current_org.get_org_users()
 
-        nodes_field = self.fields['nodes']
-        nodes_field.choices = ((n.id, n.full_value) for n in Node.get_queryset())
-
+        if self.data:
+            return
         # 前端渲染优化, 防止过多资产
-        if not self.data:
-            instance = kwargs.get('instance')
-            assets_field = self.fields['assets']
-            if instance:
-                assets_field.queryset = instance.assets.all()
-            else:
-                assets_field.queryset = Asset.objects.none()
+        assets_field = self.fields['assets']
+        nodes_field = self.fields['nodes']
+        if self.instance:
+            assets_field.queryset = self.instance.assets.all()
+            nodes_field.queryset = self.instance.nodes.all()
+        else:
+            assets_field.queryset = Asset.objects.none()
+            nodes_field.queryset = Node.objects.none()
 
     class Meta:
         model = AssetPermission
@@ -69,7 +69,7 @@ class AssetPermissionForm(OrgModelForm):
                 attrs={'class': 'select2', 'data-placeholder': _("Asset")}
             ),
             'nodes': forms.SelectMultiple(
-                attrs={'class': 'select2', 'data-placeholder': _("Node")}
+                attrs={'class': 'nodes-select2', 'data-placeholder': _("Node")}
             ),
             'system_users': forms.SelectMultiple(
                 attrs={'class': 'select2', 'data-placeholder': _('System user')}

apps/perms/templates/perms/asset_permission_create_update.html

@@ -115,6 +115,27 @@ $(document).ready(function () {
     $('.select2').select2({
         closeOnSelect: false
     });
+    $(".nodes-select2").select2({
+        closeOnSelect: false,
+        ajax: {
+            url: '{% url "api-assets:node-list" %}',
+            data: function (params) {
+                var page = params.page || 1;
+                var query = {
+                   search: params.term,
+                   offset: (page -1) * 10,
+                   limit: 10
+                };
+                return query
+            },
+            processResults: function (data) {
+                var results = $.map(data.results, function (v, i) {
+                    return {id: v.id, text: v.full_value}
+                });
+                return {results: results, pagination: {"more": true}}
+            }
+        },
+    });
     $('#date_start').daterangepicker(dateOptions);
     $('#date_expired').daterangepicker(dateOptions);
     $("#id_assets").parent().find(".select2-selection").on('click', function (e) {