@@ -44,6 +44,123 @@ def __call__(self, *args, **kwargs):
4444app .autodiscover_tasks (lambda : settings .INSTALLED_APPS )
4545
4646
47+ class DojoAsyncTask (Task ):
48+
49+ """
50+ Base task class that provides dojo_async_task functionality without using a decorator.
51+
52+ This class:
53+ - Injects user context into task kwargs
54+ - Tracks task calls for performance testing
55+ - Handles sync/async execution based on user settings
56+ - Supports all Celery features (signatures, chords, groups, chains)
57+ """
58+
59+ def apply_async (self , args = None , kwargs = None , ** options ):
60+ """Override apply_async to inject user context and track tasks."""
61+ from dojo .decorators import dojo_async_task_counter # noqa: PLC0415 circular import
62+ from dojo .utils import get_current_user # noqa: PLC0415 circular import
63+
64+ if kwargs is None :
65+ kwargs = {}
66+
67+ # Inject user context if not already present
68+ if "async_user" not in kwargs :
69+ kwargs ["async_user" ] = get_current_user ()
70+
71+ # Track task call (only if not already tracked by __call__)
72+ # Check if this is a direct call to apply_async (not from __call__)
73+ # by checking if _dojo_tracked is not set
74+ if not getattr (self , "_dojo_tracked" , False ):
75+ dojo_async_task_counter .incr (
76+ self .name ,
77+ args = args ,
78+ kwargs = kwargs ,
79+ )
80+
81+ # Call parent to execute async
82+ return super ().apply_async (args = args , kwargs = kwargs , ** options )
83+
84+ def s (self , * args , ** kwargs ):
85+ """Create a mutable signature with injected user context."""
86+ from dojo .decorators import dojo_async_task_counter # noqa: PLC0415 circular import
87+ from dojo .utils import get_current_user # noqa: PLC0415 circular import
88+
89+ if "async_user" not in kwargs :
90+ kwargs ["async_user" ] = get_current_user ()
91+
92+ # Track task call
93+ dojo_async_task_counter .incr (
94+ self .name ,
95+ args = args ,
96+ kwargs = kwargs ,
97+ )
98+
99+ return super ().s (* args , ** kwargs )
100+
101+ def si (self , * args , ** kwargs ):
102+ """Create an immutable signature with injected user context."""
103+ from dojo .decorators import dojo_async_task_counter # noqa: PLC0415 circular import
104+ from dojo .utils import get_current_user # noqa: PLC0415 circular import
105+
106+ if "async_user" not in kwargs :
107+ kwargs ["async_user" ] = get_current_user ()
108+
109+ # Track task call
110+ dojo_async_task_counter .incr (
111+ self .name ,
112+ args = args ,
113+ kwargs = kwargs ,
114+ )
115+
116+ return super ().si (* args , ** kwargs )
117+
118+ def __call__ (self , * args , ** kwargs ):
119+ """
120+ Override __call__ to handle direct task calls with sync/async logic.
121+
122+ This replicates the behavior of the dojo_async_task decorator wrapper.
123+ """
124+ # In Celery worker execution, __call__ is how tasks actually run.
125+ # We only want the sync/async decision when tasks are called directly
126+ # from application code (task(...)), not when the worker is executing a message.
127+ if not getattr (self .request , "called_directly" , True ):
128+ return super ().__call__ (* args , ** kwargs )
129+
130+ from dojo .decorators import dojo_async_task_counter , we_want_async # noqa: PLC0415 circular import
131+ from dojo .utils import get_current_user # noqa: PLC0415 circular import
132+
133+ # Inject user context if not already present
134+ if "async_user" not in kwargs :
135+ kwargs ["async_user" ] = get_current_user ()
136+
137+ # Track task call
138+ dojo_async_task_counter .incr (
139+ self .name ,
140+ args = args ,
141+ kwargs = kwargs ,
142+ )
143+
144+ # Extract countdown if present (don't pass to sync execution)
145+ countdown = kwargs .pop ("countdown" , 0 )
146+
147+ # Check if we should run async or sync
148+ if we_want_async (* args , func = self , ** kwargs ):
149+ # Mark as tracked to avoid double tracking in apply_async
150+ self ._dojo_tracked = True
151+ try :
152+ # Run asynchronously
153+ return self .apply_async (args = args , kwargs = kwargs , countdown = countdown )
154+ finally :
155+ # Clean up the flag
156+ delattr (self , "_dojo_tracked" )
157+ else :
158+ # Run synchronously in-process, matching the original decorator behavior: func(*args, **kwargs)
159+ # Remove sync from kwargs as it's a control flag, not a task argument.
160+ kwargs .pop ("sync" , None )
161+ return self .run (* args , ** kwargs )
162+
163+
47164@app .task (bind = True )
48165def debug_task (self ):
49166 logger .info (f"Request: { self .request !r} )
0 commit comments