bugfix: use subquery for (finding) counts (#12784)

* view_test: use subquery for finding counts

* fix two more group by errors

* change all counts to subqueries

* fix query

Author: valentijnscholten

.gitignore

@@ -147,3 +147,4 @@ docs/.devcontainer/devcontainer.json
 docs/.devcontainer/Dockerfile
 docs/LICENSE
 docs/.hugo_build.lock
+.cursor-rules

dojo/endpoint/views.py

@@ -8,7 +8,8 @@
 from django.contrib.admin.utils import NestedObjects
 from django.core.exceptions import PermissionDenied
 from django.db import DEFAULT_DB_ALIAS
-from django.db.models import Count, Q, QuerySet
+from django.db.models import OuterRef, QuerySet, Value
+from django.db.models.functions import Coalesce
 from django.http import HttpResponseRedirect
 from django.shortcuts import get_object_or_404, render
 from django.urls import reverse
@@ -22,6 +23,7 @@
 from dojo.filters import EndpointFilter, EndpointFilterWithoutObjectLookups
 from dojo.forms import AddEndpointForm, DeleteEndpointForm, DojoMetaDataForm, EditEndpointForm, ImportEndpointMetaForm
 from dojo.models import DojoMeta, Endpoint, Endpoint_Status, Finding, Product
+from dojo.query_utils import build_count_subquery
 from dojo.utils import (
     Product_Tab,
     add_breadcrumb,
@@ -454,7 +456,11 @@ def endpoint_status_bulk_update(request, fid):
 def prefetch_for_endpoints(endpoints):
     if isinstance(endpoints, QuerySet):
         endpoints = endpoints.prefetch_related("product", "tags", "product__tags")
-        endpoints = endpoints.annotate(active_finding_count=Count("finding__id", filter=Q(finding__active=True)))
+        active_finding_subquery = build_count_subquery(
+            Finding.objects.filter(endpoints=OuterRef("pk"), active=True),
+            group_field="endpoints",
+        )
+        endpoints = endpoints.annotate(active_finding_count=Coalesce(active_finding_subquery, Value(0)))
     else:
         logger.debug("unable to prefetch because query was already executed")
 

dojo/engagement/views.py

@@ -16,7 +16,7 @@
 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, OuterRef, Q, Value
+from django.db.models import 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
@@ -95,6 +95,7 @@
     LargeScanSizeProductAnnouncement,
     ScanTypeProductAnnouncement,
 )
+from dojo.query_utils import build_count_subquery
 from dojo.risk_acceptance.helper import prefetch_for_expiration
 from dojo.tools.factory import get_scan_types_sorted
 from dojo.user.queries import get_authorized_users
@@ -105,7 +106,6 @@
     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,
@@ -217,7 +217,11 @@ def engagements_all(request):
     # count using prefetch instead of just using 'engagement__set_test_test` to avoid loading all test in memory just to count them
     filter_string_matching = get_system_setting("filter_string_matching", False)
     products_filter_class = ProductEngagementsFilterWithoutObjectLookups if filter_string_matching else ProductEngagementsFilter
-    engagement_query = Engagement.objects.annotate(test_count=Count("test__id"))
+    test_count_subquery = build_count_subquery(
+        Test.objects.filter(engagement=OuterRef("pk")),
+        group_field="engagement_id",
+    )
+    engagement_query = Engagement.objects.annotate(test_count=Coalesce(test_count_subquery, Value(0)))
     filter_qs = products_with_engagements.prefetch_related(
         Prefetch("engagement_set", queryset=products_filter_class(request.GET, engagement_query).qs),
     )
@@ -417,7 +421,13 @@ def get_template(self):
         return "dojo/view_eng.html"
 
     def get_risks_accepted(self, eng):
-        return eng.risk_acceptance.all().select_related("owner").annotate(accepted_findings_count=Count("accepted_findings__id"))
+        accepted_findings_subquery = build_count_subquery(
+            Finding.objects.filter(risk_acceptance=OuterRef("pk")),
+            group_field="risk_acceptance",
+        )
+        return eng.risk_acceptance.all().select_related("owner").annotate(
+            accepted_findings_count=Coalesce(accepted_findings_subquery, Value(0)),
+        )
 
     def get_filtered_tests(
         self,

dojo/finding/queries.py

@@ -1,16 +1,27 @@
+import logging
+from functools import partial
+
 from crum import get_current_user
-from django.db.models import Exists, OuterRef, Q
+from django.db.models import Exists, OuterRef, Q, Value
+from django.db.models.functions import Coalesce
+from django.db.models.query import Prefetch, QuerySet
 
 from dojo.authorization.authorization import get_roles_for_permission, user_has_global_permission
 from dojo.models import (
+    IMPORT_UNTOUCHED_FINDING,
+    Endpoint_Status,
     Finding,
     Product_Group,
     Product_Member,
     Product_Type_Group,
     Product_Type_Member,
     Stub_Finding,
+    Test_Import_Finding_Action,
     Vulnerability_Id,
 )
+from dojo.query_utils import build_count_subquery
+
+logger = logging.getLogger(__name__)
 
 
 def get_authorized_groups(permission, user=None):
@@ -146,3 +157,70 @@ def get_authorized_vulnerability_ids(permission, queryset=None, user=None):
         | Q(finding__test__engagement__product__member=True)
         | Q(finding__test__engagement__product__prod_type__authorized_group=True)
         | Q(finding__test__engagement__product__authorized_group=True))
+
+
+def prefetch_for_findings(findings, prefetch_type="all", *, exclude_untouched=True):
+    """
+    Unified prefetch function for findings across the application.
+
+    Args:
+        findings: QuerySet of findings to prefetch
+        prefetch_type: "all" or "open" - controls risk acceptance prefetching
+        exclude_untouched: Whether to exclude untouched import actions
+
+    """
+    if not isinstance(findings, QuerySet):
+        logger.debug("unable to prefetch because query was already executed")
+        return findings
+
+    # Base prefetches - always needed
+    prefetched_findings = findings.prefetch_related(
+        "reviewers",
+        "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",
+        "reporter",
+    )
+
+    # Conditional prefetches for non-open findings
+    if prefetch_type != "open":
+        prefetched_findings = prefetched_findings.prefetch_related(
+            "risk_acceptance_set",
+            "risk_acceptance_set__accepted_findings",
+            "original_finding",
+            "duplicate_finding",
+        )
+
+    # Import actions - configurable filtering
+    if exclude_untouched:
+        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")
+
+    # Standard prefetches
+    prefetched_findings = prefetched_findings.prefetch_related(
+        "notes",
+        "tags",
+        "endpoints",
+        "status_finding",
+        "finding_group_set",
+        "finding_group_set__jira_issue",  # Include both variants
+        "test__engagement__product__members",
+        "test__engagement__product__prod_type__members",
+        "vulnerability_id_set",
+    )
+
+    # Endpoint counts using optimized subqueries
+    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)),
+    )

dojo/finding/views.py

@@ -6,7 +6,6 @@
 import logging
 import mimetypes
 from collections import OrderedDict, defaultdict
-from functools import partial
 from itertools import chain
 from pathlib import Path
 
@@ -15,8 +14,8 @@
 from django.core import serializers
 from django.core.exceptions import PermissionDenied, ValidationError
 from django.db import models
-from django.db.models import OuterRef, QuerySet, Value
-from django.db.models.functions import Coalesce, Length
+from django.db.models import QuerySet
+from django.db.models.functions import 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
@@ -50,7 +49,7 @@
     TestImportFilter,
     TestImportFindingActionFilter,
 )
-from dojo.finding.queries import get_authorized_findings
+from dojo.finding.queries import get_authorized_findings, prefetch_for_findings
 from dojo.forms import (
     ApplyFindingTemplateForm,
     ClearFindingReviewForm,
@@ -111,7 +110,6 @@
     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,
@@ -133,61 +131,6 @@
 logger = logging.getLogger(__name__)
 
 
-def prefetch_for_findings(findings, prefetch_type="all", *, exclude_untouched=True):
-    # 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
-
-    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",
-    )
-
-    # for open/active findings, the following 4 prefetches are not needed
-    if prefetch_type != "open":
-        prefetched_findings = prefetched_findings.prefetch_related(
-            "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(
-            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")
-
-    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",
-    )
-
-    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):
     prefetched_findings = findings
     if isinstance(

dojo/finding_group/views.py

@@ -13,7 +13,7 @@
 from dojo.authorization.authorization_decorators import user_is_authorized
 from dojo.authorization.roles_permissions import Permissions
 from dojo.filters import FindingFilter, FindingFilterWithoutObjectLookups
-from dojo.finding.views import prefetch_for_findings
+from dojo.finding.queries import prefetch_for_findings
 from dojo.forms import DeleteFindingGroupForm, EditFindingGroupForm, FindingBulkUpdateForm
 from dojo.models import Engagement, Finding, Finding_Group, GITHUB_PKey, Product
 from dojo.utils import Product_Tab, add_breadcrumb, get_page_items, get_setting, get_system_setting, get_words_for_field

dojo/product/views.py

@@ -107,6 +107,7 @@
     get_authorized_members_for_product_type,
     get_authorized_product_types,
 )
+from dojo.query_utils import build_count_subquery
 from dojo.templatetags.display_tags import asvs_calc_level
 from dojo.tool_config.factory import create_API
 from dojo.tools.factory import get_api_scan_configuration_hints
@@ -117,7 +118,6 @@
     add_external_issue,
     add_field_errors_to_response,
     async_delete,
-    build_count_subquery,
     calculate_finding_age,
     get_enabled_notifications_list,
     get_open_findings_burndown,
@@ -701,7 +701,11 @@ def view_product_metrics(request, pid):
                 accepted_objs_by_severity[finding.get("severity")] += 1
 
     tests = Test.objects.filter(engagement__product=prod).prefetch_related("finding_set", "test_type")
-    tests = tests.annotate(verified_finding_count=Count("finding__id", filter=Q(finding__verified=True)))
+    verified_finding_subquery = build_count_subquery(
+        Finding.objects.filter(test=OuterRef("pk"), verified=True),
+        group_field="test_id",
+    )
+    tests = tests.annotate(verified_finding_count=Coalesce(verified_finding_subquery, Value(0)))
 
     test_data = {}
     for t in tests:
@@ -828,9 +832,7 @@ def view_engagements(request, pid):
 
 
 def prefetch_for_view_engagements(engagements, recent_test_day_count):
-    engagements = engagements.select_related(
-        "lead",
-    ).prefetch_related(
+    engagements = engagements.prefetch_related(
         Prefetch("test_set", queryset=Test.objects.filter(
             id__in=Subquery(
                 Test.objects.filter(
@@ -840,15 +842,41 @@ def prefetch_for_view_engagements(engagements, recent_test_day_count):
             )),
                  ),
         "test_set__test_type",
-    ).annotate(
-        count_tests=Count("test", distinct=True),
-        count_findings_all=Count("test__finding__id"),
-        count_findings_open=Count("test__finding__id", filter=Q(test__finding__active=True)),
-        count_findings_open_verified=Count("test__finding__id",
-                                           filter=Q(test__finding__active=True) & Q(test__finding__verified=True)),
-        count_findings_close=Count("test__finding__id", filter=Q(test__finding__is_mitigated=True)),
-        count_findings_duplicate=Count("test__finding__id", filter=Q(test__finding__duplicate=True)),
-        count_findings_accepted=Count("test__finding__id", filter=Q(test__finding__risk_accepted=True)),
+    ).select_related(
+        "lead",
+    )
+
+    # Use subqueries to avoid GROUP BY issues
+    test_subquery = build_count_subquery(
+        Test.objects.filter(engagement=OuterRef("pk")), group_field="engagement_id",
+    )
+    finding_subquery = build_count_subquery(
+        Finding.objects.filter(test__engagement=OuterRef("pk")), group_field="test__engagement_id",
+    )
+    finding_open_subquery = build_count_subquery(
+        Finding.objects.filter(test__engagement=OuterRef("pk"), active=True), group_field="test__engagement_id",
+    )
+    finding_open_verified_subquery = build_count_subquery(
+        Finding.objects.filter(test__engagement=OuterRef("pk"), active=True, verified=True), group_field="test__engagement_id",
+    )
+    finding_close_subquery = build_count_subquery(
+        Finding.objects.filter(test__engagement=OuterRef("pk"), is_mitigated=True), group_field="test__engagement_id",
+    )
+    finding_duplicate_subquery = build_count_subquery(
+        Finding.objects.filter(test__engagement=OuterRef("pk"), duplicate=True), group_field="test__engagement_id",
+    )
+    finding_accepted_subquery = build_count_subquery(
+        Finding.objects.filter(test__engagement=OuterRef("pk"), risk_accepted=True), group_field="test__engagement_id",
+    )
+
+    engagements = engagements.annotate(
+        count_tests=Coalesce(test_subquery, Value(0)),
+        count_findings_all=Coalesce(finding_subquery, Value(0)),
+        count_findings_open=Coalesce(finding_open_subquery, Value(0)),
+        count_findings_open_verified=Coalesce(finding_open_verified_subquery, Value(0)),
+        count_findings_close=Coalesce(finding_close_subquery, Value(0)),
+        count_findings_duplicate=Coalesce(finding_duplicate_subquery, Value(0)),
+        count_findings_accepted=Coalesce(finding_accepted_subquery, Value(0)),
     )
 
     if System_Settings.objects.get().enable_jira:

dojo/product_type/views.py

@@ -35,10 +35,10 @@
     get_authorized_members_for_product_type,
     get_authorized_product_types,
 )
+from dojo.query_utils import build_count_subquery
 from dojo.utils import (
     add_breadcrumb,
     async_delete,
-    build_count_subquery,
     get_page_items,
     get_setting,
     get_system_setting,