Add tests

Signed-off-by: Kipchirchir Sigei <[email protected]>

onadata/apps/api/tests/viewsets/test_briefcase_viewset.py

@@ -12,9 +12,11 @@
 from rest_framework.test import APIRequestFactory
 
 from onadata.apps.api.tests.viewsets import test_abstract_viewset
-from onadata.apps.api.viewsets.briefcase_viewset import BriefcaseViewset
-from onadata.apps.api.viewsets.xform_submission_viewset import \
-    XFormSubmissionViewSet
+from onadata.apps.api.viewsets.briefcase_viewset import (
+    BriefcaseViewset,
+    _query_optimization_fence,
+)
+from onadata.apps.api.viewsets.xform_submission_viewset import XFormSubmissionViewSet
 from onadata.apps.api.viewsets.xform_viewset import XFormViewSet
 from onadata.apps.logger.models import Instance
 from onadata.apps.logger.models import XForm
@@ -662,6 +664,22 @@ def test_view_downloadSubmission_multiple_nodes(self, mock_get_object):
                 text = text.replace(*var)
             self.assertContains(response, instanceId, status_code=200)
 
+    def test_query_optimization_fence(self):
+        self._publish_xml_form()
+        self._make_submissions()
+        params = {"formId": self.xform.id_string}
+        params["numEntries"] = 2
+        instances = ordered_instances(self.xform)
+        optimized_instances = _query_optimization_fence(instances, 4)
+        self.assertEqual(instances.count(), optimized_instances.count())
+        op_sql_query = (
+            'SELECT "logger_instance"."id", "logger_instance"."uuid" FROM "logger_instance"'
+            f' WHERE "logger_instance"."id" IN ({optimized_instances[0].get("pk")},'
+            f' {optimized_instances[1].get("pk")}, {optimized_instances[2].get("pk")},'
+            f' {optimized_instances[3].get("pk")})'
+        )
+        self.assertEqual(str(optimized_instances.query), op_sql_query)
+
     def tearDown(self):
         # remove media files
         if self.user:

onadata/apps/api/viewsets/briefcase_viewset.py

@@ -71,6 +71,44 @@ def _parse_int(num):
     except ValueError:
         return None
 
+def _query_optimization_fence(instances, num_entries):
+    """
+    Enhances query performance by using an optimization fence.
+
+    This utility function creates an optimization fence around the provided queryset instances.
+    It encapsulates the original query within a SELECT statement with an ORDER BY and LIMIT clause,
+    optimizing the database query for improved performance.
+
+    Parameters:
+    - instances: QuerySet
+        The input QuerySet of instances to be optimized.
+    - num_entries: int
+        The number of instances to be included in the optimized result set.
+
+    Returns:
+    QuerySet
+        An optimized QuerySet containing selected fields ('pk' and 'uuid') based on the provided instances.
+    """
+    inner_raw_sql = str(instances.query)
+
+    # Create the outer query with the LIMIT clause
+    outer_query = (
+        f"SELECT id, uuid FROM ({inner_raw_sql}) AS items "  # nosec
+        "ORDER BY id ASC LIMIT %s"  # nosec
+    )
+    raw_queryset = Instance.objects.raw(outer_query, [num_entries])
+    # convert raw queryset to queryset
+    instances_data = [
+        {"pk": item.id, "uuid": item.uuid}
+        for item in raw_queryset.iterator()
+    ]
+    # Create QuerySet from the instances dict
+    instances = Instance.objects.filter(
+        pk__in=[item["pk"] for item in instances_data]
+    ).values("pk", "uuid")
+
+    return instances
+
 
 # pylint: disable=too-many-ancestors
 class BriefcaseViewset(
@@ -192,25 +230,7 @@ def filter_queryset(self, queryset):
             try:
                 instances = instances[:num_entries]
             except OperationalError:
-                # Create an optimization fence
-                # Define the base query
-                inner_raw_sql = str(instances.query)
-
-                # Create the outer query with the LIMIT clause
-                outer_query = (
-                    f"SELECT id, uuid FROM ({inner_raw_sql}) AS items "  # nosec
-                    "ORDER BY id ASC LIMIT %s"  # nosec
-                )
-                raw_queryset = Instance.objects.raw(outer_query, [num_entries])
-                # convert raw queryset to queryset
-                instances_data = [
-                    {"pk": item.id, "uuid": item.uuid}
-                    for item in raw_queryset.iterator()
-                ]
-                # Create QuerySet from the instances dict
-                instances = Instance.objects.filter(
-                    pk__in=[item["pk"] for item in instances_data]
-                ).values("pk", "uuid")
+                instances = _query_optimization_fence(instances, num_entries)
 
         # Using len() instead of .count() to prevent an extra
         # database call; len() will load the instances in memory allowing