replace dojo_async_task decorator with class+helper

Author: valentijnscholten

dojo/api_v2/views.py

@@ -46,6 +46,7 @@
 )
 from dojo.api_v2.prefetch.prefetcher import _Prefetcher
 from dojo.authorization.roles_permissions import Permissions
+from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.cred.queries import get_authorized_cred_mappings
 from dojo.endpoint.queries import (
     get_authorized_endpoint_status,
@@ -678,13 +679,13 @@ def update_jira_epic(self, request, pk=None):
         try:
 
             if engagement.has_jira_issue:
-                jira_helper.update_epic(engagement.id, **request.data)
+                dojo_dispatch_task(jira_helper.update_epic, engagement.id, **request.data)
                 response = Response(
                     {"info": "Jira Epic update query sent"},
                     status=status.HTTP_200_OK,
                 )
             else:
-                jira_helper.add_epic(engagement.id, **request.data)
+                dojo_dispatch_task(jira_helper.add_epic, engagement.id, **request.data)
                 response = Response(
                     {"info": "Jira Epic create query sent"},
                     status=status.HTTP_200_OK,

dojo/celery.py

@@ -52,7 +52,6 @@ class DojoAsyncTask(Task):
     This class:
     - Injects user context into task kwargs
     - Tracks task calls for performance testing
-    - Handles sync/async execution based on user settings
     - Supports all Celery features (signatures, chords, groups, chains)
     """
 
@@ -68,97 +67,18 @@ def apply_async(self, args=None, kwargs=None, **options):
         if "async_user" not in kwargs:
             kwargs["async_user"] = get_current_user()
 
-        # Track task call (only if not already tracked by __call__)
-        # Check if this is a direct call to apply_async (not from __call__)
-        # by checking if _dojo_tracked is not set
-        if not getattr(self, "_dojo_tracked", False):
-            dojo_async_task_counter.incr(
-                self.name,
-                args=args,
-                kwargs=kwargs,
-            )
+        # Control flag used for sync/async decision; never pass into the task itself
+        kwargs.pop("sync", None)
 
-        # Call parent to execute async
-        return super().apply_async(args=args, kwargs=kwargs, **options)
-
-    def s(self, *args, **kwargs):
-        """Create a mutable signature with injected user context."""
-        from dojo.decorators import dojo_async_task_counter  # noqa: PLC0415 circular import
-        from dojo.utils import get_current_user  # noqa: PLC0415 circular import
-
-        if "async_user" not in kwargs:
-            kwargs["async_user"] = get_current_user()
-
-        # Track task call
-        dojo_async_task_counter.incr(
-            self.name,
-            args=args,
-            kwargs=kwargs,
-        )
-
-        return super().s(*args, **kwargs)
-
-    def si(self, *args, **kwargs):
-        """Create an immutable signature with injected user context."""
-        from dojo.decorators import dojo_async_task_counter  # noqa: PLC0415 circular import
-        from dojo.utils import get_current_user  # noqa: PLC0415 circular import
-
-        if "async_user" not in kwargs:
-            kwargs["async_user"] = get_current_user()
-
-        # Track task call
-        dojo_async_task_counter.incr(
-            self.name,
-            args=args,
-            kwargs=kwargs,
-        )
-
-        return super().si(*args, **kwargs)
-
-    def __call__(self, *args, **kwargs):
-        """
-        Override __call__ to handle direct task calls with sync/async logic.
-
-        This replicates the behavior of the dojo_async_task decorator wrapper.
-        """
-        # In Celery worker execution, __call__ is how tasks actually run.
-        # We only want the sync/async decision when tasks are called directly
-        # from application code (task(...)), not when the worker is executing a message.
-        if not getattr(self.request, "called_directly", True):
-            return super().__call__(*args, **kwargs)
-
-        from dojo.decorators import dojo_async_task_counter, we_want_async  # noqa: PLC0415 circular import
-        from dojo.utils import get_current_user  # noqa: PLC0415 circular import
-
-        # Inject user context if not already present
-        if "async_user" not in kwargs:
-            kwargs["async_user"] = get_current_user()
-
-        # Track task call
+        # Track dispatch
         dojo_async_task_counter.incr(
             self.name,
             args=args,
             kwargs=kwargs,
         )
 
-        # Extract countdown if present (don't pass to sync execution)
-        countdown = kwargs.pop("countdown", 0)
-
-        # Check if we should run async or sync
-        if we_want_async(*args, func=self, **kwargs):
-            # Mark as tracked to avoid double tracking in apply_async
-            self._dojo_tracked = True
-            try:
-                # Run asynchronously
-                return self.apply_async(args=args, kwargs=kwargs, countdown=countdown)
-            finally:
-                # Clean up the flag
-                delattr(self, "_dojo_tracked")
-        else:
-            # Run synchronously in-process, matching the original decorator behavior: func(*args, **kwargs)
-            # Remove sync from kwargs as it's a control flag, not a task argument.
-            kwargs.pop("sync", None)
-            return self.run(*args, **kwargs)
+        # Call parent to execute async
+        return super().apply_async(args=args, kwargs=kwargs, **options)
 
 
 @app.task(bind=True)

dojo/celery_dispatch.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol, cast
+
+from celery.canvas import Signature
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+class _SupportsSi(Protocol):
+    def si(self, *args: Any, **kwargs: Any) -> Signature: ...
+
+
+class _SupportsApplyAsync(Protocol):
+    def apply_async(self, args: Any | None = None, kwargs: Any | None = None, **options: Any) -> Any: ...
+
+
+def _inject_async_user(kwargs: Mapping[str, Any] | None) -> dict[str, Any]:
+    result: dict[str, Any] = dict(kwargs or {})
+    if "async_user" not in result:
+        from dojo.utils import get_current_user  # noqa: PLC0415 circular import
+
+        result["async_user"] = get_current_user()
+    return result
+
+
+def dojo_create_signature(task_or_sig: _SupportsSi | Signature, *args: Any, **kwargs: Any) -> Signature:
+    """
+    Build a Celery signature with DefectDojo user context injected.
+
+    - If passed a task, returns `task_or_sig.si(*args, **kwargs)`.
+    - If passed an existing signature, returns a cloned signature with merged kwargs.
+    """
+    injected = _inject_async_user(kwargs)
+    injected.pop("countdown", None)
+
+    if isinstance(task_or_sig, Signature):
+        merged_kwargs = {**(task_or_sig.kwargs or {}), **injected}
+        return task_or_sig.clone(kwargs=merged_kwargs)
+
+    return task_or_sig.si(*args, **injected)
+
+
+def dojo_dispatch_task(task_or_sig: _SupportsSi | _SupportsApplyAsync | Signature, *args: Any, **kwargs: Any) -> Any:
+    """
+    Dispatch a task/signature using DefectDojo semantics.
+
+    - Inject `async_user` if missing.
+    - Respect `sync=True` (foreground execution) and user `block_execution`.
+    - Support `countdown=<seconds>` for async dispatch.
+
+    Returns:
+    - async: AsyncResult-like return from Celery
+    - sync: underlying return value of the task
+
+    """
+    from dojo.decorators import dojo_async_task_counter, we_want_async  # noqa: PLC0415 circular import
+
+    countdown = cast("int", kwargs.pop("countdown", 0))
+    injected = _inject_async_user(kwargs)
+
+    sig = dojo_create_signature(task_or_sig if isinstance(task_or_sig, Signature) else cast("_SupportsSi", task_or_sig), *args, **injected)
+    sig_kwargs = dict(sig.kwargs or {})
+
+    if we_want_async(*sig.args, func=getattr(sig, "type", None), **sig_kwargs):
+        # DojoAsyncTask.apply_async tracks async dispatch. Avoid double-counting here.
+        return sig.apply_async(countdown=countdown)
+
+    # Track foreground execution as a "created task" as well (matches historical dojo_async_task behavior)
+    dojo_async_task_counter.incr(str(sig.task), args=sig.args, kwargs=sig_kwargs)
+
+    sig_kwargs.pop("sync", None)
+    sig = sig.clone(kwargs=sig_kwargs)
+    eager = sig.apply()
+    return eager.get(propagate=True)

dojo/endpoint/views.py

@@ -18,6 +18,7 @@
 from dojo.authorization.authorization import user_has_permission_or_403
 from dojo.authorization.authorization_decorators import user_is_authorized
 from dojo.authorization.roles_permissions import Permissions
+from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.endpoint.queries import get_authorized_endpoints
 from dojo.endpoint.utils import clean_hosts_run, endpoint_meta_import
 from dojo.filters import EndpointFilter, EndpointFilterWithoutObjectLookups
@@ -373,7 +374,7 @@ def endpoint_bulk_update_all(request, pid=None):
             product_calc = list(Product.objects.filter(endpoint__id__in=endpoints_to_update).distinct())
             endpoints.delete()
             for prod in product_calc:
-                calculate_grade(prod.id)
+                dojo_dispatch_task(calculate_grade, prod.id)
 
             if skipped_endpoint_count > 0:
                 add_error_message_to_response(f"Skipped deletion of {skipped_endpoint_count} endpoints because you are not authorized.")

dojo/engagement/services.py

@@ -5,6 +5,7 @@
 from django.dispatch import receiver
 
 import dojo.jira_link.helper as jira_helper
+from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.models import Engagement
 
 logger = logging.getLogger(__name__)
@@ -16,7 +17,7 @@ def close_engagement(eng):
     eng.save()
 
     if jira_helper.get_jira_project(eng):
-        jira_helper.close_epic(eng.id, push_to_jira=True)
+        dojo_dispatch_task(jira_helper.close_epic, eng.id, push_to_jira=True)
 
 
 def reopen_engagement(eng):

dojo/engagement/views.py

@@ -37,6 +37,7 @@
 from dojo.authorization.authorization import user_has_permission_or_403
 from dojo.authorization.authorization_decorators import user_is_authorized
 from dojo.authorization.roles_permissions import Permissions
+from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.endpoint.utils import save_endpoints_to_add
 from dojo.engagement.queries import get_authorized_engagements
 from dojo.engagement.services import close_engagement, reopen_engagement
@@ -390,7 +391,7 @@ def copy_engagement(request, eid):
         form = DoneForm(request.POST)
         if form.is_valid():
             engagement_copy = engagement.copy()
-            calculate_grade(product.id)
+            dojo_dispatch_task(calculate_grade, product.id)
             messages.add_message(
                 request,
                 messages.SUCCESS,

dojo/finding/helper.py

@@ -434,7 +434,9 @@ def post_process_finding_save_internal(finding, dedupe_option=True, rules_option
 
     if product_grading_option:
         if system_settings.enable_product_grade:
-            calculate_grade(finding.test.engagement.product.id)
+            from dojo.celery_dispatch import dojo_dispatch_task  # noqa: PLC0415 circular import
+
+            dojo_dispatch_task(calculate_grade, finding.test.engagement.product.id)
         else:
             deduplicationLogger.debug("skipping product grading because it's disabled in system settings")
 
@@ -493,7 +495,9 @@ def post_process_findings_batch(finding_ids, *args, dedupe_option=True, rules_op
             tool_issue_updater.async_tool_issue_update(finding)
 
     if product_grading_option and system_settings.enable_product_grade:
-        calculate_grade(findings[0].test.engagement.product.id)
+        from dojo.celery_dispatch import dojo_dispatch_task  # noqa: PLC0415 circular import
+
+        dojo_dispatch_task(calculate_grade, findings[0].test.engagement.product.id)
 
     if push_to_jira:
         for finding in findings:

dojo/finding/views.py

@@ -38,6 +38,7 @@
     user_is_authorized,
 )
 from dojo.authorization.roles_permissions import Permissions
+from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.filters import (
     AcceptedFindingFilter,
     AcceptedFindingFilterWithoutObjectLookups,
@@ -1082,7 +1083,7 @@ def process_form(self, request: HttpRequest, finding: Finding, context: dict):
             product = finding.test.engagement.product
             finding.delete()
             # Update the grade of the product async
-            calculate_grade(product.id)
+            dojo_dispatch_task(calculate_grade, product.id)
             # Add a message to the request that the finding was successfully deleted
             messages.add_message(
                 request,
@@ -1353,7 +1354,7 @@ def copy_finding(request, fid):
             test = form.cleaned_data.get("test")
             product = finding.test.engagement.product
             finding_copy = finding.copy(test=test)
-            calculate_grade(product.id)
+            dojo_dispatch_task(calculate_grade, product.id)
             messages.add_message(
                 request,
                 messages.SUCCESS,

dojo/finding_group/views.py

@@ -16,6 +16,7 @@
 from dojo.authorization.authorization import user_has_permission_or_403
 from dojo.authorization.authorization_decorators import user_is_authorized
 from dojo.authorization.roles_permissions import Permissions
+from dojo.celery_dispatch import dojo_dispatch_task
 from dojo.filters import (
     FindingFilter,
     FindingFilterWithoutObjectLookups,
@@ -100,7 +101,7 @@ def view_finding_group(request, fgid):
                 elif not finding_group.has_jira_issue:
                     jira_helper.finding_group_link_jira(request, finding_group, jira_issue)
             elif push_to_jira:
-                jira_helper.push_to_jira(finding_group, sync=True)
+                dojo_dispatch_task(jira_helper.push_to_jira, finding_group, sync=True)
 
             finding_group.save()
             return HttpResponseRedirect(reverse("view_test", args=(finding_group.test.id,)))
@@ -200,7 +201,7 @@ def push_to_jira(request, fgid):
 
         # it may look like success here, but the push_to_jira are swallowing exceptions
         # but cant't change too much now without having a test suite, so leave as is for now with the addition warning message to check alerts for background errors.
-        if jira_helper.push_to_jira(group, sync=True):
+        if dojo_dispatch_task(jira_helper.push_to_jira, group, sync=True):
             messages.add_message(
                 request,
                 messages.SUCCESS,