From 9ed822bb3e4e27ff1d786f4d0be1256028bb1b9f Mon Sep 17 00:00:00 2001 From: ibuler Date: Wed, 26 Mar 2025 18:25:55 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E4=BC=98=E5=8C=96=E8=8E=B7=E5=8F=96=20?= =?UTF-8?q?labels?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/accounts/serializers/account/account.py | 2 +- apps/assets/api/platform.py | 2 +- apps/assets/serializers/asset/common.py | 4 --- apps/assets/serializers/domain.py | 5 ---- apps/common/api/mixin.py | 26 ++++++++++++++------ apps/common/serializers/mixin.py | 15 ++++++++--- apps/perms/serializers/permission.py | 3 +-- apps/perms/serializers/user_permission.py | 2 +- apps/users/serializers/group.py | 3 +-- 9 files changed, 36 insertions(+), 26 deletions(-) diff --git a/apps/accounts/serializers/account/account.py b/apps/accounts/serializers/account/account.py index 11251f3a7..78066c846 100644 --- a/apps/accounts/serializers/account/account.py +++ b/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 diff --git a/apps/assets/api/platform.py b/apps/assets/api/platform.py index bb23abf2d..2bb9f84b2 100644 --- a/apps/assets/api/platform.py +++ b/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 diff --git a/apps/assets/serializers/asset/common.py b/apps/assets/serializers/asset/common.py index 05a03fb0e..5f407267e 100644 --- a/apps/assets/serializers/asset/common.py +++ b/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 diff --git a/apps/assets/serializers/domain.py b/apps/assets/serializers/domain.py index 5bbb1c8bc..ccd970edf 100644 --- a/apps/assets/serializers/domain.py +++ b/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): diff --git a/apps/common/api/mixin.py b/apps/common/api/mixin.py index 830e7b571..ea51dd68e 100644 --- a/apps/common/api/mixin.py +++ b/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 diff --git a/apps/common/serializers/mixin.py b/apps/common/serializers/mixin.py index c3abbcd3d..01eef8182 100644 --- a/apps/common/serializers/mixin.py +++ b/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) diff --git a/apps/perms/serializers/permission.py b/apps/perms/serializers/permission.py index 821a260e7..82914b35d 100644 --- a/apps/perms/serializers/permission.py +++ b/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), diff --git a/apps/perms/serializers/user_permission.py b/apps/perms/serializers/user_permission.py index 1d108bb92..d00e16d18 100644 --- a/apps/perms/serializers/user_permission.py +++ b/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")) diff --git a/apps/users/serializers/group.py b/apps/users/serializers/group.py index 6e422414c..3b112f4cc 100644 --- a/apps/users/serializers/group.py +++ b/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