From 67f334131034312ab4831fab74a191def716eb41 Mon Sep 17 00:00:00 2001 From: fit2bot <68588906+fit2bot@users.noreply.github.com> Date: Wed, 16 Apr 2025 13:48:12 +0800 Subject: [PATCH] perf: change db prefetch (#15215) --- apps/assets/api/asset/asset.py | 5 +++ .../migrations/0016_directory_service.py | 13 +++----- .../0018_asset_directory_services.py | 22 ------------- .../migrations/0019_remove_platform_ds.py | 17 ---------- apps/assets/models/asset/common.py | 31 +++++++++++++++++-- apps/assets/serializers/asset/common.py | 2 +- apps/common/api/mixin.py | 7 +++-- 7 files changed, 44 insertions(+), 53 deletions(-) delete mode 100644 apps/assets/migrations/0018_asset_directory_services.py delete mode 100644 apps/assets/migrations/0019_remove_platform_ds.py diff --git a/apps/assets/api/asset/asset.py b/apps/assets/api/asset/asset.py index 441650205..ff542e1aa 100644 --- a/apps/assets/api/asset/asset.py +++ b/apps/assets/api/asset/asset.py @@ -144,6 +144,11 @@ class BaseAssetViewSet(OrgBulkModelViewSet): return retrieve_cls return cls + def paginate_queryset(self, queryset): + page = super().paginate_queryset(queryset) + page = Asset.compute_accounts_amount(page) + return page + def create(self, request, *args, **kwargs): if request.path.find('/api/v1/assets/assets/') > -1: error = _('Cannot create asset directly, you should create a host or other') diff --git a/apps/assets/migrations/0016_directory_service.py b/apps/assets/migrations/0016_directory_service.py index f62c593e5..39de22fff 100644 --- a/apps/assets/migrations/0016_directory_service.py +++ b/apps/assets/migrations/0016_directory_service.py @@ -46,15 +46,12 @@ class Migration(migrations.Migration): field=models.BooleanField(default=False, verbose_name="DS enabled"), ), migrations.AddField( - model_name="platform", - name="ds", - field=models.ForeignKey( - blank=True, - null=True, - on_delete=django.db.models.deletion.SET_NULL, - related_name="ds_platforms", + model_name="asset", + name="directory_services", + field=models.ManyToManyField( + related_name="assets", to="assets.directoryservice", - verbose_name="Directory service", + verbose_name="Directory services", ), ), ] diff --git a/apps/assets/migrations/0018_asset_directory_services.py b/apps/assets/migrations/0018_asset_directory_services.py deleted file mode 100644 index 5b325e5a2..000000000 --- a/apps/assets/migrations/0018_asset_directory_services.py +++ /dev/null @@ -1,22 +0,0 @@ -# Generated by Django 4.1.13 on 2025-04-15 11:09 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - - dependencies = [ - ("assets", "0017_auto_20250407_1124"), - ] - - operations = [ - migrations.AddField( - model_name="asset", - name="directory_services", - field=models.ManyToManyField( - related_name="assets", - to="assets.directoryservice", - verbose_name="Directory services", - ), - ), - ] diff --git a/apps/assets/migrations/0019_remove_platform_ds.py b/apps/assets/migrations/0019_remove_platform_ds.py deleted file mode 100644 index 8e80d57c6..000000000 --- a/apps/assets/migrations/0019_remove_platform_ds.py +++ /dev/null @@ -1,17 +0,0 @@ -# Generated by Django 4.1.13 on 2025-04-15 11:33 - -from django.db import migrations - - -class Migration(migrations.Migration): - - dependencies = [ - ("assets", "0018_asset_directory_services"), - ] - - operations = [ - migrations.RemoveField( - model_name="platform", - name="ds", - ), - ] diff --git a/apps/assets/models/asset/common.py b/apps/assets/models/asset/common.py index 4df4d2d61..fe0a8ba82 100644 --- a/apps/assets/models/asset/common.py +++ b/apps/assets/models/asset/common.py @@ -6,7 +6,7 @@ import logging from collections import defaultdict from django.db import models -from django.db.models import Q +from django.db.models import Q, Count from django.forms import model_to_dict from django.utils.translation import gettext_lazy as _ @@ -286,7 +286,7 @@ class Asset(NodesRelationMixin, LabeledMixin, AbsConnectivity, JSONFilterMixin, @property def joined_dir_svc_ids(self): - return self.directory_services.values_list('id', flat=True) + return self.directory_services.all() def is_joined_ad(self): if self.joined_dir_svc_ids: @@ -294,6 +294,33 @@ class Asset(NodesRelationMixin, LabeledMixin, AbsConnectivity, JSONFilterMixin, else: return False + @classmethod + def compute_accounts_amount(cls, assets): + from .ds import DirectoryService + asset_ids = [asset.id for asset in assets] + asset_id_dc_ids_mapper = defaultdict(list) + dc_ids = set() + relations = ( + Asset.directory_services.through.objects + .filter(asset_id__in=asset_ids) + .values_list('asset_id', 'directoryservice_id') + ) + for asset_id, ds_id in relations: + dc_ids.add(ds_id) + asset_id_dc_ids_mapper[asset_id].append(ds_id) + + directory_services = ( + DirectoryService.objects.filter(id__in=dc_ids) + .annotate(accounts_amount=Count('accounts')) + ) + ds_accounts_mapper = {ds.id: ds.accounts_amount for ds in directory_services} + for asset in assets: + asset_dc_ids = asset_id_dc_ids_mapper.get(asset.id, []) + for dc_id in asset_dc_ids: + ds_accounts = ds_accounts_mapper.get(dc_id, 0) + asset.accounts_amount += ds_accounts + return assets + @property def is_valid(self): warning = '' diff --git a/apps/assets/serializers/asset/common.py b/apps/assets/serializers/asset/common.py index d0a6b6e79..e608ada51 100644 --- a/apps/assets/serializers/asset/common.py +++ b/apps/assets/serializers/asset/common.py @@ -229,7 +229,7 @@ class AssetSerializer(BulkOrgResourceModelSerializer, ResourceLabelsMixin, Writa @classmethod def setup_eager_loading(cls, queryset): """ Perform necessary eager loading of data. """ - queryset = queryset.prefetch_related('domain', 'nodes', 'protocols', ) \ + queryset = queryset.prefetch_related('domain', 'nodes', 'protocols', 'directory_services') \ .prefetch_related('platform', 'platform__automation') \ .annotate(category=F("platform__category")) \ .annotate(type=F("platform__type")) \ diff --git a/apps/common/api/mixin.py b/apps/common/api/mixin.py index 09b9bd1aa..c7e08657a 100644 --- a/apps/common/api/mixin.py +++ b/apps/common/api/mixin.py @@ -99,6 +99,7 @@ class QuerySetMixin: return super().get_queryset() def filter_queryset(self, queryset): + queryset = super().filter_queryset(queryset) if not hasattr(self, 'action'): return queryset if self.action == 'metadata': @@ -127,7 +128,7 @@ class QuerySetMixin: def paginate_queryset(self, queryset): page = super().paginate_queryset(queryset) model = getattr(queryset, 'model', None) - if not model or not hasattr(page, 'objects'): + if not model or not hasattr(model, 'objects'): return page serializer_class = self.get_serializer_class() @@ -234,8 +235,8 @@ class OrderingFielderFieldsMixin: class CommonApiMixin( - SerializerMixin, ExtraFilterFieldsMixin, OrderingFielderFieldsMixin, - QuerySetMixin, RenderToJsonMixin, PaginatedResponseMixin + SerializerMixin, QuerySetMixin, ExtraFilterFieldsMixin, + OrderingFielderFieldsMixin, RenderToJsonMixin, PaginatedResponseMixin ): def is_swagger_request(self): return getattr(self, 'swagger_fake_view', False) or \