Support enqueue_on_commit option

matiasb · matiasb · commit deee9ea9aa0f · 2024-12-02T22:18:21.000-03:00
diff --git a/django_tasks/backends/celery/backend.py b/django_tasks/backends/celery/backend.py
@@ -1,18 +1,19 @@
-from dataclasses import dataclass
+from functools import partial
 from typing import Any, Iterable, TypeVar
 
 from celery import shared_task
 from celery.app import default_app
 from celery.local import Proxy as CeleryTaskProxy
 from django.apps import apps
 from django.core.checks import ERROR, CheckMessage
+from django.db import transaction
 from django.utils import timezone
+from kombu.utils.uuid import uuid
 from typing_extensions import ParamSpec
 
 from django_tasks.backends.base import BaseTaskBackend
 from django_tasks.task import MAX_PRIORITY, MIN_PRIORITY, ResultStatus, TaskResult
 from django_tasks.task import Task as BaseTask
-from django_tasks.utils import json_normalize
 
 if not default_app:
     from django_tasks.backends.celery.app import app as celery_app
@@ -44,7 +45,6 @@ def _map_priority(value: int) -> int:
     return mapped_value
 
 
-@dataclass
 class Task(BaseTask[P, T]):
     celery_task: CeleryTaskProxy = None
     """Celery proxy to the task in the current celery app task registry."""
@@ -77,19 +77,35 @@ def enqueue(
             priority = _map_priority(task.priority)
             apply_async_kwargs["priority"] = priority
 
+        task_id = uuid()
+        apply_async_kwargs["task_id"] = task_id
+
+        if self._get_enqueue_on_commit_for_task(task):
+            transaction.on_commit(
+                partial(
+                    task.celery_task.apply_async,
+                    args,
+                    kwargs=kwargs,
+                    **apply_async_kwargs,
+                )
+            )
+        else:
+            task.celery_task.apply_async(args, kwargs=kwargs, **apply_async_kwargs)
+
+        # TODO: send task_enqueued signal
+        # TODO: link a task to trigger the task_finished signal?
+        # TODO: consider using DBTaskResult for results?
+
         # TODO: a Celery result backend is required to get additional information
-        async_result = task.celery_task.apply_async(
-            args, kwargs=kwargs, **apply_async_kwargs
-        )
         task_result = TaskResult[T](
             task=task,
-            id=async_result.id,
+            id=task_id,
             status=ResultStatus.NEW,
             enqueued_at=timezone.now(),
             started_at=None,
             finished_at=None,
-            args=json_normalize(args),
-            kwargs=json_normalize(kwargs),
+            args=args,
+            kwargs=kwargs,
             backend=self.alias,
         )
         return task_result
diff --git a/tests/tests/test_celery_backend.py b/tests/tests/test_celery_backend.py
@@ -3,7 +3,8 @@
 
 from celery import Celery
 from celery.result import AsyncResult
-from django.test import TestCase, override_settings
+from django.db import transaction
+from django.test import TestCase, TransactionTestCase, override_settings
 from django.utils import timezone
 
 from django_tasks import ResultStatus, default_task_backend, task, tasks
@@ -16,6 +17,10 @@ def noop_task(*args: tuple, **kwargs: dict) -> None:
     return None
 
 
+def enqueue_on_commit_task(*args: tuple, **kwargs: dict) -> None:
+    pass
+
+
 @override_settings(
     TASKS={
         "default": {
@@ -24,10 +29,13 @@ def noop_task(*args: tuple, **kwargs: dict) -> None:
         }
     }
 )
-class CeleryBackendTestCase(TestCase):
+class CeleryBackendTestCase(TransactionTestCase):
     def setUp(self) -> None:
         # register task during setup so it is registered as a Celery task
         self.task = task()(noop_task)
+        self.enqueue_on_commit_task = task(enqueue_on_commit=True)(
+            enqueue_on_commit_task
+        )
 
     def test_using_correct_backend(self) -> None:
         self.assertEqual(default_task_backend, tasks["default"])
@@ -43,7 +51,7 @@ def test_celery_backend_app_missing(self) -> None:
         errors = list(default_task_backend.check())
 
         self.assertEqual(len(errors), 1)
-        self.assertIn("django_tasks.backends.celery", errors[0].hint)
+        self.assertIn("django_tasks.backends.celery", errors[0].hint)  # type:ignore[arg-type]
 
     def test_enqueue_task(self) -> None:
         task = self.task
@@ -53,52 +61,114 @@ def test_enqueue_task(self) -> None:
         from django_tasks.backends.celery.app import app as celery_app
 
         self.assertEqual(task.celery_task.app, celery_app)  # type: ignore[attr-defined]
-        with patch("celery.app.task.Task.apply_async") as mock_apply_async:
-            mock_apply_async.return_value = AsyncResult(id="123")
-            result = default_task_backend.enqueue(task, (1,), {"two": 3})
+        task_id = "123"
+        with patch("django_tasks.backends.celery.backend.uuid", return_value=task_id):
+            with patch("celery.app.task.Task.apply_async") as mock_apply_async:
+                mock_apply_async.return_value = AsyncResult(id=task_id)
+                result = default_task_backend.enqueue(task, (1,), {"two": 3})
 
-        self.assertEqual(result.id, "123")
+        self.assertEqual(result.id, task_id)
         self.assertEqual(result.status, ResultStatus.NEW)
         self.assertIsNone(result.started_at)
         self.assertIsNone(result.finished_at)
         with self.assertRaisesMessage(ValueError, "Task has not finished yet"):
-            result.result  # noqa:B018
+            result.return_value  # noqa:B018
         self.assertEqual(result.task, task)
-        self.assertEqual(result.args, [1])
+        self.assertEqual(result.args, (1,))
         self.assertEqual(result.kwargs, {"two": 3})
         expected_priority = _map_priority(DEFAULT_PRIORITY)
         mock_apply_async.assert_called_once_with(
             (1,),
             kwargs={"two": 3},
+            task_id=task_id,
             eta=None,
             priority=expected_priority,
             queue=DEFAULT_QUEUE_NAME,
         )
 
     def test_using_additional_params(self) -> None:
-        with patch("celery.app.task.Task.apply_async") as mock_apply_async:
-            mock_apply_async.return_value = AsyncResult(id="123")
-            run_after = timezone.now() + timedelta(hours=10)
-            result = self.task.using(
-                run_after=run_after, priority=75, queue_name="queue-1"
-            ).enqueue()
+        task_id = "123"
+        with patch("django_tasks.backends.celery.backend.uuid", return_value=task_id):
+            with patch("celery.app.task.Task.apply_async") as mock_apply_async:
+                mock_apply_async.return_value = AsyncResult(id=task_id)
+                run_after = timezone.now() + timedelta(hours=10)
+                result = self.task.using(
+                    run_after=run_after, priority=75, queue_name="queue-1"
+                ).enqueue()
 
-        self.assertEqual(result.id, "123")
+        self.assertEqual(result.id, task_id)
         self.assertEqual(result.status, ResultStatus.NEW)
         mock_apply_async.assert_called_once_with(
-            (), kwargs={}, eta=run_after, priority=7, queue="queue-1"
+            [], kwargs={}, task_id=task_id, eta=run_after, priority=7, queue="queue-1"
         )
 
     def test_priority_mapping(self) -> None:
         for priority, expected in [(-100, 0), (-50, 2), (0, 4), (75, 7), (100, 9)]:
-            with patch("celery.app.task.Task.apply_async") as mock_apply_async:
-                mock_apply_async.return_value = AsyncResult(id="123")
-                self.task.using(priority=priority).enqueue()
+            task_id = "123"
+            with patch(
+                "django_tasks.backends.celery.backend.uuid", return_value=task_id
+            ):
+                with patch("celery.app.task.Task.apply_async") as mock_apply_async:
+                    mock_apply_async.return_value = AsyncResult(id=task_id)
+                    self.task.using(priority=priority).enqueue()
 
             mock_apply_async.assert_called_with(
-                (), kwargs={}, eta=None, priority=expected, queue=DEFAULT_QUEUE_NAME
+                [],
+                kwargs={},
+                task_id=task_id,
+                eta=None,
+                priority=expected,
+                queue=DEFAULT_QUEUE_NAME,
             )
 
+    @override_settings(
+        TASKS={
+            "default": {
+                "BACKEND": "django_tasks.backends.celery.CeleryBackend",
+                "ENQUEUE_ON_COMMIT": True,
+            }
+        }
+    )
+    def test_wait_until_transaction_commit(self) -> None:
+        self.assertTrue(default_task_backend.enqueue_on_commit)
+        self.assertTrue(default_task_backend._get_enqueue_on_commit_for_task(self.task))
+
+        with patch("celery.app.task.Task.apply_async") as mock_apply_async:
+            mock_apply_async.return_value = AsyncResult(id="task_id")
+            with transaction.atomic():
+                self.task.enqueue()
+                assert not mock_apply_async.called
+
+        mock_apply_async.assert_called_once()
+
+    @override_settings(
+        TASKS={
+            "default": {
+                "BACKEND": "django_tasks.backends.celery.CeleryBackend",
+            }
+        }
+    )
+    def test_wait_until_transaction_by_default(self) -> None:
+        self.assertTrue(default_task_backend.enqueue_on_commit)
+        self.assertTrue(default_task_backend._get_enqueue_on_commit_for_task(self.task))
+
+    @override_settings(
+        TASKS={
+            "default": {
+                "BACKEND": "django_tasks.backends.celery.CeleryBackend",
+                "ENQUEUE_ON_COMMIT": False,
+            }
+        }
+    )
+    def test_task_specific_enqueue_on_commit(self) -> None:
+        self.assertFalse(default_task_backend.enqueue_on_commit)
+        self.assertTrue(self.enqueue_on_commit_task.enqueue_on_commit)
+        self.assertTrue(
+            default_task_backend._get_enqueue_on_commit_for_task(
+                self.enqueue_on_commit_task
+            )
+        )
+
 
 @override_settings(
     TASKS={
