Remove n+1 queries

Author: FrankApiyo

onadata/apps/api/models/temp_token.py

@@ -2,6 +2,7 @@
 """
 Temporary token authorization model class
 """
+
 import binascii
 import os
 
@@ -10,14 +11,13 @@
 
 
 class TempToken(models.Model):
-
     """
     The temporary authorization token model.
     """
 
     key = models.CharField(max_length=40, primary_key=True)
     user = models.OneToOneField(
-        get_user_model(), related_name="_user", on_delete=models.CASCADE
+        get_user_model(), related_name="temptoken", on_delete=models.CASCADE
     )
     created = models.DateTimeField(auto_now_add=True)
 

onadata/apps/api/tools.py

@@ -295,13 +295,15 @@ def get_organization_members(organization):
     """Get members team user queryset"""
     team = get_organization_members_team(organization)
 
-    return team.user_set.filter(is_active=True)
+    # Optimize: Prefetch profile to avoid N+1
+    return team.user_set.filter(is_active=True).select_related("profile")
 
 
 def get_organization_owners(organization):
     """Get owners team user queryset"""
     team = get_or_create_organization_owners_team(organization)
-    return team.user_set.filter(is_active=True)
+    # Optimize: Prefetch profile to avoid N+1
+    return team.user_set.filter(is_active=True).select_related("profile")
 
 
 def _get_owners(organization):
@@ -655,18 +657,20 @@ def check_inherit_permission_from_project(xform_id, user):
         return
 
     # get the project_xform
+    # Only fetch the fields we need: xform.id, xform.project_id, and
+    # project.id for permissions
     xform = (
         XForm.objects.filter(pk=xform_id)
         .select_related("project")
-        .only("project_id", "id")
+        .only("project_id", "id", "project__id")
         .first()
     )
 
     if not xform:
         return
 
     # ignore if forms has meta perms set
-    if xform.metadata_set.filter(data_type=XFORM_META_PERMS):
+    if xform.metadata_set.filter(data_type=XFORM_META_PERMS).exists():
         return
 
     # get and compare the project role to the xform role

onadata/apps/api/viewsets/attachment_viewset.py

@@ -59,7 +59,7 @@ class AttachmentViewSet(
     lookup_field = "pk"
     queryset = Attachment.objects.filter(
         instance__deleted_at__isnull=True, deleted_at__isnull=True
-    )
+    ).select_related("instance__xform", "xform")
     permission_classes = (AttachmentObjectPermissions,)
     serializer_class = AttachmentSerializer
     pagination_class = StandardPageNumberPagination

onadata/apps/api/viewsets/connect_viewset.py

@@ -52,7 +52,9 @@ def user_profile_w_token_response(request, status_code):
             session.set_expiry(DEFAULT_SESSION_EXPIRY_TIME)
 
     try:
-        user_profile = request.user.profile
+        user_profile = UserProfile.objects.select_related(
+            "user", "user__auth_token", "user__temptoken"
+        ).get(user=request.user)
     except UserProfile.DoesNotExist:
         user_profile = safe_cache_get(f"{USER_PROFILE_PREFIX}{request.user.username}")
         if not user_profile:

onadata/apps/api/viewsets/data_viewset.py

@@ -156,6 +156,10 @@ class DataViewSet(
 
     queryset = XForm.objects.filter(deleted_at__isnull=True)
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._cached_object = None
+
     def get_serializer_class(self):
         """Returns appropriate serializer class based on context."""
         pk_lookup, dataid_lookup = self.lookup_fields
@@ -187,6 +191,10 @@ def get_serializer_class(self):
     # pylint: disable=unused-argument
     def get_object(self, queryset=None):
         """Returns the appropriate object based on context."""
+        # Cache the object to avoid multiple queries in the same request
+        if self._cached_object is None:
+            self._cached_object = super().get_object()
+        obj = self._cached_object
         pk_lookup, dataid_lookup = self.lookup_fields
         form_pk = self.kwargs.get(pk_lookup)
         dataid = self.kwargs.get(dataid_lookup)
@@ -227,21 +235,33 @@ def _get_public_forms_queryset(self):
 
     def _filtered_or_shared_queryset(self, queryset, form_pk):
         filter_kwargs = {self.lookup_field: form_pk}
-        queryset = queryset.filter(**filter_kwargs).only("id", "shared")
+        # Don't override .only() - let the queryset keep its field restrictions
+        queryset = queryset.filter(**filter_kwargs)
 
-        if not queryset:
+        # Use .exists() to avoid executing the full query twice
+        if not queryset.exists():
             filter_kwargs["shared_data"] = True
-            queryset = XForm.objects.filter(**filter_kwargs).only("id", "shared")
+            # Use the same field restrictions as the parent filter_queryset
+            queryset = XForm.objects.filter(**filter_kwargs).only(
+                "id", "shared", "num_of_submissions", "json", "is_merged_dataset"
+            )
 
-            if not queryset:
+            if not queryset.exists():
                 raise Http404(_("No data matches with given query."))
 
         return queryset
 
     # pylint: disable=unused-argument
     def filter_queryset(self, queryset, view=None):
         """Returns and filters queryset based on context and query params."""
-        queryset = super().filter_queryset(queryset.only("id", "shared"))
+        # Fetch only the fields we need: id, shared (for permissions),
+        # num_of_submissions (for pagination), json (for survey structure),
+        # and is_merged_dataset (for query building)
+        queryset = super().filter_queryset(
+            queryset.only(
+                "id", "shared", "num_of_submissions", "json", "is_merged_dataset"
+            )
+        )
         form_pk = self.kwargs.get(self.lookup_field)
 
         if form_pk:
@@ -578,10 +598,12 @@ def list(self, request, *args, **kwargs):
             # pylint: disable=attribute-defined-outside-init
             self.object_list = self._get_public_forms_queryset()
         elif lookup:
-            queryset = self.filter_queryset(self.get_queryset()).values_list(
-                "pk", "is_merged_dataset"
-            )
-            xform_id, is_merged_dataset = queryset[0] if queryset else (lookup, False)
+            # Use get_object() to leverage caching and maintain permission
+            # checks. This replaces the separate filter_queryset() call to
+            # avoid duplicate queries
+            xform = self.get_object()
+            xform_id = xform.id
+            is_merged_dataset = xform.is_merged_dataset
             pks = [xform_id]
             if is_merged_dataset:
                 merged_form = MergedXForm.objects.get(pk=xform_id)
@@ -594,9 +616,9 @@ def list(self, request, *args, **kwargs):
                 except ValueError:
                     pks, num_of_submissions = [], 0
             else:
-                num_of_submissions = XForm.objects.get(id=xform_id).num_of_submissions
+                num_of_submissions = xform.num_of_submissions
             # pylint: disable=attribute-defined-outside-init
-            # Include geom and xform_id fields for geojson format to avoid N+1 queries
+            # For GeoJSON, we need id, json, xform_id(for xform relationship), and geom
             if export_type == "geojson":
                 self.object_list = Instance.objects.filter(
                     xform_id__in=pks, deleted_at=None

onadata/apps/api/viewsets/merged_xform_viewset.py

@@ -42,6 +42,7 @@ class MergedXFormViewSet(
     queryset = (
         MergedXForm.objects.filter(deleted_at__isnull=True)
         .annotate(number_of_submissions=Sum("xforms__num_of_submissions"))
+        .prefetch_related("xforms__user", "xforms__project")
         .all()
     )
     pagination_class = StandardPageNumberPagination
@@ -51,6 +52,29 @@ class MergedXFormViewSet(
         renderers.GeoJsonRenderer,
     ]
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._cached_object = None
+
+    def get_object(self):
+        """
+        Get object - cache to avoid multiple queries during request
+        """
+        if self._cached_object is None:
+            self._cached_object = super().get_object()
+        return self._cached_object
+
+    def get_queryset(self):
+        """
+        Get queryset - optimize for specific actions
+        """
+        queryset = super().get_queryset()
+
+        if self.action == "data":
+            queryset = queryset.only("id", "shared")
+
+        return queryset
+
     def get_serializer_class(self):
         """
         Get appropriate serializer class

onadata/apps/api/viewsets/organization_profile_viewset.py

@@ -68,7 +68,9 @@ class OrganizationProfileViewSet(
     List, Retrieve, Update, Create/Register Organizations.
     """
 
-    queryset = OrganizationProfile.objects.filter(user__is_active=True)
+    queryset = OrganizationProfile.objects.filter(user__is_active=True).select_related(
+        "user", "creator", "created_by"
+    )
     serializer_class = serializer_from_settings()
     lookup_field = "user"
     permission_classes = [permissions.OrganizationProfilePermissions]

onadata/apps/api/viewsets/project_viewset.py

@@ -6,6 +6,7 @@
 import logging
 
 from django.core.mail import send_mail
+from django.db.models import Prefetch
 from django.shortcuts import get_object_or_404
 from django.utils.translation import gettext as _
 
@@ -82,6 +83,27 @@ class ProjectViewSet(
         Project.objects.filter(deleted_at__isnull=True)
         .order_by("-date_created")
         .select_related()
+        .prefetch_related(
+            Prefetch(
+                "xform_set",
+                queryset=XForm.objects.filter(deleted_at__isnull=True)
+                .only(
+                    "id",
+                    "title",
+                    "id_string",
+                    "is_merged_dataset",
+                    "encrypted",
+                    "last_submission_time",
+                    "project_id",
+                )
+                .prefetch_related(
+                    "registration_forms__entity_list",
+                    "follow_up_forms__entity_list",
+                    "metadata_set",
+                ),
+                to_attr="xforms_prefetch",
+            )
+        )
     )
     serializer_class = ProjectSerializer
     lookup_field = "pk"
@@ -106,6 +128,19 @@ def get_serializer_class(self):
 
         return super().get_serializer_class()
 
+    def get_serializer_context(self):
+        """Pass starred projects to serializer context."""
+        context = super().get_serializer_context()
+        if (
+            self.action == "list"
+            and self.request
+            and not self.request.user.is_anonymous
+        ):
+            context["starred_projects"] = self.request.user.project_stars.values_list(
+                "pk", flat=True
+            )
+        return context
+
     def get_queryset(self):
         """Use 'prepared' prefetched queryset for GET requests."""
         if self.request.method.upper() in ["GET", "OPTIONS"]:

onadata/apps/api/viewsets/team_viewset.py

@@ -2,6 +2,7 @@
 """
 The /teams API endpoint implementation.
 """
+
 from django.contrib.auth import get_user_model
 from django.utils.translation import gettext as _
 
@@ -44,7 +45,9 @@ class TeamViewSet(
     This endpoint allows you to create, update and view team information.
     """
 
-    queryset = Team.objects.all()
+    queryset = Team.objects.select_related(
+        "organization", "created_by"
+    ).prefetch_related("user_set", "projects")
     serializer_class = TeamSerializer
     lookup_field = "pk"
     extra_lookup_fields = None

onadata/apps/logger/models/project.py

@@ -33,6 +33,9 @@ def get_queryset(self):
         # pylint: disable=invalid-name
         Team = apps.get_model("api", "Team")  # noqa N806
         XForm = apps.get_model("logger", "XForm")  # noqa N806
+        EntityList = apps.get_model("logger", "EntityList")  # noqa N806
+        RegistrationForm = apps.get_model("logger", "RegistrationForm")  # noqa N806
+        FollowUpForm = apps.get_model("logger", "FollowUpForm")  # noqa N806
 
         # pylint: disable=no-member
         return (
@@ -44,24 +47,59 @@ def get_queryset(self):
                     "xform_set",
                     queryset=XForm.objects.filter(deleted_at__isnull=True)
                     .select_related("user")
-                    .prefetch_related("user")
                     .prefetch_related("dataview_set")
                     .prefetch_related("metadata_set")
+                    .prefetch_related(
+                        Prefetch(
+                            "registration_forms",
+                            queryset=RegistrationForm.objects.filter(
+                                entity_list__deleted_at__isnull=True
+                            ).select_related("entity_list"),
+                        )
+                    )
+                    .prefetch_related(
+                        Prefetch(
+                            "follow_up_forms",
+                            queryset=FollowUpForm.objects.filter(
+                                entity_list__deleted_at__isnull=True
+                            ).select_related("entity_list"),
+                        )
+                    )
                     .only(
                         "id",
                         "user",
                         "project",
                         "title",
                         "date_created",
                         "last_submission_time",
+                        "last_updated_at",
                         "num_of_submissions",
                         "downloadable",
                         "id_string",
                         "is_merged_dataset",
+                        "encrypted",
                     ),
                     to_attr="xforms_prefetch",
                 )
             )
+            .prefetch_related(
+                Prefetch(
+                    "entity_lists",
+                    queryset=EntityList.objects.filter(deleted_at__isnull=True)
+                    .prefetch_related(
+                        Prefetch(
+                            "registrationform_set",
+                            queryset=RegistrationForm.objects.select_related("xform"),
+                        )
+                    )
+                    .prefetch_related(
+                        Prefetch(
+                            "followupform_set",
+                            queryset=FollowUpForm.objects.select_related("xform"),
+                        )
+                    ),
+                )
+            )
             .prefetch_related("tags")
             .prefetch_related(
                 Prefetch(
@@ -86,6 +124,15 @@ def get_queryset(self):
                     queryset=Team.objects.all().prefetch_related("user_set"),
                 )
             )
+            .prefetch_related(
+                Prefetch(
+                    "dataview_set",
+                    queryset=apps.get_model("logger", "DataView").objects.filter(
+                        deleted_at__isnull=True
+                    ),
+                    to_attr="dataview_prefetch",
+                )
+            )
         )