Don't dispose pools when using internal api (apache#38552)

Needed for AIP-44
astronomer · Mar 28, 2024 · 5de9075 · 5de9075
1 parent 4e070ef
commit 5de9075
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/airflow/task/task_runner/standard_task_runner.py b/airflow/task/task_runner/standard_task_runner.py
@@ -26,6 +26,7 @@
 import psutil
 from setproctitle import setproctitle
 
+from airflow.api_internal.internal_api_call import InternalApiConfig
 from airflow.models.taskinstance import TaskReturnCode
 from airflow.settings import CAN_FORK
 from airflow.task.task_runner.base_task_runner import BaseTaskRunner
@@ -74,11 +75,12 @@ def _start_by_fork(self):
             from airflow.cli.cli_parser import get_parser
             from airflow.sentry import Sentry
 
-            # Force a new SQLAlchemy session. We can't share open DB handles
-            # between process. The cli code will re-create this as part of its
-            # normal startup
-            settings.engine.pool.dispose()
-            settings.engine.dispose()
+            if not InternalApiConfig.get_use_internal_api():
+                # Force a new SQLAlchemy session. We can't share open DB handles
+                # between process. The cli code will re-create this as part of its
+                # normal startup
+                settings.engine.pool.dispose()
+                settings.engine.dispose()
 
             parser = get_parser()
             # [1:] - remove "airflow" from the start of the command
@@ -107,7 +109,7 @@ def _start_by_fork(self):
             except Exception as exc:
                 return_code = 1
 
-                self.log.error(
+                self.log.exception(
                     "Failed to execute job %s for task %s (%s; %r)",
                     job_id,
                     self._task_instance.task_id,