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

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,
@@ -592,23 +594,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")

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):

dojo/product/views.py

@@ -4,6 +4,7 @@
 import logging
 from collections import OrderedDict
 from datetime import date, datetime, timedelta
+from functools import partial
 from math import ceil
 
 from dateutil.relativedelta import relativedelta
@@ -12,8 +13,9 @@
 from django.contrib.postgres.aggregates import StringAgg
 from django.core.exceptions import PermissionDenied, ValidationError
 from django.db import DEFAULT_DB_ALIAS, connection
-from django.db.models import Count, F, Max, OuterRef, Prefetch, Q, Subquery, Sum
+from django.db.models import Count, DateField, F, OuterRef, Prefetch, Q, Subquery, Sum
 from django.db.models.expressions import Value
+from django.db.models.functions import Coalesce
 from django.db.models.query import QuerySet
 from django.http import Http404, HttpRequest, HttpResponseRedirect, JsonResponse
 from django.shortcuts import get_object_or_404, render
@@ -89,6 +91,7 @@
     Product_Type,
     System_Settings,
     Test,
+    Test_Import,
     Test_Type,
 )
 from dojo.product.queries import (
@@ -113,6 +116,7 @@
     add_external_issue,
     add_field_errors_to_response,
     async_delete,
+    build_count_subquery,
     calculate_finding_age,
     get_enabled_notifications_list,
     get_open_findings_burndown,
@@ -134,10 +138,15 @@ def product(request):
     # perform all stuff for filtering and pagination first, before annotation/prefetching
     # otherwise the paginator will perform all the annotations/prefetching already only to count the total number of records
     # see https://code.djangoproject.com/ticket/23771 and https://code.djangoproject.com/ticket/25375
+
     name_words = prods.values_list("name", flat=True)
+    base_findings = Finding.objects.filter(test__engagement__product_id=OuterRef("pk"), active=True)
     prods = prods.annotate(
-        findings_count=Count("engagement__test__finding", filter=Q(engagement__test__finding__active=True)),
+        findings_count=Coalesce(
+            build_count_subquery(base_findings, group_field="test__engagement__product_id"), Value(0),
+        ),
     )
+
     filter_string_matching = get_system_setting("filter_string_matching", False)
     filter_class = ProductFilterWithoutObjectLookups if filter_string_matching else ProductFilter
     prod_filter = filter_class(request.GET, queryset=prods, user=request.user)
@@ -157,48 +166,63 @@ def product(request):
 
 
 def prefetch_for_product(prods):
-    prefetched_prods = prods
-    if isinstance(prods,
-                  QuerySet):  # old code can arrive here with prods being a list because the query was already executed
-
-        prefetched_prods = prefetched_prods.prefetch_related("team_manager")
-        prefetched_prods = prefetched_prods.prefetch_related("product_manager")
-        prefetched_prods = prefetched_prods.prefetch_related("technical_contact")
-
-        prefetched_prods = prefetched_prods.annotate(
-            active_engagement_count=Count("engagement__id", filter=Q(engagement__active=True)))
-        prefetched_prods = prefetched_prods.annotate(
-            closed_engagement_count=Count("engagement__id", filter=Q(engagement__active=False)))
-        prefetched_prods = prefetched_prods.annotate(last_engagement_date=Max("engagement__target_start"))
-        prefetched_prods = prefetched_prods.annotate(active_finding_count=Count("engagement__test__finding__id",
-                                                                                filter=Q(
-                                                                                    engagement__test__finding__active=True)))
-        prefetched_prods = prefetched_prods.annotate(
-            active_verified_finding_count=Count("engagement__test__finding__id",
-                                                filter=Q(
-                                                    engagement__test__finding__active=True,
-                                                    engagement__test__finding__verified=True)))
-        prefetched_prods = prefetched_prods.prefetch_related("jira_project_set__jira_instance")
-        prefetched_prods = prefetched_prods.prefetch_related("members")
-        prefetched_prods = prefetched_prods.prefetch_related("prod_type__members")
-        active_endpoint_query = Endpoint.objects.filter(
-            status_endpoint__mitigated=False,
-            status_endpoint__false_positive=False,
-            status_endpoint__out_of_scope=False,
-            status_endpoint__risk_accepted=False,
-        ).distinct()
-        prefetched_prods = prefetched_prods.prefetch_related(
-            Prefetch("endpoint_set", queryset=active_endpoint_query, to_attr="active_endpoints"))
-        prefetched_prods = prefetched_prods.prefetch_related("tags")
+    # old code can arrive here with prods being a list because the query was already executed
+    if not isinstance(prods, QuerySet):
+        logger.debug("unable to prefetch because query was already executed")
+        return prods
 
-        if get_system_setting("enable_github"):
-            prefetched_prods = prefetched_prods.prefetch_related(
-                Prefetch("github_pkey_set", queryset=GITHUB_PKey.objects.all().select_related("git_conf"),
-                         to_attr="github_confs"))
+    prefetched_prods = prods.select_related("team_manager", "product_manager", "technical_contact").prefetch_related(
+        "tags",
+        "members",
+        "prod_type__members",
+        "jira_project_set__jira_instance",
+    )
 
-    else:
-        logger.debug("unable to prefetch because query was already executed")
+    engagements = Engagement.objects.filter(product_id=OuterRef("pk"))
+    count_subquery = partial(build_count_subquery, group_field="product_id")
+    prefetched_prods = prefetched_prods.annotate(
+        active_engagement_count=Coalesce(count_subquery(engagements.filter(active=True)), Value(0)),
+        closed_engagement_count=Coalesce(count_subquery(engagements.filter(active=False)), Value(0)),
+        last_engagement_date=Subquery(
+            engagements.order_by("-target_start").values("target_start")[:1], output_field=DateField(),
+        ),
+    )
 
+    base_findings = Finding.objects.filter(test__engagement__product_id=OuterRef("pk"))
+    count_subquery = partial(build_count_subquery, group_field="test__engagement__product_id")
+    prefetched_prods = prefetched_prods.annotate(
+        active_finding_count=Coalesce(count_subquery(base_findings.filter(active=True)), Value(0)),
+        active_verified_finding_count=Coalesce(
+            count_subquery(base_findings.filter(active=True, verified=True)),
+            Value(0),
+        ),
+    )
+    prefetched_prods = prefetched_prods.annotate(
+        total_reimport_count=Coalesce(
+            count_subquery(
+                Test_Import.objects.filter(test__engagement__product_id=OuterRef("pk"), type=Test_Import.REIMPORT_TYPE),
+            ),
+            Value(0),
+        ),
+    )
+
+    active_endpoint_qs = Endpoint.objects.filter(
+        status_endpoint__mitigated=False,
+        status_endpoint__false_positive=False,
+        status_endpoint__out_of_scope=False,
+        status_endpoint__risk_accepted=False,
+    ).distinct()
+
+    prefetched_prods = prefetched_prods.prefetch_related(
+        Prefetch("endpoint_set", queryset=active_endpoint_qs, to_attr="active_endpoints"),
+    )
+
+    if get_system_setting("enable_github"):
+        prefetched_prods = prefetched_prods.prefetch_related(
+            Prefetch(
+                "github_pkey_set", queryset=GITHUB_PKey.objects.all().select_related("git_conf"), to_attr="github_confs",
+            ),
+        )
     return prefetched_prods