perf(importers): add defer_product_grading option + harden watson index-flush hook

dojo/importers/default_importer.py

@@ -302,7 +302,10 @@ def _process_findings_internal(
                     finding_ids_batch,
                     dedupe_option=True,
                     rules_option=True,
-                    product_grading_option=True,
+                    # Callers may defer grading to a single end-of-import pass
+                    # (e.g. a large chunked import) to avoid one grade
+                    # recalculation per dedupe batch; default keeps per-batch grading.
+                    product_grading_option=not self.defer_product_grading,
                     issue_updater_option=True,
                     push_to_jira=push_to_jira,
                     # 'async_wait' joins on this dispatch via AsyncResult.get(), so its
@@ -333,8 +336,10 @@ def _process_findings_internal(
 
         # Note: All chord batching is now handled within the loop above
 
-        # Always perform an initial grading, even though it might get overwritten later.
-        perform_product_grading(self.test.engagement.product)
+        # Perform an initial grading (it may be overwritten later) unless the caller defers
+        # grading to a single end-of-import pass (see defer_product_grading).
+        if not self.defer_product_grading:
+            perform_product_grading(self.test.engagement.product)
 
         return new_findings
 
@@ -438,7 +443,7 @@ def close_old_findings(
                     jira_services.push(finding_group)
 
         # Calculate grade once after all findings have been closed
-        if old_findings:
+        if old_findings and not self.defer_product_grading:
             perform_product_grading(self.test.engagement.product)
 
         return old_findings

dojo/importers/default_reimporter.py

@@ -296,7 +296,8 @@ def _process_findings_internal(
 
         logger.debug(f"original_findings_qyer: {original_findings.query}")
         self.original_items = list(original_findings)
-        logger.debug(f"original_items: {[(item.id, item.hash_code) for item in self.original_items]}")
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug("original_items: %s", [(item.id, item.hash_code) for item in self.original_items])
         self.new_items = []
         self.reactivated_items = []
         self.unchanged_items = []
@@ -467,7 +468,10 @@ def _process_findings_internal(
                             finding_ids_batch,
                             dedupe_option=True,
                             rules_option=True,
-                            product_grading_option=True,
+                            # Callers may defer grading to a single end-of-import pass
+                            # (e.g. a large chunked reimport) to avoid one grade
+                            # recalculation per dedupe batch; default keeps per-batch grading.
+                            product_grading_option=not self.defer_product_grading,
                             issue_updater_option=True,
                             push_to_jira=push_to_jira,
                             jira_instance_id=getattr(self.jira_instance, "id", None),
@@ -494,8 +498,10 @@ def _process_findings_internal(
 
         # Note: All chord batching is now handled within the loop above
 
-        # Synchronous tasks were already executed during processing, just calculate grade
-        perform_product_grading(self.test.engagement.product)
+        # Synchronous tasks were already executed during processing, just calculate grade.
+        # Callers may defer grading to a single end-of-import pass (see defer_product_grading).
+        if not self.defer_product_grading:
+            perform_product_grading(self.test.engagement.product)
 
         return self.new_items, self.reactivated_items, self.to_mitigate, self.untouched
 
@@ -584,7 +590,7 @@ def close_old_findings(
                 if self.push_to_jira or jira_services.is_keep_in_sync(finding_group, prefetched_jira_instance=self.jira_instance):
                     jira_services.push(finding_group)
         # Calculate grade once after all findings have been closed
-        if mitigated_findings:
+        if mitigated_findings and not self.defer_product_grading:
             perform_product_grading(self.test.engagement.product)
 
         return mitigated_findings

dojo/importers/options.py

@@ -65,6 +65,7 @@ def load_base_options(
         self.close_old_findings_toggle: bool = self.validate_close_old_findings(*args, **kwargs)
         self.close_old_findings_product_scope: bool = self.validate_close_old_findings_product_scope(*args, **kwargs)
         self.do_not_reactivate: bool = self.validate_do_not_reactivate(*args, **kwargs)
+        self.defer_product_grading: bool = self.validate_defer_product_grading(*args, **kwargs)
         self.commit_hash: str = self.validate_commit_hash(*args, **kwargs)
         self.create_finding_groups_for_all_findings: bool = self.validate_create_finding_groups_for_all_findings(*args, **kwargs)
         self.endpoints_to_add: list[Endpoint] | list[AbstractLocation] | None = self.validate_endpoints_to_add(*args, **kwargs)
@@ -349,6 +350,23 @@ def validate_do_not_reactivate(
             **kwargs,
         )
 
+    def validate_defer_product_grading(
+        self,
+        *args: list,
+        **kwargs: dict,
+    ) -> bool:
+        # Defer per-batch/end/close-old product grade dispatches so callers (e.g. a large
+        # chunked reimport) can grade once at the end instead of once per dedupe batch.
+        # Caveat: all-or-nothing -- when True the caller MUST run perform_product_grading()
+        # itself after the import; nothing re-grades automatically, so a stale grade goes unnoticed.
+        return self.validate(
+            "defer_product_grading",
+            expected_types=[bool],
+            required=False,
+            default=False,
+            **kwargs,
+        )
+
     def validate_deduplication_execution_mode(
         self,
         *args: list,

dojo/middleware.py

@@ -1,7 +1,7 @@
 import logging
 import re
 import time
-from contextlib import suppress
+from contextlib import contextmanager, suppress
 from threading import local
 from urllib.parse import quote
 
@@ -291,6 +291,31 @@ def _drain_search_context_to_async(objects, source):
         objects.discard(entry)
 
 
+# Thread-local kill-switch for the intermediate-flush hook (below). Bulk operations that run
+# under watson.skip_index_update() set this so the hook no-ops on their thread -- see
+# suppress_intermediate_flush().
+_intermediate_flush_state = local()
+
+
+@contextmanager
+def suppress_intermediate_flush():
+    """
+    Disable the watson intermediate-flush hook on the current thread.
+
+    watson.skip_index_update() only marks its context invalid at __exit__, so while a skip block
+    is still accumulating, the intermediate-flush hook can't tell it's a skip context and would
+    drain it -- dispatching async index tasks for objects meant to be discarded, defeating the
+    skip. Wrap a skip_index_update() block in this so the hook returns early (no accumulate, no
+    drain) and the skip actually holds. Thread-local, so it only affects the wrapping thread.
+    """
+    previous = getattr(_intermediate_flush_state, "suppressed", False)
+    _intermediate_flush_state.suppressed = True
+    try:
+        yield
+    finally:
+        _intermediate_flush_state.suppressed = previous
+
+
 def install_intermediate_flush_hook():
     """
     Wrap `watson.search.search_context_manager.add_to_context` with a
@@ -311,7 +336,23 @@ def install_intermediate_flush_hook():
     original_add = cls.add_to_context
 
     def add_to_context_with_flush(self, engine, obj):
+        # Bulk ops under skip_index_update() (e.g. a large chunked reimport) set this to make the
+        # hook a full no-op on this thread. skip_index_update() only invalidates its context at
+        # __exit__, so mid-accumulation the hook would otherwise drain (index) objects meant to be
+        # discarded -- defeating the skip. Returning before original_add means nothing is even
+        # accumulated, so the skip holds and the caller's own end-of-op reindex does the indexing.
+        if getattr(_intermediate_flush_state, "suppressed", False):
+            return
         original_add(self, engine, obj)
+        # The intermediate flush is a request-path optimization on the global singleton
+        # context (AsyncSearchContextMiddleware). The async reindex task
+        # update_watson_search_index_for_model() builds its OWN SearchContextManager and
+        # IS the drain target -- if it re-drained its batch it would dispatch a clone of
+        # itself, discard those pks unindexed, and loop forever (queue ~0, worker pegged,
+        # nothing indexed). Only the singleton intermediate-flushes; any ad-hoc context
+        # manager indexes its own batch on end().
+        if self is not search_context_manager:
+            return
         threshold = getattr(settings, "WATSON_ASYNC_INDEX_UPDATE_BATCH_SIZE", 1000)
         if threshold <= 0 or not self._stack:
             return

unittests/test_importers_closeold.py

@@ -1,9 +1,11 @@
 import logging
+from unittest.mock import patch
 
 from django.utils import timezone
 
 import dojo.risk_acceptance.helper as ra_helper
 from dojo.importers.default_importer import DefaultImporter
+from dojo.importers.default_reimporter import DefaultReImporter
 from dojo.models import Development_Environment, Engagement, Product, Product_Type, User
 
 from .dojo_test_case import DojoTestCase, get_unit_tests_scans_path
@@ -59,6 +61,103 @@ def test_close_old_same_engagement(self):
             # If this behavior changes, or dedupe is on, the number of closed findings will be 4
             self.assertEqual(8, len_closed_findings)
 
+    def test_defer_product_grading_skips_close_old_grade(self):
+        """
+        With defer_product_grading=True, an import that closes old findings must skip the
+        close-old grade in close_old_findings (as well as the end-of-process grade), so the
+        importer never calls perform_product_grading -- the caller grades once itself. Pins the
+        close_old_findings deferral guard; the per-batch/end guards are covered in
+        test_importers_importer.py.
+        """
+        scan_type = "Acunetix Scan"
+        user, _ = User.objects.get_or_create(username="admin")
+        product_type, _ = Product_Type.objects.get_or_create(name="closeold-defer")
+        environment, _ = Development_Environment.objects.get_or_create(name="Development")
+        product, _ = Product.objects.get_or_create(
+            name="TestDojoCloseOldDefer",
+            description="Test",
+            prod_type=product_type,
+        )
+        engagement, _ = Engagement.objects.get_or_create(
+            name="Close Old Defer Grading",
+            product=product,
+            target_start=timezone.now(),
+            target_end=timezone.now(),
+        )
+        import_options = {
+            "user": user,
+            "lead": user,
+            "scan_date": None,
+            "environment": environment,
+            "active": True,
+            "verified": False,
+            "engagement": engagement,
+            "scan_type": scan_type,
+        }
+        # Seed the engagement with several findings (no close-old on this first import).
+        with (get_unit_tests_scans_path("acunetix") / "many_findings.xml").open(encoding="utf-8") as many_findings_scan:
+            DefaultImporter(close_old_findings=False, **import_options).process_scan(many_findings_scan)
+        # Import a single finding with close-old + deferral: the other findings close, but the
+        # importer must not grade (neither the end pass nor the close-old pass). perform_product_grading
+        # is only referenced by those two direct call sites in default_importer, so a call here would
+        # mean a guard regressed.
+        with (get_unit_tests_scans_path("acunetix") / "one_finding.xml").open(encoding="utf-8") as single_finding_scan:
+            importer = DefaultImporter(close_old_findings=True, defer_product_grading=True, **import_options)
+            with patch("dojo.importers.default_importer.perform_product_grading") as grade:
+                _, _, _, len_closed_findings, _, _, _ = importer.process_scan(single_finding_scan)
+            self.assertGreater(len_closed_findings, 0, "old findings must actually close, or the close-old guard isn't exercised")
+            grade.assert_not_called()
+
+    def test_reimport_defer_product_grading_skips_close_old_and_end_grade(self):
+        """
+        Reimport counterpart to test_defer_product_grading_skips_close_old_grade: with
+        defer_product_grading=True, a reimport that mitigates old findings must skip both the
+        end-of-process grade and the close-old grade in DefaultReImporter, so the reimporter never
+        calls perform_product_grading -- the caller grades once. Pins the reimporter's end and
+        close-old deferral guards (its per-batch guard is covered by the Pro connectors tests).
+        """
+        scan_type = "Acunetix Scan"
+        user, _ = User.objects.get_or_create(username="admin")
+        product_type, _ = Product_Type.objects.get_or_create(name="closeold-defer-reimport")
+        environment, _ = Development_Environment.objects.get_or_create(name="Development")
+        product, _ = Product.objects.get_or_create(
+            name="TestDojoCloseOldDeferReimport",
+            description="Test",
+            prod_type=product_type,
+        )
+        engagement, _ = Engagement.objects.get_or_create(
+            name="Close Old Defer Grading Reimport",
+            product=product,
+            target_start=timezone.now(),
+            target_end=timezone.now(),
+        )
+        common = {
+            "user": user,
+            "lead": user,
+            "scan_date": None,
+            "environment": environment,
+            "active": True,
+            "verified": False,
+            "scan_type": scan_type,
+        }
+        # Seed a test with several findings via an initial import.
+        with (get_unit_tests_scans_path("acunetix") / "many_findings.xml").open(encoding="utf-8") as many_findings_scan:
+            importer = DefaultImporter(engagement=engagement, close_old_findings=False, **common)
+            test, _, _, _, _, _, _ = importer.process_scan(many_findings_scan)
+        # Reimport a single finding with close-old + deferral: the absent findings mitigate, but the
+        # reimporter must not grade. perform_product_grading is only referenced by the end and
+        # close-old call sites in default_reimporter, so a call here means a guard regressed.
+        with (get_unit_tests_scans_path("acunetix") / "one_finding.xml").open(encoding="utf-8") as single_finding_scan:
+            reimporter = DefaultReImporter(test=test, close_old_findings=True, defer_product_grading=True, **common)
+            with patch("dojo.importers.default_reimporter.perform_product_grading") as grade:
+                reimporter.process_scan(single_finding_scan)
+            self.assertGreater(
+                test.finding_set.filter(is_mitigated=True).count(),
+                0,
+                "old findings must actually mitigate, or the close-old guard isn't exercised",
+            )
+            grade.assert_not_called()
+
     def test_close_old_same_product_scan(self):
         scan_type = "Acunetix Scan"
         user, _ = User.objects.get_or_create(username="admin")

unittests/test_importers_importer.py

@@ -91,6 +91,91 @@ def test_parse_findings(self):
             for finding in new_findings:
                 self.assertIn(finding.numerical_severity, ["S0", "S1", "S2", "S3", "S4"])
 
+    def _grading_import_options(self, scan_type, name_suffix, **overrides):
+        """Minimal DefaultImporter options for the product-grading-deferral tests."""
+        user, _ = User.objects.get_or_create(username="admin")
+        product_type, _ = Product_Type.objects.get_or_create(name=f"grade-defer-pt-{name_suffix}")
+        product, _ = Product.objects.get_or_create(
+            name=f"grade-defer-prod-{name_suffix}",
+            description="test product",
+            prod_type=product_type,
+        )
+        engagement, _ = Engagement.objects.get_or_create(
+            name=f"grade-defer-eng-{name_suffix}",
+            product=product,
+            target_start=timezone.now(),
+            target_end=timezone.now(),
+        )
+        environment, _ = Development_Environment.objects.get_or_create(name="Development")
+        options = {
+            "user": user,
+            "lead": user,
+            "scan_date": None,
+            "environment": environment,
+            "minimum_severity": "Info",
+            "active": True,
+            "verified": True,
+            "scan_type": scan_type,
+            "engagement": engagement,
+        }
+        options.update(overrides)
+        return options
+
+    def _process_findings_capturing_grading(self, importer, parsed_findings):
+        """Run process_findings with dispatch + grading patched; return (per_batch_options, grade_mock)."""
+        per_batch_options = []
+
+        def _spy(_task, *args, **kwargs):
+            per_batch_options.append(kwargs.get("product_grading_option"))
+
+        with (
+            patch("dojo.importers.default_importer.dojo_dispatch_task", side_effect=_spy),
+            patch("dojo.importers.default_importer.perform_product_grading") as grade,
+        ):
+            importer.process_findings(parsed_findings)
+        return per_batch_options, grade
+
+    def test_defer_product_grading_skips_per_batch_and_end_grade(self):
+        """
+        defer_product_grading=True makes DefaultImporter dispatch each dedupe batch with
+        product_grading_option=False and skip the end-of-import perform_product_grading pass,
+        leaving the caller to grade once. Mirrors DefaultReImporter's deferral.
+        """
+        with ACUNETIX_AUDIT_ONE_VULN_FILENAME.open(encoding="utf-8") as scan:
+            scan_type = "Acunetix Scan"
+            importer = DefaultImporter(**self._grading_import_options(scan_type, "defer", defer_product_grading=True))
+            test = importer.create_test(scan_type)
+            parsed_findings = importer.get_parser().get_findings(scan, test)
+
+            per_batch_options, grade = self._process_findings_capturing_grading(importer, parsed_findings)
+
+            self.assertTrue(per_batch_options, "at least one post-process batch should be dispatched")
+            self.assertTrue(
+                all(opt is False for opt in per_batch_options),
+                f"deferred import must dispatch product_grading_option=False, got {per_batch_options}",
+            )
+            grade.assert_not_called()
+
+    def test_default_import_grades_per_batch_and_at_end(self):
+        """
+        Control for the deferral test: without defer_product_grading, DefaultImporter grades per
+        dedupe batch (product_grading_option=True) and once at the end, as it always has.
+        """
+        with ACUNETIX_AUDIT_ONE_VULN_FILENAME.open(encoding="utf-8") as scan:
+            scan_type = "Acunetix Scan"
+            importer = DefaultImporter(**self._grading_import_options(scan_type, "default"))
+            test = importer.create_test(scan_type)
+            parsed_findings = importer.get_parser().get_findings(scan, test)
+
+            per_batch_options, grade = self._process_findings_capturing_grading(importer, parsed_findings)
+
+            self.assertTrue(per_batch_options, "at least one post-process batch should be dispatched")
+            self.assertTrue(
+                all(opt is True for opt in per_batch_options),
+                f"default import must dispatch product_grading_option=True, got {per_batch_options}",
+            )
+            grade.assert_called_once()
+
     def test_import_scan(self):
         with (get_unit_tests_path() / "scans" / "sarif" / "spotbugs.sarif").open(encoding="utf-8") as scan:
             scan_type = SarifParser().get_scan_types()[0]  # SARIF format implement the new method