Add support for receiving task context (#174)

Author: RealOrangeOne

README.md

@@ -78,6 +78,26 @@ modified_task = calculate_meaning_of_life.using(priority=10)
 
 In addition to the above attributes, `run_after` can be passed to specify a specific time the task should run.
 
+#### Task context
+
+Sometimes the running task may need to know context about how it was enqueued. To receive the task context as an argument to your task function, pass `takes_context` to the decorator and ensure the task takes a `context` as the first argument.
+
+```python
+from django_tasks import task, TaskContext
+
+
+@task(takes_context=True)
+def calculate_meaning_of_life(context: TaskContext) -> int:
+    return 42
+```
+
+The task context has the following attributes:
+
+- `task_result`: The running task result
+- `attempt`: The current attempt number for the task
+
+This API will be extended with additional features in future.
+
 ### Enqueueing tasks
 
 To execute a task, call the `enqueue` method on it:

django_tasks/__init__.py

@@ -16,6 +16,7 @@
     DEFAULT_TASK_BACKEND_ALIAS,
     ResultStatus,
     Task,
+    TaskContext,
     TaskResult,
     task,
 )
@@ -31,6 +32,7 @@
     "ResultStatus",
     "Task",
     "TaskResult",
+    "TaskContext",
 ]
 
 

django_tasks/backends/base.py

@@ -7,6 +7,7 @@
 from django.core.checks import messages
 from django.db import connections
 from django.utils import timezone
+from django.utils.inspect import get_func_args
 from typing_extensions import ParamSpec
 
 from django_tasks.exceptions import InvalidTaskError
@@ -61,6 +62,15 @@ def validate_task(self, task: Task) -> None:
         if not self.supports_async_task and iscoroutinefunction(task.func):
             raise InvalidTaskError("Backend does not support async tasks")
 
+        task_func_args = get_func_args(task.func)
+
+        if task.takes_context and (
+            not task_func_args or task_func_args[0] != "context"
+        ):
+            raise InvalidTaskError(
+                "Task takes context but does not have a first argument of 'context'"
+            )
+
         if (
             task.priority < MIN_PRIORITY
             or task.priority > MAX_PRIORITY

django_tasks/backends/database/management/commands/db_worker.py

@@ -22,7 +22,7 @@
 from django_tasks.backends.database.utils import exclusive_transaction
 from django_tasks.exceptions import InvalidTaskBackendError
 from django_tasks.signals import task_finished, task_started
-from django_tasks.task import DEFAULT_QUEUE_NAME
+from django_tasks.task import DEFAULT_QUEUE_NAME, TaskContext
 from django_tasks.utils import get_random_id
 
 package_logger = logging.getLogger("django_tasks")
@@ -160,8 +160,14 @@ def run_task(self, db_task_result: DBTaskResult) -> None:
             backend_type = task.get_backend()
 
             task_started.send(sender=backend_type, task_result=task_result)
-
-            return_value = task.call(*task_result.args, **task_result.kwargs)
+            if task.takes_context:
+                return_value = task.call(
+                    TaskContext(task_result=task_result),
+                    *task_result.args,
+                    **task_result.kwargs,
+                )
+            else:
+                return_value = task.call(*task_result.args, **task_result.kwargs)
 
             # Setting the return and success value inside the error handling,
             # So errors setting it (eg JSON encode) can still be recorded

django_tasks/backends/immediate.py

@@ -1,15 +1,13 @@
 import logging
 from functools import partial
-from inspect import iscoroutinefunction
 from typing import TypeVar
 
-from asgiref.sync import async_to_sync
 from django.db import transaction
 from django.utils import timezone
 from typing_extensions import ParamSpec
 
 from django_tasks.signals import task_enqueued, task_finished, task_started
-from django_tasks.task import ResultStatus, Task, TaskError, TaskResult
+from django_tasks.task import ResultStatus, Task, TaskContext, TaskError, TaskResult
 from django_tasks.utils import (
     get_exception_traceback,
     get_module_path,
@@ -43,23 +41,26 @@ def _execute_task(self, task_result: TaskResult) -> None:
 
         task = task_result.task
 
-        calling_task_func = (
-            async_to_sync(task.func) if iscoroutinefunction(task.func) else task.func
-        )
-
         object.__setattr__(task_result, "status", ResultStatus.RUNNING)
         object.__setattr__(task_result, "started_at", timezone.now())
         object.__setattr__(task_result, "last_attempted_at", timezone.now())
         task_result.worker_ids.append(self.worker_id)
         task_started.send(sender=type(self), task_result=task_result)
 
         try:
+            if task.takes_context:
+                raw_return_value = task.call(
+                    TaskContext(task_result=task_result),
+                    *task_result.args,
+                    **task_result.kwargs,
+                )
+            else:
+                raw_return_value = task.call(*task_result.args, **task_result.kwargs)
+
             object.__setattr__(
                 task_result,
                 "_return_value",
-                json_normalize(
-                    calling_task_func(*task_result.args, **task_result.kwargs)
-                ),
+                json_normalize(raw_return_value),
             )
         except BaseException as e:
             # If the user tried to terminate, let them

django_tasks/backends/rq.py

@@ -8,6 +8,7 @@
 from django.core.checks import messages
 from django.core.exceptions import SuspiciousOperation
 from django.db import transaction
+from django.utils.functional import cached_property
 from redis.client import Redis
 from rq.job import Callback, JobStatus
 from rq.job import Job as BaseJob
@@ -23,6 +24,7 @@
     MAX_PRIORITY,
     ResultStatus,
     Task,
+    TaskContext,
     TaskError,
 )
 from django_tasks.task import TaskResult as BaseTaskResult
@@ -51,21 +53,30 @@ class TaskResult(BaseTaskResult[T]):
 
 class Job(BaseJob):
     def perform(self) -> Any:
-        task_result = self.into_task_result()
-
         assert self.worker_name is not None
         self.meta.setdefault("_django_tasks_worker_ids", []).append(self.worker_name)
         self.save_meta()  # type: ignore[no-untyped-call]
 
-        task_started.send(type(task_result.task.get_backend()), task_result=task_result)
+        task_started.send(
+            type(self.task_result.task.get_backend()), task_result=self.task_result
+        )
 
         return super().perform()
 
     def _execute(self) -> Any:
         """
         Shim RQ's `Job` to call the underlying `Task` function.
         """
-        return self.func.call(*self.args, **self.kwargs)
+        try:
+            if self.func.takes_context:
+                return self.func.call(
+                    TaskContext(task_result=self.task_result), *self.args, **self.kwargs
+                )
+            return self.func.call(*self.args, **self.kwargs)
+        finally:
+            # Clear the task result cache, as it's changed now
+            self.__dict__.pop("task_result", None)
+            pass
 
     @property
     def func(self) -> Task:
@@ -78,7 +89,8 @@ def func(self) -> Task:
 
         return func
 
-    def into_task_result(self) -> TaskResult:
+    @cached_property
+    def task_result(self) -> TaskResult:
         task: Task = self.func
 
         scheduled_job_registry = ScheduledJobRegistry(  # type: ignore[no-untyped-call]
@@ -145,15 +157,15 @@ def failed_callback(
     )
     job.save_meta()  # type: ignore[no-untyped-call]
 
-    task_result = job.into_task_result()
+    task_result = job.task_result
 
     object.__setattr__(task_result, "status", ResultStatus.FAILED)
 
     task_finished.send(type(task_result.task.get_backend()), task_result=task_result)
 
 
 def success_callback(job: Job, connection: Optional[Redis], result: Any) -> None:
-    task_result = job.into_task_result()
+    task_result = job.task_result
 
     object.__setattr__(task_result, "status", ResultStatus.SUCCEEDED)
 
@@ -241,7 +253,7 @@ def get_result(self, result_id: str) -> TaskResult:
         if job is None:
             raise ResultDoesNotExist(result_id)
 
-        return job.into_task_result()
+        return job.task_result
 
     def check(self, **kwargs: Any) -> Iterable[messages.CheckMessage]:
         yield from super().check(**kwargs)

django_tasks/task.py

@@ -6,6 +6,7 @@
     Any,
     Callable,
     Generic,
+    Literal,
     Optional,
     TypeVar,
     Union,
@@ -17,7 +18,7 @@
 from django.db.models.enums import TextChoices
 from django.utils.module_loading import import_string
 from django.utils.translation import gettext_lazy as _
-from typing_extensions import ParamSpec, Self
+from typing_extensions import Concatenate, ParamSpec, Self
 
 from .exceptions import ResultDoesNotExist
 from .utils import (
@@ -87,6 +88,11 @@ class Task(Generic[P, T]):
     immediately, or whatever the backend decides
     """
 
+    takes_context: bool = False
+    """
+    Whether the task receives the task context when executed.
+    """
+
     def __post_init__(self) -> None:
         self.get_backend().validate_task(self)
 
@@ -197,18 +203,37 @@ def task(
     queue_name: str = DEFAULT_QUEUE_NAME,
     backend: str = DEFAULT_TASK_BACKEND_ALIAS,
     enqueue_on_commit: Optional[bool] = None,
+    takes_context: Literal[False] = False,
 ) -> Callable[[Callable[P, T]], Task[P, T]]: ...
 
 
-# Implementation
+# Decorator with context and arguments
+# e.g. @task(takes_context=True, ...)
+@overload
 def task(
+    *,
+    priority: int = DEFAULT_PRIORITY,
+    queue_name: str = DEFAULT_QUEUE_NAME,
+    backend: str = DEFAULT_TASK_BACKEND_ALIAS,
+    enqueue_on_commit: Optional[bool] = None,
+    takes_context: Literal[True],
+) -> Callable[[Callable[Concatenate["TaskContext", P], T]], Task[P, T]]: ...
+
+
+# Implementation
+def task(  # type: ignore[misc]
     function: Optional[Callable[P, T]] = None,
     *,
     priority: int = DEFAULT_PRIORITY,
     queue_name: str = DEFAULT_QUEUE_NAME,
     backend: str = DEFAULT_TASK_BACKEND_ALIAS,
     enqueue_on_commit: Optional[bool] = None,
-) -> Union[Task[P, T], Callable[[Callable[P, T]], Task[P, T]]]:
+    takes_context: bool = False,
+) -> Union[
+    Task[P, T],
+    Callable[[Callable[P, T]], Task[P, T]],
+    Callable[[Callable[Concatenate["TaskContext", P], T]], Task[P, T]],
+]:
     """
     A decorator used to create a task.
     """
@@ -221,6 +246,7 @@ def wrapper(f: Callable[P, T]) -> Task[P, T]:
             queue_name=queue_name,
             backend=backend,
             enqueue_on_commit=enqueue_on_commit,
+            takes_context=takes_context,
         )
 
     if function:
@@ -330,3 +356,12 @@ async def arefresh(self) -> None:
 
         for attr in TASK_REFRESH_ATTRS:
             object.__setattr__(self, attr, getattr(refreshed_task, attr))
+
+
+@dataclass(frozen=True)
+class TaskContext:
+    task_result: TaskResult
+
+    @property
+    def attempt(self) -> int:
+        return self.task_result.attempts

tests/tasks.py

@@ -1,7 +1,7 @@
 import time
 from typing import Any
 
-from django_tasks import task
+from django_tasks import TaskContext, task
 
 
 @task()
@@ -70,3 +70,14 @@ def hang() -> None:
 @task()
 def sleep_for(seconds: float) -> None:
     time.sleep(seconds)
+
+
+@task(takes_context=True)
+def get_task_id(context: TaskContext) -> str:
+    return context.task_result.id
+
+
+@task(takes_context=True)
+def test_context(context: TaskContext, attempt: int) -> None:
+    assert isinstance(context, TaskContext)
+    assert context.attempt == attempt

tests/tests/test_database_backend.py

@@ -862,6 +862,23 @@ def test_max_tasks(self) -> None:
         self.assertEqual(statuses[ResultStatus.SUCCEEDED], 2)
         self.assertEqual(statuses[ResultStatus.READY], 3)
 
+    def test_takes_context(self) -> None:
+        result = test_tasks.get_task_id.enqueue()
+
+        self.run_worker()
+
+        result.refresh()
+
+        self.assertEqual(result.return_value, result.id)
+
+    def test_context(self) -> None:
+        result = test_tasks.test_context.enqueue(1)
+
+        self.run_worker()
+        result.refresh()
+
+        self.assertEqual(result.status, ResultStatus.SUCCEEDED)
+
 
 @override_settings(
     TASKS={

tests/tests/test_dummy_backend.py

@@ -195,6 +195,10 @@ def test_enqueue_on_commit_with_no_databases(self) -> None:
         self.assertEqual(len(errors), 1)
         self.assertIn("Set `ENQUEUE_ON_COMMIT` to False", errors[0].hint)  # type:ignore[arg-type]
 
+    def test_takes_context(self) -> None:
+        result = test_tasks.get_task_id.enqueue()
+        self.assertEqual(result.status, ResultStatus.READY)
+
 
 class DummyBackendTransactionTestCase(TransactionTestCase):
     @override_settings(