Subclass DjangoTask instead of Task (#138)

Since we require `django-tenants` to be there, any customer will have
django installed. We can subclass `DjangoTask` directly and use make any
customizations available.

Ref #133

Co-authored-by: Maciej Gol <[email protected]>

README.md

@@ -31,6 +31,12 @@ Usage
    app = TenantAwareCeleryApp()
    app.config_from_object('django.conf:settings')
    app.autodiscover_tasks(lambda: settings.INSTALLED_APPS)
+
+   # On Celery >= 5.4, if you want to use DjangoTask
+   from tenant_schemas_celery.app import CeleryApp
+
+   class TenantAwareCeleryApp(CeleryApp):
+      task_cls = 'tenant_schemas_celery.task:TenantDjangoTask'
 ```
 
 This assumes a fresh Celery 5.2.0 application. For previous versions, the key is to create a new `CeleryApp` instance that will be used to access task decorator from.

run-celery-worker

@@ -1,5 +1,5 @@
 #!/bin/bash
-. .env/bin/activate
+pyenv activate tenant-schemas-celery
 cd test_app
 export DJANGO_SETTINGS_MODULE=test_app.settings
 

tenant_schemas_celery/task.py

@@ -11,7 +11,7 @@ def __init__(self):
         super().__init__(storage=_shared_storage)
 
 
-class TenantTask(Task):
+class TenantTaskMixin:
     """ Custom Task class that injects db schema currently used to the task's
         keywords so that the worker can use the same schema.
     """
@@ -66,3 +66,16 @@ def _add_current_schema(self, kwds):
     def apply(self, args=None, kwargs=None, *arg, **kw):
         self._update_headers(kw)
         return super().apply(args, kwargs, *arg, **kw)
+
+class TenantTask(TenantTaskMixin, Task):
+    ...
+
+
+# DjangoTask requires Celery 5.4. Before then, we can't make it the default.
+try:
+    from celery.contrib.django.task import DjangoTask
+
+    class TenantDjangoTask(TenantTaskMixin, DjangoTask):
+        ...
+except ImportError:
+    pass

tenant_schemas_celery/task_test.py

@@ -1,6 +1,8 @@
 from datetime import datetime, timedelta
 
+import celery
 from freezegun import freeze_time
+import pytest
 
 from tenant_schemas_celery.task import TenantTask
 from tenant_schemas_celery.app import CeleryApp
@@ -145,3 +147,32 @@ class CustomTask(TenantTask):
 
     # Should prioritize the task class's preference over Celery app config
     assert CustomTask.get_tenant_databases() == ("customdb",)
+
+
+@pytest.mark.skipif(celery.VERSION >= (5, 4, 0), reason="DjangoTask is only available on Celery 5.4+")
+def test_should_not_expose_django_task_before_celery_54():
+    class App(CeleryApp):
+        task_cls = 'tenant_schemas_celery.task:TenantDjangoTask'
+
+    app = App('testapp')
+    @app.task
+    def my_task():
+        pass
+
+    with pytest.raises(AttributeError, match="module 'tenant_schemas_celery.task' has no attribute 'TenantDjangoTask'"):
+        # Force evaluation.
+        app.tasks
+
+
+@pytest.mark.skipif(celery.VERSION < (5, 4, 0), reason="DjangoTask is only available on Celery 5.4+")
+def test_should_expose_django_task_on_celery_54():
+    class App(CeleryApp):
+        task_cls = 'tenant_schemas_celery.task:TenantDjangoTask'
+
+    app = App('testapp')
+    @app.task
+    def my_task():
+        pass
+
+    # Force evaluation.
+    app.tasks