Repository: apache/airflow

Issue:
Celery executor should support custom database backend models
Currently, the celery executor imports the `Task` model for the Celery database backend directly:

https://github.com/apache/airflow/blob/2dd4e96045d4a7f45cc8c06df3d25c4f1479392c/airflow/executors/celery_executor.py#L38

However, Celery, being highly customisable, uses `self.task_cls` in the backend implementation; this defaults to either `Task` or `TaskExtended` depending on the Celery configuration, see [`celery.backends.database:66`](https://github.com/celery/celery/blob/406f04a082949ac42ec7a4af94fed896c515aaa4/celery/backends/database/__init__.py#L66):

> ```python
>     task_cls = Task
> ```

and [`celery.backends.database:76,77`](https://github.com/celery/celery/blob/406f04a082949ac42ec7a4af94fed896c515aaa4/celery/backends/database/__init__.py#L76-L77):

> ```python
>         if self.extended_result:
>             self.task_cls = TaskExtended
> ```

Airflow should support this flexibility, and instead of importing a fixed class just use `backend.task_cls`, so in `_get_many_from_db_backend` just use:

```python
        Task = app.backend.task_cls
        with session_cleanup(session):
            tasks = session.query(Task).filter(Task.task_id.in_(task_ids)).all()
```

I suppose there is one downside to this: `TaskExtended` fetches a few more columns per task, but given that the result of `task.to_dict()` is passed to `backend.meta_from_decoded()`, which could also have been customised to make use of information contained in a custom model, that's neither here nor there. It would be the price anyone using the `result_extended` option would need to pay elsewhere anyway.


Discussion:
This looks like problems that we can get to quickly. Can I assign you to this ticket? I am happy to help with the review.

Resolving patch:
diff --git a/airflow/executors/celery_executor.py b/airflow/executors/celery_executor.py
--- a/airflow/executors/celery_executor.py
+++ b/airflow/executors/celery_executor.py
@@ -35,7 +35,7 @@
 
 from celery import Celery, Task, states as celery_states
 from celery.backends.base import BaseKeyValueStoreBackend
-from celery.backends.database import DatabaseBackend, Task as TaskDb, session_cleanup
+from celery.backends.database import DatabaseBackend, session_cleanup
 from celery.result import AsyncResult
 from celery.signals import import_modules as celery_import_modules
 from setproctitle import setproctitle  # pylint: disable=no-name-in-module
@@ -562,8 +562,9 @@ def _get_many_from_kv_backend(self, async_tasks) -> Mapping[str, EventBufferValu
     def _get_many_from_db_backend(self, async_tasks) -> Mapping[str, EventBufferValueType]:
         task_ids = _tasks_list_to_task_ids(async_tasks)
         session = app.backend.ResultSession()
+        task_cls = app.backend.task_cls
         with session_cleanup(session):
-            tasks = session.query(TaskDb).filter(TaskDb.task_id.in_(task_ids)).all()
+            tasks = session.query(task_cls).filter(task_cls.task_id.in_(task_ids)).all()
 
         task_results = [app.backend.meta_from_decoded(task.to_dict()) for task in tasks]
         task_results_by_task_id = {task_result["task_id"]: task_result for task_result in task_results}


Test patch: