perf: 优化获取 labels

apps/accounts/serializers/account/account.py

@@ -258,7 +258,7 @@ class AccountSerializer(AccountCreateUpdateSerializerMixin, BaseAccountSerialize
         queryset = queryset.prefetch_related(
             'asset', 'asset__platform',
             'asset__platform__automation'
-        ).prefetch_related('labels', 'labels__label')
+        )
         return queryset
 
 

apps/assets/api/platform.py

@@ -52,7 +52,7 @@ class AssetPlatformViewSet(JMSModelViewSet):
         queryset = (
             super().get_queryset()
             .annotate(assets_amount=Coalesce(Subquery(asset_count_subquery), Value(0)))
-            .prefetch_related('protocols', 'automation', 'labels', 'labels__label')
+            .prefetch_related('protocols', 'automation')
         )
         queryset = queryset.filter(type__in=AllTypes.get_types_values())
         return queryset

apps/assets/serializers/asset/common.py

@@ -231,10 +231,6 @@ class AssetSerializer(BulkOrgResourceModelSerializer, ResourceLabelsMixin, Writa
             .annotate(category=F("platform__category")) \
             .annotate(type=F("platform__type")) \
             .annotate(accounts_amount=Count('accounts'))
-        if queryset.model is Asset:
-            queryset = queryset.prefetch_related('labels__label', 'labels')
-        else:
-            queryset = queryset.prefetch_related('asset_ptr__labels__label', 'asset_ptr__labels')
         return queryset
 
     @staticmethod

apps/assets/serializers/domain.py

@@ -55,11 +55,6 @@ class DomainSerializer(ResourceLabelsMixin, BulkOrgResourceModelSerializer):
         validated_data['assets'] = assets + gateways
         return super().update(instance, validated_data)
 
-    @classmethod
-    def setup_eager_loading(cls, queryset):
-        queryset = queryset.prefetch_related('labels', 'labels__label')
-        return queryset
-
 
 class DomainListSerializer(DomainSerializer):
     class Meta(DomainSerializer.Meta):

apps/common/api/mixin.py

@@ -103,22 +103,34 @@ class QuerySetMixin:
         queryset = self.setup_eager_loading(queryset)
         return queryset
 
-    # Todo: 未来考虑自定义 pagination
-    def setup_eager_loading(self, queryset):
-        if self.request.query_params.get('format') not in ['csv', 'xlsx']:
+    def setup_eager_loading(self, queryset, is_paginated=False):
+        is_export_request = self.request.query_params.get('format') in ['csv', 'xlsx']
+        # 不分页不走一般这个，是因为会消耗多余的 sql 查询, 不如分页的时候查询一次
+        if not is_export_request and not is_paginated:
             return queryset
+
         serializer_class = self.get_serializer_class()
-        if not serializer_class or not hasattr(serializer_class, 'setup_eager_loading'):
+        if not serializer_class:
             return queryset
-        return serializer_class.setup_eager_loading(queryset)
+
+        if hasattr(serializer_class, 'setup_eager_loading'):
+            queryset = serializer_class.setup_eager_loading(queryset)
+
+        if hasattr(serializer_class, 'setup_eager_labels'):
+            queryset = serializer_class.setup_eager_labels(queryset)
+        return queryset
 
     def paginate_queryset(self, queryset):
         page = super().paginate_queryset(queryset)
+        model = getattr(queryset, 'model', None)
+        if not model:
+            return page
+
         serializer_class = self.get_serializer_class()
-        if page and serializer_class and hasattr(serializer_class, 'setup_eager_loading'):
+        if page and serializer_class:
             ids = [str(obj.id) for obj in page]
             page = self.get_queryset().filter(id__in=ids)
-            page = serializer_class.setup_eager_loading(page)
+            page = self.setup_eager_loading(page, is_paginated=True)
             page_mapper = {str(obj.id): obj for obj in page}
             page = [page_mapper.get(_id) for _id in ids if _id in page_mapper]
         return page

apps/common/serializers/mixin.py

@@ -7,7 +7,7 @@ else:
     from collections import Iterable
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
-from django.db.models import NOT_PROVIDED
+from django.db.models import NOT_PROVIDED, OneToOneField
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
 from rest_framework.exceptions import ValidationError
@@ -466,5 +466,14 @@ class ResourceLabelsMixin(serializers.Serializer):
         return instance
 
     @classmethod
-    def setup_eager_loading(cls, queryset):
-        return queryset.prefetch_related("labels")
+    def setup_eager_labels(cls, queryset):
+        if not hasattr(queryset, 'model'):
+            return queryset
+
+        fields = ['labels', 'labels__label']
+        model = queryset.model
+        pk_field = model._meta.pk
+
+        if isinstance(pk_field, OneToOneField):
+            fields = ['{}__{}'.format(pk_field.name, f) for f in fields]
+        return queryset.prefetch_related(*fields)

apps/perms/serializers/permission.py

@@ -155,7 +155,7 @@ class AssetPermissionSerializer(ResourceLabelsMixin, BulkOrgResourceModelSeriali
         """Perform necessary eager loading of data."""
         queryset = queryset.prefetch_related(
             "users", "user_groups", "assets", "nodes",
-        ).prefetch_related('labels', 'labels__label')
+        )
         return queryset
 
     @staticmethod
@@ -212,7 +212,6 @@ class AssetPermissionListSerializer(AssetPermissionSerializer):
     def setup_eager_loading(cls, queryset):
         """Perform necessary eager loading of data."""
         queryset = queryset \
-            .prefetch_related('labels', 'labels__label') \
             .annotate(users_amount=Count("users", distinct=True),
                       user_groups_amount=Count("user_groups", distinct=True),
                       assets_amount=Count("assets", distinct=True),

apps/perms/serializers/user_permission.py

@@ -40,7 +40,7 @@ class AssetPermedSerializer(OrgResourceModelSerializerMixin, ResourceLabelsMixin
     @classmethod
     def setup_eager_loading(cls, queryset):
         """ Perform necessary eager loading of data. """
-        queryset = queryset.prefetch_related('domain', 'nodes', 'labels') \
+        queryset = queryset.prefetch_related('domain', 'nodes') \
             .prefetch_related('platform') \
             .annotate(category=F("platform__category")) \
             .annotate(type=F("platform__type"))

apps/users/serializers/group.py

@@ -45,8 +45,7 @@ class UserGroupSerializer(ResourceLabelsMixin, BulkOrgResourceModelSerializer):
     @classmethod
     def setup_eager_loading(cls, queryset):
         """ Perform necessary eager loading of data. """
-        queryset = queryset.prefetch_related('labels', 'labels__label') \
-            .annotate(users_amount=Count('users', distinct=True, filter=Q(users__is_service_account=False)))
+        queryset = queryset.annotate(users_amount=Count('users', distinct=True, filter=Q(users__is_service_account=False)))
         return queryset