perf: change db prefetch (#15215)

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')

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",
             ),
         ),
     ]

apps/assets/migrations/0018_asset_directory_services.py

@@ -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",
-            ),
-        ),
-    ]

apps/assets/migrations/0019_remove_platform_ds.py

@@ -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",
-        ),
-    ]

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 = ''

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")) \

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 \