perf(tags): wrap importer in tag_inheritance.batch() + bulk flush (Phase B Stage 2)

dojo/importers/default_importer.py

@@ -5,6 +5,7 @@
 from django.db.models.query_utils import Q
 from django.urls import reverse
 
+from dojo import tag_inheritance
 from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.finding import helper as finding_helper
 from dojo.importers.base_importer import BaseImporter, Parser
@@ -114,29 +115,36 @@ def process_scan(
         # Note: for fresh imports, parse_findings() calls create_test() internally,
         # so self.test is guaranteed to be set after this call.
         parsed_findings = self.parse_findings(scan, parser) or []
-        new_findings = self.process_findings(parsed_findings, **kwargs)
-        # Close any old findings in the processed list if the the user specified for that
-        # to occur in the form that is then passed to the kwargs
-        closed_findings = self.close_old_findings(self.test.finding_set.all(), **kwargs)
-        # Update the timestamps of the test object by looking at the findings imported
-        self.update_timestamps()
-        # Update the test meta
-        self.update_test_meta()
-        # Save the test and engagement for changes to take affect
-        self.test.save()
-        self.test.engagement.save()
-        # Create a test import history object to record the flags sent to the importer
-        # This operation will return None if the user does not have the import history
-        # feature enabled
-        test_import_history = self.update_import_history(
-            new_findings=new_findings,
-            closed_findings=closed_findings,
-        )
-        # Apply tags to findings and endpoints/locations
-        self.apply_import_tags(
-            new_findings=new_findings,
-            closed_findings=closed_findings,
-        )
+        # Open a tag-inheritance context. Signal handlers register touched
+        # instances into the context; the context auto-flushes (bulk-applies
+        # inherited tags) on exit. Mid-context `ctx.flush()` calls drain the
+        # accumulated set early — used before per-batch post-process dispatch
+        # so JIRA labels reflect the full tag state on first push.
+        with tag_inheritance.batch() as tag_ctx:
+            new_findings = self.process_findings(parsed_findings, **kwargs)
+            # Close any old findings in the processed list if the the user specified for that
+            # to occur in the form that is then passed to the kwargs
+            closed_findings = self.close_old_findings(self.test.finding_set.all(), **kwargs)
+            # Update the timestamps of the test object by looking at the findings imported
+            self.update_timestamps()
+            # Update the test meta
+            self.update_test_meta()
+            # Save the test and engagement for changes to take affect
+            self.test.save()
+            self.test.engagement.save()
+            # Create a test import history object to record the flags sent to the importer
+            # This operation will return None if the user does not have the import history
+            # feature enabled
+            test_import_history = self.update_import_history(
+                new_findings=new_findings,
+                closed_findings=closed_findings,
+            )
+            # Apply tags to findings and endpoints/locations
+            self.apply_import_tags(
+                new_findings=new_findings,
+                closed_findings=closed_findings,
+            )
+        # Inheritance auto-flushes on context exit above.
         # Send out some notifications to the user
         logger.debug("IMPORT_SCAN: Generating notifications")
         dojo_dispatch_task(
@@ -269,6 +277,11 @@ def process_findings(
                 findings_with_parser_tags.clear()
                 finding_ids_batch = list(batch_finding_ids)
                 batch_finding_ids.clear()
+                # Drain the inheritance context BEFORE dispatching post-process
+                # so the JIRA push inside that task sees inherited tags on the
+                # findings (otherwise inheritance lands later, on context exit).
+                if (ctx := tag_inheritance.current()) is not None:
+                    ctx.flush()
                 logger.debug("process_findings: dispatching batch with push_to_jira=%s (batch_size=%d, is_final=%s)",
                              push_to_jira, len(finding_ids_batch), is_final_finding)
                 dojo_dispatch_task(

dojo/importers/default_reimporter.py

@@ -6,6 +6,7 @@
 from django.db.models.query_utils import Q
 
 import dojo.finding.helper as finding_helper
+from dojo import tag_inheritance
 from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.finding.deduplication import (
     find_candidates_for_deduplication_hash,
@@ -107,43 +108,48 @@ def process_scan(
         # Get the findings from the parser based on what methods the parser supplies
         # This could either mean traditional file parsing, or API pull parsing
         parsed_findings = self.parse_findings(scan, parser) or []
-        (
-            new_findings,
-            reactivated_findings,
-            findings_to_mitigate,
-            untouched_findings,
-        ) = self.process_findings(parsed_findings, **kwargs)
-        # Close any old findings in the processed list if the the user specified for that
-        # to occur in the form that is then passed to the kwargs
-        closed_findings = self.close_old_findings(findings_to_mitigate, **kwargs)
-        # Update the timestamps of the test object by looking at the findings imported
-        logger.debug("REIMPORT_SCAN: Updating test/engagement timestamps")
-        # Update the timestamps of the test object by looking at the findings imported
-        self.update_timestamps()
-        # Update the test meta
-        self.update_test_meta()
-        # Update the test tags
-        self.update_test_tags()
-        # Save the test and engagement for changes to take affect
-        self.test.save()
-        self.test.engagement.save()
-        logger.debug("REIMPORT_SCAN: Updating test tags")
-        # Create a test import history object to record the flags sent to the importer
-        # This operation will return None if the user does not have the import history
-        # feature enabled
-        test_import_history = self.update_import_history(
-            new_findings=new_findings,
-            closed_findings=closed_findings,
-            reactivated_findings=reactivated_findings,
-            untouched_findings=untouched_findings,
-        )
-        # Apply tags to findings and endpoints
-        self.apply_import_tags(
-            new_findings=new_findings,
-            closed_findings=closed_findings,
-            reactivated_findings=reactivated_findings,
-            untouched_findings=untouched_findings,
-        )
+        # Open a tag-inheritance context (auto-flushes on exit). Signals
+        # register touched instances; `ctx.flush()` mid-loop drains them
+        # before each post-process dispatch so JIRA labels are correct.
+        with tag_inheritance.batch() as tag_ctx:
+            (
+                new_findings,
+                reactivated_findings,
+                findings_to_mitigate,
+                untouched_findings,
+            ) = self.process_findings(parsed_findings, **kwargs)
+            # Close any old findings in the processed list if the the user specified for that
+            # to occur in the form that is then passed to the kwargs
+            closed_findings = self.close_old_findings(findings_to_mitigate, **kwargs)
+            # Update the timestamps of the test object by looking at the findings imported
+            logger.debug("REIMPORT_SCAN: Updating test/engagement timestamps")
+            # Update the timestamps of the test object by looking at the findings imported
+            self.update_timestamps()
+            # Update the test meta
+            self.update_test_meta()
+            # Update the test tags
+            self.update_test_tags()
+            # Save the test and engagement for changes to take affect
+            self.test.save()
+            self.test.engagement.save()
+            logger.debug("REIMPORT_SCAN: Updating test tags")
+            # Create a test import history object to record the flags sent to the importer
+            # This operation will return None if the user does not have the import history
+            # feature enabled
+            test_import_history = self.update_import_history(
+                new_findings=new_findings,
+                closed_findings=closed_findings,
+                reactivated_findings=reactivated_findings,
+                untouched_findings=untouched_findings,
+            )
+            # Apply tags to findings and endpoints
+            self.apply_import_tags(
+                new_findings=new_findings,
+                closed_findings=closed_findings,
+                reactivated_findings=reactivated_findings,
+                untouched_findings=untouched_findings,
+            )
+        # Inheritance auto-flushes on context exit above.
         # Send out som notifications to the user
         logger.debug("REIMPORT_SCAN: Generating notifications")
         updated_count = (
@@ -427,6 +433,11 @@ def process_findings(
                         findings_with_parser_tags.clear()
                         finding_ids_batch = list(batch_finding_ids)
                         batch_finding_ids.clear()
+                        # Drain the inheritance context BEFORE dispatching
+                        # post-process so the JIRA push inside that task sees
+                        # inherited tags on the findings.
+                        if (ctx := tag_inheritance.current()) is not None:
+                            ctx.flush()
                         dojo_dispatch_task(
                             finding_helper.post_process_findings_batch,
                             finding_ids_batch,

dojo/importers/location_manager.py

@@ -7,14 +7,13 @@
 
 from django.core.exceptions import ValidationError
 from django.db import transaction
-from django.db.models import signals
 from django.utils import timezone
 
+from dojo import tag_inheritance
 from dojo.importers.base_location_manager import BaseLocationManager
 from dojo.location.models import AbstractLocation, Location, LocationFindingReference, LocationProductReference
 from dojo.location.status import FindingLocationStatus, ProductLocationStatus
 from dojo.models import Product, _manage_inherited_tags
-from dojo.tags_signals import make_inherited_tags_sticky
 from dojo.tools.locations import LocationData
 from dojo.url.models import URL
 from dojo.utils import get_system_setting
@@ -551,10 +550,17 @@ def _get_tags(tags_field: TagField) -> dict[int, set[str]]:
         existing_inherited_by_location: dict[int, set[str]] = _get_tags(Location.inherited_tags)
         existing_tags_by_location: dict[int, set[str]] = _get_tags(Location.tags)
 
-        # Perform the bulk updates. First, though, disconnect the make_inherited_tags_sticky signal on Location.tags
-        # while updating, otherwise each (inherited_)tags.set() will trigger, defeating the purpose of this bulk update.
-        disconnected = signals.m2m_changed.disconnect(make_inherited_tags_sticky, sender=Location.tags.through)
-        try:
+        # Perform the bulk updates inside a `tag_inheritance.batch()` context.
+        # While the batch is active, signal handlers in `dojo/tags_signals.py`
+        # short-circuit per-row inheritance work that would otherwise fire on
+        # every `(inherited_)tags.set()` and defeat the bulk update.
+        #
+        # This replaces a previous `signals.m2m_changed.disconnect(...)` /
+        # `connect(...)` dance which was process-global and therefore unsafe
+        # under threaded gunicorn / Celery thread pools / ASGI threadpools:
+        # while disconnected, every thread in the process lost sticky
+        # enforcement. Thread-local batch state avoids that hazard.
+        with tag_inheritance.batch():
             for location in locations:
                 target_tag_names: set[str] = set()
                 for pid in product_ids_by_location[location.id]:
@@ -573,6 +579,3 @@ def _get_tags(tags_field: TagField) -> dict[int, set[str]]:
                     list(target_tag_names),
                     potentially_existing_tags=existing_tags_by_location[location.id],
                 )
-        finally:
-            if disconnected:
-                signals.m2m_changed.connect(make_inherited_tags_sticky, sender=Location.tags.through)