[Update] 优化asset api, 获取时增加prefetch

apps/assets/api/asset.py

@@ -40,7 +40,9 @@ class AssetViewSet(IDInFilterMixin, LabelFilter, BulkModelViewSet):
     permission_classes = (IsSuperUserOrAppUser,)
 
     def get_queryset(self):
-        queryset = super().get_queryset()
+        queryset = super().get_queryset()\
+            .prefetch_related('labels', 'nodes')\
+            .select_related('admin_user')
         admin_user_id = self.request.query_params.get('admin_user_id')
         node_id = self.request.query_params.get("node_id")
         show_current_asset = self.request.query_params.get("show_current_asset")
@@ -66,7 +68,6 @@ class AssetViewSet(IDInFilterMixin, LabelFilter, BulkModelViewSet):
                 queryset = queryset.filter(
                     nodes__key__regex='^{}(:[0-9]+)*$'.format(node.key),
                 ).distinct()
-
         return queryset
 
 

apps/assets/models/asset.py

@@ -59,42 +59,70 @@ class Asset(models.Model):
         ('Other', 'Other'),
     )
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
-    ip = models.GenericIPAddressField(max_length=32, verbose_name=_('IP'), db_index=True)
-    hostname = models.CharField(max_length=128, unique=True, verbose_name=_('Hostname'))
+    ip = models.GenericIPAddressField(max_length=32, verbose_name=_('IP'),
+                                      db_index=True)
+    hostname = models.CharField(max_length=128, unique=True,
+                                verbose_name=_('Hostname'))
     port = models.IntegerField(default=22, verbose_name=_('Port'))
-    platform = models.CharField(max_length=128, choices=PLATFORM_CHOICES, default='Linux', verbose_name=_('Platform'))
-    domain = models.ForeignKey("assets.Domain", null=True, blank=True, related_name='assets', verbose_name=_("Domain"), on_delete=models.SET_NULL)
-    nodes = models.ManyToManyField('assets.Node', default=default_node, related_name='assets', verbose_name=_("Nodes"))
+    platform = models.CharField(max_length=128, choices=PLATFORM_CHOICES,
+                                default='Linux', verbose_name=_('Platform'))
+    domain = models.ForeignKey("assets.Domain", null=True, blank=True,
+                               related_name='assets', verbose_name=_("Domain"),
+                               on_delete=models.SET_NULL)
+    nodes = models.ManyToManyField('assets.Node', default=default_node,
+                                   related_name='assets',
+                                   verbose_name=_("Nodes"))
     is_active = models.BooleanField(default=True, verbose_name=_('Is active'))
 
     # Auth
-    admin_user = models.ForeignKey('assets.AdminUser', on_delete=models.PROTECT, null=True, verbose_name=_("Admin user"))
+    admin_user = models.ForeignKey('assets.AdminUser', on_delete=models.PROTECT,
+                                   null=True, verbose_name=_("Admin user"))
 
     # Some information
-    public_ip = models.GenericIPAddressField(max_length=32, blank=True, null=True, verbose_name=_('Public IP'))
-    number = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Asset number'))
+    public_ip = models.GenericIPAddressField(max_length=32, blank=True,
+                                             null=True,
+                                             verbose_name=_('Public IP'))
+    number = models.CharField(max_length=32, null=True, blank=True,
+                              verbose_name=_('Asset number'))
 
     # Collect
-    vendor = models.CharField(max_length=64, null=True, blank=True, verbose_name=_('Vendor'))
-    model = models.CharField(max_length=54, null=True, blank=True, verbose_name=_('Model'))
-    sn = models.CharField(max_length=128, null=True, blank=True, verbose_name=_('Serial number'))
-
-    cpu_model = models.CharField(max_length=64, null=True, blank=True, verbose_name=_('CPU model'))
+    vendor = models.CharField(max_length=64, null=True, blank=True,
+                              verbose_name=_('Vendor'))
+    model = models.CharField(max_length=54, null=True, blank=True,
+                             verbose_name=_('Model'))
+    sn = models.CharField(max_length=128, null=True, blank=True,
+                          verbose_name=_('Serial number'))
+
+    cpu_model = models.CharField(max_length=64, null=True, blank=True,
+                                 verbose_name=_('CPU model'))
     cpu_count = models.IntegerField(null=True, verbose_name=_('CPU count'))
     cpu_cores = models.IntegerField(null=True, verbose_name=_('CPU cores'))
-    memory = models.CharField(max_length=64, null=True, blank=True, verbose_name=_('Memory'))
-    disk_total = models.CharField(max_length=1024, null=True, blank=True, verbose_name=_('Disk total'))
-    disk_info = models.CharField(max_length=1024, null=True, blank=True, verbose_name=_('Disk info'))
-
-    os = models.CharField(max_length=128, null=True, blank=True, verbose_name=_('OS'))
-    os_version = models.CharField(max_length=16, null=True, blank=True, verbose_name=_('OS version'))
-    os_arch = models.CharField(max_length=16, blank=True, null=True, verbose_name=_('OS arch'))
-    hostname_raw = models.CharField(max_length=128, blank=True, null=True, verbose_name=_('Hostname raw'))
-
-    labels = models.ManyToManyField('assets.Label', blank=True, related_name='assets', verbose_name=_("Labels"))
-    created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by'))
-    date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created'))
-    comment = models.TextField(max_length=128, default='', blank=True, verbose_name=_('Comment'))
+    memory = models.CharField(max_length=64, null=True, blank=True,
+                              verbose_name=_('Memory'))
+    disk_total = models.CharField(max_length=1024, null=True, blank=True,
+                                  verbose_name=_('Disk total'))
+    disk_info = models.CharField(max_length=1024, null=True, blank=True,
+                                 verbose_name=_('Disk info'))
+
+    os = models.CharField(max_length=128, null=True, blank=True,
+                          verbose_name=_('OS'))
+    os_version = models.CharField(max_length=16, null=True, blank=True,
+                                  verbose_name=_('OS version'))
+    os_arch = models.CharField(max_length=16, blank=True, null=True,
+                               verbose_name=_('OS arch'))
+    hostname_raw = models.CharField(max_length=128, blank=True, null=True,
+                                    verbose_name=_('Hostname raw'))
+
+    labels = models.ManyToManyField('assets.Label', blank=True,
+                                    related_name='assets',
+                                    verbose_name=_("Labels"))
+    created_by = models.CharField(max_length=32, null=True, blank=True,
+                                  verbose_name=_('Created by'))
+    date_created = models.DateTimeField(auto_now_add=True, null=True,
+                                        blank=True,
+                                        verbose_name=_('Date created'))
+    comment = models.TextField(max_length=128, default='', blank=True,
+                               verbose_name=_('Comment'))
 
     objects = AssetManager()
 
@@ -121,6 +149,22 @@ class Asset(models.Model):
         nodes = self.nodes.all() or [Node.root()]
         return nodes
 
+    @property
+    def nodes_cache_key(self):
+        key = "NODES_OF_{}".format(str(self.id))
+        return key
+
+    def get_nodes_or_cache(self):
+        cached = cache.get(self.nodes_cache_key)
+        if cached is not None:
+            return cached
+        nodes = list(self.get_nodes())
+        cache.set(self.nodes_cache_key, nodes, 3600)
+        return nodes
+
+    def expire_nodes_cache(self):
+        cache.delete(self.nodes_cache_key)
+
     @property
     def hardware_info(self):
         if self.cpu_count:

apps/assets/models/node.py

@@ -13,9 +13,6 @@ __all__ = ['Node']
 class Node(models.Model):
     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, unique=True, verbose_name=_("Value")
-    # )
     value = models.CharField(max_length=128, verbose_name=_("Value"))
     child_mark = models.IntegerField(default=0)
     date_create = models.DateTimeField(auto_now_add=True)
@@ -31,10 +28,11 @@ class Node(models.Model):
 
     @property
     def full_value(self):
-        if self == self.__class__.root():
+        ancestor = [a.value for a in self.ancestor]
+        if self.is_root():
             return self.value
-        else:
-            return '{} / {}'.format(self.parent.full_value, self.value)
+        ancestor.append(self.value)
+        return ' / '.join(ancestor)
 
     @property
     def level(self):
@@ -118,7 +116,6 @@ class Node(models.Model):
     def parent(self):
         if self.key == "0" or not self.key.startswith("0"):
             return self.__class__.root()
-
         parent_key = ":".join(self.key.split(":")[:-1])
         try:
             parent = self.__class__.objects.get(key=parent_key)
@@ -132,16 +129,17 @@ class Node(models.Model):
 
     @property
     def ancestor(self):
-        _key = self.key.split(':')
-        ancestor_keys = []
-
         if self.is_root():
-            return [self.__class__.root()]
-
-        for i in range(len(_key)-1):
-            _key.pop()
-            ancestor_keys.append(':'.join(_key))
-        return self.__class__.objects.filter(key__in=ancestor_keys)
+            ancestor = self.__class__.objects.filter(key='0')
+        else:
+            _key = self.key.split(':')
+            ancestor_keys = []
+            for i in range(len(_key)-1):
+                _key.pop()
+                ancestor_keys.append(':'.join(_key))
+            ancestor = self.__class__.objects.filter(key__in=ancestor_keys)
+        ancestor = list(ancestor)
+        return ancestor
 
     @property
     def ancestor_with_self(self):

apps/assets/serializers/asset.py

@@ -12,34 +12,11 @@ __all__ = [
 ]
 
 
-class NodeTMPSerializer(serializers.ModelSerializer):
-    parent = serializers.SerializerMethodField()
-    assets_amount = serializers.SerializerMethodField()
-
-    class Meta:
-        model = Node
-        fields = ['id', 'key', 'value', 'parent', 'assets_amount', 'is_node']
-        list_serializer_class = BulkListSerializer
-
-    @staticmethod
-    def get_parent(obj):
-        return obj.parent.id
-
-    @staticmethod
-    def get_assets_amount(obj):
-        return obj.get_all_assets().count()
-
-    def get_fields(self):
-        fields = super().get_fields()
-        field = fields["key"]
-        field.required = False
-        return fields
-
-
 class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer):
     """
     资产的数据结构
     """
+    nodes = serializers.SerializerMethodField()
 
     class Meta:
         model = Asset
@@ -54,6 +31,10 @@ class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer):
         ])
         return fields
 
+    @staticmethod
+    def get_nodes(obj):
+        return [n.id for n in obj.get_nodes_or_cache()]
+
 
 class AssetGrantedSerializer(serializers.ModelSerializer):
     """

apps/assets/signals_handler.py

@@ -63,22 +63,31 @@ def on_system_user_assets_change(sender, instance=None, **kwargs):
 
 @receiver(m2m_changed, sender=Asset.nodes.through)
 def on_asset_node_changed(sender, instance=None, **kwargs):
-    if isinstance(instance, Asset) and kwargs['action'] == 'post_add':
-        logger.debug("Asset node change signal received")
-        nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
-        system_users_assets = defaultdict(set)
-        system_users = SystemUser.objects.filter(nodes__in=nodes)
-        for system_user in system_users:
-            system_users_assets[system_user].update({instance})
-        for system_user, assets in system_users_assets.items():
-            system_user.assets.add(*tuple(assets))
+    if isinstance(instance, Asset):
+        instance.expire_nodes_cache()
+        if kwargs['action'] == 'post_add':
+            logger.debug("Asset node change signal received")
+            nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
+            system_users_assets = defaultdict(set)
+            system_users = SystemUser.objects.filter(nodes__in=nodes)
+            # 清理节点缓存
+            for system_user in system_users:
+                system_users_assets[system_user].update({instance})
+            for system_user, assets in system_users_assets.items():
+                system_user.assets.add(*tuple(assets))
 
 
 @receiver(m2m_changed, sender=Asset.nodes.through)
 def on_node_assets_changed(sender, instance=None, **kwargs):
-    if isinstance(instance, Node) and kwargs['action'] == 'post_add':
-        logger.debug("Node assets change signal received")
+    if isinstance(instance, Node):
         assets = kwargs['model'].objects.filter(pk__in=kwargs['pk_set'])
-        system_users = SystemUser.objects.filter(nodes=instance)
-        for system_user in system_users:
-            system_user.assets.add(*tuple(assets))
+        # 清理资产节点缓存
+        for asset in assets:
+            asset.expire_nodes_cache()
+
+        if kwargs['action'] == 'post_add':
+            logger.debug("Node assets change signal received")
+            # 重新关联系统用户和资产的关系
+            system_users = SystemUser.objects.filter(nodes=instance)
+            for system_user in system_users:
+                system_user.assets.add(*tuple(assets))

apps/perms/utils.py

@@ -15,7 +15,7 @@ logger = get_logger(__file__)
 
 class Tree:
     def __init__(self):
-        self.__all_nodes = list(Node.objects.all())
+        self.__all_nodes = list(Node.objects.all().prefetch_related('assets'))
         self.__node_asset_map = defaultdict(set)
         self.nodes = defaultdict(dict)
         self.root = Node.root()
@@ -134,7 +134,7 @@ class AssetPermissionUtil:
         _assets = cls.get_user_group_assets(group)
         tree = Tree()
         for asset, _system_users in _assets.items():
-            _nodes = asset.get_nodes()
+            _nodes = asset.get_nodes_or_cache()
             tree.add_nodes(_nodes)
             for node in _nodes:
                 tree.nodes[node][asset].update(_system_users)

jms

@@ -123,6 +123,7 @@ def start_gunicorn():
         'gunicorn', 'jumpserver.wsgi',
         '-b', bind,
         '-w', str(WORKERS),
+        '-k', 'eventlet',
         '--access-logformat', log_format,
         '-p', pid_file,
     ]