Add configurable transaction management

Author: RealOrangeOne

README.md

@@ -79,6 +79,23 @@ The returned `TaskResult` can be interrogated to query the current state of the
 
 If the task takes arguments, these can be passed as-is to `enqueue`.
 
+#### Transactions
+
+By default, it's up to the backend to determine whether the task should be enqueued immediately (after calling `.enqueue`) or wait until the end of the current database transaction (if there is one).
+
+This can be configured using the `ENQUEUE_ON_COMMIT` setting. `True` and `False` force the behaviour, and `None` is used to let the backend decide.
+
+```python
+TASKS = {
+    "default": {
+        "BACKEND": "django_tasks.backends.immediate.ImmediateBackend",
+        "ENQUEUE_ON_COMMIT": False
+    }
+}
+```
+
+All built-in backends default to waiting for the end of the transaction (`"ENQUEUE_ON_COMMIT": True`).
+
 ### Executing tasks with the database backend
 
 First, you'll need to add `django_tasks.backends.database`  to `INSTALLED_APPS`, and run `manage.py migrate`.

django_tasks/backends/base.py

@@ -1,9 +1,9 @@
 from abc import ABCMeta, abstractmethod
 from inspect import iscoroutinefunction
-from typing import Any, List, TypeVar
+from typing import Any, Iterable, Optional, TypeVar
 
 from asgiref.sync import sync_to_async
-from django.core.checks.messages import CheckMessage
+from django.core.checks import messages
 from django.utils import timezone
 from typing_extensions import ParamSpec
 
@@ -16,6 +16,9 @@
 
 
 class BaseTaskBackend(metaclass=ABCMeta):
+    alias: str
+    enqueue_on_commit: Optional[bool]
+
     task_class = Task
 
     supports_defer = False
@@ -29,6 +32,18 @@ class BaseTaskBackend(metaclass=ABCMeta):
 
     def __init__(self, options: dict) -> None:
         self.alias = options["ALIAS"]
+        self.enqueue_on_commit = options.get("ENQUEUE_ON_COMMIT", None)
+
+    def _get_enqueue_on_commit_for_task(self, task: Task) -> Optional[bool]:
+        """
+        Determine the correct `enqueue_on_commit` setting to use for a given task.
+
+        If the task defines it, use that, otherwise, fall back to the backend.
+        """
+        if isinstance(task.enqueue_on_commit, bool):
+            return task.enqueue_on_commit
+
+        return self.enqueue_on_commit
 
     def validate_task(self, task: Task) -> None:
         """
@@ -94,8 +109,8 @@ def close(self) -> None:
         # HACK: `close` isn't abstract, but should do nothing by default
         return None
 
-    def check(self, **kwargs: Any) -> List[CheckMessage]:
-        raise NotImplementedError(
-            "subclasses may provide a check() method to verify that task "
-            "backend is configured correctly."
-        )
+    def check(self, **kwargs: Any) -> Iterable[messages.CheckMessage]:
+        if self.enqueue_on_commit not in {True, False, None}:
+            yield messages.CheckMessage(
+                messages.ERROR, "`ENQUEUE_ON_COMMIT` must be a bool or None"
+            )

django_tasks/backends/database/backend.py

@@ -1,13 +1,13 @@
 from dataclasses import asdict, dataclass
-from typing import TYPE_CHECKING, Any, List, TypeVar
+from typing import TYPE_CHECKING, Any, Iterable, TypeVar
 
 from django.apps import apps
-from django.core.checks import ERROR, CheckMessage
+from django.core.checks import messages
 from django.core.exceptions import ValidationError
 from typing_extensions import ParamSpec
 
 from django_tasks.backends.base import BaseTaskBackend
-from django_tasks.exceptions import ResultDoesNotExist
+from django_tasks.exceptions import InvalidTaskError, ResultDoesNotExist
 from django_tasks.task import Task
 from django_tasks.task import TaskResult as BaseTaskResult
 from django_tasks.utils import json_normalize
@@ -39,6 +39,14 @@ class DatabaseBackend(BaseTaskBackend):
     supports_get_result = True
     supports_defer = True
 
+    def validate_task(self, task: Task[P, T]) -> None:
+        super().validate_task(task)
+
+        if task.enqueue_on_commit is False:
+            raise InvalidTaskError(
+                "enqueue_on_commit must be True or None when using database backend"
+            )
+
     def _task_to_db_task(
         self, task: Task[P, T], args: P.args, kwargs: P.kwargs
     ) -> "DBTaskResult":
@@ -91,16 +99,21 @@ async def aget_result(self, result_id: str) -> TaskResult:
         except (DBTaskResult.DoesNotExist, ValidationError) as e:
             raise ResultDoesNotExist(result_id) from e
 
-    def check(self, **kwargs: Any) -> List[CheckMessage]:
-        if not apps.is_installed("django_tasks.backends.database"):
-            backend_name = self.__class__.__name__
+    def check(self, **kwargs: Any) -> Iterable[messages.CheckMessage]:
+        yield from super().check(**kwargs)
 
-            return [
-                CheckMessage(
-                    ERROR,
-                    f"{backend_name} configured as django_tasks backend, but database app not installed",
-                    "Insert 'django_tasks.backends.database' in INSTALLED_APPS",
-                )
-            ]
+        backend_name = self.__class__.__name__
 
-        return []
+        if not apps.is_installed("django_tasks.backends.database"):
+            yield messages.CheckMessage(
+                messages.ERROR,
+                f"{backend_name} configured as django_tasks backend, but database app not installed",
+                "Insert 'django_tasks.backends.database' in INSTALLED_APPS",
+            )
+
+        if self.enqueue_on_commit is False:
+            yield messages.CheckMessage(
+                messages.WARNING,
+                f"{backend_name} must enqueue tasks at the end of a transaction",
+                "Ensure ENQUEUE_ON_COMMIT is True or None",
+            )

django_tasks/backends/dummy.py

@@ -1,7 +1,9 @@
 from copy import deepcopy
+from functools import partial
 from typing import List, TypeVar
 from uuid import uuid4
 
+from django.db import transaction
 from django.utils import timezone
 from typing_extensions import ParamSpec
 
@@ -41,8 +43,11 @@ def enqueue(
             backend=self.alias,
         )
 
-        # Copy the task to prevent mutation issues
-        self.results.append(deepcopy(result))
+        if self._get_enqueue_on_commit_for_task(task) is not False:
+            # Copy the task to prevent mutation issues
+            transaction.on_commit(partial(self.results.append, deepcopy(result)))
+        else:
+            self.results.append(deepcopy(result))
 
         return result
 

django_tasks/backends/immediate.py

@@ -1,8 +1,10 @@
+from functools import partial
 from inspect import iscoroutinefunction
 from typing import TypeVar
 from uuid import uuid4
 
 from asgiref.sync import async_to_sync
+from django.db import transaction
 from django.utils import timezone
 from typing_extensions import ParamSpec
 
@@ -18,34 +20,46 @@
 class ImmediateBackend(BaseTaskBackend):
     supports_async_task = True
 
-    def enqueue(
-        self, task: Task[P, T], args: P.args, kwargs: P.kwargs
-    ) -> TaskResult[T]:
-        self.validate_task(task)
-
+    def _execute_task(self, task_result: TaskResult) -> None:
+        """
+        Execute the task for the given `TaskResult`, mutating it with the outcome
+        """
         calling_task_func = (
-            async_to_sync(task.func) if iscoroutinefunction(task.func) else task.func
+            async_to_sync(task_result.task.func)
+            if iscoroutinefunction(task_result.task.func)
+            else task_result.task.func
         )
 
-        enqueued_at = timezone.now()
         try:
-            result = json_normalize(calling_task_func(*args, **kwargs))
-            status = ResultStatus.COMPLETE
+            task_result._result = json_normalize(
+                calling_task_func(*task_result.args, **task_result.kwargs)
+            )
+            task_result.status = ResultStatus.COMPLETE
         except Exception:
-            result = None
-            status = ResultStatus.FAILED
+            task_result._result = None
+            task_result.status = ResultStatus.FAILED
+
+        task_result.finished_at = timezone.now()
+
+    def enqueue(
+        self, task: Task[P, T], args: P.args, kwargs: P.kwargs
+    ) -> TaskResult[T]:
+        self.validate_task(task)
 
         task_result = TaskResult[T](
             task=task,
             id=str(uuid4()),
-            status=status,
-            enqueued_at=enqueued_at,
-            finished_at=timezone.now(),
+            status=ResultStatus.NEW,
+            enqueued_at=timezone.now(),
+            finished_at=None,
             args=json_normalize(args),
             kwargs=json_normalize(kwargs),
             backend=self.alias,
         )
 
-        task_result._result = result
+        if self._get_enqueue_on_commit_for_task(task) is not False:
+            transaction.on_commit(partial(self._execute_task, task_result))
+        else:
+            self._execute_task(task_result)
 
         return task_result

django_tasks/checks.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Sequence
+from typing import Any, Iterable, Sequence
 
 from django.apps.config import AppConfig
 from django.core.checks.messages import CheckMessage
@@ -8,16 +8,11 @@
 
 def check_tasks(
     app_configs: Sequence[AppConfig] = None, **kwargs: Any
-) -> List[CheckMessage]:
+) -> Iterable[CheckMessage]:
     """Checks all registered task backends."""
 
-    errors = []
     for backend in tasks.all():
         try:
-            backend_errors = backend.check()
+            yield from backend.check()
         except NotImplementedError:
             pass
-        else:
-            errors.extend(backend_errors)
-
-    return errors

django_tasks/task.py

@@ -56,6 +56,8 @@ class Task(Generic[P, T]):
     run_after: Optional[datetime] = None
     """The earliest this task will run"""
 
+    enqueue_on_commit: Optional[bool] = None
+
     def __post_init__(self) -> None:
         self.get_backend().validate_task(self)
 
@@ -164,6 +166,7 @@ def task(
     priority: int = 0,
     queue_name: str = DEFAULT_QUEUE_NAME,
     backend: str = DEFAULT_TASK_BACKEND_ALIAS,
+    enqueue_on_commit: Optional[bool] = None,
 ) -> Callable[[Callable[P, T]], Task[P, T]]: ...
 
 
@@ -174,6 +177,7 @@ def task(
     priority: int = 0,
     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]]]:
     """
     A decorator used to create a task.
@@ -182,7 +186,11 @@ def task(
 
     def wrapper(f: Callable[P, T]) -> Task[P, T]:
         return tasks[backend].task_class(
-            priority=priority, func=f, queue_name=queue_name, backend=backend
+            priority=priority,
+            func=f,
+            queue_name=queue_name,
+            backend=backend,
+            enqueue_on_commit=enqueue_on_commit,
         )
 
     if function:

tests/settings.py

@@ -58,6 +58,6 @@
 
 USE_TZ = True
 
-if sys.argv[0] != "test":
+if sys.argv[1] == "runserver":
     DEBUG = True
     TASKS = {"default": {"BACKEND": "django_tasks.backends.database.DatabaseBackend"}}

tests/tasks.py

@@ -29,3 +29,13 @@ def failing_task() -> None:
 @task()
 def exit_task() -> None:
     exit(1)
+
+
+@task(enqueue_on_commit=True)
+def enqueue_on_commit_task() -> None:
+    pass
+
+
+@task(enqueue_on_commit=False)
+def never_enqueue_on_commit_task() -> None:
+    pass