From 53a947873a4f271ad198a0ef5825c3a96d272e2d Mon Sep 17 00:00:00 2001 From: Kipchirchir Sigei Date: Thu, 16 Nov 2023 12:09:28 +0300 Subject: [PATCH] Enhancement: Add Optimization Fence Query to gracefully handle cancel statement timeout Signed-off-by: Kipchirchir Sigei --- .../apps/api/viewsets/briefcase_viewset.py | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/onadata/apps/api/viewsets/briefcase_viewset.py b/onadata/apps/api/viewsets/briefcase_viewset.py index 7e0648a7a5..8834f67cbf 100644 --- a/onadata/apps/api/viewsets/briefcase_viewset.py +++ b/onadata/apps/api/viewsets/briefcase_viewset.py @@ -8,6 +8,7 @@ from django.contrib.auth import get_user_model from django.core.files import File from django.core.validators import ValidationError +from django.db import OperationalError from django.http import Http404 from django.utils.translation import gettext as _ @@ -121,7 +122,7 @@ def get_object(self, queryset=None): return obj - # pylint: disable=too-many-branches,too-many-statements + # pylint: disable=too-many-branches,too-many-statements,too-many-locals def filter_queryset(self, queryset): """ Filters an XForm submission instances using ODK Aggregate query parameters. @@ -188,7 +189,28 @@ def filter_queryset(self, queryset): num_entries = _parse_int(num_entries) if num_entries: - instances = instances[:num_entries] + 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") # Using len() instead of .count() to prevent an extra # database call; len() will load the instances in memory allowing