Repository: maciej-gol/tenant-schemas-celery

Issue:
Unable to override the default base Task class for tasks registered without the base option
**Context:**
1. We have overridden the `TenantTask` class provided by `tenant-schemas-celery` for our use-case - consider it to be named as `DerivedTenantTask`
2. There exists a task, defined as, for example:
```
@app.task
def my_task(arg1, arg2):
    ...
```

**Issue:**
- `tenant-schemas-celery`'s `CeleryApp.create_task_cls` hardcodes `"tenant_schemas_celery.task:TenantTask"` as the base class that all tasks will be subclassed from
- Due to this, even though we're specifying our custom `DerivedTenantTask` as the base `task_cls` for our celery app, all tasks where the base is not explicity mentioned in the `@app.task` decorator are being subclassed from `tenant-schemas-celery`'s `TenantTask` and due to this, the custom logic we've written in our `DerivedTenantTask` class is not being executed

**Request:**
- Please keep the celery app's `task_cls` attribute in `tenant-schemas-celery`'s `CeleryApp.create_task_cls`:
```
class CeleryApp:
    registry_cls = 'tenant_schemas_celery.registry:TenantTaskRegistry'
    task_cls = "tenant_schemas_celery.task:TenantTask`

    def create_task_cls(self):
        return self.subclass_with_self(
            self.task_cls,
            abstract=True,
            name="TenantTask",
            attribute="_app",
        )
```
- LMK if this was kept intentionally or if you agree that this can be fixed 
- If it can be fixed, also LMK if you're running short on bandwidth and would like me to pick this up myself

Discussion:

Patch:
diff --git a/tenant_schemas_celery/app.py b/tenant_schemas_celery/app.py
index c96e557..f6bc93c 100644
--- a/tenant_schemas_celery/app.py
+++ b/tenant_schemas_celery/app.py
@@ -76,10 +76,11 @@ task_postrun.connect(
 
 class CeleryApp(Celery):
     registry_cls = 'tenant_schemas_celery.registry:TenantTaskRegistry'
-    
+    task_cls = 'tenant_schemas_celery.task:TenantTask'
+
     def create_task_cls(self):
         return self.subclass_with_self(
-            "tenant_schemas_celery.task:TenantTask",
+            self.task_cls,
             abstract=True,
             name="TenantTask",
             attribute="_app",


Test patch:
diff --git a/tenant_schemas_celery/app_test.py b/tenant_schemas_celery/app_test.py
new file mode 100644
index 0000000..58526f8
--- /dev/null
+++ b/tenant_schemas_celery/app_test.py
@@ -0,0 +1,32 @@
+from tenant_schemas_celery.app import CeleryApp
+from tenant_schemas_celery.task import TenantTask
+
+
+class DummyTask(TenantTask):
+    ...
+
+
+def test_celery_app_should_allow_overriding_task_cls_as_object() -> None:
+    class App(CeleryApp):
+        task_cls = DummyTask
+
+    app = App(set_as_current=False)
+
+    @app.task()
+    def some_task() -> None:
+        ...
+
+    assert isinstance(some_task, DummyTask)
+
+
+def test_celery_app_should_allow_overriding_task_cls_as_string() -> None:
+    class App(CeleryApp):
+        task_cls = f"{DummyTask.__module__}:{DummyTask.__name__}"
+
+    app = App(set_as_current=False)
+
+    @app.task()
+    def some_task() -> None:
+        ...
+
+    assert isinstance(some_task, DummyTask)