extract chord orchestration into method

Author: valentijnscholten

dojo/importers/base_importer.py

@@ -1,6 +1,7 @@
 import base64
 import logging
 
+from celery import chord, group
 from django.conf import settings
 from django.core.exceptions import ValidationError
 from django.core.files.base import ContentFile
@@ -10,6 +11,7 @@
 from django.utils.timezone import make_aware
 
 import dojo.finding.helper as finding_helper
+from dojo import utils
 from dojo.importers.endpoint_manager import EndpointManager
 from dojo.importers.options import ImporterOptions
 from dojo.models import (
@@ -24,6 +26,7 @@
     Endpoint,
     FileUpload,
     Finding,
+    System_Settings,
     Test,
     Test_Import,
     Test_Import_Finding_Action,
@@ -532,6 +535,47 @@ def update_test_type_from_internal_test(self, internal_test: ParserTest) -> None
             self.test.test_type.dynamic_tool = dynamic_tool
         self.test.test_type.save()
 
+    def maybe_launch_post_processing_chord(
+        self,
+        post_processing_task_signatures,
+        current_batch_number: int,
+        max_batch_size: int,
+        *
+        is_final_batch: bool,
+    ) -> tuple[list, int, bool]:
+        """
+        Helper to optionally launch a chord of post-processing tasks with a calculate-grade callback
+        when async is desired. Uses exponential batch sizing up to the configured max batch size.
+
+        Returns a tuple of (post_processing_task_signatures, current_batch_number, launched)
+        where launched indicates whether a chord/group was dispatched and signatures were reset.
+        """
+        launched = False
+        if not post_processing_task_signatures:
+            return post_processing_task_signatures, current_batch_number, launched
+
+        current_batch_size = min(2 ** current_batch_number, max_batch_size)
+        batch_full = len(post_processing_task_signatures) >= current_batch_size
+
+        if batch_full or is_final_batch:
+            product = self.test.engagement.product
+            system_settings = System_Settings.objects.get()
+            if system_settings.enable_product_grade:
+                calculate_grade_signature = utils.calculate_grade_signature(product)
+                chord(post_processing_task_signatures)(calculate_grade_signature)
+            else:
+                group(post_processing_task_signatures).apply_async()
+
+            logger.debug(
+                f"Launched chord with {len(post_processing_task_signatures)} tasks (batch #{current_batch_number}, size: {len(post_processing_task_signatures)})",
+            )
+            post_processing_task_signatures = []
+            if not is_final_batch:
+                current_batch_number += 1
+            launched = True
+
+        return post_processing_task_signatures, current_batch_number, launched
+
     def verify_tool_configuration_from_test(self):
         """
         Verify that the Tool_Configuration supplied along with the

dojo/importers/default_importer.py

@@ -1,21 +1,18 @@
 import logging
 
-from celery import chord, group
 from django.core.files.uploadedfile import TemporaryUploadedFile
 from django.core.serializers import serialize
 from django.db.models.query_utils import Q
 from django.urls import reverse
 
 import dojo.jira_link.helper as jira_helper
-from dojo import utils
 from dojo.decorators import we_want_async
 from dojo.finding import helper as finding_helper
 from dojo.importers.base_importer import BaseImporter, Parser
 from dojo.importers.options import ImporterOptions
 from dojo.models import (
     Engagement,
     Finding,
-    System_Settings,
     Test,
     Test_Import,
 )
@@ -253,30 +250,12 @@ def process_findings(
 
             # Check if we should launch a chord (batch full or end of findings)
             if we_want_async(async_user=self.user) and post_processing_task_signatures:
-                # Calculate current batch size: 2^batch_number, capped at max_batch_size
-                # We do this because post processing only starts after all tasks have been added to the chord
-                # So we start with small batches to minmize the delay
-                current_batch_size = min(2 ** current_batch_number, max_batch_size)
-
-                batch_full = len(post_processing_task_signatures) >= current_batch_size
-
-                if batch_full or is_final_finding:
-                    # Launch chord with current batch of signatures
-                    product = self.test.engagement.product
-                    system_settings = System_Settings.objects.get()
-                    if system_settings.enable_product_grade:
-                        calculate_grade_signature = utils.calculate_grade_signature(product)
-                        chord(post_processing_task_signatures)(calculate_grade_signature)
-                    elif post_processing_task_signatures:
-                        # If product grading is disabled, just run the post-processing tasks without the grade calculation callback
-                        group(post_processing_task_signatures).apply_async()
-
-                    logger.debug(f"Launched chord with {len(post_processing_task_signatures)} tasks (batch #{current_batch_number}, size: {len(post_processing_task_signatures)})")
-
-                    # Reset for next batch (only if not final)
-                    post_processing_task_signatures = []
-                    if not is_final_finding:
-                        current_batch_number += 1
+                post_processing_task_signatures, current_batch_number, _ = self.maybe_launch_post_processing_chord(
+                    post_processing_task_signatures,
+                    current_batch_number,
+                    max_batch_size,
+                    is_final_finding,
+                )
             else:
                 # Execute task immediately for synchronous processing
                 post_processing_task_signature()

dojo/importers/default_reimporter.py

@@ -1,21 +1,18 @@
 import logging
 
-from celery import chord, group
 from django.core.files.uploadedfile import TemporaryUploadedFile
 from django.core.serializers import serialize
 from django.db.models.query_utils import Q
 
 import dojo.finding.helper as finding_helper
 import dojo.jira_link.helper as jira_helper
-from dojo import utils
 from dojo.decorators import we_want_async
 from dojo.importers.base_importer import BaseImporter, Parser
 from dojo.importers.options import ImporterOptions
 from dojo.models import (
     Development_Environment,
     Finding,
     Notes,
-    System_Settings,
     Test,
     Test_Import,
 )
@@ -271,30 +268,12 @@ def process_findings(
 
             # Check if we should launch a chord (batch full or end of findings)
             if we_want_async(async_user=self.user) and post_processing_task_signatures:
-                # Calculate current batch size: 2^batch_number, capped at max_batch_size
-                # We do this because post processing only starts after all tasks have been added to the chord
-                # So we start with small batches to minmize the delay
-                current_batch_size = min(2 ** current_batch_number, max_batch_size)
-
-                batch_full = len(post_processing_task_signatures) >= current_batch_size
-
-                if batch_full or is_final:
-                    # Launch chord with current batch of signatures
-                    product = self.test.engagement.product
-                    system_settings = System_Settings.objects.get()
-                    if system_settings.enable_product_grade:
-                        calculate_grade_signature = utils.calculate_grade_signature(product)
-                        chord(post_processing_task_signatures)(calculate_grade_signature)
-                    elif post_processing_task_signatures:
-                        # If product grading is disabled, just run the post-processing tasks without the grade calculation callback
-                        group(post_processing_task_signatures).apply_async()
-
-                    logger.debug(f"Launched chord with {len(post_processing_task_signatures)} tasks (batch #{current_batch_number}, size: {len(post_processing_task_signatures)})")
-
-                    # Reset for next batch (only if not final)
-                    post_processing_task_signatures = []
-                    if not is_final:
-                        current_batch_number += 1
+                post_processing_task_signatures, current_batch_number, _ = self.maybe_launch_post_processing_chord(
+                    post_processing_task_signatures,
+                    current_batch_number,
+                    max_batch_size,
+                    is_final,
+                )
             else:
                 post_processing_task_signature()