[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']
+        if self.instance:
             nodes_field.choices = ((n.id, n.full_value) for n in
-                               Node.get_queryset())
+                                   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/models/node.py

@@ -2,6 +2,7 @@
 #
 import uuid
 import re
+import time
 
 from django.db import models, transaction
 from django.db.models import Q
@@ -13,6 +14,7 @@ from orgs.mixins import OrgModelMixin, OrgManager
 from orgs.utils import set_current_org, get_current_org
 from orgs.models import Organization
 
+
 __all__ = ['Node']
 
 
@@ -22,76 +24,102 @@ class NodeQuerySet(models.QuerySet):
 
 
 class FamilyMixin:
-    _parents = None
-    _children = None
-    _all_children = None
+    __parents = None
+    __children = None
+    __all_children = None
     is_node = True
+    time_tree_updated = None
+    time_tree_updated_cache_key = 'NODE_TREE_CREATED_AT'
+    tree_cache_time = 3600
+    _tree_service = None
+
+    @classmethod
+    def tree(cls):
+        from ..utils import TreeService
+        cache_updated_time = cls.get_cache_time()
+        if not cls.time_tree_updated or cache_updated_time != cls.time_tree_updated:
+            t = TreeService.new()
+            cls.update_cache_tree(t)
+            return t
+        return cls._tree_service
+
+    @classmethod
+    def get_cache_time(cls):
+        return cache.get(cls.time_tree_updated_cache_key)
+
+    @classmethod
+    def update_cache_tree(cls, t):
+        cls._tree_service = t
+        now = time.time()
+        cls.time_tree_updated = now
+        cache.set(cls.time_tree_updated_cache_key, now, cls.tree_cache_time)
+
+    @classmethod
+    def expire_cache_tree(cls):
+        cache.delete(cls.time_tree_updated_cache_key)
+
+    @classmethod
+    def refresh_tree(cls):
+        cls.expire_cache_tree()
 
     @property
-    def children(self):
-        if self._children:
-            return self._children
-        pattern = r'^{0}:[0-9]+$'.format(self.key)
-        return Node.objects.filter(key__regex=pattern)
+    def _tree(self):
+        return self.__class__.tree()
 
-    @children.setter
-    def children(self, value):
-        self._children = value
+    @property
+    def children(self):
+        return self.get_children(with_self=False)
 
     @property
     def all_children(self):
-        if self._all_children:
-            return self._all_children
-        pattern = r'^{0}:'.format(self.key)
-        return Node.objects.filter(
-            key__regex=pattern
-        )
+        return self.get_all_children(with_self=False)
+
+    def _get_children(self, with_self=False):
+        pattern = r'^{0}:[0-9]+$'.format(self.key)
+        if with_self:
+            pattern += r'^{0}$'
+        return Node.objects.filter(key__regex=pattern)
 
     def get_children(self, with_self=False):
-        children = list(self.children)
+        _children = self._tree.children(self.key)
+        children = [n.data for n in _children]
         if with_self:
             children.append(self)
         return children
 
-    def get_all_children(self, with_self=False):
-        children = self.all_children
+    def _get_all_children(self, with_self=False):
+        pattern = r'^{0}:'.format(self.key)
         if with_self:
-            children = list(children)
-            children.append(self)
+            pattern += r'|^{0}$'
+        children = Node.objects.filter(key__regex=pattern)
+        return children
+
+    def get_all_children(self, with_self=False):
+        _children = self._tree.all_children(self.key, with_self=with_self)
+        children = [n.data for n in _children]
         return children
 
     @property
     def parents(self):
-        if self._parents:
-            return self._parents
-        ancestor_keys = self.get_ancestor_keys()
-        ancestor = Node.objects.filter(
-            key__in=ancestor_keys
-        ).order_by('key')
-        return ancestor
-
-    @parents.setter
-    def parents(self, value):
-        self._parents = value
+        return self.get_ancestor(with_self=False)
 
-    def get_ancestor(self, with_self=False):
+    def _get_ancestor(self, with_self=False):
         parents = self.parents
         if with_self:
             parents = list(parents)
             parents.append(self)
         return parents
 
+    def get_ancestor(self, with_self=False):
+        _ancestor = self._tree.ancestors(self.key, with_self=with_self)
+        ancestor = [n.data for n in _ancestor]
+        return ancestor
+
     @property
     def parent(self):
-        if self._parents:
-            return self._parents[0]
         if self.is_root():
             return self
-        try:
-            parent = Node.objects.get(key=self.parent_key)
-            return parent
-        except Node.DoesNotExist:
-            return Node.root()
+        return self._tree.parent(self.key).data
 
     @parent.setter
     def parent(self, parent):
@@ -107,7 +135,7 @@ class FamilyMixin:
                 child.save()
             self.save()
 
-    def get_sibling(self, with_self=False):
+    def _get_siblings(self, with_self=False):
         key = ':'.join(self.key.split(':')[:-1])
         pattern = r'^{}:[0-9]+$'.format(key)
         sibling = Node.objects.filter(
@@ -117,6 +145,13 @@ class FamilyMixin:
             sibling = sibling.exclude(key=self.key)
         return sibling
 
+    def get_siblings(self, with_self=False):
+        _siblings = self._tree.siblings(self.key)
+        siblings = [n.data for n in _siblings]
+        if with_self:
+            siblings.append(self)
+        return siblings
+
     def get_family(self):
         ancestor = self.get_ancestor()
         children = self.get_all_children()
@@ -133,12 +168,10 @@ class FamilyMixin:
         return parent_keys
 
     def is_children(self, other):
-        pattern = re.compile(r'^{0}:[0-9]+$'.format(self.key))
-        return pattern.match(other.key)
+        return other.key.startswith(self.key + ':')
 
     def is_parent(self, other):
-        pattern = re.compile(r'^{0}:[0-9]+$'.format(other.key))
-        return pattern.match(self.key)
+        return other.is_children(self)
 
     @property
     def parent_key(self):
@@ -164,40 +197,18 @@ class FullValueMixin:
 
     @property
     def full_value(self):
-        if self._full_value:
-            return self._full_value
-        key = self._full_value_cache_key.format(self.key)
-        cached = cache.get(key)
-        if cached:
-            return cached
         if self.is_root():
             return self.value
-        parent_full_value = self.parent.full_value
-        value = parent_full_value + ' / ' + self.value
-        self.full_value = value
+        value = self._tree.get_node_full_tag(self.key)
         return value
 
-    @full_value.setter
-    def full_value(self, value):
-        self._full_value = value
-        key = self._full_value_cache_key.format(self.key)
-        cache.set(key, value, 3600*24)
-
-    def expire_full_value(self):
-        key = self._full_value_cache_key.format(self.key)
-        cache.delete_pattern(key+'*')
 
-    @classmethod
-    def expire_nodes_full_value(cls, nodes=None):
-        key = cls._full_value_cache_key.format('*')
-        cache.delete_pattern(key+'*')
-
-
-class AssetsAmountMixin:
+class NodeAssetsMixin:
     _assets_amount_cache_key = '_NODE_ASSETS_AMOUNT_{}'
     _assets_amount = None
     key = ''
     cache_time = 3600 * 24 * 7
+    id = None
 
     @property
     def assets_amount(self):
@@ -212,14 +223,9 @@ class AssetsAmountMixin:
         if cached is not None:
             return cached
         assets_amount = self.get_all_assets().only('id').count()
-        self.assets_amount = assets_amount
-        return assets_amount
-
-    @assets_amount.setter
-    def assets_amount(self, value):
-        self._assets_amount = value
         cache_key = self._assets_amount_cache_key.format(self.key)
-        cache.set(cache_key, value, self.cache_time)
+        cache.set(cache_key, assets_amount, self.cache_time)
+        return assets_amount
 
     def expire_assets_amount(self):
         ancestor_keys = self.get_ancestor_keys(with_self=True)
@@ -229,18 +235,38 @@ class AssetsAmountMixin:
 
     @classmethod
     def expire_nodes_assets_amount(cls, nodes=None):
-        key = cls._assets_amount_cache_key.format('*')
-        cache.delete_pattern(key)
+        for node in nodes:
+            node.expire_assets_amount()
 
     @classmethod
-    def refresh_nodes(cls):
-        from ..utils import NodeUtil
-        util = NodeUtil(with_assets_amount=True)
-        util.set_assets_amount()
-        util.set_full_value()
+    def refresh_assets_amount(cls):
+        cache_key = cls._assets_amount_cache_key.format('*')
+        cache.delete_pattern(cache_key)
+
+    def get_all_assets(self):
+        from .asset import Asset
+        if self.is_root():
+            return Asset.objects.filter(org_id=self.org_id)
+        children = self.get_all_children(with_self=True)
+        assets = Asset.objects.filter(nodes__in=children).distinct()
+        return assets
+
+    def get_assets(self):
+        from .asset import Asset
+        if self.is_default_node():
+            assets = Asset.objects.filter(Q(nodes__id=self.id) | Q(nodes__isnull=True))
+        else:
+            assets = Asset.objects.filter(nodes__id=self.id)
+        return assets.distinct()
 
+    def get_valid_assets(self):
+        return self.get_assets().valid()
 
-class Node(OrgModelMixin, FamilyMixin, FullValueMixin, AssetsAmountMixin):
+    def get_all_valid_assets(self):
+        return self.get_all_assets().valid()
+
+
+class Node(OrgModelMixin, FamilyMixin, FullValueMixin, NodeAssetsMixin):
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
     key = models.CharField(unique=True, max_length=64, verbose_name=_("Key"))  # '1:1:1:1'
     value = models.CharField(max_length=128, verbose_name=_("Value"))
@@ -256,7 +282,7 @@ class Node(OrgModelMixin, FamilyMixin, FullValueMixin, AssetsAmountMixin):
         ordering = ['key']
 
     def __str__(self):
-        return self.full_value
+        return self.value
 
     def __eq__(self, other):
         if not other:
@@ -316,31 +342,10 @@ class Node(OrgModelMixin, FamilyMixin, FullValueMixin, AssetsAmountMixin):
             child = self.__class__.objects.create(id=_id, key=child_key, value=value)
             return child
 
-    def get_assets(self):
-        from .asset import Asset
-        if self.is_default_node():
-            assets = Asset.objects.filter(Q(nodes__id=self.id) | Q(nodes__isnull=True))
-        else:
-            assets = Asset.objects.filter(nodes__id=self.id)
-        return assets.distinct()
-
-    def get_valid_assets(self):
-        return self.get_assets().valid()
-
-    def get_all_assets(self):
-        from .asset import Asset
-        pattern = r'^{0}$|^{0}:'.format(self.key)
-        args = []
-        kwargs = {}
-        if self.is_root():
-            args.append(Q(nodes__key__regex=pattern) | Q(nodes=None))
-        else:
-            kwargs['nodes__key__regex'] = pattern
-        assets = Asset.objects.filter(*args, **kwargs).distinct()
-        return assets
-
-    def get_all_valid_assets(self):
-        return self.get_all_assets().valid()
+    @classmethod
+    def refresh_nodes(cls):
+        cls.refresh_assets_amount()
+        cls.refresh_tree()
 
     def is_default_node(self):
         return self.is_root() and self.key == '1'
@@ -384,6 +389,7 @@ class Node(OrgModelMixin, FamilyMixin, FullValueMixin, AssetsAmountMixin):
     def as_tree_node(self):
         from common.tree import TreeNode
         name = '{} ({})'.format(self.value, self.assets_amount)
+        # name = self.value
         data = {
             'id': self.key,
             'name': name,
@@ -422,7 +428,14 @@ class Node(OrgModelMixin, FamilyMixin, FullValueMixin, AssetsAmountMixin):
         org = get_current_org()
         if not org or not org.is_real():
             Organization.default().change_to()
+        i = 0
+        while i < count:
+            nodes = list(cls.objects.all())
+            if count > 100:
+                length = 100
+            else:
+                length = count
 
-        for i in range(count):
-            node = random.choice(cls.objects.all())
+            for i in range(length):
+                node = random.choice(nodes)
                 node.create_child('Node {}'.format(i))

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()
+        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) {