Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Maffooch merged 3 commits into from
Jun 30, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
63d682c
Optimize the DB queries for engagement, finding, product and product_…
DenysMoskalenko Jun 13, 2025
34b4bb1
Apply suggestions from code review
DenysMoskalenko Jun 13, 2025
f6f89c1
Add subquery for test_count for engagements getting
DenysMoskalenko Jun 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 46 additions & 46 deletions dojo/engagement/views.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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"

Expand All @@ -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,
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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():
Expand All @@ -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"
Expand All @@ -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")
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down
123 changes: 48 additions & 75 deletions dojo/finding/views.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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):
Expand Down
Loading