perf(assets): prefetch related XForm objects to reduce database queries DEV-880 (#6487)

### 📣 Summary
Improve performance of asset queries by prefetching related XForm
objects.


### 📖 Description
This change optimizes asset retrieval by prefetching related XForm
objects, reducing the number of database queries executed when listing
or accessing assets. Previously, each access to related XForm data
triggered additional queries, which impacted performance on endpoints
returning multiple assets. With prefetching, related data is loaded
efficiently in a single query, improving response times and lowering
database load without altering API behavior or response content.

Author: noliveleger

kpi/serializers/v2/asset.py

@@ -10,7 +10,6 @@
 from django.db.models import F
 from django.utils.translation import gettext as t
 from django.utils.translation import ngettext as nt
-from django_request_cache import cache_for_request
 from drf_spectacular.types import OpenApiTypes
 from drf_spectacular.utils import extend_schema_field
 from rest_framework import exceptions, serializers
@@ -1280,13 +1279,3 @@ def _get_user_detail(obj, attr: str) -> str:
     def _get_view(self) -> str:
         request = self.context['request']
         return request.parser_context['kwargs']['uid_project_view']
-
-    # FIXME Remove this method, seems to not be used anywhere
-    @cache_for_request
-    def _user_has_asset_perms(self, obj: Asset, perm: str) -> bool:
-        request = self.context.get('request')
-        user = get_database_user(request.user)
-        self._set_asset_ids_cache(obj)
-        if obj.owner == user or obj.has_perm(user, perm):
-            return True
-        return False

kpi/tests/api/v1/test_api_assets.py

@@ -4,10 +4,10 @@
 from urllib.parse import unquote_plus
 
 from django.urls import reverse
-from formpack.utils.expand_content import SCHEMA_VERSION
 from rest_framework import status
 from rest_framework.authtoken.models import Token
 
+from formpack.utils.expand_content import SCHEMA_VERSION
 from kobo.apps.kobo_auth.shortcuts import User
 from kpi.constants import ASSET_TYPE_COLLECTION
 from kpi.models import Asset, SubmissionExportTask
@@ -61,13 +61,13 @@ def test_query_counts(self):
         # expected query counts are different in v1 and v2 so override the test here
         self.create_asset()
 
-        with self.assertNumQueries(FuzzyInt(28, 37)):
+        with self.assertNumQueries(FuzzyInt(28, 40)):
             self.client.get(self.list_url)
         # test query count does not increase with more assets
         self.create_asset()
         self.create_asset()
         self.create_asset()
-        with self.assertNumQueries(FuzzyInt(28, 37)):
+        with self.assertNumQueries(FuzzyInt(28, 40)):
             self.client.get(self.list_url)
 
 

kpi/tests/api/v2/test_api_assets.py

@@ -419,18 +419,18 @@ def test_creator_permissions_on_import(self):
     def test_query_counts(self):
         self.create_asset()
         # 45 when stripe is disabled, 46 when enabled
-        with self.assertNumQueries(FuzzyInt(42, 51)):
+        with self.assertNumQueries(FuzzyInt(36, 45)):
             self.client.get(self.list_url)
         # test query count does not increase with more assets
         # add several assets so the fuzziness of the count doesn't hide an O(n) addition
         self.create_asset()
         self.create_asset()
         self.create_asset()
-        with self.assertNumQueries(FuzzyInt(42, 51)):
+        with self.assertNumQueries(FuzzyInt(36, 45)):
             self.client.get(self.list_url)
 
         # test query counts with search filter
-        with self.assertNumQueries(FuzzyInt(42, 51)):
+        with self.assertNumQueries(FuzzyInt(36, 45)):
             self.client.get(self.list_url, data={'q': 'asset_type:survey'})
 
 

kpi/views/v2/asset.py

@@ -15,6 +15,7 @@
 
 from kobo.apps.audit_log.base_views import AuditLoggedModelViewSet
 from kobo.apps.audit_log.models import AuditType
+from kobo.apps.openrosa.apps.logger.models.xform import XForm
 from kpi.constants import (
     ASSET_TYPE_ARG_NAME,
     ASSET_TYPE_SURVEY,
@@ -372,6 +373,10 @@ class AssetViewSet(
     ]
     log_type = AuditType.PROJECT_HISTORY
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._serializer_context = {}
+
     @action(
         detail=False,
         methods=['POST'],
@@ -694,8 +699,9 @@ def get_serializer_context(self):
         """
         Extra context provided to the serializer class.
         """
-
         context_ = super().get_serializer_context()
+        context_.update(self._serializer_context)
+
         if self.action == 'list':
             # To avoid making a triple join-query for each asset in the list
             # to retrieve related objects, we populated dicts key-ed by asset
@@ -705,34 +711,37 @@ def get_serializer_context(self):
             # The serializer will be able to pick what it needs from that dict
             # and narrow down data according to users' permissions.
 
-            # self.__filtered_queryset is set in the `list()` method that
+            # self._filtered_queryset is set in the `list()` method that
             # DRF automatically calls and is overridden below. This is
             # to prevent double calls to `filter_queryset()` as described in
             # the issue here: https://github.com/kobotoolbox/kpi/issues/2576
-            queryset = self.__filtered_queryset
+            queryset = self._filtered_queryset
 
             # 1) Retrieve all asset IDs of the current list
-            asset_ids = AssetPagination.get_all_asset_ids_from_queryset(
-                queryset
-            )
+            if 'asset_ids_cache' not in context_:
+                asset_ids = AssetPagination.get_all_asset_ids_from_queryset(queryset)
+                context_['asset_ids_cache'] = asset_ids
+            else:
+                asset_ids = context_['asset_ids_cache']
 
             # 2) Get object permissions per asset
             context_[
                 'object_permissions_per_asset'
             ] = self.cache_all_assets_perms(asset_ids)
-            context_['asset_ids_cache'] = asset_ids
 
             # 3) Get the collection subscriptions per asset
             subscriptions_queryset = (
                 UserAssetSubscription.objects.values('asset_id', 'user_id')
                 .distinct()
+                .filter(asset_id__in=asset_ids)
                 .order_by('asset_id')
             )
 
             user_subscriptions_per_asset = defaultdict(list)
             for record in subscriptions_queryset:
                 user_subscriptions_per_asset[record['asset_id']].append(
-                    record['user_id'])
+                    record['user_id']
+                )
 
             context_['user_subscriptions_per_asset'] = user_subscriptions_per_asset
 
@@ -770,17 +779,26 @@ def get_serializer_context(self):
     def list(self, request, *args, **kwargs):
         # assigning global filtered query set to prevent additional,
         # unnecessary calls to `filter_queryset`
-        self.__filtered_queryset = self.filter_queryset(self.get_queryset())
+        self._filtered_queryset = self.filter_queryset(self.get_queryset())
+
+        page = self.paginate_queryset(self._filtered_queryset)
 
-        page = self.paginate_queryset(self.__filtered_queryset)
         if page is not None:
-            serializer = self.get_serializer(page, many=True)
+            self._serializer_context['asset_ids_cache'] = []
+            self._serializer_context['asset_uids_cache'] = []
+            for asset in page:
+                self._serializer_context['asset_ids_cache'].append(asset.pk)
+                self._serializer_context['asset_uids_cache'].append(asset.uid)
+
+            serializer = self.get_serializer(
+                self._attach_xforms_to_assets(page), many=True
+            )
             metadata = None
             if request.GET.get('metadata') == 'on':
-                metadata = self.get_metadata(self.__filtered_queryset)
+                metadata = self.get_metadata(self._filtered_queryset)
             return self.get_paginated_response(serializer.data, metadata)
 
-        serializer = self.get_serializer(self.__filtered_queryset, many=True)
+        serializer = self.get_serializer(self._filtered_queryset, many=True)
         return Response(serializer.data)
 
     @action(detail=False, methods=['GET'])
@@ -861,8 +879,9 @@ def reports(self, request, *args, **kwargs):
         asset = self.get_object()
         if not asset.has_deployment:
             raise Http404
-        serializer = ReportsDetailSerializer(asset,
-                                             context=self.get_serializer_context())
+        serializer = ReportsDetailSerializer(
+            asset, context=self.get_serializer_context()
+        )
         return Response(serializer.data)
 
     @extend_schema(tags=['Form content'])
@@ -904,6 +923,39 @@ def xform(self, request, *args, **kwargs):
     def xls(self, request, *args, **kwargs):
         return self.table_view(self, request, *args, **kwargs)
 
+    def _attach_xforms_to_assets(self, assets: list):
+        """
+        Attach the related XForm to each Asset and yield them to
+        stay memory-efficient.
+        """
+
+        asset_uids = self._serializer_context['asset_uids_cache']
+        xform_qs = (
+            XForm.all_objects.filter(kpi_asset_uid__in=asset_uids)
+            .only(
+                'id_string',
+                'num_of_submissions',
+                'attachment_storage_bytes',
+                'require_auth',
+                'uuid',
+                'mongo_uuid',
+                'encrypted',
+                'last_submission_time',
+                'kpi_asset_uid',
+            )
+            .order_by()
+        )
+
+        xforms_by_uid = {xf.kpi_asset_uid: xf for xf in xform_qs}
+
+        # Single pass over assets: attach _xform then yield
+        for asset in assets:
+            if asset.has_deployment:
+                xf = xforms_by_uid.get(asset.uid, None)
+                xf.user = asset.owner
+                asset.deployment._xform = xf
+            yield asset
+
     def _bulk_asset_actions(self, data: dict) -> dict:
         params = {
             'data': data,