Repository: apache/airflow

Issue:
 Celery Executor cannot connect to the database to get information, resulting in a scheduler exit abnormally
### Apache Airflow version

Other Airflow 2 version (please specify below)

### What happened


We  use Celery Executor where using RabbitMQ as a broker and postgresql as a result backend

Airflow Version: 2.2.3
Celery Version: 5.2.3


apache-airflow-providers-celery==2.1.0



Below is the error message:
```
_The above exception was the direct cause of the following exception: Traceback (most recent call last):
File"/app/airflow2.2.3/airflow/airflow/jobs/schedulerjob.py”, line 672, in _execute self._run_scheduler_loop()
File"/app/airflow2.2.3/airflow/airflow/jobs/scheduler_job.py", line 754, in _run_scheduler_loop self.executor.heartbeat()
File"/app/airflow2.2.3/airflow/airflow/executors/base_executor.py”, line 168, in heartbeat self.sync()
File"/app/airflow2.2.3/airflow/airflow/executors/celery_executorpy”, line 330, in sync self.update_all_task_states()
File"/app/airflow223/airflow/airflow/executors/celery_executor.py”,line 442,in update_all_task_states state_and_info_by_celery_task_id=self.bulk_state_fetcher.get_many(self.tasks. values()) File"/app/airflow2.2.3/airflow/airflow/executors/celery_executorpy”,line 598, in get_many result = self._get many_from db backend(async_results)
File"/app/airflow2.2.3/airflow/airflow/executors/celery_executor.py”,line 618, in get_many_from_db_backend tasks-session.query(task_cls).filter(task_cls.task_id.in(task_ids)).all()
File“/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/orm/query.py”, line 3373, in all return list(self)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/orm/query.py”, line 3535, in iter return self._execute_and_instances(context)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/orm/query.py”,line 3556, in _execute_and_instances conn =self._get bind args(
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/salalchemy/orm/query.py”, line 3571, in _get_bind_args return fn(
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/orm/query.py”,line 3550, in _connection_from_session conn=self.session.connection(**kw)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/orm/session.py”, line 1142, in connection return self._connection_for_bind(
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/orm/session.py”,line 1150, in _connection_for_bind return self.transaction.connection_for bind(
File“/app/airflow2.2.3/airflow2_env/Iib/python3.8/site-packages/sqlalchemy/orm/session.py”, line 433, in _connection_for_bind conn=bind._contextual_connect()
File“/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/salalchemy/engine/base.py”,line 2302, in _contextual_connect self._wrap_pool_connect(self.pool.connect,None),
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/engine/base.py”, line 2339, in _wrap_pool_connect
Tracking Connection.handle dbapi_exception_noconnection(
File "/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/engine/base.py”， line 1583,in handle_dbapi_exception_noconnection util.raise (
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/util/compat.py”, line 182, in raise
ents raise exception File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/salalchemy/engine/base.py”, line 2336, in _wrap_pool_connect
return fn()

2023-06-05 16:39:05.069 ERROR -Exception when executing SchedulerJob. run scheduler loop Traceback (most recent call last):
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/engine/base.py”, line 2336,in _wrap_pool_connect return fno
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/pool/base.py”, line 364, in connect returnConnectionFairy.checkout(self)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/pool/base.py”, line 778, in _checkout fairy=ConnectionRecordcheckout(pool)
File"/app/airflow2.2.3/airflow2_env/lib/python3.8/site-packages/sqlalchemy/pool/base.py”, line 495, in checkout rec=pool. do_get()
File“/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/pool/impl.py”, line 241, in _do_get return self._createconnection()
File"/app/airflow2.2.3/airflow2_env/lib/python3.8/site-packages/salalchemy/pool/base.py”, line 309, in _create_connection return _ConnectionRecord(self)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/sitepackages/sqlalchemy/pool/base.py”, line 440, in init self. connect(firstconnectcheck=True)
File"/app/airflow2.2.3/airflow2_env/lib/python3.8/site-packages/sqlalchemy/pool/base.py”, line 661, in connect pool.logger.debug"Error onconnect(:%s"，e)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/util/langhelpers.py”, line 68, in exit compat.raise(
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/salalchemy/util/compat.py", line 182, in raise raise exception
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/pool/base.py”, line 656, in _connect connection =pool.invoke_creator(sel f)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/engine/strategies.py”, line 114, in connect return dialect.connect(*cargs, **cparans)
File"/app/airflow2.2.3/airflow2_env/1ib/python3.8/site-packages/sqlalchemy/engine/default.py”,line 508, in connect return self.dbapi.connect(*cargs, **cparams)
File"/app/airflow2.2.3/airflow2_env/lib/python3.8/site-packages/psycopg2/init.py”, line 126, in connect conn=connect(dsn，connection_factory=connection_factory, **kwasync) psycopg2.0perationalError: could not connect to server: Connection timed out
Is the server running on host"xxxxxxxxxx”and accepting TCP/IP connections on port 5432?
```
### What you think should happen instead

I think it may be caused by network jitter issues, add retries to solve it

### How to reproduce

celeryExecutor fails to create a PG connection while retrieving metadata information, and it can be reproduced

### Operating System

NAME="RedFlag Asianux"  VERSION="7 (Lotus)"   

### Versions of Apache Airflow Providers

_No response_

### Deployment

Virtualenv installation

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Resolving patch:
diff --git a/airflow/executors/celery_executor_utils.py b/airflow/executors/celery_executor_utils.py
--- a/airflow/executors/celery_executor_utils.py
+++ b/airflow/executors/celery_executor_utils.py
@@ -31,7 +31,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, Task as TaskDb, retry, session_cleanup
 from celery.result import AsyncResult
 from celery.signals import import_modules as celery_import_modules
 from setproctitle import setproctitle
@@ -250,15 +250,19 @@ def _get_many_from_kv_backend(self, async_tasks) -> Mapping[str, EventBufferValu
 
         return self._prepare_state_and_info_by_task_dict(task_ids, task_results_by_task_id)
 
-    def _get_many_from_db_backend(self, async_tasks) -> Mapping[str, EventBufferValueType]:
-        task_ids = self._tasks_list_to_task_ids(async_tasks)
+    @retry
+    def _query_task_cls_from_db_backend(self, task_ids, **kwargs):
         session = app.backend.ResultSession()
         task_cls = getattr(app.backend, "task_cls", TaskDb)
         with session_cleanup(session):
-            tasks = session.query(task_cls).filter(task_cls.task_id.in_(task_ids)).all()
+            return session.query(task_cls).filter(task_cls.task_id.in_(task_ids)).all()
 
+    def _get_many_from_db_backend(self, async_tasks) -> Mapping[str, EventBufferValueType]:
+        task_ids = self._tasks_list_to_task_ids(async_tasks)
+        tasks = self._query_task_cls_from_db_backend(task_ids)
         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}
+
         return self._prepare_state_and_info_by_task_dict(task_ids, task_results_by_task_id)
 
     @staticmethod


Test patch:
diff --git a/tests/integration/executors/test_celery_executor.py b/tests/integration/executors/test_celery_executor.py
--- a/tests/integration/executors/test_celery_executor.py
+++ b/tests/integration/executors/test_celery_executor.py
@@ -310,6 +310,39 @@ def test_should_support_db_backend(self, mock_session, caplog):
         assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)}
         assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"]
 
+    @mock.patch("celery.backends.database.DatabaseBackend.ResultSession")
+    def test_should_retry_db_backend(self, mock_session, caplog):
+        caplog.set_level(logging.DEBUG, logger=self.bulk_state_fetcher_logger)
+        from sqlalchemy.exc import DatabaseError
+
+        with _prepare_app():
+            mock_backend = DatabaseBackend(app=celery_executor.app, url="sqlite3://")
+            with mock.patch("airflow.executors.celery_executor_utils.Celery.backend", mock_backend):
+                caplog.clear()
+                mock_session = mock_backend.ResultSession.return_value
+                mock_retry_db_result = mock_session.query.return_value.filter.return_value.all
+                mock_retry_db_result.return_value = [
+                    mock.MagicMock(**{"to_dict.return_value": {"status": "SUCCESS", "task_id": "123"}})
+                ]
+                mock_retry_db_result.side_effect = [
+                    DatabaseError("DatabaseError", "DatabaseError", "DatabaseError"),
+                    mock_retry_db_result.return_value,
+                ]
+
+                fetcher = celery_executor_utils.BulkStateFetcher()
+                result = fetcher.get_many(
+                    [
+                        mock.MagicMock(task_id="123"),
+                        mock.MagicMock(task_id="456"),
+                    ]
+                )
+        assert mock_retry_db_result.call_count == 2
+        assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)}
+        assert caplog.messages == [
+            "Failed operation _query_task_cls_from_db_backend.  Retrying 2 more times.",
+            "Fetched 2 state(s) for 2 task(s)",
+        ]
+
     def test_should_support_base_backend(self, caplog):
         caplog.set_level(logging.DEBUG, logger=self.bulk_state_fetcher_logger)
         with _prepare_app():