Optimize queryset annotations & prefetches to cut DB time for test / finding / product views (issue #12575) (#12603)

* Optimize the DB queries for engagement, finding, product and product_type views. Mostly counters fixes

* Apply suggestions from code review

Fix typo in the "build_count_subquery" docstring

Co-authored-by: valentijnscholten <valentijnscholten@gmail.com>

* Add subquery for test_count for engagements getting

---------

Co-authored-by: valentijnscholten <valentijnscholten@gmail.com>

Author: DenysMoskalenko

dojo/engagement/views.py

@@ -5,7 +5,7 @@
 import re
 import time
 from datetime import datetime
-from functools import reduce
+from functools import partial, reduce
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from time import strftime
@@ -16,7 +16,8 @@
 from django.contrib.auth.models import User
 from django.core.exceptions import PermissionDenied, ValidationError
 from django.db import DEFAULT_DB_ALIAS
-from django.db.models import Count, Q
+from django.db.models import Count, OuterRef, Q, Value
+from django.db.models.functions import Coalesce
 from django.db.models.query import Prefetch, QuerySet
 from django.http import HttpRequest, HttpResponse, HttpResponseRedirect, QueryDict, StreamingHttpResponse
 from django.shortcuts import get_object_or_404, render
@@ -104,6 +105,7 @@
     add_error_message_to_response,
     add_success_message_to_response,
     async_delete,
+    build_count_subquery,
     calculate_grade,
     generate_file_response_from_file_path,
     get_cal_event,
@@ -151,7 +153,6 @@ def engagement_calendar(request):
 
 
 def get_filtered_engagements(request, view):
-
     if view not in {"all", "active"}:
         msg = f"View {view} is not allowed"
         raise ValidationError(msg)
@@ -161,37 +162,29 @@ def get_filtered_engagements(request, view):
     if view == "active":
         engagements = engagements.filter(active=True)
 
-    engagements = engagements.select_related("product", "product__prod_type") \
+    engagements = (
+        engagements
+        .select_related("product", "product__prod_type")
         .prefetch_related("lead", "tags", "product__tags")
+    )
 
     if System_Settings.objects.get().enable_jira:
         engagements = engagements.prefetch_related(
             "jira_project__jira_instance",
             "product__jira_project_set__jira_instance",
         )
 
+    test_count_subquery = build_count_subquery(
+        Test.objects.filter(engagement=OuterRef("pk")), group_field="engagement_id",
+    )
+    engagements = engagements.annotate(test_count=Coalesce(test_count_subquery, Value(0)))
+
     filter_string_matching = get_system_setting("filter_string_matching", False)
     filter_class = EngagementDirectFilterWithoutObjectLookups if filter_string_matching else EngagementDirectFilter
     return filter_class(request.GET, queryset=engagements)
 
 
-def get_test_counts(engagements):
-    # Get the test counts per engagement. As a separate query, this is much
-    # faster than annotating the above `engagements` query.
-    return {
-        test["engagement"]: test["test_count"]
-        for test in Test.objects.filter(
-            engagement__in=engagements,
-        ).values(
-            "engagement",
-        ).annotate(
-            test_count=Count("engagement"),
-        )
-    }
-
-
 def engagements(request, view):
-
     if not view:
         view = "active"
 
@@ -209,7 +202,6 @@ def engagements(request, view):
     return render(
         request, "dojo/engagement.html", {
             "engagements": engs,
-            "engagement_test_counts": get_test_counts(filtered_engagements.qs),
             "filter_form": filtered_engagements.form,
             "product_name_words": product_name_words,
             "engagement_name_words": engagement_name_words,
@@ -592,23 +584,27 @@ def post(self, request, eid, *args, **kwargs):
 
 
 def prefetch_for_view_tests(tests):
-    prefetched = tests
-    if isinstance(tests,
-                  QuerySet):  # old code can arrive here with prods being a list because the query was already executed
-
-        prefetched = prefetched.select_related("lead")
-        prefetched = prefetched.prefetch_related("tags", "test_type", "notes")
-        prefetched = prefetched.annotate(count_findings_test_all=Count("finding__id", distinct=True))
-        prefetched = prefetched.annotate(count_findings_test_active=Count("finding__id", filter=Q(finding__active=True), distinct=True))
-        prefetched = prefetched.annotate(count_findings_test_active_verified=Count("finding__id", filter=Q(finding__active=True) & Q(finding__verified=True), distinct=True))
-        prefetched = prefetched.annotate(count_findings_test_mitigated=Count("finding__id", filter=Q(finding__is_mitigated=True), distinct=True))
-        prefetched = prefetched.annotate(count_findings_test_dups=Count("finding__id", filter=Q(finding__duplicate=True), distinct=True))
-        prefetched = prefetched.annotate(total_reimport_count=Count("test_import__id", filter=Q(test_import__type=Test_Import.REIMPORT_TYPE), distinct=True))
-
-    else:
+    # old code can arrive here with prods being a list because the query was already executed
+    if not isinstance(tests, QuerySet):
         logger.warning("unable to prefetch because query was already executed")
-
-    return prefetched
+        return tests
+
+    prefetched = tests.select_related("lead", "test_type").prefetch_related("tags", "notes")
+    base_findings = Finding.objects.filter(test_id=OuterRef("pk"))
+    count_subquery = partial(build_count_subquery, group_field="test_id")
+    return prefetched.annotate(
+        count_findings_test_all=Coalesce(count_subquery(base_findings), Value(0)),
+        count_findings_test_active=Coalesce(count_subquery(base_findings.filter(active=True)), Value(0)),
+        count_findings_test_active_verified=Coalesce(
+            count_subquery(base_findings.filter(active=True, verified=True)), Value(0),
+        ),
+        count_findings_test_mitigated=Coalesce(count_subquery(base_findings.filter(is_mitigated=True)), Value(0)),
+        count_findings_test_dups=Coalesce(count_subquery(base_findings.filter(duplicate=True)), Value(0)),
+        total_reimport_count=Coalesce(
+            count_subquery(Test_Import.objects.filter(test_id=OuterRef("pk"), type=Test_Import.REIMPORT_TYPE)),
+            Value(0),
+        ),
+    )
 
 
 @user_is_authorized(Engagement, Permissions.Test_Add, "eid")
@@ -1583,14 +1579,18 @@ def get_engagements(request):
         query = get_list_index(path_items, 1)
 
     request.GET = QueryDict(query)
-    engagements = get_filtered_engagements(request, view).qs
-    test_counts = get_test_counts(engagements)
-
-    return engagements, test_counts
+    return get_filtered_engagements(request, view).qs
 
 
 def get_excludes():
-    return ["is_ci_cd", "jira_issue", "jira_project", "objects", "unaccepted_open_findings"]
+    return [
+        "is_ci_cd",
+        "jira_issue",
+        "jira_project",
+        "objects",
+        "unaccepted_open_findings",
+        "test_count",  # already exported separately as “tests”
+    ]
 
 
 def get_foreign_keys():
@@ -1600,7 +1600,7 @@ def get_foreign_keys():
 
 def csv_export(request):
     logger.debug("starting csv export")
-    engagements, test_counts = get_engagements(request)
+    engagements = get_engagements(request)
 
     response = HttpResponse(content_type="text/csv")
     response["Content-Disposition"] = "attachment; filename=engagements.csv"
@@ -1627,7 +1627,7 @@ def csv_export(request):
                     if value and isinstance(value, str):
                         value = value.replace("\n", " NEWLINE ").replace("\r", "")
                     fields.append(value)
-            fields.append(test_counts.get(engagement.id, 0))
+            fields.append(getattr(engagement, "test_count", 0))
 
             writer.writerow(fields)
     logger.debug("done with csv export")
@@ -1636,7 +1636,7 @@ def csv_export(request):
 
 def excel_export(request):
     logger.debug("starting excel export")
-    engagements, test_counts = get_engagements(request)
+    engagements = get_engagements(request)
 
     workbook = Workbook()
     workbook.iso_dates = True
@@ -1668,7 +1668,7 @@ def excel_export(request):
                         value = value.replace(tzinfo=None)
                     worksheet.cell(row=row_num, column=col_num, value=value)
                     col_num += 1
-            worksheet.cell(row=row_num, column=col_num, value=test_counts.get(engagement.id, 0))
+            worksheet.cell(row=row_num, column=col_num, value=getattr(engagement, "test_count", 0))
         row_num += 1
 
     with NamedTemporaryFile() as tmp:

dojo/finding/views.py

@@ -6,6 +6,7 @@
 import logging
 import mimetypes
 from collections import OrderedDict, defaultdict
+from functools import partial
 from itertools import chain
 from pathlib import Path
 
@@ -14,8 +15,8 @@
 from django.core import serializers
 from django.core.exceptions import PermissionDenied, ValidationError
 from django.db import models
-from django.db.models import Count, Q, QuerySet
-from django.db.models.functions import Length
+from django.db.models import OuterRef, QuerySet, Value
+from django.db.models.functions import Coalesce, Length
 from django.db.models.query import Prefetch
 from django.http import Http404, HttpRequest, HttpResponse, HttpResponseRedirect, JsonResponse, StreamingHttpResponse
 from django.shortcuts import get_object_or_404, render
@@ -110,6 +111,7 @@
     add_field_errors_to_response,
     add_success_message_to_response,
     apply_cwe_to_template,
+    build_count_subquery,
     calculate_grade,
     close_external_issue,
     do_false_positive_history,
@@ -132,87 +134,58 @@
 
 
 def prefetch_for_findings(findings, prefetch_type="all", *, exclude_untouched=True):
-    prefetched_findings = findings
-    if isinstance(
-        findings, QuerySet,
-    ):  # old code can arrive here with prods being a list because the query was already executed
-        prefetched_findings = prefetched_findings.prefetch_related(
-            "reviewers",
-        )
-        prefetched_findings = prefetched_findings.prefetch_related("reporter")
-        prefetched_findings = prefetched_findings.prefetch_related(
-            "jira_issue__jira_project__jira_instance",
-        )
-        prefetched_findings = prefetched_findings.prefetch_related("test__test_type")
-        prefetched_findings = prefetched_findings.prefetch_related(
-            "test__engagement__jira_project__jira_instance",
-        )
-        prefetched_findings = prefetched_findings.prefetch_related(
-            "test__engagement__product__jira_project_set__jira_instance",
-        )
-        prefetched_findings = prefetched_findings.prefetch_related("found_by")
+    # old code can arrive here with prods being a list because the query was already executed
+    if not isinstance(findings, QuerySet):
+        logger.debug("unable to prefetch because query was already executed")
+        return findings
 
-        # for open/active findings the following 4 prefetches are not needed
-        if prefetch_type != "open":
-            prefetched_findings = prefetched_findings.prefetch_related(
-                "risk_acceptance_set",
-            )
-            prefetched_findings = prefetched_findings.prefetch_related(
-                "risk_acceptance_set__accepted_findings",
-            )
-            prefetched_findings = prefetched_findings.prefetch_related(
-                "original_finding",
-            )
-            prefetched_findings = prefetched_findings.prefetch_related(
-                "duplicate_finding",
-            )
+    prefetched_findings = findings.prefetch_related(
+        "reviewers",
+        "reporter",
+        "jira_issue__jira_project__jira_instance",
+        "test__test_type",
+        "test__engagement__jira_project__jira_instance",
+        "test__engagement__product__jira_project_set__jira_instance",
+        "found_by",
+    )
 
-        if exclude_untouched:
-            # filter out noop reimport actions from finding status history
-            prefetched_findings = prefetched_findings.prefetch_related(
-                Prefetch(
-                    "test_import_finding_action_set",
-                    queryset=Test_Import_Finding_Action.objects.exclude(
-                        action=IMPORT_UNTOUCHED_FINDING,
-                    ),
-                ),
-            )
-        else:
-            prefetched_findings = prefetched_findings.prefetch_related(
-                "test_import_finding_action_set",
-            )
-        """
-        we could try to prefetch only the latest note with SubQuery and OuterRef,
-        but I'm getting that MySql doesn't support limits in subqueries.
-        """
-        prefetched_findings = prefetched_findings.prefetch_related("notes")
-        prefetched_findings = prefetched_findings.prefetch_related("tags")
-        prefetched_findings = prefetched_findings.prefetch_related("endpoints")
-        prefetched_findings = prefetched_findings.prefetch_related("status_finding")
-        prefetched_findings = prefetched_findings.annotate(
-            active_endpoint_count=Count(
-                "status_finding__id", filter=Q(status_finding__mitigated=False),
-            ),
-        )
-        prefetched_findings = prefetched_findings.annotate(
-            mitigated_endpoint_count=Count(
-                "status_finding__id", filter=Q(status_finding__mitigated=True),
-            ),
-        )
-        prefetched_findings = prefetched_findings.prefetch_related("finding_group_set")
+    # for open/active findings, the following 4 prefetches are not needed
+    if prefetch_type != "open":
         prefetched_findings = prefetched_findings.prefetch_related(
-            "test__engagement__product__members",
-        )
-        prefetched_findings = prefetched_findings.prefetch_related(
-            "test__engagement__product__prod_type__members",
+            "risk_acceptance_set",
+            "risk_acceptance_set__accepted_findings",
+            "original_finding",
+            "duplicate_finding",
         )
+
+    if exclude_untouched:
+        # filter out noop reimport actions from finding status history
         prefetched_findings = prefetched_findings.prefetch_related(
-            "vulnerability_id_set",
+            Prefetch(
+                "test_import_finding_action_set",
+                queryset=Test_Import_Finding_Action.objects.exclude(action=IMPORT_UNTOUCHED_FINDING),
+            ),
         )
     else:
-        logger.debug("unable to prefetch because query was already executed")
+        prefetched_findings = prefetched_findings.prefetch_related("test_import_finding_action_set")
+
+    prefetched_findings = prefetched_findings.prefetch_related(
+        "notes",
+        "tags",
+        "endpoints",
+        "status_finding",
+        "finding_group_set",
+        "test__engagement__product__members",
+        "test__engagement__product__prod_type__members",
+        "vulnerability_id_set",
+    )
 
-    return prefetched_findings
+    base_status = Endpoint_Status.objects.filter(finding_id=OuterRef("pk"))
+    count_subquery = partial(build_count_subquery, group_field="finding_id")
+    return prefetched_findings.annotate(
+        active_endpoint_count=Coalesce(count_subquery(base_status.filter(mitigated=False)), Value(0)),
+        mitigated_endpoint_count=Coalesce(count_subquery(base_status.filter(mitigated=True)), Value(0)),
+    )
 
 
 def prefetch_for_similar_findings(findings):