From a1f754d4a83d682c59d8928977b75b27654e33e4 Mon Sep 17 00:00:00 2001
From: Joshua <joshua.chyeung@gmail.com>
Date: Wed, 28 Sep 2022 12:13:16 +0800
Subject: [PATCH 01/22] Added IstioKubernetesPodOperator and IstioPodManager

---
 .../operators/istio_kubernetes_pod.py         | 58 +++++++++++++++++
 .../kubernetes/utils/istio_pod_manager.py     | 65 +++++++++++++++++++
 .../cncf/kubernetes/utils/pod_manager.py      | 30 +++++++++
 3 files changed, 153 insertions(+)
 create mode 100644 airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
 create mode 100644 airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py

diff --git a/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
new file mode 100644
index 0000000000000..ab6401f65cd3a
--- /dev/null
+++ b/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Executes task in a Kubernetes POD with Istio enabled"""
+from __future__ import annotations
+
+from kubernetes.client import models as k8s
+
+from airflow import AirflowException
+from airflow.compat.functools import cached_property
+from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator, _suppress
+from airflow.providers.cncf.kubernetes.utils.istio_pod_manager import IstioPodManager
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
+
+
+class IstioKubernetesPodOperator(KubernetesPodOperator):
+    """
+    Execute a task in a Kubernetes Pod with Istio enabled
+    All parameters are the same with KubernetesPodOperator, the only difference is the cleanup part
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KubernetesPodOperator`
+    """
+
+    @cached_property
+    def pod_manager(self) -> IstioPodManager:
+        return IstioPodManager(kube_client=self.client)
+
+    def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
+        pod_phase = remote_pod.status.phase if hasattr(remote_pod, 'status') else None
+        if pod_phase != PodPhase.SUCCEEDED and not self.pod_manager.container_is_succeeded(pod, 'base'):
+            if self.log_events_on_failure:
+                with _suppress(Exception):
+                    for event in self.pod_manager.read_pod_events(pod).items:
+                        self.log.error("Pod Event: %s - %s", event.reason, event.message)
+            if not self.is_delete_operator_pod:
+                with _suppress(Exception):
+                    self.patch_already_checked(pod)
+            with _suppress(Exception):
+                self.process_pod_deletion(pod)
+            raise AirflowException(f'Pod {pod and pod.metadata.name} returned a failure: {remote_pod}')
+        else:
+            with _suppress(Exception):
+                self.process_pod_deletion(pod)
diff --git a/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py b/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
new file mode 100644
index 0000000000000..2248d565181f0
--- /dev/null
+++ b/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import time
+
+from kubernetes.client.models.v1_pod import V1Pod
+
+from airflow.providers.cncf.kubernetes.utils.pod_manager import (
+    PodManager,
+    PodPhase,
+    container_is_completed,
+    container_is_succeeded,
+)
+
+
+class IstioPodManager(PodManager):
+    """
+    Helper class for checking the status of base container in Kubernetes pods
+    for use with the IstioKubernetesPodOperator
+    """
+
+    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is completed"""
+        remote_pod = self.read_pod(pod)
+        return container_is_completed(pod=remote_pod, container_name=container_name)
+
+    def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is succeeded"""
+        remote_pod = self.read_pod(pod)
+        return container_is_succeeded(pod=remote_pod, container_name=container_name)
+
+    def await_pod_completion(self, pod: V1Pod) -> V1Pod:
+        """
+        Monitors a pod and returns the final state
+        (neglect sidecar state e.g. istio-proxy, vault-agent)
+        :param pod: pod spec that will be monitored
+        :return:  Tuple[State, Optional[str]]
+        """
+        while True:
+            remote_pod = self.read_pod(pod)
+            if remote_pod.status.phase in PodPhase.terminal_states:
+                break
+            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
+                remote_pod, 'base'
+            ):
+                break
+            self.log.info('Pod %s has phase %s', pod.metadata.name, remote_pod.status.phase)
+            time.sleep(2)
+        return remote_pod
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index bd431e078b9fc..d89ff978dcaa5 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -144,6 +144,36 @@ def container_is_terminated(pod: V1Pod, container_name: str) -> bool:
     return container_status.state.terminated is not None
 
 
+def container_is_completed(pod: V1Pod, container_name: str) -> bool:
+    """
+    Examines V1Pod ``pod`` to determine whether ``container_name`` is completed.
+    If that container is present and completed, returns True.  Returns False otherwise.
+    """
+    container_statuses = pod.status.container_statuses if pod and pod.status else None
+    if not container_statuses:
+        return False
+    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+    if not container_status:
+        return False
+    return container_status.state.terminated is not None
+
+
+def container_is_succeeded(pod: V1Pod, container_name: str) -> bool:
+    """
+    Examines V1Pod ``pod`` to determine whether ``container_name`` is completed and succeeded.
+    If that container is present and completed and succeeded, returns True.  Returns False otherwise.
+    """
+    if not container_is_completed(pod, container_name):
+        return False
+    container_statuses = pod.status.container_statuses if pod and pod.status else None
+    if not container_statuses:
+        return False
+    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+    if not container_status:
+        return False
+    return container_status.state.terminated.exit_code == 0
+
+
 def get_container_termination_message(pod: V1Pod, container_name: str):
     with suppress(AttributeError, TypeError):
         container_statuses = pod.status.container_statuses

From adbffcb54dec5c4254441bbd6f8194454897f1ee Mon Sep 17 00:00:00 2001
From: Joshua <joshua.chyeung@gmail.com>
Date: Fri, 21 Oct 2022 14:31:03 +0800
Subject: [PATCH 02/22] Merged into KubernetesPodOperator(istio_enabled=True)

---
 .../operators/istio_kubernetes_pod.py         | 58 -----------------
 .../kubernetes/utils/istio_pod_manager.py     | 65 -------------------
 2 files changed, 123 deletions(-)
 delete mode 100644 airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
 delete mode 100644 airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py

diff --git a/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
deleted file mode 100644
index ab6401f65cd3a..0000000000000
--- a/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Executes task in a Kubernetes POD with Istio enabled"""
-from __future__ import annotations
-
-from kubernetes.client import models as k8s
-
-from airflow import AirflowException
-from airflow.compat.functools import cached_property
-from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator, _suppress
-from airflow.providers.cncf.kubernetes.utils.istio_pod_manager import IstioPodManager
-from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
-
-
-class IstioKubernetesPodOperator(KubernetesPodOperator):
-    """
-    Execute a task in a Kubernetes Pod with Istio enabled
-    All parameters are the same with KubernetesPodOperator, the only difference is the cleanup part
-
-    .. seealso::
-        For more information on how to use this operator, take a look at the guide:
-        :ref:`howto/operator:KubernetesPodOperator`
-    """
-
-    @cached_property
-    def pod_manager(self) -> IstioPodManager:
-        return IstioPodManager(kube_client=self.client)
-
-    def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
-        pod_phase = remote_pod.status.phase if hasattr(remote_pod, 'status') else None
-        if pod_phase != PodPhase.SUCCEEDED and not self.pod_manager.container_is_succeeded(pod, 'base'):
-            if self.log_events_on_failure:
-                with _suppress(Exception):
-                    for event in self.pod_manager.read_pod_events(pod).items:
-                        self.log.error("Pod Event: %s - %s", event.reason, event.message)
-            if not self.is_delete_operator_pod:
-                with _suppress(Exception):
-                    self.patch_already_checked(pod)
-            with _suppress(Exception):
-                self.process_pod_deletion(pod)
-            raise AirflowException(f'Pod {pod and pod.metadata.name} returned a failure: {remote_pod}')
-        else:
-            with _suppress(Exception):
-                self.process_pod_deletion(pod)
diff --git a/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py b/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
deleted file mode 100644
index 2248d565181f0..0000000000000
--- a/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-from __future__ import annotations
-
-import time
-
-from kubernetes.client.models.v1_pod import V1Pod
-
-from airflow.providers.cncf.kubernetes.utils.pod_manager import (
-    PodManager,
-    PodPhase,
-    container_is_completed,
-    container_is_succeeded,
-)
-
-
-class IstioPodManager(PodManager):
-    """
-    Helper class for checking the status of base container in Kubernetes pods
-    for use with the IstioKubernetesPodOperator
-    """
-
-    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is completed"""
-        remote_pod = self.read_pod(pod)
-        return container_is_completed(pod=remote_pod, container_name=container_name)
-
-    def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is succeeded"""
-        remote_pod = self.read_pod(pod)
-        return container_is_succeeded(pod=remote_pod, container_name=container_name)
-
-    def await_pod_completion(self, pod: V1Pod) -> V1Pod:
-        """
-        Monitors a pod and returns the final state
-        (neglect sidecar state e.g. istio-proxy, vault-agent)
-        :param pod: pod spec that will be monitored
-        :return:  Tuple[State, Optional[str]]
-        """
-        while True:
-            remote_pod = self.read_pod(pod)
-            if remote_pod.status.phase in PodPhase.terminal_states:
-                break
-            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
-                remote_pod, 'base'
-            ):
-                break
-            self.log.info('Pod %s has phase %s', pod.metadata.name, remote_pod.status.phase)
-            time.sleep(2)
-        return remote_pod

From 8b831dea97dc41428b7e4f202e0b048f232f94bb Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Thu, 18 May 2023 22:19:17 +0800
Subject: [PATCH 03/22] Add support to istio

Continuing from joshua's work to integrate istio support into KubernetesPodOperator
---
 .../cncf/kubernetes/operators/pod.py          |  18 ++-
 .../cncf/kubernetes/utils/pod_manager.py      |  30 +++--
 .../cncf/kubernetes/utils/test_pod_manager.py | 117 ++++++++++++++++++
 3 files changed, 153 insertions(+), 12 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index e3f363e7c2319..19ca759b1cd2c 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -227,6 +227,7 @@ class KubernetesPodOperator(BaseOperator):
         container name to use.
     :param deferrable: Run operator in the deferrable mode.
     :param poll_interval: Polling period in seconds to check for the status. Used only in deferrable mode.
+    :param istio_enabled: Whether Istio is enabled in k8s cluster. False by default.
     """
 
     # This field can be overloaded at the instance level via base_container_name
@@ -303,6 +304,7 @@ def __init__(
         base_container_name: str | None = None,
         deferrable: bool = False,
         poll_interval: float = 2,
+        istio_enabled: bool = False,
         **kwargs,
     ) -> None:
         # TODO: remove in provider 6.0.0 release. This is a mitigate step to advise users to switch to the
@@ -384,6 +386,7 @@ def __init__(
         self.poll_interval = poll_interval
         self.remote_pod: k8s.V1Pod | None = None
         self._config_dict: dict | None = None  # TODO: remove it when removing convert_config_file_to_dict
+        self.istio_enabled = istio_enabled
 
     @cached_property
     def _incluster_namespace(self):
@@ -461,7 +464,7 @@ def _get_ti_pod_labels(context: Context | None = None, include_try_number: bool
 
     @cached_property
     def pod_manager(self) -> PodManager:
-        return PodManager(kube_client=self.client)
+        return PodManager(kube_client=self.client, istio_enabled=self.istio_enabled)
 
     @cached_property
     def hook(self) -> PodOperatorHookProtocol:
@@ -687,13 +690,20 @@ def post_complete_action(self, *, pod, remote_pod, **kwargs):
             remote_pod=remote_pod,
         )
 
-    def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
+    def time_to_clean_up(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod) -> bool:
         pod_phase = remote_pod.status.phase if hasattr(remote_pod, "status") else None
+        if self.istio_enabled:
+            return pod_phase != PodPhase.SUCCEEDED and not self.pod_manager.container_is_succeeded(
+                pod=pod, container_name="base"
+            )
+        else:
+            return pod_phase != PodPhase.SUCCEEDED
 
-        if pod_phase != PodPhase.SUCCEEDED or not self.is_delete_operator_pod:
+    def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
+        if (self.time_to_clean_up(pod=pod, remote_pod=remote_pod)) or (not self.is_delete_operator_pod):
             self.patch_already_checked(remote_pod, reraise=False)
 
-        if pod_phase != PodPhase.SUCCEEDED:
+        if self.time_to_clean_up(pod=pod, remote_pod=remote_pod):
             if self.log_events_on_failure:
                 self._read_pod_events(pod, reraise=False)
 
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index d89ff978dcaa5..17ed954d73a41 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -149,10 +149,7 @@ def container_is_completed(pod: V1Pod, container_name: str) -> bool:
     Examines V1Pod ``pod`` to determine whether ``container_name`` is completed.
     If that container is present and completed, returns True.  Returns False otherwise.
     """
-    container_statuses = pod.status.container_statuses if pod and pod.status else None
-    if not container_statuses:
-        return False
-    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+    container_status = get_container_status(pod, container_name)
     if not container_status:
         return False
     return container_status.state.terminated is not None
@@ -165,10 +162,8 @@ def container_is_succeeded(pod: V1Pod, container_name: str) -> bool:
     """
     if not container_is_completed(pod, container_name):
         return False
-    container_statuses = pod.status.container_statuses if pod and pod.status else None
-    if not container_statuses:
-        return False
-    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+
+    container_status = get_container_status(pod, container_name)
     if not container_status:
         return False
     return container_status.state.terminated.exit_code == 0
@@ -284,6 +279,7 @@ class PodManager(LoggingMixin):
     def __init__(
         self,
         kube_client: client.CoreV1Api,
+        istio_enabled: bool = False,
     ):
         """
         Creates the launcher.
@@ -293,6 +289,7 @@ def __init__(
         super().__init__()
         self._client = kube_client
         self._watch = watch.Watch()
+        self._istio_enabled = istio_enabled
 
     def run_pod_async(self, pod: V1Pod, **kwargs) -> V1Pod:
         """Runs POD asynchronously"""
@@ -459,6 +456,13 @@ def await_pod_completion(self, pod: V1Pod) -> V1Pod:
             remote_pod = self.read_pod(pod)
             if remote_pod.status.phase in PodPhase.terminal_states:
                 break
+
+            if (
+                self._istio_enabled
+                and remote_pod.status.phase == PodPhase.RUNNING
+                and self.container_is_completed(remote_pod, "base")
+            ):
+                break
             self.log.info("Pod %s has phase %s", pod.metadata.name, remote_pod.status.phase)
             time.sleep(2)
         return remote_pod
@@ -497,6 +501,16 @@ def container_is_terminated(self, pod: V1Pod, container_name: str) -> bool:
         remote_pod = self.read_pod(pod)
         return container_is_terminated(pod=remote_pod, container_name=container_name)
 
+    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is completed"""
+        remote_pod = self.read_pod(pod)
+        return container_is_completed(pod=remote_pod, container_name=container_name)
+
+    def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is succeeded"""
+        remote_pod = self.read_pod(pod)
+        return container_is_succeeded(pod=remote_pod, container_name=container_name)
+
     @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_exponential(), reraise=True)
     def read_pod_logs(
         self,
diff --git a/tests/providers/cncf/kubernetes/utils/test_pod_manager.py b/tests/providers/cncf/kubernetes/utils/test_pod_manager.py
index a55be38a5e227..4bdffde707deb 100644
--- a/tests/providers/cncf/kubernetes/utils/test_pod_manager.py
+++ b/tests/providers/cncf/kubernetes/utils/test_pod_manager.py
@@ -18,6 +18,7 @@
 
 import logging
 from datetime import datetime
+from types import SimpleNamespace
 from unittest import mock
 from unittest.mock import MagicMock
 
@@ -34,7 +35,9 @@
     PodLogsConsumer,
     PodManager,
     PodPhase,
+    container_is_completed,
     container_is_running,
+    container_is_succeeded,
     container_is_terminated,
 )
 from airflow.utils.timezone import utc
@@ -371,6 +374,120 @@ def test_container_is_terminated_with_waiting_state(self, container_state, expec
         assert container_is_terminated(pod_info, "base") == expected_is_terminated
 
 
+def params_for_test_container_is_completed():
+    """The `container_is_completed` method is designed to handle an assortment of bad objects
+    returned from `read_pod`.  E.g. a None object, an object `e` such that `e.status` is None,
+    an object `e` such that `e.status.container_statuses` is None, and so on.  This function
+    emits params used in `test_container_is_completed` to verify this behavior.
+    We create mock classes not derived from MagicMock because with an instance `e` of MagicMock,
+    tests like `e.hello is not None` are always True.
+    """
+
+    class RemotePodMock:
+        pass
+
+    class ContainerStatusMock:
+        def __init__(self, name):
+            self.name = name
+
+    def remote_pod(completed=None, not_completed=None):
+        e = RemotePodMock()
+        e.status = RemotePodMock()
+        e.status.container_statuses = []
+        for r in not_completed or []:
+            e.status.container_statuses.append(container(r, False))
+        for r in completed or []:
+            e.status.container_statuses.append(container(r, True))
+        return e
+
+    def container(name, completed):
+        c = ContainerStatusMock(name)
+        c.state = RemotePodMock()
+        c.state.terminated = {"a": "b"} if completed else None
+        return c
+
+    pod_mock_list = []
+    pod_mock_list.append(pytest.param(None, False, id="None remote_pod"))
+    p = RemotePodMock()
+    p.status = None
+    pod_mock_list.append(pytest.param(p, False, id="None remote_pod.status"))
+    p = RemotePodMock()
+    p.status = RemotePodMock()
+    p.status.container_statuses = []
+    pod_mock_list.append(pytest.param(p, False, id="empty remote_pod.status.container_statuses"))
+    pod_mock_list.append(pytest.param(remote_pod(), False, id="filter empty"))
+    pod_mock_list.append(pytest.param(remote_pod(None, ["base"]), False, id="filter 0 completed"))
+    pod_mock_list.append(pytest.param(remote_pod(["hello"], ["base"]), False, id="filter 1 not completed"))
+    pod_mock_list.append(pytest.param(remote_pod(["base"], ["hello"]), True, id="filter 1 completed"))
+    return pod_mock_list
+
+
+@pytest.mark.parametrize("remote_pod, result", params_for_test_container_is_completed())
+def test_container_is_completed(remote_pod, result):
+    """The `container_is_completed` function is designed to handle an assortment of bad objects
+    returned from `read_pod`.  E.g. a None object, an object `e` such that `e.status` is None,
+    an object `e` such that `e.status.container_statuses` is None, and so on.  This test
+    verifies the expected behavior."""
+    assert container_is_completed(remote_pod, "base") is result
+
+
+def params_for_test_container_is_succeeded():
+    """The `container_is_succeeded` method is designed to handle an assortment of bad objects
+    returned from `read_pod`.  E.g. a None object, an object `e` such that `e.status` is None,
+    an object `e` such that `e.status.container_statuses` is None, and so on.  This function
+    emits params used in `test_container_is_succeeded` to verify this behavior.
+    We create mock classes not derived from MagicMock because with an instance `e` of MagicMock,
+    tests like `e.hello is not None` are always True.
+    """
+
+    class RemotePodMock:
+        pass
+
+    class ContainerStatusMock:
+        def __init__(self, name):
+            self.name = name
+
+    def remote_pod(succeeded=None, not_succeeded=None):
+        e = RemotePodMock()
+        e.status = RemotePodMock()
+        e.status.container_statuses = []
+        for r in not_succeeded or []:
+            e.status.container_statuses.append(container(r, False))
+        for r in succeeded or []:
+            e.status.container_statuses.append(container(r, True))
+        return e
+
+    def container(name, succeeded):
+        c = ContainerStatusMock(name)
+        c.state = RemotePodMock()
+        c.state.terminated = SimpleNamespace(**{"exit_code": 0}) if succeeded else None
+        return c
+
+    pod_mock_list = []
+    pod_mock_list.append(pytest.param(None, False, id="None remote_pod"))
+    p = RemotePodMock()
+    p.status = None
+    pod_mock_list.append(pytest.param(p, False, id="None remote_pod.status"))
+    p = RemotePodMock()
+    p.status = RemotePodMock()
+    p.status.container_statuses = []
+    pod_mock_list.append(pytest.param(p, False, id="empty remote_pod.status.container_statuses"))
+    pod_mock_list.append(pytest.param(remote_pod(), False, id="filter empty"))
+    pod_mock_list.append(pytest.param(remote_pod(None, ["base"]), False, id="filter 0 succeeded"))
+    pod_mock_list.append(pytest.param(remote_pod(["hello"], ["base"]), False, id="filter 1 not succeeded"))
+    pod_mock_list.append(pytest.param(remote_pod(["base"], ["hello"]), True, id="filter 1 succeeded"))
+    return pod_mock_list
+
+
+@pytest.mark.parametrize("remote_pod, result", params_for_test_container_is_succeeded())
+def test_container_is_succeeded(remote_pod, result):
+    """The `container_is_succeeded` function is designed to handle an assortment of bad objects
+    returned from `read_pod`.  E.g. a None object, an object `e` such that `e.status` is None,
+    an object `e` such that `e.status.container_statuses` is None, and so on.  This test
+    verifies the expected behavior."""
+    assert container_is_succeeded(remote_pod, "base") is result
+
+
 def params_for_test_container_is_running():
     """The `container_is_running` method is designed to handle an assortment of bad objects
     returned from `read_pod`.  E.g. a None object, an object `e` such that `e.status` is None,

From 1aa7685491a00dda841767265f28a56922875b46 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Fri, 19 May 2023 01:01:51 +0800
Subject: [PATCH 04/22] Add istio into spelling word

Fix failing spell check
---
 airflow/providers/cncf/kubernetes/operators/pod.py | 2 +-
 docs/spelling_wordlist.txt                         | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index 19ca759b1cd2c..b9b5eafa9d906 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -227,7 +227,7 @@ class KubernetesPodOperator(BaseOperator):
         container name to use.
     :param deferrable: Run operator in the deferrable mode.
     :param poll_interval: Polling period in seconds to check for the status. Used only in deferrable mode.
-    :param istio_enabled: Whether Istio is enabled in k8s cluster. False by default.
+    :param istio_enabled: Whether istio is enabled in k8s cluster. False by default.
     """
 
     # This field can be overloaded at the instance level via base_container_name
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index dc4da70a26524..a3f834c8e7a82 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -772,6 +772,7 @@ isfile
 ish
 isn
 isort
+istio
 iterable
 iterables
 iteratively

From 53bd8122056e9f887a01025c548f83335d4d5a12 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Sat, 20 May 2023 20:52:26 +0800
Subject: [PATCH 05/22] add system test dag with istio enabled

Add a new system test dag with istio enabled
---
 .../kubernetes/example_kubernetes_istio.py    | 181 ++++++++++++++++++
 1 file changed, 181 insertions(+)
 create mode 100644 tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py

diff --git a/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
new file mode 100644
index 0000000000000..c0abe2f845462
--- /dev/null
+++ b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
@@ -0,0 +1,181 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+This is an example dag for using the KubernetesPodOperator.
+"""
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+from kubernetes.client import models as k8s
+
+from airflow import DAG
+from airflow.kubernetes.secret import Secret
+from airflow.operators.bash import BashOperator
+from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
+
+# [START howto_operator_k8s_cluster_resources]
+secret_file = Secret("volume", "/etc/sql_conn", "airflow-secrets", "sql_alchemy_conn")
+secret_env = Secret("env", "SQL_CONN", "airflow-secrets", "sql_alchemy_conn")
+secret_all_keys = Secret("env", None, "airflow-secrets-2")
+volume_mount = k8s.V1VolumeMount(
+    name="test-volume", mount_path="/root/mount_file", sub_path=None, read_only=True
+)
+
+configmaps = [
+    k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name="test-configmap-1")),
+    k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name="test-configmap-2")),
+]
+
+volume = k8s.V1Volume(
+    name="test-volume",
+    persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name="test-volume"),
+)
+
+port = k8s.V1ContainerPort(name="http", container_port=80)
+
+init_container_volume_mounts = [
+    k8s.V1VolumeMount(mount_path="/etc/foo", name="test-volume", sub_path=None, read_only=True)
+]
+
+init_environments = [k8s.V1EnvVar(name="key1", value="value1"), k8s.V1EnvVar(name="key2", value="value2")]
+
+init_container = k8s.V1Container(
+    name="init-container",
+    image="ubuntu:16.04",
+    env=init_environments,
+    volume_mounts=init_container_volume_mounts,
+    command=["bash", "-cx"],
+    args=["echo 10"],
+)
+
+affinity = k8s.V1Affinity(
+    node_affinity=k8s.V1NodeAffinity(
+        preferred_during_scheduling_ignored_during_execution=[
+            k8s.V1PreferredSchedulingTerm(
+                weight=1,
+                preference=k8s.V1NodeSelectorTerm(
+                    match_expressions=[
+                        k8s.V1NodeSelectorRequirement(key="disktype", operator="In", values=["ssd"])
+                    ]
+                ),
+            )
+        ]
+    ),
+    pod_affinity=k8s.V1PodAffinity(
+        required_during_scheduling_ignored_during_execution=[
+            k8s.V1WeightedPodAffinityTerm(
+                weight=1,
+                pod_affinity_term=k8s.V1PodAffinityTerm(
+                    label_selector=k8s.V1LabelSelector(
+                        match_expressions=[
+                            k8s.V1LabelSelectorRequirement(key="security", operator="In", values="S1")
+                        ]
+                    ),
+                    topology_key="failure-domain.beta.kubernetes.io/zone",
+                ),
+            )
+        ]
+    ),
+)
+
+tolerations = [k8s.V1Toleration(key="key", operator="Equal", value="value")]
+
+# [END howto_operator_k8s_cluster_resources]
+
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_kubernetes_operator_istio"
+
+with DAG(
+    dag_id="example_kubernetes_operator_istio",
+    schedule=None,
+    start_date=datetime(2021, 1, 1),
+    tags=["example"],
+) as dag:
+    k = KubernetesPodOperator(
+        namespace="default",
+        image="ubuntu:16.04",
+        cmds=["bash", "-cx"],
+        arguments=["echo", "10"],
+        labels={"foo": "bar"},
+        secrets=[secret_file, secret_env, secret_all_keys],
+        ports=[port],
+        volumes=[volume],
+        volume_mounts=[volume_mount],
+        env_from=configmaps,
+        name="airflow-test-pod",
+        task_id="task",
+        affinity=affinity,
+        is_delete_operator_pod=True,
+        hostnetwork=False,
+        tolerations=tolerations,
+        init_containers=[init_container],
+        priority_class_name="medium",
+        istio_enabled=True,
+    )
+
+    # [START howto_operator_k8s_private_image]
+    quay_k8s = KubernetesPodOperator(
+        namespace="default",
+        image="quay.io/apache/bash",
+        image_pull_secrets=[k8s.V1LocalObjectReference("testquay")],
+        cmds=["bash", "-cx"],
+        arguments=["echo", "10", "echo pwd"],
+        labels={"foo": "bar"},
+        name="airflow-private-image-pod",
+        is_delete_operator_pod=True,
+        in_cluster=True,
+        task_id="task-two",
+        get_logs=True,
+        istio_enabled=True,
+    )
+    # [END howto_operator_k8s_private_image]
+
+    # [START howto_operator_k8s_write_xcom]
+    write_xcom = KubernetesPodOperator(
+        namespace="default",
+        image="alpine",
+        cmds=["sh", "-c", "mkdir -p /airflow/xcom/;echo '[1,2,3,4]' > /airflow/xcom/return.json"],
+        name="write-xcom",
+        do_xcom_push=True,
+        is_delete_operator_pod=True,
+        in_cluster=True,
+        task_id="write-xcom",
+        get_logs=True,
+        istio_enabled=True,
+    )
+
+    pod_task_xcom_result = BashOperator(
+        bash_command="echo \"{{ task_instance.xcom_pull('write-xcom')[0] }}\"",
+        task_id="pod_task_xcom_result",
+    )
+
+    write_xcom >> pod_task_xcom_result
+    # [END howto_operator_k8s_write_xcom]
+
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)

From f82acc1579ed8e56c32cd7c42bab10f8db5cc024 Mon Sep 17 00:00:00 2001
From: eladkal <45845474+eladkal@users.noreply.github.com>
Date: Wed, 9 Aug 2023 15:25:00 +0300
Subject: [PATCH 06/22] Prepare docs for RC3 providers (google) (#33253)

---
 airflow/providers/google/CHANGELOG.rst           |  7 +++++++
 docs/apache-airflow-providers-google/commits.rst | 11 +++++++----
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/google/CHANGELOG.rst b/airflow/providers/google/CHANGELOG.rst
index 998e10fe5b7b3..8383d41374ff4 100644
--- a/airflow/providers/google/CHANGELOG.rst
+++ b/airflow/providers/google/CHANGELOG.rst
@@ -48,6 +48,12 @@ Bug Fixes
 * ``Fix 'DataFusionAsyncHook' catch 404 (#32855)``
 * ``Fix system test for MetastoreHivePartitionSensor (#32861)``
 * ``Fix catching 409 error (#33173)``
+* ``make 'sql' a cached property in 'BigQueryInsertJobOperator' (#33218)``
+
+Misc
+~~~~
+
+* ``refactor(providers.google): use module level __getattr__ for DATAPROC_JOB_LOG_LINK to DATAPROC_JOB_LINK and add deprecation warning (#33189)``
 
 .. Below changes are excluded from the changelog. Move them to
    appropriate section above if needed. Do not delete the lines(!):
@@ -58,6 +64,7 @@ Bug Fixes
    * ``Fix DataflowStartSqlJobOperator system test (#32823)``
    * ``Alias 'DATAPROC_JOB_LOG_LINK' to 'DATAPROC_JOB_LINK' (#33148)``
    * ``Prepare docs for Aug 2023 1st wave of Providers (#33128)``
+   * ``Prepare docs for RC2 providers (google, redis) (#33185)``
 
 10.5.0
 ......
diff --git a/docs/apache-airflow-providers-google/commits.rst b/docs/apache-airflow-providers-google/commits.rst
index bfbafa5fbd9bb..15bd2a90d8f3d 100644
--- a/docs/apache-airflow-providers-google/commits.rst
+++ b/docs/apache-airflow-providers-google/commits.rst
@@ -38,11 +38,14 @@ For high-level changelog, see :doc:`package information including changelog <ind
 10.6.0
 ......
 
-Latest change: 2023-08-07
+Latest change: 2023-08-09
 
-=================================================================================================  ===========  ===================================================================================================
+=================================================================================================  ===========  ================================================================================================================================================
 Commit                                                                                             Committed    Subject
-=================================================================================================  ===========  ===================================================================================================
+=================================================================================================  ===========  ================================================================================================================================================
+`e1af008711 <https://github.com/apache/airflow/commit/e1af008711039f7f53b2a21cad70a08b514cd8f6>`_  2023-08-09   ``make 'sql' a cached property in 'BigQueryInsertJobOperator' (#33218)``
+`e43206eb2e <https://github.com/apache/airflow/commit/e43206eb2e055a78814fcff7e8c35c6fd9c11e85>`_  2023-08-08   ``refactor(providers.google): use module level __getattr__ for DATAPROC_JOB_LOG_LINK to DATAPROC_JOB_LINK and add deprecation warning (#33189)``
+`8f27d3ff47 <https://github.com/apache/airflow/commit/8f27d3ff472e6ddd52c08a948d51901adb8a1f1a>`_  2023-08-08   ``Prepare docs for RC2 providers (google, redis) (#33185)``
 `15ede4a1a1 <https://github.com/apache/airflow/commit/15ede4a1a1af8bebb52799420bf2f328406c126e>`_  2023-08-07   ``Fix DataflowStartSqlJobOperator system test (#32823)``
 `9736143468 <https://github.com/apache/airflow/commit/9736143468cfe034e65afb3df3031ab3626f0f6d>`_  2023-08-07   ``Add a new parameter to SQL operators to specify conn id field (#30784)``
 `0e076dceaf <https://github.com/apache/airflow/commit/0e076dceaf169174bbd4d8ee46911b60505eb098>`_  2023-08-07   ``Fix catching 409 error (#33173)``
@@ -57,7 +60,7 @@ Commit
 `f9cddf332d <https://github.com/apache/airflow/commit/f9cddf332db712ce2121d355dd94c78b392a7bb9>`_  2023-07-31   ``openlineage,gcs: use proper name for openlineage methods (#32956)``
 `ce5eebd004 <https://github.com/apache/airflow/commit/ce5eebd00403beabc23b4f0b4bedba5b5c397c42>`_  2023-07-31   ``Fix system test for MetastoreHivePartitionSensor (#32861)``
 `c422920c6e <https://github.com/apache/airflow/commit/c422920c6ef9ac2cceb0369cd0ec0989887688d9>`_  2023-07-31   ``Add system test and docs for CloudDataTransferServiceGCSToGCSOperator (#32960)``
-=================================================================================================  ===========  ===================================================================================================
+=================================================================================================  ===========  ================================================================================================================================================
 
 10.5.0
 ......

From 70a050b1c3a46e3eabae7f6052c95718800dc604 Mon Sep 17 00:00:00 2001
From: Usiel Riedl <usiel.riedl@gmail.com>
Date: Wed, 9 Aug 2023 20:34:37 +0800
Subject: [PATCH 07/22] Enable pools to consider deferred tasks (#32709)

* Makes pools respect deferrable tasks (with extra setting)

See https://github.com/apache/airflow/discussions/21243

This commit makes pools consider deferred tasks if the `include_deferred` flag is set. By default a pool will not consider deferred tasks as occupied slots, but still show the number of deferred tasks in its stats.

---------

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/api/client/api_client.py              |   3 +-
 airflow/api/client/json_client.py             |   9 +-
 airflow/api/client/local_client.py            |  10 +-
 airflow/api/common/experimental/pool.py       |   2 +-
 .../api_connexion/endpoints/pool_endpoint.py  |   4 +-
 airflow/api_connexion/openapi/v1.yaml         |  15 ++-
 airflow/api_connexion/schemas/pool_schema.py  |   8 ++
 airflow/cli/cli_config.py                     |  16 ++-
 airflow/cli/commands/pool_command.py          |  18 ++-
 airflow/jobs/scheduler_job_runner.py          |   2 +
 ...7_0_add_include_deferred_column_to_pool.py |  50 ++++++++
 airflow/models/pool.py                        |  58 +++++++--
 .../ti_deps/deps/pool_slots_available_dep.py  |  11 +-
 airflow/utils/db.py                           |   1 +
 airflow/www/static/js/types/api-generated.ts  |  14 +-
 airflow/www/views.py                          |  29 ++++-
 .../logging-monitoring/metrics.rst            |   1 +
 .../administration-and-deployment/pools.rst   |   2 +-
 .../authoring-and-scheduling/deferring.rst    |   2 +-
 docs/apache-airflow/img/airflow_erd.sha256    |   2 +-
 docs/apache-airflow/img/airflow_erd.svg       |  86 +++++++------
 docs/apache-airflow/migrations-ref.rst        |   4 +-
 .../endpoints/test_pool_endpoint.py           | 120 ++++++++++++++----
 .../schemas/test_pool_schemas.py              |  14 +-
 tests/api_connexion/test_auth.py              |   4 +
 .../client/test_local_client.py               |  20 ++-
 .../common/experimental/test_pool.py          |   7 +-
 tests/cli/commands/test_pool_command.py       |  67 +++++++---
 tests/cli/commands/test_task_command.py       |   2 +-
 tests/jobs/test_backfill_job.py               |   3 +-
 tests/jobs/test_scheduler_job.py              |  31 +++--
 tests/models/test_pool.py                     |  70 +++++++++-
 tests/models/test_taskinstance.py             |   2 +-
 tests/operators/test_subdag_operator.py       |   8 +-
 .../deps/test_pool_slots_available_dep.py     |  13 +-
 tests/www/api/experimental/test_endpoints.py  |   2 +
 tests/www/views/test_views_pool.py            |   5 +
 37 files changed, 549 insertions(+), 166 deletions(-)
 create mode 100644 airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py

diff --git a/airflow/api/client/api_client.py b/airflow/api/client/api_client.py
index 1fdbc7027b3f7..d4a6d2357a8bf 100644
--- a/airflow/api/client/api_client.py
+++ b/airflow/api/client/api_client.py
@@ -60,12 +60,13 @@ def get_pools(self):
         """Get all pools."""
         raise NotImplementedError()
 
-    def create_pool(self, name, slots, description):
+    def create_pool(self, name, slots, description, include_deferred):
         """Create a pool.
 
         :param name: pool name
         :param slots: pool slots amount
         :param description: pool description
+        :param include_deferred: include deferred tasks in pool calculations
         """
         raise NotImplementedError()
 
diff --git a/airflow/api/client/json_client.py b/airflow/api/client/json_client.py
index 2eb91b3a422f7..ac73813136544 100644
--- a/airflow/api/client/json_client.py
+++ b/airflow/api/client/json_client.py
@@ -108,23 +108,26 @@ def get_pools(self):
         pools = self._request(url)
         return [(p["pool"], p["slots"], p["description"]) for p in pools]
 
-    def create_pool(self, name: str, slots: int, description: str):
+    def create_pool(self, name: str, slots: int, description: str, include_deferred: bool):
         """Create a new pool.
 
         :param name: The name of the pool to create.
         :param slots: The number of slots in the pool.
         :param description: A description of the pool.
+        :param include_deferred: include deferred tasks in pool calculations
+
         :return: A tuple containing the name of the pool, the number of slots in the pool,
-            and a description of the pool.
+            a description of the pool and the include_deferred flag.
         """
         endpoint = "/api/experimental/pools"
         data = {
             "name": name,
             "slots": slots,
             "description": description,
+            "include_deferred": include_deferred,
         }
         response = self._request(urljoin(self._api_base_url, endpoint), method="POST", json=data)
-        return response["pool"], response["slots"], response["description"]
+        return response["pool"], response["slots"], response["description"], response["include_deferred"]
 
     def delete_pool(self, name: str):
         """Delete a pool.
diff --git a/airflow/api/client/local_client.py b/airflow/api/client/local_client.py
index b90039608f4ca..d8b95650998e1 100644
--- a/airflow/api/client/local_client.py
+++ b/airflow/api/client/local_client.py
@@ -63,12 +63,12 @@ def get_pool(self, name):
         pool = Pool.get_pool(pool_name=name)
         if not pool:
             raise PoolNotFound(f"Pool {name} not found")
-        return pool.pool, pool.slots, pool.description
+        return pool.pool, pool.slots, pool.description, pool.include_deferred
 
     def get_pools(self):
-        return [(p.pool, p.slots, p.description) for p in Pool.get_pools()]
+        return [(p.pool, p.slots, p.description, p.include_deferred) for p in Pool.get_pools()]
 
-    def create_pool(self, name, slots, description):
+    def create_pool(self, name, slots, description, include_deferred):
         if not (name and name.strip()):
             raise AirflowBadRequest("Pool name shouldn't be empty")
         pool_name_length = Pool.pool.property.columns[0].type.length
@@ -78,7 +78,9 @@ def create_pool(self, name, slots, description):
             slots = int(slots)
         except ValueError:
             raise AirflowBadRequest(f"Bad value for `slots`: {slots}")
-        pool = Pool.create_or_update_pool(name=name, slots=slots, description=description)
+        pool = Pool.create_or_update_pool(
+            name=name, slots=slots, description=description, include_deferred=include_deferred
+        )
         return pool.pool, pool.slots, pool.description
 
     def delete_pool(self, name):
diff --git a/airflow/api/common/experimental/pool.py b/airflow/api/common/experimental/pool.py
index cff21d4d56392..34e35cd43556f 100644
--- a/airflow/api/common/experimental/pool.py
+++ b/airflow/api/common/experimental/pool.py
@@ -68,7 +68,7 @@ def create_pool(name, slots, description, session: Session = NEW_SESSION):
     session.expire_on_commit = False
     pool = session.scalar(select(Pool).filter_by(pool=name).limit(1))
     if pool is None:
-        pool = Pool(pool=name, slots=slots, description=description)
+        pool = Pool(pool=name, slots=slots, description=description, include_deferred=False)
         session.add(pool)
     else:
         pool.slots = slots
diff --git a/airflow/api_connexion/endpoints/pool_endpoint.py b/airflow/api_connexion/endpoints/pool_endpoint.py
index 6503705d85b10..fb59ab5ea24d2 100644
--- a/airflow/api_connexion/endpoints/pool_endpoint.py
+++ b/airflow/api_connexion/endpoints/pool_endpoint.py
@@ -88,10 +88,10 @@ def patch_pool(
 ) -> APIResponse:
     """Update a pool."""
     request_dict = get_json_request_dict()
-    # Only slots can be modified in 'default_pool'
+    # Only slots and include_deferred can be modified in 'default_pool'
     try:
         if pool_name == Pool.DEFAULT_POOL_NAME and request_dict["name"] != Pool.DEFAULT_POOL_NAME:
-            if update_mask and len(update_mask) == 1 and update_mask[0].strip() == "slots":
+            if update_mask and all(mask.strip() in {"slots", "include_deferred"} for mask in update_mask):
                 pass
             else:
                 raise BadRequest(detail="Default Pool's name can't be modified")
diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index da7dca04f64c5..72f65438b0200 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -3158,7 +3158,7 @@ components:
         occupied_slots:
           type: integer
           readOnly: true
-          description: The number of slots used by running/queued tasks at the moment.
+          description: The number of slots used by running/queued tasks at the moment. May include deferred tasks if 'include_deferred' is set to true.
         running_slots:
           type: integer
           readOnly: true
@@ -3175,6 +3175,13 @@ components:
           type: integer
           readOnly: true
           description: The number of slots used by scheduled tasks at the moment.
+        deferred_slots:
+          type: integer
+          readOnly: true
+          description: |
+            The number of slots used by deferred tasks at the moment. Relevant if 'include_deferred' is set to true.
+
+            *New in version 2.7.0*
         description:
           type: string
           description: |
@@ -3182,6 +3189,12 @@ components:
 
             *New in version 2.3.0*
           nullable: true
+        include_deferred:
+          type: boolean
+          description: |
+            If set to true, deferred tasks are considered when calculating open pool slots.
+
+            *New in version 2.7.0*
 
     PoolCollection:
       type: object
diff --git a/airflow/api_connexion/schemas/pool_schema.py b/airflow/api_connexion/schemas/pool_schema.py
index 4e25287d1d357..e18548cee844d 100644
--- a/airflow/api_connexion/schemas/pool_schema.py
+++ b/airflow/api_connexion/schemas/pool_schema.py
@@ -39,7 +39,10 @@ class Meta:
     queued_slots = fields.Method("get_queued_slots", dump_only=True)
     scheduled_slots = fields.Method("get_scheduled_slots", dump_only=True)
     open_slots = fields.Method("get_open_slots", dump_only=True)
+    deferred_slots = fields.Method("get_deferred_slots", dump_only=True)
     description = auto_field()
+    # we skip auto_field() here to be compatible with the manual validation in the pool_endpoint module
+    include_deferred = fields.Boolean(load_default=False)
 
     @staticmethod
     def get_occupied_slots(obj: Pool) -> int:
@@ -61,6 +64,11 @@ def get_scheduled_slots(obj: Pool) -> int:
         """Returns the scheduled slots of the pool."""
         return obj.scheduled_slots()
 
+    @staticmethod
+    def get_deferred_slots(obj: Pool) -> int:
+        """Returns the deferred slots of the pool."""
+        return obj.deferred_slots()
+
     @staticmethod
     def get_open_slots(obj: Pool) -> float:
         """Returns the open slots of the pool."""
diff --git a/airflow/cli/cli_config.py b/airflow/cli/cli_config.py
index 7569e2ff581b6..f2248efe5e635 100644
--- a/airflow/cli/cli_config.py
+++ b/airflow/cli/cli_config.py
@@ -513,6 +513,9 @@ def string_lower_type(val):
 ARG_POOL_NAME = Arg(("pool",), metavar="NAME", help="Pool name")
 ARG_POOL_SLOTS = Arg(("slots",), type=int, help="Pool slots")
 ARG_POOL_DESCRIPTION = Arg(("description",), help="Pool description")
+ARG_POOL_INCLUDE_DEFERRED = Arg(
+    ("--include-deferred",), help="Include deferred tasks in calculations for Pool", action="store_true"
+)
 ARG_POOL_IMPORT = Arg(
     ("file",),
     metavar="FILEPATH",
@@ -521,8 +524,8 @@ def string_lower_type(val):
         textwrap.dedent(
             """
             {
-                "pool_1": {"slots": 5, "description": ""},
-                "pool_2": {"slots": 10, "description": "test"}
+                "pool_1": {"slots": 5, "description": "", "include_deferred": true},
+                "pool_2": {"slots": 10, "description": "test", "include_deferred": false}
             }"""
         ),
         " " * 4,
@@ -1456,7 +1459,14 @@ class GroupCommand(NamedTuple):
         name="set",
         help="Configure pool",
         func=lazy_load_command("airflow.cli.commands.pool_command.pool_set"),
-        args=(ARG_POOL_NAME, ARG_POOL_SLOTS, ARG_POOL_DESCRIPTION, ARG_OUTPUT, ARG_VERBOSE),
+        args=(
+            ARG_POOL_NAME,
+            ARG_POOL_SLOTS,
+            ARG_POOL_DESCRIPTION,
+            ARG_POOL_INCLUDE_DEFERRED,
+            ARG_OUTPUT,
+            ARG_VERBOSE,
+        ),
     ),
     ActionCommand(
         name="delete",
diff --git a/airflow/cli/commands/pool_command.py b/airflow/cli/commands/pool_command.py
index 8d9e206f1bb2d..7e7c63640cb6e 100644
--- a/airflow/cli/commands/pool_command.py
+++ b/airflow/cli/commands/pool_command.py
@@ -38,6 +38,7 @@ def _show_pools(pools, output):
             "pool": x[0],
             "slots": x[1],
             "description": x[2],
+            "include_deferred": x[3],
         },
     )
 
@@ -69,7 +70,9 @@ def pool_get(args):
 def pool_set(args):
     """Creates new pool with a given name and slots."""
     api_client = get_current_api_client()
-    api_client.create_pool(name=args.pool, slots=args.slots, description=args.description)
+    api_client.create_pool(
+        name=args.pool, slots=args.slots, description=args.description, include_deferred=args.include_deferred
+    )
     print(f"Pool {args.pool} created")
 
 
@@ -119,8 +122,15 @@ def pool_import_helper(filepath):
     pools = []
     failed = []
     for k, v in pools_json.items():
-        if isinstance(v, dict) and len(v) == 2:
-            pools.append(api_client.create_pool(name=k, slots=v["slots"], description=v["description"]))
+        if isinstance(v, dict) and "slots" in v and "description" in v:
+            pools.append(
+                api_client.create_pool(
+                    name=k,
+                    slots=v["slots"],
+                    description=v["description"],
+                    include_deferred=v.get("include_deferred", False),
+                )
+            )
         else:
             failed.append(k)
     return pools, failed
@@ -132,7 +142,7 @@ def pool_export_helper(filepath):
     pool_dict = {}
     pools = api_client.get_pools()
     for pool in pools:
-        pool_dict[pool[0]] = {"slots": pool[1], "description": pool[2]}
+        pool_dict[pool[0]] = {"slots": pool[1], "description": pool[2], "include_deferred": pool[3]}
     with open(filepath, "w") as poolfile:
         poolfile.write(json.dumps(pool_dict, sort_keys=True, indent=4))
     return pools
diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py
index e47fc7bd86558..0ed7aae5ceb21 100644
--- a/airflow/jobs/scheduler_job_runner.py
+++ b/airflow/jobs/scheduler_job_runner.py
@@ -1571,10 +1571,12 @@ def _emit_pool_metrics(self, session: Session = NEW_SESSION) -> None:
             Stats.gauge(f"pool.open_slots.{pool_name}", slot_stats["open"])
             Stats.gauge(f"pool.queued_slots.{pool_name}", slot_stats["queued"])
             Stats.gauge(f"pool.running_slots.{pool_name}", slot_stats["running"])
+            Stats.gauge(f"pool.deferred_slots.{pool_name}", slot_stats["deferred"])
             # Same metrics with tagging
             Stats.gauge("pool.open_slots", slot_stats["open"], tags={"pool_name": pool_name})
             Stats.gauge("pool.queued_slots", slot_stats["queued"], tags={"pool_name": pool_name})
             Stats.gauge("pool.running_slots", slot_stats["running"], tags={"pool_name": pool_name})
+            Stats.gauge("pool.deferred_slots", slot_stats["deferred"], tags={"pool_name": pool_name})
 
     @provide_session
     def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int:
diff --git a/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py b/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
new file mode 100644
index 0000000000000..ea103204801b0
--- /dev/null
+++ b/airflow/migrations/versions/0128_2_7_0_add_include_deferred_column_to_pool.py
@@ -0,0 +1,50 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""add include_deferred column to pool
+
+Revision ID: 405de8318b3a
+Revises: 788397e78828
+Create Date: 2023-07-20 04:22:21.007342
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "405de8318b3a"
+down_revision = "788397e78828"
+branch_labels = None
+depends_on = None
+airflow_version = "2.7.0"
+
+
+def upgrade():
+    """Apply add include_deferred column to pool"""
+    with op.batch_alter_table("slot_pool") as batch_op:
+        batch_op.add_column(
+            sa.Column("include_deferred", sa.Boolean, nullable=False, server_default=sa.false())
+        )
+
+
+def downgrade():
+    """Unapply add include_deferred column to pool"""
+    with op.batch_alter_table("slot_pool") as batch_op:
+        batch_op.drop_column("include_deferred")
diff --git a/airflow/models/pool.py b/airflow/models/pool.py
index 3b5708b121298..13a2e15e66bec 100644
--- a/airflow/models/pool.py
+++ b/airflow/models/pool.py
@@ -19,7 +19,7 @@
 
 from typing import Any
 
-from sqlalchemy import Column, Integer, String, Text, func, select
+from sqlalchemy import Boolean, Column, Integer, String, Text, func, select
 from sqlalchemy.orm.session import Session
 
 from airflow.exceptions import AirflowException, PoolNotFound
@@ -37,6 +37,7 @@ class PoolStats(TypedDict):
 
     total: int
     running: int
+    deferred: int
     queued: int
     open: int
 
@@ -51,6 +52,7 @@ class Pool(Base):
     # -1 for infinite
     slots = Column(Integer, default=0)
     description = Column(Text)
+    include_deferred = Column(Boolean, nullable=False)
 
     DEFAULT_POOL_NAME = "default_pool"
 
@@ -108,6 +110,7 @@ def create_or_update_pool(
         name: str,
         slots: int,
         description: str,
+        include_deferred: bool,
         session: Session = NEW_SESSION,
     ) -> Pool:
         """Create a pool with given parameters or update it if it already exists."""
@@ -116,11 +119,12 @@ def create_or_update_pool(
 
         pool = session.scalar(select(Pool).filter_by(pool=name))
         if pool is None:
-            pool = Pool(pool=name, slots=slots, description=description)
+            pool = Pool(pool=name, slots=slots, description=description, include_deferred=include_deferred)
             session.add(pool)
         else:
             pool.slots = slots
             pool.description = description
+            pool.include_deferred = include_deferred
 
         session.commit()
         return pool
@@ -161,21 +165,26 @@ def slots_stats(
         from airflow.models.taskinstance import TaskInstance  # Avoid circular import
 
         pools: dict[str, PoolStats] = {}
+        pool_includes_deferred: dict[str, bool] = {}
 
-        query = select(Pool.pool, Pool.slots)
+        query = select(Pool.pool, Pool.slots, Pool.include_deferred)
 
         if lock_rows:
             query = with_row_locks(query, session=session, **nowait(session))
 
         pool_rows = session.execute(query)
-        for (pool_name, total_slots) in pool_rows:
+        for (pool_name, total_slots, include_deferred) in pool_rows:
             if total_slots == -1:
                 total_slots = float("inf")  # type: ignore
-            pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0, open=0)
+            pools[pool_name] = PoolStats(total=total_slots, running=0, queued=0, open=0, deferred=0)
+            pool_includes_deferred[pool_name] = include_deferred
 
+        allowed_execution_states = EXECUTION_STATES | {
+            TaskInstanceState.DEFERRED,
+        }
         state_count_by_pool = session.execute(
             select(TaskInstance.pool, TaskInstance.state, func.sum(TaskInstance.pool_slots))
-            .filter(TaskInstance.state.in_(EXECUTION_STATES))
+            .filter(TaskInstance.state.in_(allowed_execution_states))
             .group_by(TaskInstance.pool, TaskInstance.state)
         )
 
@@ -192,12 +201,16 @@ def slots_stats(
                 stats_dict["running"] = count
             elif state == TaskInstanceState.QUEUED:
                 stats_dict["queued"] = count
+            elif state == TaskInstanceState.DEFERRED:
+                stats_dict["deferred"] = count
             else:
-                raise AirflowException(f"Unexpected state. Expected values: {EXECUTION_STATES}.")
+                raise AirflowException(f"Unexpected state. Expected values: {allowed_execution_states}.")
 
         # calculate open metric
         for pool_name, stats_dict in pools.items():
             stats_dict["open"] = stats_dict["total"] - stats_dict["running"] - stats_dict["queued"]
+            if pool_includes_deferred[pool_name]:
+                stats_dict["open"] -= stats_dict["deferred"]
 
         return pools
 
@@ -212,6 +225,7 @@ def to_json(self) -> dict[str, Any]:
             "pool": self.pool,
             "slots": self.slots,
             "description": self.description,
+            "include_deferred": self.include_deferred,
         }
 
     @provide_session
@@ -224,15 +238,24 @@ def occupied_slots(self, session: Session = NEW_SESSION) -> int:
         """
         from airflow.models.taskinstance import TaskInstance  # Avoid circular import
 
+        occupied_states = self.get_occupied_states()
+
         return int(
             session.scalar(
                 select(func.sum(TaskInstance.pool_slots))
                 .filter(TaskInstance.pool == self.pool)
-                .filter(TaskInstance.state.in_(EXECUTION_STATES))
+                .filter(TaskInstance.state.in_(occupied_states))
             )
             or 0
         )
 
+    def get_occupied_states(self):
+        if self.include_deferred:
+            return EXECUTION_STATES | {
+                TaskInstanceState.DEFERRED,
+            }
+        return EXECUTION_STATES
+
     @provide_session
     def running_slots(self, session: Session = NEW_SESSION) -> int:
         """
@@ -290,6 +313,25 @@ def scheduled_slots(self, session: Session = NEW_SESSION) -> int:
             or 0
         )
 
+    @provide_session
+    def deferred_slots(self, session: Session = NEW_SESSION) -> int:
+        """
+        Get the number of slots deferred at the moment.
+
+        :param session: SQLAlchemy ORM Session
+        :return: the number of deferred slots
+        """
+        from airflow.models.taskinstance import TaskInstance  # Avoid circular import
+
+        return int(
+            session.scalar(
+                select(func.sum(TaskInstance.pool_slots)).where(
+                    TaskInstance.pool == self.pool, TaskInstance.state == TaskInstanceState.DEFERRED
+                )
+            )
+            or 0
+        )
+
     @provide_session
     def open_slots(self, session: Session = NEW_SESSION) -> float:
         """
diff --git a/airflow/ti_deps/deps/pool_slots_available_dep.py b/airflow/ti_deps/deps/pool_slots_available_dep.py
index 38bfe31936679..e095cfb68e084 100644
--- a/airflow/ti_deps/deps/pool_slots_available_dep.py
+++ b/airflow/ti_deps/deps/pool_slots_available_dep.py
@@ -18,7 +18,6 @@
 """This module defines dep for pool slots availability."""
 from __future__ import annotations
 
-from airflow.ti_deps.dependencies_states import EXECUTION_STATES
 from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
 from airflow.utils.session import provide_session
 
@@ -49,12 +48,12 @@ def _get_dep_statuses(self, ti, session, dep_context=None):
                 reason=f"Tasks using non-existent pool '{pool_name}' will not be scheduled"
             )
             return
-        else:
-            # Controlled by UNIQUE key in slot_pool table,
-            # only one result can be returned.
-            open_slots = pools[0].open_slots(session=session)
+        # Controlled by UNIQUE key in slot_pool table,
+        # only one result can be returned.
+        open_slots = pools[0].open_slots(session=session)
 
-        if ti.state in EXECUTION_STATES:
+        occupied_states = pools[0].get_occupied_states()
+        if ti.state in occupied_states:
             open_slots += ti.pool_slots
 
         if open_slots <= (ti.pool_slots - 1):
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 1a4a5db6ae482..92ad989800635 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -124,6 +124,7 @@ def add_default_pool_if_not_exists(session: Session = NEW_SESSION):
             pool=Pool.DEFAULT_POOL_NAME,
             slots=conf.getint(section="core", key="default_pool_task_slot_count"),
             description="Default pool",
+            include_deferred=False,
         )
         session.add(default_pool)
         session.commit()
diff --git a/airflow/www/static/js/types/api-generated.ts b/airflow/www/static/js/types/api-generated.ts
index e2d38d072597e..8dead88fc7ed8 100644
--- a/airflow/www/static/js/types/api-generated.ts
+++ b/airflow/www/static/js/types/api-generated.ts
@@ -1239,7 +1239,7 @@ export interface components {
       name?: string;
       /** @description The maximum number of slots that can be assigned to tasks. One job may occupy one or more slots. */
       slots?: number;
-      /** @description The number of slots used by running/queued tasks at the moment. */
+      /** @description The number of slots used by running/queued tasks at the moment. May include deferred tasks if 'include_deferred' is set to true. */
       occupied_slots?: number;
       /** @description The number of slots used by running tasks at the moment. */
       running_slots?: number;
@@ -1249,12 +1249,24 @@ export interface components {
       open_slots?: number;
       /** @description The number of slots used by scheduled tasks at the moment. */
       scheduled_slots?: number;
+      /**
+       * @description The number of slots used by deferred tasks at the moment. Relevant if 'include_deferred' is set to true.
+       *
+       * *New in version 2.7.0*
+       */
+      deferred_slots?: number;
       /**
        * @description The description of the pool.
        *
        * *New in version 2.3.0*
        */
       description?: string | null;
+      /**
+       * @description If set to true, deferred tasks are considered when calculating open pool slots.
+       *
+       * *New in version 2.7.0*
+       */
+      include_deferred?: boolean;
     };
     /**
      * @description Collection of pools.
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 3faedd3d8b15f..5f12a94969148 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -70,7 +70,7 @@
 from sqlalchemy import Date, and_, case, desc, func, inspect, or_, select, union_all
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session, joinedload
-from wtforms import validators
+from wtforms import BooleanField, validators
 
 import airflow
 from airflow import models, plugins_manager, settings
@@ -4874,9 +4874,17 @@ class PoolModelView(AirflowModelView):
         permissions.ACTION_CAN_ACCESS_MENU,
     ]
 
-    list_columns = ["pool", "slots", "running_slots", "queued_slots", "scheduled_slots"]
-    add_columns = ["pool", "slots", "description"]
-    edit_columns = ["pool", "slots", "description"]
+    list_columns = ["pool", "slots", "running_slots", "queued_slots", "scheduled_slots", "deferred_slots"]
+    add_columns = ["pool", "slots", "description", "include_deferred"]
+    edit_columns = ["pool", "slots", "description", "include_deferred"]
+
+    # include_deferred is non-nullable, but as a checkbox in the resulting form we want to allow it unchecked
+    include_deferred_field = BooleanField(
+        validators=[validators.Optional()],
+        description="Check to include deferred tasks when calculating open pool slots.",
+    )
+    edit_form_extra_fields = {"include_deferred": include_deferred_field}
+    add_form_extra_fields = {"include_deferred": include_deferred_field}
 
     base_order = ("pool", "asc")
 
@@ -4943,11 +4951,24 @@ def fscheduled_slots(self):
         else:
             return Markup('<span class="label label-danger">Invalid</span>')
 
+    def fdeferred_slots(self):
+        """Deferred slots rendering."""
+        pool_id = self.get("pool")
+        deferred_slots = self.get("deferred_slots")
+        if pool_id is not None and deferred_slots is not None:
+            url = url_for("TaskInstanceModelView.list", _flt_3_pool=pool_id, _flt_3_state="deferred")
+            return Markup("<a href='{url}'>{deferred_slots}</a>").format(
+                url=url, deferred_slots=deferred_slots
+            )
+        else:
+            return Markup('<span class="label label-danger">Invalid</span>')
+
     formatters_columns = {
         "pool": pool_link,
         "running_slots": frunning_slots,
         "queued_slots": fqueued_slots,
         "scheduled_slots": fscheduled_slots,
+        "deferred_slots": fdeferred_slots,
     }
 
     validators_columns = {"pool": [validators.DataRequired()], "slots": [validators.NumberRange(min=-1)]}
diff --git a/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst b/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst
index e4bb1da68a26e..b18d05a87768f 100644
--- a/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst
+++ b/docs/apache-airflow/administration-and-deployment/logging-monitoring/metrics.rst
@@ -214,6 +214,7 @@ Name                                                Description
 ``pool.open_slots.<pool_name>``                     Number of open slots in the pool
 ``pool.queued_slots.<pool_name>``                   Number of queued slots in the pool
 ``pool.running_slots.<pool_name>``                  Number of running slots in the pool
+``pool.deferred_slots.<pool_name>``                 Number of deferred slots in the pool
 ``pool.starving_tasks.<pool_name>``                 Number of starving tasks in the pool
 ``triggers.running.<hostname>``                     Number of triggers currently running for a triggerer (described by hostname)
 =================================================== ========================================================================
diff --git a/docs/apache-airflow/administration-and-deployment/pools.rst b/docs/apache-airflow/administration-and-deployment/pools.rst
index 5395ccfe3866b..5a906d3f6fed2 100644
--- a/docs/apache-airflow/administration-and-deployment/pools.rst
+++ b/docs/apache-airflow/administration-and-deployment/pools.rst
@@ -22,7 +22,7 @@ Pools
 
 Some systems can get overwhelmed when too many processes hit them at the same time. Airflow pools can be used to
 **limit the execution parallelism** on arbitrary sets of tasks. The list of pools is managed in the UI
-(``Menu -> Admin -> Pools``) by giving the pools a name and assigning it a number of worker slots.
+(``Menu -> Admin -> Pools``) by giving the pools a name and assigning it a number of worker slots. There you can also decide whether the pool should include :doc:`deferred tasks </authoring-and-scheduling/deferring>` in its calculation of occupied slots.
 
 Tasks can then be associated with one of the existing pools by using the ``pool`` parameter when creating tasks:
 
diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index 6b5f124e8f621..c44ac2c1ddc38 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -20,7 +20,7 @@ Deferrable Operators & Triggers
 
 Standard :doc:`Operators </core-concepts/operators>` and :doc:`Sensors <../core-concepts/sensors>` take up a full *worker slot* for the entire time they are running, even if they are idle; for example, if you only have 100 worker slots available to run Tasks, and you have 100 DAGs waiting on a Sensor that's currently running but idle, then you *cannot run anything else* - even though your entire Airflow cluster is essentially idle. ``reschedule`` mode for Sensors solves some of this, allowing Sensors to only run at fixed intervals, but it is inflexible and only allows using time as the reason to resume, not anything else.
 
-This is where *Deferrable Operators* come in. A deferrable operator is one that is written with the ability to suspend itself and free up the worker when it knows it has to wait, and hand off the job of resuming it to something called a *Trigger*. As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a lot less resources wasted on idle Operators or Sensors.
+This is where *Deferrable Operators* come in. A deferrable operator is one that is written with the ability to suspend itself and free up the worker when it knows it has to wait, and hand off the job of resuming it to something called a *Trigger*. As a result, while it is suspended (deferred), it is not taking up a worker slot and your cluster will have a lot less resources wasted on idle Operators or Sensors. Note that by default deferred tasks will not use up pool slots, if you would like them to, you can change this by editing the pool in question.
 
 *Triggers* are small, asynchronous pieces of Python code designed to be run all together in a single Python process; because they are asynchronous, they are able to all co-exist efficiently. As an overview of how this process works:
 
diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256
index f10af90f5bb5b..5f9515fc19e5c 100644
--- a/docs/apache-airflow/img/airflow_erd.sha256
+++ b/docs/apache-airflow/img/airflow_erd.sha256
@@ -1 +1 @@
-7ff18b1eafa528dbdfb62d75151a526de280f73e1edf7fea18da7c64644f0da9
\ No newline at end of file
+3420f98a9dfa230d76be4da57bc4a3641cb929b33b9eafe7ffa41fb4fcc4d283
\ No newline at end of file
diff --git a/docs/apache-airflow/img/airflow_erd.svg b/docs/apache-airflow/img/airflow_erd.svg
index 18e5f83b117c1..ba548c60fdddf 100644
--- a/docs/apache-airflow/img/airflow_erd.svg
+++ b/docs/apache-airflow/img/airflow_erd.svg
@@ -4,11 +4,11 @@
 <!-- Generated by graphviz version 2.43.0 (0)
  -->
 <!-- Title: %3 Pages: 1 -->
-<svg width="1559pt" height="5959pt"
- viewBox="0.00 0.00 1559.00 5958.50" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
-<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 5954.5)">
+<svg width="1559pt" height="5983pt"
+ viewBox="0.00 0.00 1559.00 5982.50" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
+<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 5978.5)">
 <title>%3</title>
-<polygon fill="white" stroke="transparent" points="-4,4 -4,-5954.5 1555,-5954.5 1555,4 -4,4"/>
+<polygon fill="white" stroke="transparent" points="-4,4 -4,-5978.5 1555,-5978.5 1555,4 -4,4"/>
 <!-- ab_permission -->
 <g id="node1" class="node">
 <title>ab_permission</title>
@@ -1228,28 +1228,28 @@
 <g id="edge48" class="edge">
 <title>task_instance&#45;&#45;xcom</title>
 <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-482.51C1190.54,-462.01 1215.43,-441.43 1239,-422 1245.31,-416.8 1251.76,-411.48 1258.29,-406.14"/>
-<text text-anchor="start" x="1248.29" y="-394.94" font-family="Times,serif" font-size="14.00">1</text>
+<text text-anchor="start" x="1227.29" y="-394.94" font-family="Times,serif" font-size="14.00">0..N</text>
 <text text-anchor="start" x="1166.17" y="-471.31" font-family="Times,serif" font-size="14.00">1</text>
 </g>
 <!-- task_instance&#45;&#45;xcom -->
 <g id="edge49" class="edge">
 <title>task_instance&#45;&#45;xcom</title>
 <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-500.3C1190.54,-480.01 1215.43,-459.43 1239,-440 1245.31,-434.8 1251.76,-429.48 1258.29,-424.12"/>
-<text text-anchor="start" x="1227.29" y="-412.92" font-family="Times,serif" font-size="14.00">0..N</text>
+<text text-anchor="start" x="1248.29" y="-412.92" font-family="Times,serif" font-size="14.00">1</text>
 <text text-anchor="start" x="1166.17" y="-489.1" font-family="Times,serif" font-size="14.00">1</text>
 </g>
 <!-- task_instance&#45;&#45;xcom -->
 <g id="edge50" class="edge">
 <title>task_instance&#45;&#45;xcom</title>
 <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-518.08C1190.54,-498.01 1215.43,-477.43 1239,-458 1246.11,-452.14 1253.42,-446.12 1260.79,-440.05"/>
-<text text-anchor="start" x="1250.79" y="-443.85" font-family="Times,serif" font-size="14.00">1</text>
+<text text-anchor="start" x="1229.79" y="-443.85" font-family="Times,serif" font-size="14.00">0..N</text>
 <text text-anchor="start" x="1166.17" y="-506.88" font-family="Times,serif" font-size="14.00">1</text>
 </g>
 <!-- task_instance&#45;&#45;xcom -->
 <g id="edge51" class="edge">
 <title>task_instance&#45;&#45;xcom</title>
 <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1166.17,-535.87C1190.54,-516.01 1215.43,-495.43 1239,-476 1252.96,-464.49 1267.64,-452.42 1282.18,-440.26"/>
-<text text-anchor="start" x="1282.18" y="-444.06" font-family="Times,serif" font-size="14.00">0..N</text>
+<text text-anchor="start" x="1272.18" y="-444.06" font-family="Times,serif" font-size="14.00">1</text>
 <text text-anchor="start" x="1166.17" y="-524.67" font-family="Times,serif" font-size="14.00">1</text>
 </g>
 <!-- log_template -->
@@ -1582,43 +1582,47 @@
 <!-- slot_pool -->
 <g id="node41" class="node">
 <title>slot_pool</title>
-<polygon fill="none" stroke="black" points="103.5,-5739 103.5,-5767 279.5,-5767 279.5,-5739 103.5,-5739"/>
-<text text-anchor="start" x="150.5" y="-5750.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">slot_pool</text>
-<polygon fill="none" stroke="black" points="103.5,-5714 103.5,-5739 279.5,-5739 279.5,-5714 103.5,-5714"/>
-<text text-anchor="start" x="108.5" y="-5723.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
-<text text-anchor="start" x="121.5" y="-5723.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
-<text text-anchor="start" x="198.5" y="-5723.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
-<polygon fill="none" stroke="black" points="103.5,-5689 103.5,-5714 279.5,-5714 279.5,-5689 103.5,-5689"/>
-<text text-anchor="start" x="108.5" y="-5698.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
-<text text-anchor="start" x="186.5" y="-5698.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
-<polygon fill="none" stroke="black" points="103.5,-5664 103.5,-5689 279.5,-5689 279.5,-5664 103.5,-5664"/>
-<text text-anchor="start" x="108.5" y="-5673.8" font-family="Helvetica,sans-Serif" font-size="14.00">pool</text>
-<text text-anchor="start" x="138.5" y="-5673.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(256)]</text>
-<polygon fill="none" stroke="black" points="103.5,-5639 103.5,-5664 279.5,-5664 279.5,-5639 103.5,-5639"/>
-<text text-anchor="start" x="108.5" y="-5648.8" font-family="Helvetica,sans-Serif" font-size="14.00">slots</text>
-<text text-anchor="start" x="141.5" y="-5648.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
+<polygon fill="none" stroke="black" points="48.5,-5763 48.5,-5791 334.5,-5791 334.5,-5763 48.5,-5763"/>
+<text text-anchor="start" x="150.5" y="-5774.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">slot_pool</text>
+<polygon fill="none" stroke="black" points="48.5,-5738 48.5,-5763 334.5,-5763 334.5,-5738 48.5,-5738"/>
+<text text-anchor="start" x="53.5" y="-5747.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
+<text text-anchor="start" x="66.5" y="-5747.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
+<text text-anchor="start" x="143.5" y="-5747.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
+<polygon fill="none" stroke="black" points="48.5,-5713 48.5,-5738 334.5,-5738 334.5,-5713 48.5,-5713"/>
+<text text-anchor="start" x="53.5" y="-5722.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
+<text text-anchor="start" x="131.5" y="-5722.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
+<polygon fill="none" stroke="black" points="48.5,-5688 48.5,-5713 334.5,-5713 334.5,-5688 48.5,-5688"/>
+<text text-anchor="start" x="53.5" y="-5697.8" font-family="Helvetica,sans-Serif" font-size="14.00">include_deferred</text>
+<text text-anchor="start" x="169.5" y="-5697.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [BOOLEAN]</text>
+<text text-anchor="start" x="253.5" y="-5697.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
+<polygon fill="none" stroke="black" points="48.5,-5663 48.5,-5688 334.5,-5688 334.5,-5663 48.5,-5663"/>
+<text text-anchor="start" x="53.5" y="-5672.8" font-family="Helvetica,sans-Serif" font-size="14.00">pool</text>
+<text text-anchor="start" x="83.5" y="-5672.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(256)]</text>
+<polygon fill="none" stroke="black" points="48.5,-5638 48.5,-5663 334.5,-5663 334.5,-5638 48.5,-5638"/>
+<text text-anchor="start" x="53.5" y="-5647.8" font-family="Helvetica,sans-Serif" font-size="14.00">slots</text>
+<text text-anchor="start" x="86.5" y="-5647.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
 </g>
 <!-- variable -->
 <g id="node42" class="node">
 <title>variable</title>
-<polygon fill="none" stroke="black" points="100.5,-5918 100.5,-5946 283.5,-5946 283.5,-5918 100.5,-5918"/>
-<text text-anchor="start" x="155.5" y="-5929.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">variable</text>
-<polygon fill="none" stroke="black" points="100.5,-5893 100.5,-5918 283.5,-5918 283.5,-5893 100.5,-5893"/>
-<text text-anchor="start" x="105.5" y="-5902.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
-<text text-anchor="start" x="118.5" y="-5902.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
-<text text-anchor="start" x="195.5" y="-5902.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
-<polygon fill="none" stroke="black" points="100.5,-5868 100.5,-5893 283.5,-5893 283.5,-5868 100.5,-5868"/>
-<text text-anchor="start" x="105.5" y="-5877.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
-<text text-anchor="start" x="183.5" y="-5877.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
-<polygon fill="none" stroke="black" points="100.5,-5843 100.5,-5868 283.5,-5868 283.5,-5843 100.5,-5843"/>
-<text text-anchor="start" x="105.5" y="-5852.8" font-family="Helvetica,sans-Serif" font-size="14.00">is_encrypted</text>
-<text text-anchor="start" x="194.5" y="-5852.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [BOOLEAN]</text>
-<polygon fill="none" stroke="black" points="100.5,-5818 100.5,-5843 283.5,-5843 283.5,-5818 100.5,-5818"/>
-<text text-anchor="start" x="105.5" y="-5827.8" font-family="Helvetica,sans-Serif" font-size="14.00">key</text>
-<text text-anchor="start" x="130.5" y="-5827.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(250)]</text>
-<polygon fill="none" stroke="black" points="100.5,-5793 100.5,-5818 283.5,-5818 283.5,-5793 100.5,-5793"/>
-<text text-anchor="start" x="105.5" y="-5802.8" font-family="Helvetica,sans-Serif" font-size="14.00">val</text>
-<text text-anchor="start" x="126.5" y="-5802.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
+<polygon fill="none" stroke="black" points="100.5,-5942 100.5,-5970 283.5,-5970 283.5,-5942 100.5,-5942"/>
+<text text-anchor="start" x="155.5" y="-5953.2" font-family="Helvetica,sans-Serif" font-weight="bold" font-size="16.00">variable</text>
+<polygon fill="none" stroke="black" points="100.5,-5917 100.5,-5942 283.5,-5942 283.5,-5917 100.5,-5917"/>
+<text text-anchor="start" x="105.5" y="-5926.8" font-family="Helvetica,sans-Serif" text-decoration="underline" font-size="14.00">id</text>
+<text text-anchor="start" x="118.5" y="-5926.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [INTEGER]</text>
+<text text-anchor="start" x="195.5" y="-5926.8" font-family="Helvetica,sans-Serif" font-size="14.00"> NOT NULL</text>
+<polygon fill="none" stroke="black" points="100.5,-5892 100.5,-5917 283.5,-5917 283.5,-5892 100.5,-5892"/>
+<text text-anchor="start" x="105.5" y="-5901.8" font-family="Helvetica,sans-Serif" font-size="14.00">description</text>
+<text text-anchor="start" x="183.5" y="-5901.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
+<polygon fill="none" stroke="black" points="100.5,-5867 100.5,-5892 283.5,-5892 283.5,-5867 100.5,-5867"/>
+<text text-anchor="start" x="105.5" y="-5876.8" font-family="Helvetica,sans-Serif" font-size="14.00">is_encrypted</text>
+<text text-anchor="start" x="194.5" y="-5876.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [BOOLEAN]</text>
+<polygon fill="none" stroke="black" points="100.5,-5842 100.5,-5867 283.5,-5867 283.5,-5842 100.5,-5842"/>
+<text text-anchor="start" x="105.5" y="-5851.8" font-family="Helvetica,sans-Serif" font-size="14.00">key</text>
+<text text-anchor="start" x="130.5" y="-5851.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [VARCHAR(250)]</text>
+<polygon fill="none" stroke="black" points="100.5,-5817 100.5,-5842 283.5,-5842 283.5,-5817 100.5,-5817"/>
+<text text-anchor="start" x="105.5" y="-5826.8" font-family="Helvetica,sans-Serif" font-size="14.00">val</text>
+<text text-anchor="start" x="126.5" y="-5826.8" font-family="Helvetica,sans-Serif" font-size="14.00"> [TEXT]</text>
 </g>
 </g>
 </svg>
diff --git a/docs/apache-airflow/migrations-ref.rst b/docs/apache-airflow/migrations-ref.rst
index be796354da6f8..a3a42654f6a76 100644
--- a/docs/apache-airflow/migrations-ref.rst
+++ b/docs/apache-airflow/migrations-ref.rst
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
 | Revision ID                     | Revises ID        | Airflow Version   | Description                                                  |
 +=================================+===================+===================+==============================================================+
-| ``788397e78828`` (head)         | ``937cbd173ca1``  | ``2.7.0``         | Add custom_operator_name column                              |
+| ``405de8318b3a`` (head)         | ``788397e78828``  | ``2.7.0``         | add include_deferred column to pool                          |
++---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
+| ``788397e78828``                | ``937cbd173ca1``  | ``2.7.0``         | Add custom_operator_name column                              |
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
 | ``937cbd173ca1``                | ``c804e5c76e3e``  | ``2.7.0``         | Add index to task_instance table                             |
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
diff --git a/tests/api_connexion/endpoints/test_pool_endpoint.py b/tests/api_connexion/endpoints/test_pool_endpoint.py
index c12e2473f2c8d..a0121fccd185c 100644
--- a/tests/api_connexion/endpoints/test_pool_endpoint.py
+++ b/tests/api_connexion/endpoints/test_pool_endpoint.py
@@ -63,7 +63,7 @@ def teardown_method(self) -> None:
 
 class TestGetPools(TestBasePoolEndpoints):
     def test_response_200(self, session):
-        pool_model = Pool(pool="test_pool_a", slots=3)
+        pool_model = Pool(pool="test_pool_a", slots=3, include_deferred=True)
         session.add(pool_model)
         session.commit()
         result = session.query(Pool).all()
@@ -79,8 +79,10 @@ def test_response_200(self, session):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 128,
                     "description": "Default pool",
+                    "include_deferred": False,
                 },
                 {
                     "name": "test_pool_a",
@@ -89,15 +91,17 @@ def test_response_200(self, session):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 3,
                     "description": None,
+                    "include_deferred": True,
                 },
             ],
             "total_entries": 2,
         } == response.json
 
     def test_response_200_with_order_by(self, session):
-        pool_model = Pool(pool="test_pool_a", slots=3)
+        pool_model = Pool(pool="test_pool_a", slots=3, include_deferred=True)
         session.add(pool_model)
         session.commit()
         result = session.query(Pool).all()
@@ -113,8 +117,10 @@ def test_response_200_with_order_by(self, session):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 3,
                     "description": None,
+                    "include_deferred": True,
                 },
                 {
                     "name": "default_pool",
@@ -123,8 +129,10 @@ def test_response_200_with_order_by(self, session):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 128,
                     "description": "Default pool",
+                    "include_deferred": False,
                 },
             ],
             "total_entries": 2,
@@ -164,7 +172,7 @@ class TestGetPoolsPagination(TestBasePoolEndpoints):
     )
     @provide_session
     def test_limit_and_offset(self, url, expected_pool_ids, session):
-        pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 121)]
+        pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 121)]
         session.add_all(pools)
         session.commit()
         result = session.query(Pool).count()
@@ -175,7 +183,7 @@ def test_limit_and_offset(self, url, expected_pool_ids, session):
         assert pool_ids == expected_pool_ids
 
     def test_should_respect_page_size_limit_default(self, session):
-        pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 121)]
+        pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 121)]
         session.add_all(pools)
         session.commit()
         result = session.query(Pool).count()
@@ -185,7 +193,7 @@ def test_should_respect_page_size_limit_default(self, session):
         assert len(response.json["pools"]) == 100
 
     def test_should_raise_400_for_invalid_orderby(self, session):
-        pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 121)]
+        pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 121)]
         session.add_all(pools)
         session.commit()
         result = session.query(Pool).count()
@@ -199,7 +207,7 @@ def test_should_raise_400_for_invalid_orderby(self, session):
 
     @conf_vars({("api", "maximum_page_limit"): "150"})
     def test_should_return_conf_max_if_req_max_above_conf(self, session):
-        pools = [Pool(pool=f"test_pool{i}", slots=1) for i in range(1, 200)]
+        pools = [Pool(pool=f"test_pool{i}", slots=1, include_deferred=False) for i in range(1, 200)]
         session.add_all(pools)
         session.commit()
         result = session.query(Pool).count()
@@ -211,7 +219,7 @@ def test_should_return_conf_max_if_req_max_above_conf(self, session):
 
 class TestGetPool(TestBasePoolEndpoints):
     def test_response_200(self, session):
-        pool_model = Pool(pool="test_pool_a", slots=3)
+        pool_model = Pool(pool="test_pool_a", slots=3, include_deferred=True)
         session.add(pool_model)
         session.commit()
         response = self.client.get("/api/v1/pools/test_pool_a", environ_overrides={"REMOTE_USER": "test"})
@@ -223,8 +231,10 @@ def test_response_200(self, session):
             "running_slots": 0,
             "queued_slots": 0,
             "scheduled_slots": 0,
+            "deferred_slots": 0,
             "open_slots": 3,
             "description": None,
+            "include_deferred": True,
         } == response.json
 
     def test_response_404(self):
@@ -246,7 +256,7 @@ def test_should_raises_401_unauthenticated(self):
 class TestDeletePool(TestBasePoolEndpoints):
     def test_response_204(self, session):
         pool_name = "test_pool"
-        pool_instance = Pool(pool=pool_name, slots=3)
+        pool_instance = Pool(pool=pool_name, slots=3, include_deferred=False)
         session.add(pool_instance)
         session.commit()
 
@@ -268,7 +278,7 @@ def test_response_404(self):
 
     def test_should_raises_401_unauthenticated(self, session):
         pool_name = "test_pool"
-        pool_instance = Pool(pool=pool_name, slots=3)
+        pool_instance = Pool(pool=pool_name, slots=3, include_deferred=False)
         session.add(pool_instance)
         session.commit()
 
@@ -285,7 +295,7 @@ class TestPostPool(TestBasePoolEndpoints):
     def test_response_200(self):
         response = self.client.post(
             "api/v1/pools",
-            json={"name": "test_pool_a", "slots": 3, "description": "test pool"},
+            json={"name": "test_pool_a", "slots": 3, "description": "test pool", "include_deferred": True},
             environ_overrides={"REMOTE_USER": "test"},
         )
         assert response.status_code == 200
@@ -296,18 +306,20 @@ def test_response_200(self):
             "running_slots": 0,
             "queued_slots": 0,
             "scheduled_slots": 0,
+            "deferred_slots": 0,
             "open_slots": 3,
             "description": "test pool",
+            "include_deferred": True,
         } == response.json
 
     def test_response_409(self, session):
         pool_name = "test_pool_a"
-        pool_instance = Pool(pool=pool_name, slots=3)
+        pool_instance = Pool(pool=pool_name, slots=3, include_deferred=False)
         session.add(pool_instance)
         session.commit()
         response = self.client.post(
             "api/v1/pools",
-            json={"name": "test_pool_a", "slots": 3},
+            json={"name": "test_pool_a", "slots": 3, "include_deferred": False},
             environ_overrides={"REMOTE_USER": "test"},
         )
         assert response.status_code == 409
@@ -334,7 +346,7 @@ def test_response_409(self, session):
             pytest.param(
                 {},
                 "Missing required property(ies): ['name', 'slots']",
-                id="for missing pool name AND slots",
+                id="for missing pool name AND slots AND include_deferred",
             ),
             pytest.param(
                 {"name": "invalid_pool", "slots": 3, "extra_field_1": "extra"},
@@ -363,12 +375,12 @@ def test_should_raises_401_unauthenticated(self):
 
 class TestPatchPool(TestBasePoolEndpoints):
     def test_response_200(self, session):
-        pool = Pool(pool="test_pool", slots=2)
+        pool = Pool(pool="test_pool", slots=2, include_deferred=True)
         session.add(pool)
         session.commit()
         response = self.client.patch(
             "api/v1/pools/test_pool",
-            json={"name": "test_pool_a", "slots": 3},
+            json={"name": "test_pool_a", "slots": 3, "include_deferred": False},
             environ_overrides={"REMOTE_USER": "test"},
         )
         assert response.status_code == 200
@@ -379,8 +391,10 @@ def test_response_200(self, session):
             "open_slots": 3,
             "running_slots": 0,
             "scheduled_slots": 0,
+            "deferred_slots": 0,
             "slots": 3,
             "description": None,
+            "include_deferred": False,
         } == response.json
 
     @pytest.mark.parametrize(
@@ -393,13 +407,13 @@ def test_response_200(self, session):
             # Extra properties
             (
                 "{'extra_field': ['Unknown field.']}",
-                {"name": "test_pool_a", "slots": 3, "extra_field": "extra"},
+                {"name": "test_pool_a", "slots": 3, "include_deferred": True, "extra_field": "extra"},
             ),
         ],
     )
     @provide_session
     def test_response_400(self, error_detail, request_json, session):
-        pool = Pool(pool="test_pool", slots=2)
+        pool = Pool(pool="test_pool", slots=2, include_deferred=False)
         session.add(pool)
         session.commit()
         response = self.client.patch(
@@ -414,7 +428,7 @@ def test_response_400(self, error_detail, request_json, session):
         } == response.json
 
     def test_should_raises_401_unauthenticated(self, session):
-        pool = Pool(pool="test_pool", slots=2)
+        pool = Pool(pool="test_pool", slots=2, include_deferred=False)
         session.add(pool)
         session.commit()
 
@@ -443,7 +457,7 @@ def test_delete_400(self):
             pytest.param(
                 400,
                 "api/v1/pools/default_pool",
-                {"name": "test_pool_a", "slots": 3},
+                {"name": "test_pool_a", "slots": 3, "include_deferred": False},
                 {
                     "detail": "Default Pool's name can't be modified",
                     "status": 400,
@@ -455,7 +469,7 @@ def test_delete_400(self):
             pytest.param(
                 400,
                 "api/v1/pools/default_pool?update_mask=name, slots",
-                {"name": "test_pool_a", "slots": 3},
+                {"name": "test_pool_a", "slots": 3, "include_deferred": False},
                 {
                     "detail": "Default Pool's name can't be modified",
                     "status": 400,
@@ -475,11 +489,49 @@ def test_delete_400(self):
                     "open_slots": 3,
                     "running_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "slots": 3,
                     "description": "Default pool",
+                    "include_deferred": False,
                 },
                 id="200 Update mask with slots",
             ),
+            pytest.param(
+                200,
+                "api/v1/pools/default_pool?update_mask=include_deferred",
+                {"name": "test_pool_a", "include_deferred": True},
+                {
+                    "occupied_slots": 0,
+                    "queued_slots": 0,
+                    "name": "default_pool",
+                    "open_slots": 128,
+                    "running_slots": 0,
+                    "scheduled_slots": 0,
+                    "deferred_slots": 0,
+                    "slots": 128,
+                    "description": "Default pool",
+                    "include_deferred": True,
+                },
+                id="200 Update mask with include_deferred",
+            ),
+            pytest.param(
+                200,
+                "api/v1/pools/default_pool?update_mask=slots,include_deferred",
+                {"name": "test_pool_a", "slots": 3, "include_deferred": True},
+                {
+                    "occupied_slots": 0,
+                    "queued_slots": 0,
+                    "name": "default_pool",
+                    "open_slots": 3,
+                    "running_slots": 0,
+                    "scheduled_slots": 0,
+                    "deferred_slots": 0,
+                    "slots": 3,
+                    "description": "Default pool",
+                    "include_deferred": True,
+                },
+                id="200 Update mask with slots AND include_deferred",
+            ),
             pytest.param(
                 200,
                 "api/v1/pools/default_pool?update_mask=name,slots",
@@ -491,8 +543,10 @@ def test_delete_400(self):
                     "open_slots": 3,
                     "running_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "slots": 3,
                     "description": "Default pool",
+                    "include_deferred": False,
                 },
                 id="200 Update mask with slots and name",
             ),
@@ -502,6 +556,7 @@ def test_delete_400(self):
                 {
                     "name": "default_pool",
                     "slots": 3,
+                    "include_deferred": True,
                 },
                 {
                     "occupied_slots": 0,
@@ -510,8 +565,10 @@ def test_delete_400(self):
                     "open_slots": 3,
                     "running_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "slots": 3,
                     "description": "Default pool",
+                    "include_deferred": True,
                 },
                 id="200 no update mask",
             ),
@@ -525,37 +582,50 @@ def test_patch(self, status_code, url, json, expected_response):
 
 class TestPatchPoolWithUpdateMask(TestBasePoolEndpoints):
     @pytest.mark.parametrize(
-        "url, patch_json, expected_name, expected_slots",
+        "url, patch_json, expected_name, expected_slots, expected_include_deferred",
         [
             (
                 "api/v1/pools/test_pool?update_mask=name, slots",
                 {"name": "test_pool_a", "slots": 2},
                 "test_pool_a",
                 2,
+                False,
             ),
             (
                 "api/v1/pools/test_pool?update_mask=name",
                 {"name": "test_pool_a", "slots": 2},
                 "test_pool_a",
                 3,
+                False,
             ),
             (
                 "api/v1/pools/test_pool?update_mask=slots",
                 {"name": "test_pool_a", "slots": 2},
                 "test_pool",
                 2,
+                False,
             ),
             (
                 "api/v1/pools/test_pool?update_mask=slots",
                 {"slots": 2},
                 "test_pool",
                 2,
+                False,
+            ),
+            (
+                "api/v1/pools/test_pool?update_mask=include_deferred",
+                {"include_deferred": True},
+                "test_pool",
+                3,
+                True,
             ),
         ],
     )
     @provide_session
-    def test_response_200(self, url, patch_json, expected_name, expected_slots, session):
-        pool = Pool(pool="test_pool", slots=3)
+    def test_response_200(
+        self, url, patch_json, expected_name, expected_slots, expected_include_deferred, session
+    ):
+        pool = Pool(pool="test_pool", slots=3, include_deferred=False)
         session.add(pool)
         session.commit()
         response = self.client.patch(url, json=patch_json, environ_overrides={"REMOTE_USER": "test"})
@@ -567,8 +637,10 @@ def test_response_200(self, url, patch_json, expected_name, expected_slots, sess
             "running_slots": 0,
             "queued_slots": 0,
             "scheduled_slots": 0,
+            "deferred_slots": 0,
             "open_slots": expected_slots,
             "description": None,
+            "include_deferred": expected_include_deferred,
         } == response.json
 
     @pytest.mark.parametrize(
@@ -602,7 +674,7 @@ def test_response_200(self, url, patch_json, expected_name, expected_slots, sess
     )
     @provide_session
     def test_response_400(self, error_detail, url, patch_json, session):
-        pool = Pool(pool="test_pool", slots=3)
+        pool = Pool(pool="test_pool", slots=3, include_deferred=False)
         session.add(pool)
         session.commit()
         response = self.client.patch(url, json=patch_json, environ_overrides={"REMOTE_USER": "test"})
diff --git a/tests/api_connexion/schemas/test_pool_schemas.py b/tests/api_connexion/schemas/test_pool_schemas.py
index f0eb0c0a492f5..9c5a78994db43 100644
--- a/tests/api_connexion/schemas/test_pool_schemas.py
+++ b/tests/api_connexion/schemas/test_pool_schemas.py
@@ -31,7 +31,7 @@ def teardown_method(self) -> None:
 
     @provide_session
     def test_serialize(self, session):
-        pool_model = Pool(pool="test_pool", slots=2)
+        pool_model = Pool(pool="test_pool", slots=2, include_deferred=False)
         session.add(pool_model)
         session.commit()
         pool_instance = session.query(Pool).filter(Pool.pool == pool_model.pool).first()
@@ -43,13 +43,15 @@ def test_serialize(self, session):
             "running_slots": 0,
             "queued_slots": 0,
             "scheduled_slots": 0,
+            "deferred_slots": 0,
             "open_slots": 2,
             "description": None,
+            "include_deferred": False,
         }
 
     @provide_session
     def test_deserialize(self, session):
-        pool_dict = {"name": "test_pool", "slots": 3}
+        pool_dict = {"name": "test_pool", "slots": 3, "include_deferred": True}
         deserialized_pool = pool_schema.load(pool_dict, session=session)
         assert not isinstance(deserialized_pool, Pool)  # Checks if load_instance is set to True
 
@@ -62,8 +64,8 @@ def teardown_method(self) -> None:
         clear_db_pools()
 
     def test_serialize(self):
-        pool_model_a = Pool(pool="test_pool_a", slots=3)
-        pool_model_b = Pool(pool="test_pool_b", slots=3)
+        pool_model_a = Pool(pool="test_pool_a", slots=3, include_deferred=False)
+        pool_model_b = Pool(pool="test_pool_b", slots=3, include_deferred=True)
         instance = PoolCollection(pools=[pool_model_a, pool_model_b], total_entries=2)
         assert {
             "pools": [
@@ -74,8 +76,10 @@ def test_serialize(self):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 3,
                     "description": None,
+                    "include_deferred": False,
                 },
                 {
                     "name": "test_pool_b",
@@ -84,8 +88,10 @@ def test_serialize(self):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 3,
                     "description": None,
+                    "include_deferred": True,
                 },
             ],
             "total_entries": 2,
diff --git a/tests/api_connexion/test_auth.py b/tests/api_connexion/test_auth.py
index 35491ad7b2adf..b4cdc4a2bc7e6 100644
--- a/tests/api_connexion/test_auth.py
+++ b/tests/api_connexion/test_auth.py
@@ -78,8 +78,10 @@ def test_success(self):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 128,
                     "description": "Default pool",
+                    "include_deferred": False,
                 },
             ],
             "total_entries": 1,
@@ -153,8 +155,10 @@ def test_success(self):
                     "running_slots": 0,
                     "queued_slots": 0,
                     "scheduled_slots": 0,
+                    "deferred_slots": 0,
                     "open_slots": 128,
                     "description": "Default pool",
+                    "include_deferred": False,
                 },
             ],
             "total_entries": 1,
diff --git a/tests/api_experimental/client/test_local_client.py b/tests/api_experimental/client/test_local_client.py
index 575400a115de0..67246106709e6 100644
--- a/tests/api_experimental/client/test_local_client.py
+++ b/tests/api_experimental/client/test_local_client.py
@@ -181,22 +181,26 @@ def test_delete_dag(self):
             assert session.query(DagModel).filter(DagModel.dag_id == key).count() == 0
 
     def test_get_pool(self):
-        self.client.create_pool(name="foo", slots=1, description="")
+        self.client.create_pool(name="foo", slots=1, description="", include_deferred=False)
         pool = self.client.get_pool(name="foo")
-        assert pool == ("foo", 1, "")
+        assert pool == ("foo", 1, "", False)
 
     def test_get_pool_non_existing_raises(self):
         with pytest.raises(PoolNotFound):
             self.client.get_pool(name="foo")
 
     def test_get_pools(self):
-        self.client.create_pool(name="foo1", slots=1, description="")
-        self.client.create_pool(name="foo2", slots=2, description="")
+        self.client.create_pool(name="foo1", slots=1, description="", include_deferred=False)
+        self.client.create_pool(name="foo2", slots=2, description="", include_deferred=True)
         pools = sorted(self.client.get_pools(), key=lambda p: p[0])
-        assert pools == [("default_pool", 128, "Default pool"), ("foo1", 1, ""), ("foo2", 2, "")]
+        assert pools == [
+            ("default_pool", 128, "Default pool", False),
+            ("foo1", 1, "", False),
+            ("foo2", 2, "", True),
+        ]
 
     def test_create_pool(self):
-        pool = self.client.create_pool(name="foo", slots=1, description="")
+        pool = self.client.create_pool(name="foo", slots=1, description="", include_deferred=False)
         assert pool == ("foo", 1, "")
         with create_session() as session:
             assert session.query(Pool).count() == 2
@@ -207,6 +211,7 @@ def test_create_pool_bad_slots(self):
                 name="foo",
                 slots="foo",
                 description="",
+                include_deferred=True,
             )
 
     def test_create_pool_name_too_long(self):
@@ -219,10 +224,11 @@ def test_create_pool_name_too_long(self):
                 name=long_name,
                 slots=5,
                 description="",
+                include_deferred=False,
             )
 
     def test_delete_pool(self):
-        self.client.create_pool(name="foo", slots=1, description="")
+        self.client.create_pool(name="foo", slots=1, description="", include_deferred=False)
         with create_session() as session:
             assert session.query(Pool).count() == 2
         self.client.delete_pool(name="foo")
diff --git a/tests/api_experimental/common/experimental/test_pool.py b/tests/api_experimental/common/experimental/test_pool.py
index 4e6bdda7ac34b..d00f318ffac9e 100644
--- a/tests/api_experimental/common/experimental/test_pool.py
+++ b/tests/api_experimental/common/experimental/test_pool.py
@@ -45,6 +45,7 @@ def setup_method(self):
                 pool=name,
                 slots=i,
                 description=name,
+                include_deferred=True,
             )
             self.pools.append(pool)
         with create_session() as session:
@@ -77,10 +78,12 @@ def test_create_pool(self):
             assert session.query(models.Pool).count() == self.TOTAL_POOL_COUNT + 1
 
     def test_create_pool_existing(self):
-        pool = pool_api.create_pool(name=self.pools[0].pool, slots=5, description="")
-        assert pool.pool == self.pools[0].pool
+        pool = pool_api.create_pool(name=self.pools[1].pool, slots=5, description="")
+        assert pool.pool == self.pools[1].pool
         assert pool.slots == 5
         assert pool.description == ""
+        # do not change include_deferred on an existing pool
+        assert pool.include_deferred is True
         with create_session() as session:
             assert session.query(models.Pool).count() == self.TOTAL_POOL_COUNT
 
diff --git a/tests/cli/commands/test_pool_command.py b/tests/cli/commands/test_pool_command.py
index ec175912d8503..ca76a965145c3 100644
--- a/tests/cli/commands/test_pool_command.py
+++ b/tests/cli/commands/test_pool_command.py
@@ -19,7 +19,6 @@
 
 import io
 import json
-import os
 from contextlib import redirect_stdout
 
 import pytest
@@ -67,6 +66,18 @@ def test_pool_create(self):
         pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
         assert self.session.query(Pool).count() == 2
 
+    def test_pool_update_deferred(self):
+        pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
+        assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is False
+
+        pool_command.pool_set(
+            self.parser.parse_args(["pools", "set", "foo", "1", "test", "--include-deferred"])
+        )
+        assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is True
+
+        pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
+        assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is False
+
     def test_pool_get(self):
         pool_command.pool_set(self.parser.parse_args(["pools", "set", "foo", "1", "test"]))
         pool_command.pool_get(self.parser.parse_args(["pools", "get", "foo"]))
@@ -80,39 +91,57 @@ def test_pool_import_nonexistent(self):
         with pytest.raises(SystemExit):
             pool_command.pool_import(self.parser.parse_args(["pools", "import", "nonexistent.json"]))
 
-    def test_pool_import_invalid_json(self):
-        with open("pools_import_invalid.json", mode="w") as file:
+    def test_pool_import_invalid_json(self, tmp_path):
+        invalid_pool_import_file_path = tmp_path / "pools_import_invalid.json"
+        with open(invalid_pool_import_file_path, mode="w") as file:
             file.write("not valid json")
 
         with pytest.raises(SystemExit):
-            pool_command.pool_import(self.parser.parse_args(["pools", "import", "pools_import_invalid.json"]))
-
-    def test_pool_import_invalid_pools(self):
-        pool_config_input = {"foo": {"description": "foo_test"}}
-        with open("pools_import_invalid.json", mode="w") as file:
+            pool_command.pool_import(
+                self.parser.parse_args(["pools", "import", str(invalid_pool_import_file_path)])
+            )
+
+    def test_pool_import_invalid_pools(self, tmp_path):
+        invalid_pool_import_file_path = tmp_path / "pools_import_invalid.json"
+        pool_config_input = {"foo": {"description": "foo_test", "include_deferred": False}}
+        with open(invalid_pool_import_file_path, mode="w") as file:
             json.dump(pool_config_input, file)
 
         with pytest.raises(SystemExit):
-            pool_command.pool_import(self.parser.parse_args(["pools", "import", "pools_import_invalid.json"]))
+            pool_command.pool_import(
+                self.parser.parse_args(["pools", "import", str(invalid_pool_import_file_path)])
+            )
 
-    def test_pool_import_export(self):
-        # Create two pools first
+    def test_pool_import_backwards_compatibility(self, tmp_path):
+        pool_import_file_path = tmp_path / "pools_import.json"
         pool_config_input = {
+            # JSON before version 2.7.0 does not contain `include_deferred`
             "foo": {"description": "foo_test", "slots": 1},
-            "default_pool": {"description": "Default pool", "slots": 128},
-            "baz": {"description": "baz_test", "slots": 2},
         }
-        with open("pools_import.json", mode="w") as file:
+        with open(pool_import_file_path, mode="w") as file:
+            json.dump(pool_config_input, file)
+
+        pool_command.pool_import(self.parser.parse_args(["pools", "import", str(pool_import_file_path)]))
+
+        assert self.session.query(Pool).filter(Pool.pool == "foo").first().include_deferred is False
+
+    def test_pool_import_export(self, tmp_path):
+        pool_import_file_path = tmp_path / "pools_import.json"
+        pool_export_file_path = tmp_path / "pools_export.json"
+        pool_config_input = {
+            "foo": {"description": "foo_test", "slots": 1, "include_deferred": True},
+            "default_pool": {"description": "Default pool", "slots": 128, "include_deferred": False},
+            "baz": {"description": "baz_test", "slots": 2, "include_deferred": False},
+        }
+        with open(pool_import_file_path, mode="w") as file:
             json.dump(pool_config_input, file)
 
         # Import json
-        pool_command.pool_import(self.parser.parse_args(["pools", "import", "pools_import.json"]))
+        pool_command.pool_import(self.parser.parse_args(["pools", "import", str(pool_import_file_path)]))
 
         # Export json
-        pool_command.pool_export(self.parser.parse_args(["pools", "export", "pools_export.json"]))
+        pool_command.pool_export(self.parser.parse_args(["pools", "export", str(pool_export_file_path)]))
 
-        with open("pools_export.json") as file:
+        with open(pool_export_file_path) as file:
             pool_config_output = json.load(file)
             assert pool_config_input == pool_config_output, "Input and output pool files are not same"
-        os.remove("pools_import.json")
-        os.remove("pools_export.json")
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
index a96c5cbddfed2..eac5fe6802523 100644
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -813,7 +813,7 @@ def test_run_task_with_pool(self):
 
         clear_db_pools()
         with create_session() as session:
-            pool = Pool(pool=pool_name, slots=1)
+            pool = Pool(pool=pool_name, slots=1, include_deferred=False)
             session.add(pool)
             session.commit()
 
diff --git a/tests/jobs/test_backfill_job.py b/tests/jobs/test_backfill_job.py
index a11b8d56886d6..e4efd8682ea98 100644
--- a/tests/jobs/test_backfill_job.py
+++ b/tests/jobs/test_backfill_job.py
@@ -598,6 +598,7 @@ def test_backfill_respect_pool_limit(self, mock_log, dag_maker):
         pool = Pool(
             pool="pool_with_two_slots",
             slots=slots,
+            include_deferred=False,
         )
         session.add(pool)
         session.commit()
@@ -961,7 +962,7 @@ def test_backfill_pooled_tasks(self):
         Test that queued tasks are executed by BackfillJobRunner
         """
         session = settings.Session()
-        pool = Pool(pool="test_backfill_pooled_task_pool", slots=1)
+        pool = Pool(pool="test_backfill_pooled_task_pool", slots=1, include_deferred=False)
         session.add(pool)
         session.commit()
         session.close()
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 3a3d3f55b4b42..ed09d9cd18201 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -605,8 +605,8 @@ def test_find_executable_task_instances_pool(self, dag_maker):
         for ti in tis:
             ti.state = State.SCHEDULED
             session.merge(ti)
-        pool = Pool(pool="a", slots=1, description="haha")
-        pool2 = Pool(pool="b", slots=100, description="haha")
+        pool = Pool(pool="a", slots=1, description="haha", include_deferred=False)
+        pool2 = Pool(pool="b", slots=100, description="haha", include_deferred=False)
         session.add(pool)
         session.add(pool2)
         session.flush()
@@ -732,8 +732,8 @@ def test_find_executable_task_instances_order_priority_with_pools(self, dag_make
 
         dag_id = "SchedulerJobTest.test_find_executable_task_instances_order_priority_with_pools"
 
-        session.add(Pool(pool="pool1", slots=32))
-        session.add(Pool(pool="pool2", slots=32))
+        session.add(Pool(pool="pool1", slots=32, include_deferred=False))
+        session.add(Pool(pool="pool2", slots=32, include_deferred=False))
 
         with dag_maker(dag_id=dag_id, max_active_tasks=2):
             op1 = EmptyOperator(task_id="dummy1", priority_weight=1, pool="pool1")
@@ -888,7 +888,12 @@ def test_infinite_pool(self, dag_maker):
         ti = dr.task_instances[0]
         ti.state = State.SCHEDULED
         session.merge(ti)
-        infinite_pool = Pool(pool="infinite_pool", slots=-1, description="infinite pool")
+        infinite_pool = Pool(
+            pool="infinite_pool",
+            slots=-1,
+            description="infinite pool",
+            include_deferred=False,
+        )
         session.add(infinite_pool)
         session.commit()
 
@@ -913,7 +918,7 @@ def test_not_enough_pool_slots(self, caplog, dag_maker):
         ti = dr.task_instances[1]
         ti.state = State.SCHEDULED
         session.merge(ti)
-        some_pool = Pool(pool="some_pool", slots=2, description="my pool")
+        some_pool = Pool(pool="some_pool", slots=2, description="my pool", include_deferred=False)
         session.add(some_pool)
         session.commit()
         with caplog.at_level(logging.WARNING):
@@ -1336,8 +1341,8 @@ def test_find_executable_task_instances_negative_open_pool_slots(self, dag_maker
         self.job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull)
         session = settings.Session()
 
-        pool1 = Pool(pool="pool1", slots=1)
-        pool2 = Pool(pool="pool2", slots=1)
+        pool1 = Pool(pool="pool1", slots=1, include_deferred=False)
+        pool2 = Pool(pool="pool2", slots=1, include_deferred=False)
 
         session.add(pool1)
         session.add(pool2)
@@ -2578,7 +2583,7 @@ def test_scheduler_verify_pool_full(self, dag_maker, configs):
                 )
 
             session = settings.Session()
-            pool = Pool(pool="test_scheduler_verify_pool_full", slots=1)
+            pool = Pool(pool="test_scheduler_verify_pool_full", slots=1, include_deferred=False)
             session.add(pool)
             session.flush()
 
@@ -2619,7 +2624,7 @@ def test_scheduler_verify_pool_full_2_slots_per_task(self, dag_maker, session):
                 bash_command="echo hi",
             )
 
-        pool = Pool(pool="test_scheduler_verify_pool_full_2_slots_per_task", slots=6)
+        pool = Pool(pool="test_scheduler_verify_pool_full_2_slots_per_task", slots=6, include_deferred=False)
         session.add(pool)
         session.flush()
 
@@ -2673,8 +2678,8 @@ def test_scheduler_keeps_scheduling_pool_full(self, dag_maker):
         dag_d2 = dag_maker.dag
 
         session = settings.Session()
-        pool_p1 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p1", slots=1)
-        pool_p2 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p2", slots=10)
+        pool_p1 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p1", slots=1, include_deferred=False)
+        pool_p2 = Pool(pool="test_scheduler_keeps_scheduling_pool_full_p2", slots=10, include_deferred=False)
         session.add(pool_p1)
         session.add(pool_p2)
         session.flush()
@@ -2751,7 +2756,7 @@ def test_scheduler_verify_priority_and_slots(self, dag_maker):
             )
 
         session = settings.Session()
-        pool = Pool(pool="test_scheduler_verify_priority_and_slots", slots=2)
+        pool = Pool(pool="test_scheduler_verify_priority_and_slots", slots=2, include_deferred=False)
         session.add(pool)
         session.flush()
 
diff --git a/tests/models/test_pool.py b/tests/models/test_pool.py
index 2e142814c6d72..32d30b99bb90f 100644
--- a/tests/models/test_pool.py
+++ b/tests/models/test_pool.py
@@ -55,6 +55,7 @@ def add_pools(self):
                 pool=name,
                 slots=i,
                 description=name,
+                include_deferred=False,
             )
             self.pools.append(pool)
         with create_session() as session:
@@ -64,23 +65,27 @@ def teardown_method(self):
         self.clean_db()
 
     def test_open_slots(self, dag_maker):
-        pool = Pool(pool="test_pool", slots=5)
+        pool = Pool(pool="test_pool", slots=5, include_deferred=False)
         with dag_maker(
             dag_id="test_open_slots",
             start_date=DEFAULT_DATE,
         ):
             op1 = EmptyOperator(task_id="dummy1", pool="test_pool")
             op2 = EmptyOperator(task_id="dummy2", pool="test_pool")
+            op3 = EmptyOperator(task_id="dummy3", pool="test_pool")
         dag_maker.create_dagrun()
         ti1 = TI(task=op1, execution_date=DEFAULT_DATE)
         ti2 = TI(task=op2, execution_date=DEFAULT_DATE)
+        ti3 = TI(task=op3, execution_date=DEFAULT_DATE)
         ti1.state = State.RUNNING
         ti2.state = State.QUEUED
+        ti3.state = State.DEFERRED
 
         session = settings.Session()
         session.add(pool)
         session.merge(ti1)
         session.merge(ti2)
+        session.merge(ti3)
         session.commit()
         session.close()
 
@@ -88,23 +93,69 @@ def test_open_slots(self, dag_maker):
         assert 1 == pool.running_slots()
         assert 1 == pool.queued_slots()
         assert 2 == pool.occupied_slots()
+        assert 1 == pool.deferred_slots()
         assert {
             "default_pool": {
                 "open": 128,
                 "queued": 0,
                 "total": 128,
                 "running": 0,
+                "deferred": 0,
             },
             "test_pool": {
                 "open": 3,
                 "queued": 1,
                 "running": 1,
+                "deferred": 1,
+                "total": 5,
+            },
+        } == pool.slots_stats()
+
+    def test_open_slots_including_deferred(self, dag_maker):
+        pool = Pool(pool="test_pool", slots=5, include_deferred=True)
+        with dag_maker(
+            dag_id="test_open_slots_including_deferred",
+            start_date=DEFAULT_DATE,
+        ):
+            op1 = EmptyOperator(task_id="dummy1", pool="test_pool")
+            op2 = EmptyOperator(task_id="dummy2", pool="test_pool")
+        dag_maker.create_dagrun()
+        ti1 = TI(task=op1, execution_date=DEFAULT_DATE)
+        ti2 = TI(task=op2, execution_date=DEFAULT_DATE)
+        ti1.state = State.RUNNING
+        ti2.state = State.DEFERRED
+
+        session = settings.Session()
+        session.add(pool)
+        session.merge(ti1)
+        session.merge(ti2)
+        session.commit()
+        session.close()
+
+        assert 3 == pool.open_slots()
+        assert 1 == pool.running_slots()
+        assert 0 == pool.queued_slots()
+        assert 1 == pool.deferred_slots()
+        assert 2 == pool.occupied_slots()
+        assert {
+            "default_pool": {
+                "open": 128,
+                "queued": 0,
+                "total": 128,
+                "running": 0,
+                "deferred": 0,
+            },
+            "test_pool": {
+                "open": 3,
+                "queued": 0,
+                "running": 1,
+                "deferred": 1,
                 "total": 5,
             },
         } == pool.slots_stats()
 
     def test_infinite_slots(self, dag_maker):
-        pool = Pool(pool="test_pool", slots=-1)
+        pool = Pool(pool="test_pool", slots=-1, include_deferred=False)
         with dag_maker(
             dag_id="test_infinite_slots",
         ):
@@ -133,12 +184,14 @@ def test_infinite_slots(self, dag_maker):
                 "queued": 0,
                 "total": 128,
                 "running": 0,
+                "deferred": 0,
             },
             "test_pool": {
                 "open": float("inf"),
                 "queued": 1,
                 "running": 1,
                 "total": float("inf"),
+                "deferred": 0,
             },
         } == pool.slots_stats()
 
@@ -170,6 +223,7 @@ def test_default_pool_open_slots(self, dag_maker):
                 "queued": 2,
                 "total": 5,
                 "running": 1,
+                "deferred": 0,
             }
         } == Pool.slots_stats()
 
@@ -194,18 +248,22 @@ def test_get_pools(self):
 
     def test_create_pool(self, session):
         self.add_pools()
-        pool = Pool.create_or_update_pool(name="foo", slots=5, description="")
+        pool = Pool.create_or_update_pool(name="foo", slots=5, description="", include_deferred=True)
         assert pool.pool == "foo"
         assert pool.slots == 5
         assert pool.description == ""
+        assert pool.include_deferred is True
         assert session.query(Pool).count() == self.TOTAL_POOL_COUNT + 1
 
     def test_create_pool_existing(self, session):
         self.add_pools()
-        pool = Pool.create_or_update_pool(name=self.pools[0].pool, slots=5, description="")
+        pool = Pool.create_or_update_pool(
+            name=self.pools[0].pool, slots=5, description="", include_deferred=False
+        )
         assert pool.pool == self.pools[0].pool
         assert pool.slots == 5
         assert pool.description == ""
+        assert pool.include_deferred is False
         assert session.query(Pool).count() == self.TOTAL_POOL_COUNT
 
     def test_delete_pool(self, session):
@@ -223,7 +281,9 @@ def test_delete_default_pool_not_allowed(self):
             Pool.delete_pool(Pool.DEFAULT_POOL_NAME)
 
     def test_is_default_pool(self):
-        pool = Pool.create_or_update_pool(name="not_default_pool", slots=1, description="test")
+        pool = Pool.create_or_update_pool(
+            name="not_default_pool", slots=1, description="test", include_deferred=False
+        )
         default_pool = Pool.get_default_pool()
         assert not Pool.is_default_pool(id=pool.id)
         assert Pool.is_default_pool(str(default_pool.id))
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9e7f5eb48a17c..ac2f7f71748ef 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -96,7 +96,7 @@
 @pytest.fixture
 def test_pool():
     with create_session() as session:
-        test_pool = Pool(pool="test_pool", slots=1)
+        test_pool = Pool(pool="test_pool", slots=1, include_deferred=False)
         session.add(test_pool)
         session.flush()
         yield test_pool
diff --git a/tests/operators/test_subdag_operator.py b/tests/operators/test_subdag_operator.py
index 8df1ee9fe371b..c814e64a7e2be 100644
--- a/tests/operators/test_subdag_operator.py
+++ b/tests/operators/test_subdag_operator.py
@@ -88,8 +88,8 @@ def test_subdag_pools(self):
         subdag = DAG("parent.child", default_args=default_args)
 
         session = airflow.settings.Session()
-        pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1)
-        pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10)
+        pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1, include_deferred=False)
+        pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10, include_deferred=False)
         session.add(pool_1)
         session.add(pool_10)
         session.commit()
@@ -116,8 +116,8 @@ def test_subdag_pools_no_possible_conflict(self):
         subdag = DAG("parent.child", default_args=default_args)
 
         session = airflow.settings.Session()
-        pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1)
-        pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10)
+        pool_1 = airflow.models.Pool(pool="test_pool_1", slots=1, include_deferred=False)
+        pool_10 = airflow.models.Pool(pool="test_pool_10", slots=10, include_deferred=False)
         session.add(pool_1)
         session.add(pool_10)
         session.commit()
diff --git a/tests/ti_deps/deps/test_pool_slots_available_dep.py b/tests/ti_deps/deps/test_pool_slots_available_dep.py
index cba53ef65b2a3..08abba29e4171 100644
--- a/tests/ti_deps/deps/test_pool_slots_available_dep.py
+++ b/tests/ti_deps/deps/test_pool_slots_available_dep.py
@@ -23,6 +23,7 @@
 from airflow.ti_deps.dependencies_states import EXECUTION_STATES
 from airflow.ti_deps.deps.pool_slots_available_dep import PoolSlotsAvailableDep
 from airflow.utils.session import create_session
+from airflow.utils.state import TaskInstanceState
 from tests.test_utils import db
 
 
@@ -30,8 +31,9 @@ class TestPoolSlotsAvailableDep:
     def setup_method(self):
         db.clear_db_pools()
         with create_session() as session:
-            test_pool = Pool(pool="test_pool")
-            session.add(test_pool)
+            test_pool = Pool(pool="test_pool", include_deferred=False)
+            test_includes_deferred_pool = Pool(pool="test_includes_deferred_pool", include_deferred=True)
+            session.add_all([test_pool, test_includes_deferred_pool])
             session.commit()
 
     def teardown_method(self):
@@ -53,6 +55,13 @@ def test_running_pooled_task_pass(self, mock_open_slots):
             ti = Mock(pool="test_pool", state=state, pool_slots=1)
             assert PoolSlotsAvailableDep().is_met(ti=ti)
 
+    @patch("airflow.models.Pool.open_slots", return_value=0)
+    def test_deferred_pooled_task_pass(self, mock_open_slots):
+        ti = Mock(pool="test_includes_deferred_pool", state=TaskInstanceState.DEFERRED, pool_slots=1)
+        assert PoolSlotsAvailableDep().is_met(ti=ti)
+        ti_to_fail = Mock(pool="test_pool", state=TaskInstanceState.DEFERRED, pool_slots=1)
+        assert not PoolSlotsAvailableDep().is_met(ti=ti_to_fail)
+
     def test_task_with_nonexistent_pool(self):
         ti = Mock(pool="nonexistent_pool", pool_slots=1)
         assert not PoolSlotsAvailableDep().is_met(ti=ti)
diff --git a/tests/www/api/experimental/test_endpoints.py b/tests/www/api/experimental/test_endpoints.py
index 86a145f5b81ae..d42cfb2237004 100644
--- a/tests/www/api/experimental/test_endpoints.py
+++ b/tests/www/api/experimental/test_endpoints.py
@@ -383,6 +383,7 @@ def _setup_attrs(self, _setup_attrs_base):
                 pool=name,
                 slots=i,
                 description=name,
+                include_deferred=False,
             )
             self.session.add(pool)
             self.pools.append(pool)
@@ -434,6 +435,7 @@ def test_create_pool(self):
         assert pool["pool"] == "foo"
         assert pool["slots"] == 1
         assert pool["description"] == ""
+        assert pool["include_deferred"] is False
         assert self._get_pool_count() == self.TOTAL_POOL_COUNT + 1
 
     def test_create_pool_with_bad_name(self):
diff --git a/tests/www/views/test_views_pool.py b/tests/www/views/test_views_pool.py
index aebcd2c8444da..e79ec46c56ddf 100644
--- a/tests/www/views/test_views_pool.py
+++ b/tests/www/views/test_views_pool.py
@@ -29,6 +29,7 @@
     "pool": "test-pool",
     "slots": 777,
     "description": "test-pool-description",
+    "include_deferred": False,
 }
 
 
@@ -90,9 +91,13 @@ def test_list(app, admin_client, pool_factory):
         url = flask.url_for("TaskInstanceModelView.list", _flt_3_pool="test-pool", _flt_3_state="scheduled")
         scheduled_tag = markupsafe.Markup("<a href='{url}'>{slots}</a>").format(url=url, slots=0)
 
+        url = flask.url_for("TaskInstanceModelView.list", _flt_3_pool="test-pool", _flt_3_state="deferred")
+        deferred_tag = markupsafe.Markup("<a href='{url}'>{slots}</a>").format(url=url, slots=0)
+
     check_content_in_response(used_tag, resp)
     check_content_in_response(queued_tag, resp)
     check_content_in_response(scheduled_tag, resp)
+    check_content_in_response(deferred_tag, resp)
 
 
 def test_pool_muldelete(session, admin_client, pool_factory):

From 95187f3fe8bbbf323b59f4932de41ec03249e711 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Wed, 9 Aug 2023 20:52:56 +0800
Subject: [PATCH 08/22] Remove duplicated function

---
 .../providers/cncf/kubernetes/utils/pod_manager.py    | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index 6f39392f321ab..41518db27a221 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -153,17 +153,6 @@ def container_is_terminated(pod: V1Pod, container_name: str) -> bool:
     return container_status.state.terminated is not None
 
 
-def container_is_completed(pod: V1Pod, container_name: str) -> bool:
-    """
-    Examines V1Pod ``pod`` to determine whether ``container_name`` is completed.
-    If that container is present and completed, returns True.  Returns False otherwise.
-    """
-    container_status = get_container_status(pod, container_name)
-    if not container_status:
-        return False
-    return container_status.state.terminated is not None
-
-
 def container_is_succeeded(pod: V1Pod, container_name: str) -> bool:
     """
     Examines V1Pod ``pod`` to determine whether ``container_name`` is completed and succeeded.

From 1dcf05f1c69e52667064259308ba5950b594b268 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent@astronomer.io>
Date: Wed, 9 Aug 2023 09:37:54 -0400
Subject: [PATCH 09/22] Gantt chart: Use earliest/oldest ti dates if different
 than dag run start/end (#33215)

---
 .../js/dag/details/gantt/GanttTooltip.tsx     |  3 ++
 .../www/static/js/dag/details/gantt/Row.tsx   | 20 ++++++++----
 .../www/static/js/dag/details/gantt/index.tsx | 32 +++++++++++++++++--
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx b/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx
index 2f7d617c8cd3c..2f501209256d4 100644
--- a/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx
+++ b/airflow/www/static/js/dag/details/gantt/GanttTooltip.tsx
@@ -42,6 +42,9 @@ const GanttTooltip = ({ task, instance }: Props) => {
       <Text>
         Task{isGroup ? " Group" : ""}: {task.label}
       </Text>
+      {!!instance?.tryNumber && instance.tryNumber > 1 && (
+        <Text>Try Number: {instance.tryNumber}</Text>
+      )}
       <br />
       {instance?.queuedDttm && (
         <Text>
diff --git a/airflow/www/static/js/dag/details/gantt/Row.tsx b/airflow/www/static/js/dag/details/gantt/Row.tsx
index 96f079154e1e1..304fb83e22559 100644
--- a/airflow/www/static/js/dag/details/gantt/Row.tsx
+++ b/airflow/www/static/js/dag/details/gantt/Row.tsx
@@ -24,24 +24,31 @@ import { getDuration } from "src/datetime_utils";
 import { SimpleStatus } from "src/dag/StatusBox";
 import { useContainerRef } from "src/context/containerRef";
 import { hoverDelay } from "src/utils";
-import type { DagRun, Task } from "src/types";
+import type { Task } from "src/types";
 import GanttTooltip from "./GanttTooltip";
 
 interface Props {
   ganttWidth?: number;
   openGroupIds: string[];
-  dagRun: DagRun;
   task: Task;
+  ganttStartDate?: string | null;
+  ganttEndDate?: string | null;
 }
 
-const Row = ({ ganttWidth = 500, openGroupIds, task, dagRun }: Props) => {
+const Row = ({
+  ganttWidth = 500,
+  openGroupIds,
+  task,
+  ganttStartDate,
+  ganttEndDate,
+}: Props) => {
   const {
     selected: { runId, taskId },
     onSelect,
   } = useSelection();
   const containerRef = useContainerRef();
 
-  const runDuration = getDuration(dagRun?.startDate, dagRun?.endDate);
+  const runDuration = getDuration(ganttStartDate, ganttEndDate);
 
   const instance = task.instances.find((ti) => ti.runId === runId);
   const isSelected = taskId === instance?.taskId;
@@ -54,7 +61,7 @@ const Row = ({ ganttWidth = 500, openGroupIds, task, dagRun }: Props) => {
     ? getDuration(instance?.queuedDttm, instance?.startDate)
     : 0;
   const taskStartOffset = getDuration(
-    dagRun.startDate,
+    ganttStartDate,
     instance?.queuedDttm || instance?.startDate
   );
 
@@ -127,7 +134,8 @@ const Row = ({ ganttWidth = 500, openGroupIds, task, dagRun }: Props) => {
           <Row
             ganttWidth={ganttWidth}
             openGroupIds={openGroupIds}
-            dagRun={dagRun}
+            ganttStartDate={ganttStartDate}
+            ganttEndDate={ganttEndDate}
             task={c}
             key={`gantt-${c.id}`}
           />
diff --git a/airflow/www/static/js/dag/details/gantt/index.tsx b/airflow/www/static/js/dag/details/gantt/index.tsx
index 0e4b7adb8480f..83200ac2ff386 100644
--- a/airflow/www/static/js/dag/details/gantt/index.tsx
+++ b/airflow/www/static/js/dag/details/gantt/index.tsx
@@ -106,10 +106,35 @@ const Gantt = ({ openGroupIds, gridScrollRef, ganttScrollRef }: Props) => {
 
   const dagRun = dagRuns.find((dr) => dr.runId === runId);
 
-  const startDate = dagRun?.startDate;
+  let startDate = dagRun?.queuedAt || dagRun?.startDate;
+  let endDate = dagRun?.endDate;
+
+  // Check if any task instance dates are outside the bounds of the dag run dates and update our min start and max end
+  groups.children?.forEach((task) => {
+    const taskInstance = task.instances.find((ti) => ti.runId === runId);
+    if (
+      taskInstance?.queuedDttm &&
+      (!startDate ||
+        Date.parse(taskInstance.queuedDttm) < Date.parse(startDate))
+    ) {
+      startDate = taskInstance.queuedDttm;
+    } else if (
+      taskInstance?.startDate &&
+      (!startDate || Date.parse(taskInstance.startDate) < Date.parse(startDate))
+    ) {
+      startDate = taskInstance.startDate;
+    }
+
+    if (
+      taskInstance?.endDate &&
+      (!endDate || Date.parse(taskInstance.endDate) > Date.parse(endDate))
+    ) {
+      endDate = taskInstance.endDate;
+    }
+  });
 
   const numBars = Math.round(width / 100);
-  const runDuration = getDuration(dagRun?.startDate, dagRun?.endDate);
+  const runDuration = getDuration(startDate, endDate);
   const intervals = runDuration / numBars;
 
   return (
@@ -166,8 +191,9 @@ const Gantt = ({ openGroupIds, gridScrollRef, ganttScrollRef }: Props) => {
               <Row
                 ganttWidth={width}
                 openGroupIds={openGroupIds}
-                dagRun={dagRun}
                 task={c}
+                ganttStartDate={startDate}
+                ganttEndDate={endDate}
                 key={`gantt-${c.id}`}
               />
             ))}

From 5e484c3e8ab4382d8fd5e1b27b27c98496c74c5d Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Wed, 9 Aug 2023 22:15:00 +0800
Subject: [PATCH 10/22] Add new enum in OnFinishAction, and revise test

---
 .../providers/cncf/kubernetes/operators/pod.py   | 16 ++++++++++++----
 .../cncf/kubernetes/utils/pod_manager.py         | 14 ++++----------
 .../cncf/kubernetes/example_kubernetes_istio.py  |  8 +++-----
 3 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index d72b2fec0c283..6cb07654e9fd4 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -233,7 +233,8 @@ class KubernetesPodOperator(BaseOperator):
     :param log_pod_spec_on_failure: Log the pod's specification if a failure occurs
     :param on_finish_action: What to do when the pod reaches its final state, or the execution is interrupted.
         If "delete_pod", the pod will be deleted regardless it's state; if "delete_succeeded_pod",
-        only succeeded pod will be deleted. You can set to "keep_pod" to keep the pod.
+        only succeeded pod will be deleted. Use "delete_succeeded_pod_with_istio" if istio sidecar injection
+        is enabled in your environment. You can set to "keep_pod" to keep the pod.
     :param is_delete_operator_pod: What to do when the pod reaches its final
         state, or the execution is interrupted. If True (default), delete the
         pod; if False, leave the pod.
@@ -751,9 +752,16 @@ def _read_pod_events(self, pod, *, reraise=True):
     def process_pod_deletion(self, pod: k8s.V1Pod, *, reraise=True):
         with _optionally_suppress(reraise=reraise):
             if pod is not None:
-                should_delete_pod = (self.on_finish_action == OnFinishAction.DELETE_POD) or (
-                    self.on_finish_action == OnFinishAction.DELETE_SUCCEEDED_POD
-                    and pod.status.phase == PodPhase.SUCCEEDED
+                should_delete_pod = (
+                    (self.on_finish_action == OnFinishAction.DELETE_POD)
+                    or (
+                        self.on_finish_action == OnFinishAction.DELETE_SUCCEEDED_POD
+                        and pod.status.phase == PodPhase.SUCCEEDED
+                    )
+                    or (
+                        self.on_finish_action == OnFinishAction.DELETE_SUCCEEDED_POD_WITH_ISTIO
+                        and self.pod_manager.container_is_succeeded(pod, self.base_container_name)
+                    )
                 )
                 if should_delete_pod:
                     self.log.info("Deleting pod: %s", pod.metadata.name)
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index 137d8efeb5390..acd6edfdeb8ad 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -278,11 +278,7 @@ class PodLoggingStatus:
 class PodManager(LoggingMixin):
     """Create, monitor, and otherwise interact with Kubernetes pods for use with the KubernetesPodOperator."""
 
-    def __init__(
-        self,
-        kube_client: client.CoreV1Api,
-        istio_enabled: bool = False,
-    ):
+    def __init__(self, kube_client: client.CoreV1Api):
         """
         Creates the launcher.
 
@@ -291,7 +287,6 @@ def __init__(
         super().__init__()
         self._client = kube_client
         self._watch = watch.Watch()
-        self._istio_enabled = istio_enabled
 
     def run_pod_async(self, pod: V1Pod, **kwargs) -> V1Pod:
         """Runs POD asynchronously."""
@@ -536,10 +531,8 @@ def await_pod_completion(self, pod: V1Pod) -> V1Pod:
             if remote_pod.status.phase in PodPhase.terminal_states:
                 break
 
-            if (
-                self._istio_enabled
-                and remote_pod.status.phase == PodPhase.RUNNING
-                and self.container_is_completed(remote_pod, "base")
+            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
+                remote_pod, "base"
             ):
                 break
             self.log.info("Pod %s has phase %s", pod.metadata.name, remote_pod.status.phase)
@@ -760,3 +753,4 @@ class OnFinishAction(enum.Enum):
     KEEP_POD = "keep_pod"
     DELETE_POD = "delete_pod"
     DELETE_SUCCEEDED_POD = "delete_succeeded_pod"
+    DELETE_SUCCEEDED_POD_WITH_ISTIO = "delete_succeeded_pod_with_istio"
diff --git a/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
index c0abe2f845462..eb17b582b1b9a 100644
--- a/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
+++ b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
@@ -122,12 +122,11 @@
         name="airflow-test-pod",
         task_id="task",
         affinity=affinity,
-        is_delete_operator_pod=True,
         hostnetwork=False,
         tolerations=tolerations,
         init_containers=[init_container],
         priority_class_name="medium",
-        istio_enabled=True,
+        on_finish_action="delete_succeeded_pod_with_istio",
     )
 
     # [START howto_operator_k8s_private_image]
@@ -139,11 +138,10 @@
         arguments=["echo", "10", "echo pwd"],
         labels={"foo": "bar"},
         name="airflow-private-image-pod",
-        is_delete_operator_pod=True,
         in_cluster=True,
         task_id="task-two",
         get_logs=True,
-        istio_enabled=True,
+        on_finish_action="delete_succeeded_pod_with_istio",
     )
     # [END howto_operator_k8s_private_image]
 
@@ -158,7 +156,7 @@
         in_cluster=True,
         task_id="write-xcom",
         get_logs=True,
-        istio_enabled=True,
+        on_finish_action="delete_succeeded_pod_with_istio",
     )
 
     pod_task_xcom_result = BashOperator(

From 83bd60fd97d4ca622adcbd7898d88880fee43054 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?=
 <6774676+eumiro@users.noreply.github.com>
Date: Wed, 9 Aug 2023 14:19:00 +0000
Subject: [PATCH 11/22] Refactor: Simplify code in providers/amazon (#33222)

---
 .../providers/amazon/aws/hooks/base_aws.py    |  7 +-
 .../amazon/aws/hooks/batch_client.py          | 91 +++++++++----------
 .../providers/amazon/aws/hooks/datasync.py    | 28 ++----
 .../amazon/aws/hooks/redshift_data.py         |  6 +-
 airflow/providers/amazon/aws/hooks/s3.py      |  7 +-
 .../providers/amazon/aws/hooks/sagemaker.py   | 31 +++----
 .../amazon/aws/log/s3_task_handler.py         |  7 +-
 .../amazon/aws/operators/redshift_cluster.py  | 18 ++--
 .../amazon/aws/secrets/secrets_manager.py     |  5 +-
 .../amazon/aws/transfers/redshift_to_s3.py    |  2 +-
 .../amazon/aws/transfers/s3_to_redshift.py    |  2 +-
 .../amazon/aws/transfers/sql_to_s3.py         |  2 +-
 .../providers/amazon/aws/triggers/batch.py    |  8 +-
 airflow/providers/amazon/aws/triggers/ecs.py  |  4 +-
 airflow/providers/amazon/aws/triggers/emr.py  |  4 +-
 .../amazon/aws/triggers/sagemaker.py          |  4 +-
 .../amazon/aws/utils/waiter_with_logging.py   | 30 +++---
 17 files changed, 112 insertions(+), 144 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py
index 01748089c8936..a908882e89ba2 100644
--- a/airflow/providers/amazon/aws/hooks/base_aws.py
+++ b/airflow/providers/amazon/aws/hooks/base_aws.py
@@ -494,10 +494,9 @@ def _find_class_name(target_function_name: str) -> str:
         responsible with catching and handling those.
         """
         stack = inspect.stack()
-        # Find the index of the most recent frame which called the provided function name.
-        target_frame_index = [frame.function for frame in stack].index(target_function_name)
-        # Pull that frame off the stack.
-        target_frame = stack[target_frame_index][0]
+        # Find the index of the most recent frame which called the provided function name
+        # and pull that frame off the stack.
+        target_frame = next(frame for frame in stack if frame.function == target_function_name)[0]
         # Get the local variables for that frame.
         frame_variables = target_frame.f_locals["self"]
         # Get the class object for that frame.
diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py
index e585297d3a624..87637be410a19 100644
--- a/airflow/providers/amazon/aws/hooks/batch_client.py
+++ b/airflow/providers/amazon/aws/hooks/batch_client.py
@@ -26,6 +26,7 @@
 """
 from __future__ import annotations
 
+import itertools as it
 from random import uniform
 from time import sleep
 from typing import Callable
@@ -343,8 +344,17 @@ def poll_job_status(self, job_id: str, match_status: list[str]) -> bool:
 
         :raises: AirflowException
         """
-        retries = 0
-        while True:
+        for retries in range(1 + self.max_retries):
+            if retries:
+                pause = self.exponential_delay(retries)
+                self.log.info(
+                    "AWS Batch job (%s) status check (%d of %d) in the next %.2f seconds",
+                    job_id,
+                    retries,
+                    self.max_retries,
+                    pause,
+                )
+                self.delay(pause)
 
             job = self.get_job_description(job_id)
             job_status = job.get("status")
@@ -354,23 +364,10 @@ def poll_job_status(self, job_id: str, match_status: list[str]) -> bool:
                 job_status,
                 match_status,
             )
-
             if job_status in match_status:
                 return True
-
-            if retries >= self.max_retries:
-                raise AirflowException(f"AWS Batch job ({job_id}) status checks exceed max_retries")
-
-            retries += 1
-            pause = self.exponential_delay(retries)
-            self.log.info(
-                "AWS Batch job (%s) status check (%d of %d) in the next %.2f seconds",
-                job_id,
-                retries,
-                self.max_retries,
-                pause,
-            )
-            self.delay(pause)
+        else:
+            raise AirflowException(f"AWS Batch job ({job_id}) status checks exceed max_retries")
 
     def get_job_description(self, job_id: str) -> dict:
         """
@@ -382,12 +379,21 @@ def get_job_description(self, job_id: str) -> dict:
 
         :raises: AirflowException
         """
-        retries = 0
-        while True:
+        for retries in range(self.status_retries):
+            if retries:
+                pause = self.exponential_delay(retries)
+                self.log.info(
+                    "AWS Batch job (%s) description retry (%d of %d) in the next %.2f seconds",
+                    job_id,
+                    retries,
+                    self.status_retries,
+                    pause,
+                )
+                self.delay(pause)
+
             try:
                 response = self.get_conn().describe_jobs(jobs=[job_id])
                 return self.parse_job_description(job_id, response)
-
             except botocore.exceptions.ClientError as err:
                 # Allow it to retry in case of exceeded quota limit of requests to AWS API
                 if err.response.get("Error", {}).get("Code") != "TooManyRequestsException":
@@ -398,23 +404,11 @@ def get_job_description(self, job_id: str) -> dict:
                     "check Amazon Provider AWS Connection documentation for more details.",
                     str(err),
                 )
-
-            retries += 1
-            if retries >= self.status_retries:
-                raise AirflowException(
-                    f"AWS Batch job ({job_id}) description error: exceeded status_retries "
-                    f"({self.status_retries})"
-                )
-
-            pause = self.exponential_delay(retries)
-            self.log.info(
-                "AWS Batch job (%s) description retry (%d of %d) in the next %.2f seconds",
-                job_id,
-                retries,
-                self.status_retries,
-                pause,
+        else:
+            raise AirflowException(
+                f"AWS Batch job ({job_id}) description error: exceeded status_retries "
+                f"({self.status_retries})"
             )
-            self.delay(pause)
 
     @staticmethod
     def parse_job_description(job_id: str, response: dict) -> dict:
@@ -476,7 +470,7 @@ def get_job_all_awslogs_info(self, job_id: str) -> list[dict[str, str]]:
             )
 
         # If the user selected another logDriver than "awslogs", then CloudWatch logging is disabled.
-        if any([c.get("logDriver", "awslogs") != "awslogs" for c in log_configs]):
+        if any(c.get("logDriver", "awslogs") != "awslogs" for c in log_configs):
             self.log.warning(
                 f"AWS Batch job ({job_id}) uses non-aws log drivers. AWS CloudWatch logging disabled."
             )
@@ -494,18 +488,17 @@ def get_job_all_awslogs_info(self, job_id: str) -> list[dict[str, str]]:
 
         # cross stream names with options (i.e. attempts X nodes) to generate all log infos
         result = []
-        for stream in stream_names:
-            for option in log_options:
-                result.append(
-                    {
-                        "awslogs_stream_name": stream,
-                        # If the user did not specify anything, the default settings are:
-                        #   awslogs-group = /aws/batch/job
-                        #   awslogs-region = `same as AWS Batch Job region`
-                        "awslogs_group": option.get("awslogs-group", "/aws/batch/job"),
-                        "awslogs_region": option.get("awslogs-region", self.conn_region_name),
-                    }
-                )
+        for stream, option in it.product(stream_names, log_options):
+            result.append(
+                {
+                    "awslogs_stream_name": stream,
+                    # If the user did not specify anything, the default settings are:
+                    #   awslogs-group = /aws/batch/job
+                    #   awslogs-region = `same as AWS Batch Job region`
+                    "awslogs_group": option.get("awslogs-group", "/aws/batch/job"),
+                    "awslogs_region": option.get("awslogs-region", self.conn_region_name),
+                }
+            )
         return result
 
     @staticmethod
diff --git a/airflow/providers/amazon/aws/hooks/datasync.py b/airflow/providers/amazon/aws/hooks/datasync.py
index 1fd06eeefe3dd..32fc7621531f3 100644
--- a/airflow/providers/amazon/aws/hooks/datasync.py
+++ b/airflow/providers/amazon/aws/hooks/datasync.py
@@ -125,17 +125,11 @@ def get_location_arns(
 
     def _refresh_locations(self) -> None:
         """Refresh the local list of Locations."""
-        self.locations = []
-        next_token = None
-        while True:
-            if next_token:
-                locations = self.get_conn().list_locations(NextToken=next_token)
-            else:
-                locations = self.get_conn().list_locations()
+        locations = self.get_conn().list_locations()
+        self.locations = locations["Locations"]
+        while "NextToken" in locations:
+            locations = self.get_conn().list_locations(NextToken=locations["NextToken"])
             self.locations.extend(locations["Locations"])
-            if "NextToken" not in locations:
-                break
-            next_token = locations["NextToken"]
 
     def create_task(
         self, source_location_arn: str, destination_location_arn: str, **create_task_kwargs
@@ -181,17 +175,11 @@ def delete_task(self, task_arn: str) -> None:
 
     def _refresh_tasks(self) -> None:
         """Refreshes the local list of Tasks."""
-        self.tasks = []
-        next_token = None
-        while True:
-            if next_token:
-                tasks = self.get_conn().list_tasks(NextToken=next_token)
-            else:
-                tasks = self.get_conn().list_tasks()
+        tasks = self.get_conn().list_tasks()
+        self.tasks = tasks["Tasks"]
+        while "NextToken" in tasks:
+            tasks = self.get_conn().list_tasks(NextToken=tasks["NextToken"])
             self.tasks.extend(tasks["Tasks"])
-            if "NextToken" not in tasks:
-                break
-            next_token = tasks["NextToken"]
 
     def get_task_arns_for_location_arns(
         self,
diff --git a/airflow/providers/amazon/aws/hooks/redshift_data.py b/airflow/providers/amazon/aws/hooks/redshift_data.py
index a522d3e8c9688..ba7a61c28f579 100644
--- a/airflow/providers/amazon/aws/hooks/redshift_data.py
+++ b/airflow/providers/amazon/aws/hooks/redshift_data.py
@@ -195,9 +195,9 @@ def get_table_primary_key(
             # we only select a single column (that is a string),
             # so safe to assume that there is only a single col in the record
             pk_columns += [y["stringValue"] for x in response["Records"] for y in x]
-            if "NextToken" not in response.keys():
-                break
-            else:
+            if "NextToken" in response:
                 token = response["NextToken"]
+            else:
+                break
 
         return pk_columns or None
diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
index 8878f04e1b582..f0e98ed448872 100644
--- a/airflow/providers/amazon/aws/hooks/s3.py
+++ b/airflow/providers/amazon/aws/hooks/s3.py
@@ -1284,18 +1284,15 @@ def delete_bucket(self, bucket_name: str, force_delete: bool = False, max_retrie
             bucket and trying to delete the bucket.
         :return: None
         """
-        tries_remaining = max_retries + 1
         if force_delete:
-            while tries_remaining:
+            for retry in range(max_retries):
                 bucket_keys = self.list_keys(bucket_name=bucket_name)
                 if not bucket_keys:
                     break
-                if tries_remaining <= max_retries:
-                    # Avoid first loop
+                if retry:  # Avoid first loop
                     sleep(500)
 
                 self.delete_objects(bucket=bucket_name, keys=bucket_keys)
-                tries_remaining -= 1
 
         self.conn.delete_bucket(Bucket=bucket_name)
 
diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py
index 758ba7e9c5d5c..b8247d9de97bb 100644
--- a/airflow/providers/amazon/aws/hooks/sagemaker.py
+++ b/airflow/providers/amazon/aws/hooks/sagemaker.py
@@ -1143,15 +1143,12 @@ def stop_pipeline(
         if check_interval is None:
             check_interval = 10
 
-        retries = 2  # i.e. 3 calls max, 1 initial + 2 retries
-        while True:
+        for retries in (2, 1, 0):
             try:
                 self.conn.stop_pipeline_execution(PipelineExecutionArn=pipeline_exec_arn)
-                break
             except ClientError as ce:
                 # this can happen if the pipeline was transitioning between steps at that moment
-                if ce.response["Error"]["Code"] == "ConflictException" and retries > 0:
-                    retries = retries - 1
+                if ce.response["Error"]["Code"] == "ConflictException" and retries:
                     self.log.warning(
                         "Got a conflict exception when trying to stop the pipeline, "
                         "retrying %s more times. Error was: %s",
@@ -1159,18 +1156,20 @@ def stop_pipeline(
                         ce,
                     )
                     time.sleep(0.3)  # error is due to a race condition, so it should be very transient
-                    continue
-                # we have to rely on the message to catch the right error here, because its type
-                # (ValidationException) is shared with other kinds of errors (e.g. badly formatted ARN)
-                if (
-                    not fail_if_not_running
-                    and "Only pipelines with 'Executing' status can be stopped"
-                    in ce.response["Error"]["Message"]
-                ):
-                    self.log.warning("Cannot stop pipeline execution, as it was not running: %s", ce)
                 else:
-                    self.log.error(ce)
-                    raise
+                    # we have to rely on the message to catch the right error here, because its type
+                    # (ValidationException) is shared with other kinds of errors (e.g. badly formatted ARN)
+                    if (
+                        not fail_if_not_running
+                        and "Only pipelines with 'Executing' status can be stopped"
+                        in ce.response["Error"]["Message"]
+                    ):
+                        self.log.warning("Cannot stop pipeline execution, as it was not running: %s", ce)
+                        break
+                    else:
+                        self.log.error(ce)
+                        raise
+            else:
                 break
 
         res = self.describe_pipeline_exec(pipeline_exec_arn)
diff --git a/airflow/providers/amazon/aws/log/s3_task_handler.py b/airflow/providers/amazon/aws/log/s3_task_handler.py
index a353d9b07ed87..4a813715fdd81 100644
--- a/airflow/providers/amazon/aws/log/s3_task_handler.py
+++ b/airflow/providers/amazon/aws/log/s3_task_handler.py
@@ -122,9 +122,10 @@ def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[list[str], l
         bucket, prefix = self.hook.parse_s3_url(s3url=os.path.join(self.remote_base, worker_log_rel_path))
         keys = self.hook.list_keys(bucket_name=bucket, prefix=prefix)
         if keys:
-            keys = [f"s3://{bucket}/{key}" for key in keys]
-            messages.extend(["Found logs in s3:", *[f"  * {x}" for x in sorted(keys)]])
-            for key in sorted(keys):
+            keys = sorted(f"s3://{bucket}/{key}" for key in keys)
+            messages.append("Found logs in s3:")
+            messages.extend(f"  * {key}" for key in keys)
+            for key in keys:
                 logs.append(self.s3_read(key, return_error=True))
         else:
             messages.append(f"No logs found on s3 for ti={ti}")
diff --git a/airflow/providers/amazon/aws/operators/redshift_cluster.py b/airflow/providers/amazon/aws/operators/redshift_cluster.py
index ed6aa79e9f2dd..bdce036c312a4 100644
--- a/airflow/providers/amazon/aws/operators/redshift_cluster.py
+++ b/airflow/providers/amazon/aws/operators/redshift_cluster.py
@@ -492,14 +492,14 @@ def __init__(
     def execute(self, context: Context):
         redshift_hook = RedshiftHook(aws_conn_id=self.aws_conn_id)
         self.log.info("Starting resume cluster")
-        while self._remaining_attempts >= 1:
+        while self._remaining_attempts:
             try:
                 redshift_hook.get_conn().resume_cluster(ClusterIdentifier=self.cluster_identifier)
                 break
             except redshift_hook.get_conn().exceptions.InvalidClusterStateFault as error:
-                self._remaining_attempts = self._remaining_attempts - 1
+                self._remaining_attempts -= 1
 
-                if self._remaining_attempts > 0:
+                if self._remaining_attempts:
                     self.log.error(
                         "Unable to resume cluster. %d attempts remaining.", self._remaining_attempts
                     )
@@ -580,14 +580,14 @@ def __init__(
 
     def execute(self, context: Context):
         redshift_hook = RedshiftHook(aws_conn_id=self.aws_conn_id)
-        while self._remaining_attempts >= 1:
+        while self._remaining_attempts:
             try:
                 redshift_hook.get_conn().pause_cluster(ClusterIdentifier=self.cluster_identifier)
                 break
             except redshift_hook.get_conn().exceptions.InvalidClusterStateFault as error:
-                self._remaining_attempts = self._remaining_attempts - 1
+                self._remaining_attempts -= 1
 
-                if self._remaining_attempts > 0:
+                if self._remaining_attempts:
                     self.log.error(
                         "Unable to pause cluster. %d attempts remaining.", self._remaining_attempts
                     )
@@ -669,7 +669,7 @@ def __init__(
         self.max_attempts = max_attempts
 
     def execute(self, context: Context):
-        while self._attempts >= 1:
+        while self._attempts:
             try:
                 self.redshift_hook.delete_cluster(
                     cluster_identifier=self.cluster_identifier,
@@ -678,9 +678,9 @@ def execute(self, context: Context):
                 )
                 break
             except self.redshift_hook.get_conn().exceptions.InvalidClusterStateFault:
-                self._attempts = self._attempts - 1
+                self._attempts -= 1
 
-                if self._attempts > 0:
+                if self._attempts:
                     self.log.error("Unable to delete cluster. %d attempts remaining.", self._attempts)
                     time.sleep(self._attempt_interval)
                 else:
diff --git a/airflow/providers/amazon/aws/secrets/secrets_manager.py b/airflow/providers/amazon/aws/secrets/secrets_manager.py
index a5acf19a372e8..f33f328c00218 100644
--- a/airflow/providers/amazon/aws/secrets/secrets_manager.py
+++ b/airflow/providers/amazon/aws/secrets/secrets_manager.py
@@ -217,10 +217,7 @@ def _standardize_secret_keys(self, secret: dict[str, Any]) -> dict[str, Any]:
 
         conn_d: dict[str, Any] = {}
         for conn_field, possible_words in possible_words_for_conn_fields.items():
-            try:
-                conn_d[conn_field] = [v for k, v in secret.items() if k in possible_words][0]
-            except IndexError:
-                conn_d[conn_field] = None
+            conn_d[conn_field] = next((v for k, v in secret.items() if k in possible_words), None)
 
         return conn_d
 
diff --git a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
index 45f1e1783c49d..6b6b4890c093a 100644
--- a/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
+++ b/airflow/providers/amazon/aws/transfers/redshift_to_s3.py
@@ -137,7 +137,7 @@ def __init__(
 
         if self.redshift_data_api_kwargs:
             for arg in ["sql", "parameters"]:
-                if arg in self.redshift_data_api_kwargs.keys():
+                if arg in self.redshift_data_api_kwargs:
                     raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")
 
     def _build_unload_query(
diff --git a/airflow/providers/amazon/aws/transfers/s3_to_redshift.py b/airflow/providers/amazon/aws/transfers/s3_to_redshift.py
index b80b5bed3276d..6bedb092b42b5 100644
--- a/airflow/providers/amazon/aws/transfers/s3_to_redshift.py
+++ b/airflow/providers/amazon/aws/transfers/s3_to_redshift.py
@@ -116,7 +116,7 @@ def __init__(
 
         if self.redshift_data_api_kwargs:
             for arg in ["sql", "parameters"]:
-                if arg in self.redshift_data_api_kwargs.keys():
+                if arg in self.redshift_data_api_kwargs:
                     raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")
 
     def _build_copy_query(
diff --git a/airflow/providers/amazon/aws/transfers/sql_to_s3.py b/airflow/providers/amazon/aws/transfers/sql_to_s3.py
index 0324406820e2f..1b327def97c93 100644
--- a/airflow/providers/amazon/aws/transfers/sql_to_s3.py
+++ b/airflow/providers/amazon/aws/transfers/sql_to_s3.py
@@ -194,7 +194,7 @@ def _partition_dataframe(self, df: DataFrame) -> Iterable[tuple[str, DataFrame]]
             yield "", df
         else:
             grouped_df = df.groupby(**self.groupby_kwargs)
-            for group_label in grouped_df.groups.keys():
+            for group_label in grouped_df.groups:
                 yield group_label, grouped_df.get_group(group_label).reset_index(drop=True)
 
     def _get_hook(self) -> DbApiHook:
diff --git a/airflow/providers/amazon/aws/triggers/batch.py b/airflow/providers/amazon/aws/triggers/batch.py
index 8966540a3849f..40b230221fc08 100644
--- a/airflow/providers/amazon/aws/triggers/batch.py
+++ b/airflow/providers/amazon/aws/triggers/batch.py
@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import asyncio
+import itertools as it
 from functools import cached_property
 from typing import Any
 
@@ -161,9 +162,7 @@ async def run(self):
         """
         async with self.hook.async_conn as client:
             waiter = self.hook.get_waiter("batch_job_complete", deferrable=True, client=client)
-            attempt = 0
-            while True:
-                attempt = attempt + 1
+            for attempt in it.count(1):
                 try:
                     await waiter.wait(
                         jobs=[self.job_id],
@@ -172,7 +171,6 @@ async def run(self):
                             "MaxAttempts": 1,
                         },
                     )
-                    break
                 except WaiterError as error:
                     if "error" in str(error):
                         yield TriggerEvent({"status": "failure", "message": f"Job Failed: {error}"})
@@ -183,6 +181,8 @@ async def run(self):
                         attempt,
                     )
                     await asyncio.sleep(int(self.poke_interval))
+                else:
+                    break
 
             yield TriggerEvent(
                 {
diff --git a/airflow/providers/amazon/aws/triggers/ecs.py b/airflow/providers/amazon/aws/triggers/ecs.py
index cc203207d158b..4a380d23f0fa1 100644
--- a/airflow/providers/amazon/aws/triggers/ecs.py
+++ b/airflow/providers/amazon/aws/triggers/ecs.py
@@ -165,8 +165,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
             # fmt: on
             waiter = ecs_client.get_waiter("tasks_stopped")
             logs_token = None
-            while self.waiter_max_attempts >= 1:
-                self.waiter_max_attempts = self.waiter_max_attempts - 1
+            while self.waiter_max_attempts:
+                self.waiter_max_attempts -= 1
                 try:
                     await waiter.wait(
                         cluster=self.cluster, tasks=[self.task_arn], WaiterConfig={"MaxAttempts": 1}
diff --git a/airflow/providers/amazon/aws/triggers/emr.py b/airflow/providers/amazon/aws/triggers/emr.py
index d7d9844af33fb..af8fef5f7bcf3 100644
--- a/airflow/providers/amazon/aws/triggers/emr.py
+++ b/airflow/providers/amazon/aws/triggers/emr.py
@@ -70,10 +70,8 @@ async def run(self):
         self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
         async with self.hook.async_conn as client:
             for step_id in self.step_ids:
-                attempt = 0
                 waiter = client.get_waiter("step_complete")
-                while attempt < int(self.max_attempts):
-                    attempt += 1
+                for attempt in range(1, 1 + self.max_attempts):
                     try:
                         await waiter.wait(
                             ClusterId=self.job_flow_id,
diff --git a/airflow/providers/amazon/aws/triggers/sagemaker.py b/airflow/providers/amazon/aws/triggers/sagemaker.py
index ec11323d426ee..7165276c4c915 100644
--- a/airflow/providers/amazon/aws/triggers/sagemaker.py
+++ b/airflow/providers/amazon/aws/triggers/sagemaker.py
@@ -164,12 +164,10 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
     }
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
-        attempts = 0
         hook = SageMakerHook(aws_conn_id=self.aws_conn_id)
         async with hook.async_conn as conn:
             waiter = hook.get_waiter(self._waiter_name[self.waiter_type], deferrable=True, client=conn)
-            while attempts < self.waiter_max_attempts:
-                attempts = attempts + 1
+            for _ in range(self.waiter_max_attempts):
                 try:
                     await waiter.wait(
                         PipelineExecutionArn=self.pipeline_execution_arn, WaiterConfig={"MaxAttempts": 1}
diff --git a/airflow/providers/amazon/aws/utils/waiter_with_logging.py b/airflow/providers/amazon/aws/utils/waiter_with_logging.py
index 1e506927a1f43..2c0d403426530 100644
--- a/airflow/providers/amazon/aws/utils/waiter_with_logging.py
+++ b/airflow/providers/amazon/aws/utils/waiter_with_logging.py
@@ -63,22 +63,21 @@ def wait(
         status_args = ["Clusters[0].state", "Clusters[0].details"]
     """
     log = logging.getLogger(__name__)
-    attempt = 0
-    while True:
-        attempt += 1
+    for attempt in range(waiter_max_attempts):
+        if attempt:
+            time.sleep(waiter_delay)
         try:
             waiter.wait(**args, WaiterConfig={"MaxAttempts": 1})
-            break
         except WaiterError as error:
             if "terminal failure" in str(error):
                 log.error("%s: %s", failure_message, _LazyStatusFormatter(status_args, error.last_response))
                 raise AirflowException(f"{failure_message}: {error}")
 
             log.info("%s: %s", status_message, _LazyStatusFormatter(status_args, error.last_response))
-            if attempt >= waiter_max_attempts:
-                raise AirflowException("Waiter error: max attempts reached")
-
-            time.sleep(waiter_delay)
+        else:
+            break
+    else:
+        raise AirflowException("Waiter error: max attempts reached")
 
 
 async def async_wait(
@@ -115,22 +114,21 @@ async def async_wait(
         status_args = ["Clusters[0].state", "Clusters[0].details"]
     """
     log = logging.getLogger(__name__)
-    attempt = 0
-    while True:
-        attempt += 1
+    for attempt in range(waiter_max_attempts):
+        if attempt:
+            await asyncio.sleep(waiter_delay)
         try:
             await waiter.wait(**args, WaiterConfig={"MaxAttempts": 1})
-            break
         except WaiterError as error:
             if "terminal failure" in str(error):
                 log.error("%s: %s", failure_message, _LazyStatusFormatter(status_args, error.last_response))
                 raise AirflowException(f"{failure_message}: {error}")
 
             log.info("%s: %s", status_message, _LazyStatusFormatter(status_args, error.last_response))
-            if attempt >= waiter_max_attempts:
-                raise AirflowException("Waiter error: max attempts reached")
-
-            await asyncio.sleep(waiter_delay)
+        else:
+            break
+    else:
+        raise AirflowException("Waiter error: max attempts reached")
 
 
 class _LazyStatusFormatter:

From ee767a4240b78d53617157ee941e4d4d048504a8 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Wed, 9 Aug 2023 22:43:39 +0800
Subject: [PATCH 12/22] Modify await_pod_completion logic

---
 .../providers/cncf/kubernetes/operators/pod.py |  4 ++--
 .../cncf/kubernetes/utils/pod_manager.py       | 18 +++++++-----------
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index 6cb07654e9fd4..a9a80058a0130 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -601,7 +601,7 @@ def execute_sync(self, context: Context):
             if self.do_xcom_push:
                 self.pod_manager.await_xcom_sidecar_container_start(pod=self.pod)
                 result = self.extract_xcom(pod=self.pod)
-            self.remote_pod = self.pod_manager.await_pod_completion(self.pod)
+            self.remote_pod = self.pod_manager.await_pod_completion(self.pod, self.base_container_name)
         finally:
             self.cleanup(
                 pod=self.pod or self.pod_request_obj,
@@ -664,7 +664,7 @@ def execute_complete(self, context: Context, event: dict, **kwargs):
                     xcom_sidecar_output = self.extract_xcom(pod=pod)
                     return xcom_sidecar_output
         finally:
-            pod = self.pod_manager.await_pod_completion(pod)
+            pod = self.pod_manager.await_pod_completion(pod, self.base_container_name)
             if pod is not None:
                 self.post_complete_action(
                     pod=pod,
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index acd6edfdeb8ad..b9666c81a9978 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -519,7 +519,7 @@ def await_container_completion(self, pod: V1Pod, container_name: str) -> None:
             self.log.info("Waiting for container '%s' state to be completed", container_name)
             time.sleep(1)
 
-    def await_pod_completion(self, pod: V1Pod) -> V1Pod:
+    def await_pod_completion(self, pod: V1Pod, container_name: str) -> V1Pod:
         """
         Monitors a pod and returns the final state.
 
@@ -531,8 +531,8 @@ def await_pod_completion(self, pod: V1Pod) -> V1Pod:
             if remote_pod.status.phase in PodPhase.terminal_states:
                 break
 
-            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
-                remote_pod, "base"
+            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_succeeded(
+                remote_pod, container_name
             ):
                 break
             self.log.info("Pod %s has phase %s", pod.metadata.name, remote_pod.status.phase)
@@ -573,15 +573,9 @@ def container_is_terminated(self, pod: V1Pod, container_name: str) -> bool:
         remote_pod = self.read_pod(pod)
         return container_is_terminated(pod=remote_pod, container_name=container_name)
 
-    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is completed."""
-        remote_pod = self.read_pod(pod)
-        return container_is_completed(pod=remote_pod, container_name=container_name)
-
     def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
         """Reads pod and checks if container is succeeded."""
-        remote_pod = self.read_pod(pod)
-        return container_is_succeeded(pod=remote_pod, container_name=container_name)
+        return container_is_succeeded(pod=pod, container_name=container_name)
 
     @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_exponential(), reraise=True)
     def read_pod_logs(
@@ -693,7 +687,9 @@ def extract_xcom_json(self, pod: V1Pod) -> str:
         ) as resp:
             result = self._exec_pod_command(
                 resp,
-                f"if [ -s {PodDefaults.XCOM_MOUNT_PATH}/return.json ]; then cat {PodDefaults.XCOM_MOUNT_PATH}/return.json; else echo __airflow_xcom_result_empty__; fi",  # noqa
+                f"if [ -s {PodDefaults.XCOM_MOUNT_PATH}/return.json ]; "
+                f"then cat {PodDefaults.XCOM_MOUNT_PATH}/return.json; "
+                f"else echo __airflow_xcom_result_empty__; fi",
             )
             if result and result.rstrip() != "__airflow_xcom_result_empty__":
                 # Note: result string is parsed to check if its valid json.

From eb07618d29ce957ef4401a98f25449a23d0135ad Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Wed, 9 Aug 2023 17:41:06 +0200
Subject: [PATCH 13/22] Fix missing related PRs when cherry-picking (#33261)

The `airflow-github` tooling had a check if PR has been already
merged, but the PR did not work well when there was already a PR
cherry-picked that referred the PR in question in their log message.

The check was done via `git log --grep=#PR` essentially and so
it happenes that `--grep` matches not only the first line of
the message but the whole log message. So if the log message
contained reference to the cherry-picked PR it was seen as merged.

This change adds ( ) around and $ to the regexp to match not
only the PR number but also parentheses around.
---
 dev/airflow-github | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev/airflow-github b/dev/airflow-github
index 33e543a77f3b0..115ed48d19169 100755
--- a/dev/airflow-github
+++ b/dev/airflow-github
@@ -96,7 +96,7 @@ def get_commit_in_main_associated_with_pr(repo: git.Repo, issue: Issue) -> str |
 
 def is_cherrypicked(repo: git.Repo, issue: Issue, previous_version: str | None = None) -> bool:
     """Check if a given issue is cherry-picked in the current branch or not"""
-    log_args = ["--format=%H", f"--grep=#{issue.number}"]
+    log_args = ["--format=%H", f"--grep=(#{issue.number})$"]
     if previous_version:
         log_args.append(previous_version + "..")
     log = repo.git.log(*log_args)

From b5b717f8fe3b09110c139e99eb9816e665e896df Mon Sep 17 00:00:00 2001
From: Joshua <joshua.chyeung@gmail.com>
Date: Wed, 28 Sep 2022 12:13:16 +0800
Subject: [PATCH 14/22] Added IstioKubernetesPodOperator and IstioPodManager

---
 .../operators/istio_kubernetes_pod.py         | 58 +++++++++++++++++
 .../kubernetes/utils/istio_pod_manager.py     | 65 +++++++++++++++++++
 .../cncf/kubernetes/utils/pod_manager.py      | 30 +++++++++
 3 files changed, 153 insertions(+)
 create mode 100644 airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
 create mode 100644 airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py

diff --git a/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
new file mode 100644
index 0000000000000..ab6401f65cd3a
--- /dev/null
+++ b/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Executes task in a Kubernetes POD with Istio enabled"""
+from __future__ import annotations
+
+from kubernetes.client import models as k8s
+
+from airflow import AirflowException
+from airflow.compat.functools import cached_property
+from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator, _suppress
+from airflow.providers.cncf.kubernetes.utils.istio_pod_manager import IstioPodManager
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
+
+
+class IstioKubernetesPodOperator(KubernetesPodOperator):
+    """
+    Execute a task in a Kubernetes Pod with Istio enabled
+    All parameters are the same with KubernetesPodOperator, the only difference is the cleanup part
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KubernetesPodOperator`
+    """
+
+    @cached_property
+    def pod_manager(self) -> IstioPodManager:
+        return IstioPodManager(kube_client=self.client)
+
+    def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
+        pod_phase = remote_pod.status.phase if hasattr(remote_pod, 'status') else None
+        if pod_phase != PodPhase.SUCCEEDED and not self.pod_manager.container_is_succeeded(pod, 'base'):
+            if self.log_events_on_failure:
+                with _suppress(Exception):
+                    for event in self.pod_manager.read_pod_events(pod).items:
+                        self.log.error("Pod Event: %s - %s", event.reason, event.message)
+            if not self.is_delete_operator_pod:
+                with _suppress(Exception):
+                    self.patch_already_checked(pod)
+            with _suppress(Exception):
+                self.process_pod_deletion(pod)
+            raise AirflowException(f'Pod {pod and pod.metadata.name} returned a failure: {remote_pod}')
+        else:
+            with _suppress(Exception):
+                self.process_pod_deletion(pod)
diff --git a/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py b/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
new file mode 100644
index 0000000000000..2248d565181f0
--- /dev/null
+++ b/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import time
+
+from kubernetes.client.models.v1_pod import V1Pod
+
+from airflow.providers.cncf.kubernetes.utils.pod_manager import (
+    PodManager,
+    PodPhase,
+    container_is_completed,
+    container_is_succeeded,
+)
+
+
+class IstioPodManager(PodManager):
+    """
+    Helper class for checking the status of base container in Kubernetes pods
+    for use with the IstioKubernetesPodOperator
+    """
+
+    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is completed"""
+        remote_pod = self.read_pod(pod)
+        return container_is_completed(pod=remote_pod, container_name=container_name)
+
+    def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is succeeded"""
+        remote_pod = self.read_pod(pod)
+        return container_is_succeeded(pod=remote_pod, container_name=container_name)
+
+    def await_pod_completion(self, pod: V1Pod) -> V1Pod:
+        """
+        Monitors a pod and returns the final state
+        (neglect sidecar state e.g. istio-proxy, vault-agent)
+        :param pod: pod spec that will be monitored
+        :return:  Tuple[State, Optional[str]]
+        """
+        while True:
+            remote_pod = self.read_pod(pod)
+            if remote_pod.status.phase in PodPhase.terminal_states:
+                break
+            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
+                remote_pod, 'base'
+            ):
+                break
+            self.log.info('Pod %s has phase %s', pod.metadata.name, remote_pod.status.phase)
+            time.sleep(2)
+        return remote_pod
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index aff81a1f18d1e..c6d01b9055295 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -160,6 +160,36 @@ def container_is_terminated(pod: V1Pod, container_name: str) -> bool:
     return container_status.state.terminated is not None
 
 
+def container_is_completed(pod: V1Pod, container_name: str) -> bool:
+    """
+    Examines V1Pod ``pod`` to determine whether ``container_name`` is completed.
+    If that container is present and completed, returns True.  Returns False otherwise.
+    """
+    container_statuses = pod.status.container_statuses if pod and pod.status else None
+    if not container_statuses:
+        return False
+    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+    if not container_status:
+        return False
+    return container_status.state.terminated is not None
+
+
+def container_is_succeeded(pod: V1Pod, container_name: str) -> bool:
+    """
+    Examines V1Pod ``pod`` to determine whether ``container_name`` is completed and succeeded.
+    If that container is present and completed and succeeded, returns True.  Returns False otherwise.
+    """
+    if not container_is_completed(pod, container_name):
+        return False
+    container_statuses = pod.status.container_statuses if pod and pod.status else None
+    if not container_statuses:
+        return False
+    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+    if not container_status:
+        return False
+    return container_status.state.terminated.exit_code == 0
+
+
 def get_container_termination_message(pod: V1Pod, container_name: str):
     with suppress(AttributeError, TypeError):
         container_statuses = pod.status.container_statuses

From ff43f86238a95d31149096e0f977db0529c9e65e Mon Sep 17 00:00:00 2001
From: Joshua <joshua.chyeung@gmail.com>
Date: Fri, 21 Oct 2022 14:31:03 +0800
Subject: [PATCH 15/22] Merged into KubernetesPodOperator(istio_enabled=True)

---
 .../operators/istio_kubernetes_pod.py         | 58 -----------------
 .../kubernetes/utils/istio_pod_manager.py     | 65 -------------------
 2 files changed, 123 deletions(-)
 delete mode 100644 airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
 delete mode 100644 airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py

diff --git a/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
deleted file mode 100644
index ab6401f65cd3a..0000000000000
--- a/airflow/providers/cncf/kubernetes/operators/istio_kubernetes_pod.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Executes task in a Kubernetes POD with Istio enabled"""
-from __future__ import annotations
-
-from kubernetes.client import models as k8s
-
-from airflow import AirflowException
-from airflow.compat.functools import cached_property
-from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator, _suppress
-from airflow.providers.cncf.kubernetes.utils.istio_pod_manager import IstioPodManager
-from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
-
-
-class IstioKubernetesPodOperator(KubernetesPodOperator):
-    """
-    Execute a task in a Kubernetes Pod with Istio enabled
-    All parameters are the same with KubernetesPodOperator, the only difference is the cleanup part
-
-    .. seealso::
-        For more information on how to use this operator, take a look at the guide:
-        :ref:`howto/operator:KubernetesPodOperator`
-    """
-
-    @cached_property
-    def pod_manager(self) -> IstioPodManager:
-        return IstioPodManager(kube_client=self.client)
-
-    def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
-        pod_phase = remote_pod.status.phase if hasattr(remote_pod, 'status') else None
-        if pod_phase != PodPhase.SUCCEEDED and not self.pod_manager.container_is_succeeded(pod, 'base'):
-            if self.log_events_on_failure:
-                with _suppress(Exception):
-                    for event in self.pod_manager.read_pod_events(pod).items:
-                        self.log.error("Pod Event: %s - %s", event.reason, event.message)
-            if not self.is_delete_operator_pod:
-                with _suppress(Exception):
-                    self.patch_already_checked(pod)
-            with _suppress(Exception):
-                self.process_pod_deletion(pod)
-            raise AirflowException(f'Pod {pod and pod.metadata.name} returned a failure: {remote_pod}')
-        else:
-            with _suppress(Exception):
-                self.process_pod_deletion(pod)
diff --git a/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py b/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
deleted file mode 100644
index 2248d565181f0..0000000000000
--- a/airflow/providers/cncf/kubernetes/utils/istio_pod_manager.py
+++ /dev/null
@@ -1,65 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-from __future__ import annotations
-
-import time
-
-from kubernetes.client.models.v1_pod import V1Pod
-
-from airflow.providers.cncf.kubernetes.utils.pod_manager import (
-    PodManager,
-    PodPhase,
-    container_is_completed,
-    container_is_succeeded,
-)
-
-
-class IstioPodManager(PodManager):
-    """
-    Helper class for checking the status of base container in Kubernetes pods
-    for use with the IstioKubernetesPodOperator
-    """
-
-    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is completed"""
-        remote_pod = self.read_pod(pod)
-        return container_is_completed(pod=remote_pod, container_name=container_name)
-
-    def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is succeeded"""
-        remote_pod = self.read_pod(pod)
-        return container_is_succeeded(pod=remote_pod, container_name=container_name)
-
-    def await_pod_completion(self, pod: V1Pod) -> V1Pod:
-        """
-        Monitors a pod and returns the final state
-        (neglect sidecar state e.g. istio-proxy, vault-agent)
-        :param pod: pod spec that will be monitored
-        :return:  Tuple[State, Optional[str]]
-        """
-        while True:
-            remote_pod = self.read_pod(pod)
-            if remote_pod.status.phase in PodPhase.terminal_states:
-                break
-            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
-                remote_pod, 'base'
-            ):
-                break
-            self.log.info('Pod %s has phase %s', pod.metadata.name, remote_pod.status.phase)
-            time.sleep(2)
-        return remote_pod

From 2b3dfb2f846a8bda387df73b1a83db55b5e51bd8 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Thu, 18 May 2023 22:19:17 +0800
Subject: [PATCH 16/22] Add support to istio

Continuing from joshua's work to integrate istio support into KubernetesPodOperator
---
 .../cncf/kubernetes/utils/pod_manager.py      | 30 ++++++++++++++-----
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index c6d01b9055295..4a3ca40a428ae 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -165,10 +165,7 @@ def container_is_completed(pod: V1Pod, container_name: str) -> bool:
     Examines V1Pod ``pod`` to determine whether ``container_name`` is completed.
     If that container is present and completed, returns True.  Returns False otherwise.
     """
-    container_statuses = pod.status.container_statuses if pod and pod.status else None
-    if not container_statuses:
-        return False
-    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+    container_status = get_container_status(pod, container_name)
     if not container_status:
         return False
     return container_status.state.terminated is not None
@@ -181,10 +178,8 @@ def container_is_succeeded(pod: V1Pod, container_name: str) -> bool:
     """
     if not container_is_completed(pod, container_name):
         return False
-    container_statuses = pod.status.container_statuses if pod and pod.status else None
-    if not container_statuses:
-        return False
-    container_status = next(iter([x for x in container_statuses if x.name == container_name]), None)
+
+    container_status = get_container_status(pod, container_name)
     if not container_status:
         return False
     return container_status.state.terminated.exit_code == 0
@@ -297,6 +292,7 @@ class PodManager(LoggingMixin):
     def __init__(
         self,
         kube_client: client.CoreV1Api,
+        istio_enabled: bool = False,
     ):
         """
         Creates the launcher.
@@ -306,6 +302,7 @@ def __init__(
         super().__init__()
         self._client = kube_client
         self._watch = watch.Watch()
+        self._istio_enabled = istio_enabled
 
     def run_pod_async(self, pod: V1Pod, **kwargs) -> V1Pod:
         """Runs POD asynchronously."""
@@ -549,6 +546,13 @@ def await_pod_completion(self, pod: V1Pod) -> V1Pod:
             remote_pod = self.read_pod(pod)
             if remote_pod.status.phase in PodPhase.terminal_states:
                 break
+
+            if (
+                self._istio_enabled
+                and remote_pod.status.phase == PodPhase.RUNNING
+                and self.container_is_completed(remote_pod, "base")
+            ):
+                break
             self.log.info("Pod %s has phase %s", pod.metadata.name, remote_pod.status.phase)
             time.sleep(2)
         return remote_pod
@@ -587,6 +591,16 @@ def container_is_terminated(self, pod: V1Pod, container_name: str) -> bool:
         remote_pod = self.read_pod(pod)
         return container_is_terminated(pod=remote_pod, container_name=container_name)
 
+    def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is completed"""
+        remote_pod = self.read_pod(pod)
+        return container_is_completed(pod=remote_pod, container_name=container_name)
+
+    def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
+        """Reads pod and checks if container is succeeded"""
+        remote_pod = self.read_pod(pod)
+        return container_is_succeeded(pod=remote_pod, container_name=container_name)
+
     @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_exponential(), reraise=True)
     def read_pod_logs(
         self,

From a0099e56a5fde1d61db49398ca0e1a81d487daa4 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Fri, 19 May 2023 01:01:51 +0800
Subject: [PATCH 17/22] Add istio into spelling word

Fix failing spell check
---
 docs/spelling_wordlist.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index fbca728f8ed22..a2347070c16da 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -786,6 +786,7 @@ isfile
 ish
 isn
 isort
+istio
 iterable
 iterables
 iteratively

From ccdb8797106e43e9d5d47fa075999ae23035f220 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Sat, 20 May 2023 20:52:26 +0800
Subject: [PATCH 18/22] add system test dag with istio enabled

Add a new system test dag with istio enabled
---
 .../kubernetes/example_kubernetes_istio.py    | 181 ++++++++++++++++++
 1 file changed, 181 insertions(+)
 create mode 100644 tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py

diff --git a/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
new file mode 100644
index 0000000000000..c0abe2f845462
--- /dev/null
+++ b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
@@ -0,0 +1,181 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+This is an example dag for using the KubernetesPodOperator.
+"""
+from __future__ import annotations
+
+import os
+from datetime import datetime
+
+from kubernetes.client import models as k8s
+
+from airflow import DAG
+from airflow.kubernetes.secret import Secret
+from airflow.operators.bash import BashOperator
+from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
+
+# [START howto_operator_k8s_cluster_resources]
+secret_file = Secret("volume", "/etc/sql_conn", "airflow-secrets", "sql_alchemy_conn")
+secret_env = Secret("env", "SQL_CONN", "airflow-secrets", "sql_alchemy_conn")
+secret_all_keys = Secret("env", None, "airflow-secrets-2")
+volume_mount = k8s.V1VolumeMount(
+    name="test-volume", mount_path="/root/mount_file", sub_path=None, read_only=True
+)
+
+configmaps = [
+    k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name="test-configmap-1")),
+    k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource(name="test-configmap-2")),
+]
+
+volume = k8s.V1Volume(
+    name="test-volume",
+    persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name="test-volume"),
+)
+
+port = k8s.V1ContainerPort(name="http", container_port=80)
+
+init_container_volume_mounts = [
+    k8s.V1VolumeMount(mount_path="/etc/foo", name="test-volume", sub_path=None, read_only=True)
+]
+
+init_environments = [k8s.V1EnvVar(name="key1", value="value1"), k8s.V1EnvVar(name="key2", value="value2")]
+
+init_container = k8s.V1Container(
+    name="init-container",
+    image="ubuntu:16.04",
+    env=init_environments,
+    volume_mounts=init_container_volume_mounts,
+    command=["bash", "-cx"],
+    args=["echo 10"],
+)
+
+affinity = k8s.V1Affinity(
+    node_affinity=k8s.V1NodeAffinity(
+        preferred_during_scheduling_ignored_during_execution=[
+            k8s.V1PreferredSchedulingTerm(
+                weight=1,
+                preference=k8s.V1NodeSelectorTerm(
+                    match_expressions=[
+                        k8s.V1NodeSelectorRequirement(key="disktype", operator="In", values=["ssd"])
+                    ]
+                ),
+            )
+        ]
+    ),
+    pod_affinity=k8s.V1PodAffinity(
+        required_during_scheduling_ignored_during_execution=[
+            k8s.V1WeightedPodAffinityTerm(
+                weight=1,
+                pod_affinity_term=k8s.V1PodAffinityTerm(
+                    label_selector=k8s.V1LabelSelector(
+                        match_expressions=[
+                            k8s.V1LabelSelectorRequirement(key="security", operator="In", values="S1")
+                        ]
+                    ),
+                    topology_key="failure-domain.beta.kubernetes.io/zone",
+                ),
+            )
+        ]
+    ),
+)
+
+tolerations = [k8s.V1Toleration(key="key", operator="Equal", value="value")]
+
+# [END howto_operator_k8s_cluster_resources]
+
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_kubernetes_operator_istio"
+
+with DAG(
+    dag_id="example_kubernetes_operator_istio",
+    schedule=None,
+    start_date=datetime(2021, 1, 1),
+    tags=["example"],
+) as dag:
+    k = KubernetesPodOperator(
+        namespace="default",
+        image="ubuntu:16.04",
+        cmds=["bash", "-cx"],
+        arguments=["echo", "10"],
+        labels={"foo": "bar"},
+        secrets=[secret_file, secret_env, secret_all_keys],
+        ports=[port],
+        volumes=[volume],
+        volume_mounts=[volume_mount],
+        env_from=configmaps,
+        name="airflow-test-pod",
+        task_id="task",
+        affinity=affinity,
+        is_delete_operator_pod=True,
+        hostnetwork=False,
+        tolerations=tolerations,
+        init_containers=[init_container],
+        priority_class_name="medium",
+        istio_enabled=True,
+    )
+
+    # [START howto_operator_k8s_private_image]
+    quay_k8s = KubernetesPodOperator(
+        namespace="default",
+        image="quay.io/apache/bash",
+        image_pull_secrets=[k8s.V1LocalObjectReference("testquay")],
+        cmds=["bash", "-cx"],
+        arguments=["echo", "10", "echo pwd"],
+        labels={"foo": "bar"},
+        name="airflow-private-image-pod",
+        is_delete_operator_pod=True,
+        in_cluster=True,
+        task_id="task-two",
+        get_logs=True,
+        istio_enabled=True,
+    )
+    # [END howto_operator_k8s_private_image]
+
+    # [START howto_operator_k8s_write_xcom]
+    write_xcom = KubernetesPodOperator(
+        namespace="default",
+        image="alpine",
+        cmds=["sh", "-c", "mkdir -p /airflow/xcom/;echo '[1,2,3,4]' > /airflow/xcom/return.json"],
+        name="write-xcom",
+        do_xcom_push=True,
+        is_delete_operator_pod=True,
+        in_cluster=True,
+        task_id="write-xcom",
+        get_logs=True,
+        istio_enabled=True,
+    )
+
+    pod_task_xcom_result = BashOperator(
+        bash_command="echo \"{{ task_instance.xcom_pull('write-xcom')[0] }}\"",
+        task_id="pod_task_xcom_result",
+    )
+
+    write_xcom >> pod_task_xcom_result
+    # [END howto_operator_k8s_write_xcom]
+
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)

From aec6669ec2a00a267556c0cc0713d9d3c5ad75dd Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Wed, 9 Aug 2023 20:52:56 +0800
Subject: [PATCH 19/22] Remove duplicated function

---
 .../providers/cncf/kubernetes/utils/pod_manager.py    | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index 4a3ca40a428ae..9d38279924f26 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -160,17 +160,6 @@ def container_is_terminated(pod: V1Pod, container_name: str) -> bool:
     return container_status.state.terminated is not None
 
 
-def container_is_completed(pod: V1Pod, container_name: str) -> bool:
-    """
-    Examines V1Pod ``pod`` to determine whether ``container_name`` is completed.
-    If that container is present and completed, returns True.  Returns False otherwise.
-    """
-    container_status = get_container_status(pod, container_name)
-    if not container_status:
-        return False
-    return container_status.state.terminated is not None
-
-
 def container_is_succeeded(pod: V1Pod, container_name: str) -> bool:
     """
     Examines V1Pod ``pod`` to determine whether ``container_name`` is completed and succeeded.

From fff686ed4ccb3f59b1db3823c11a205e527172d6 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Wed, 9 Aug 2023 22:15:00 +0800
Subject: [PATCH 20/22] Add new enum in OnFinishAction, and revise test

---
 .../providers/cncf/kubernetes/operators/pod.py   | 16 ++++++++++++----
 .../cncf/kubernetes/utils/pod_manager.py         | 14 ++++----------
 .../cncf/kubernetes/example_kubernetes_istio.py  |  8 +++-----
 3 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index d72b2fec0c283..6cb07654e9fd4 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -233,7 +233,8 @@ class KubernetesPodOperator(BaseOperator):
     :param log_pod_spec_on_failure: Log the pod's specification if a failure occurs
     :param on_finish_action: What to do when the pod reaches its final state, or the execution is interrupted.
         If "delete_pod", the pod will be deleted regardless it's state; if "delete_succeeded_pod",
-        only succeeded pod will be deleted. You can set to "keep_pod" to keep the pod.
+        only succeeded pod will be deleted. Use "delete_succeeded_pod_with_istio" if istio sidecar injection
+        is enabled in your environment. You can set to "keep_pod" to keep the pod.
     :param is_delete_operator_pod: What to do when the pod reaches its final
         state, or the execution is interrupted. If True (default), delete the
         pod; if False, leave the pod.
@@ -751,9 +752,16 @@ def _read_pod_events(self, pod, *, reraise=True):
     def process_pod_deletion(self, pod: k8s.V1Pod, *, reraise=True):
         with _optionally_suppress(reraise=reraise):
             if pod is not None:
-                should_delete_pod = (self.on_finish_action == OnFinishAction.DELETE_POD) or (
-                    self.on_finish_action == OnFinishAction.DELETE_SUCCEEDED_POD
-                    and pod.status.phase == PodPhase.SUCCEEDED
+                should_delete_pod = (
+                    (self.on_finish_action == OnFinishAction.DELETE_POD)
+                    or (
+                        self.on_finish_action == OnFinishAction.DELETE_SUCCEEDED_POD
+                        and pod.status.phase == PodPhase.SUCCEEDED
+                    )
+                    or (
+                        self.on_finish_action == OnFinishAction.DELETE_SUCCEEDED_POD_WITH_ISTIO
+                        and self.pod_manager.container_is_succeeded(pod, self.base_container_name)
+                    )
                 )
                 if should_delete_pod:
                     self.log.info("Deleting pod: %s", pod.metadata.name)
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index 9d38279924f26..14c09bf6684d6 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -278,11 +278,7 @@ class PodLoggingStatus:
 class PodManager(LoggingMixin):
     """Create, monitor, and otherwise interact with Kubernetes pods for use with the KubernetesPodOperator."""
 
-    def __init__(
-        self,
-        kube_client: client.CoreV1Api,
-        istio_enabled: bool = False,
-    ):
+    def __init__(self, kube_client: client.CoreV1Api):
         """
         Creates the launcher.
 
@@ -291,7 +287,6 @@ def __init__(
         super().__init__()
         self._client = kube_client
         self._watch = watch.Watch()
-        self._istio_enabled = istio_enabled
 
     def run_pod_async(self, pod: V1Pod, **kwargs) -> V1Pod:
         """Runs POD asynchronously."""
@@ -536,10 +531,8 @@ def await_pod_completion(self, pod: V1Pod) -> V1Pod:
             if remote_pod.status.phase in PodPhase.terminal_states:
                 break
 
-            if (
-                self._istio_enabled
-                and remote_pod.status.phase == PodPhase.RUNNING
-                and self.container_is_completed(remote_pod, "base")
+            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
+                remote_pod, "base"
             ):
                 break
             self.log.info("Pod %s has phase %s", pod.metadata.name, remote_pod.status.phase)
@@ -760,3 +753,4 @@ class OnFinishAction(enum.Enum):
     KEEP_POD = "keep_pod"
     DELETE_POD = "delete_pod"
     DELETE_SUCCEEDED_POD = "delete_succeeded_pod"
+    DELETE_SUCCEEDED_POD_WITH_ISTIO = "delete_succeeded_pod_with_istio"
diff --git a/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
index c0abe2f845462..eb17b582b1b9a 100644
--- a/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
+++ b/tests/system/providers/cncf/kubernetes/example_kubernetes_istio.py
@@ -122,12 +122,11 @@
         name="airflow-test-pod",
         task_id="task",
         affinity=affinity,
-        is_delete_operator_pod=True,
         hostnetwork=False,
         tolerations=tolerations,
         init_containers=[init_container],
         priority_class_name="medium",
-        istio_enabled=True,
+        on_finish_action="delete_succeeded_pod_with_istio",
     )
 
     # [START howto_operator_k8s_private_image]
@@ -139,11 +138,10 @@
         arguments=["echo", "10", "echo pwd"],
         labels={"foo": "bar"},
         name="airflow-private-image-pod",
-        is_delete_operator_pod=True,
         in_cluster=True,
         task_id="task-two",
         get_logs=True,
-        istio_enabled=True,
+        on_finish_action="delete_succeeded_pod_with_istio",
     )
     # [END howto_operator_k8s_private_image]
 
@@ -158,7 +156,7 @@
         in_cluster=True,
         task_id="write-xcom",
         get_logs=True,
-        istio_enabled=True,
+        on_finish_action="delete_succeeded_pod_with_istio",
     )
 
     pod_task_xcom_result = BashOperator(

From c2e0a222ce06accbed0b9968645ab684af178b7b Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Wed, 9 Aug 2023 22:43:39 +0800
Subject: [PATCH 21/22] Modify await_pod_completion logic

---
 airflow/providers/cncf/kubernetes/operators/pod.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index 6cb07654e9fd4..a9a80058a0130 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -601,7 +601,7 @@ def execute_sync(self, context: Context):
             if self.do_xcom_push:
                 self.pod_manager.await_xcom_sidecar_container_start(pod=self.pod)
                 result = self.extract_xcom(pod=self.pod)
-            self.remote_pod = self.pod_manager.await_pod_completion(self.pod)
+            self.remote_pod = self.pod_manager.await_pod_completion(self.pod, self.base_container_name)
         finally:
             self.cleanup(
                 pod=self.pod or self.pod_request_obj,
@@ -664,7 +664,7 @@ def execute_complete(self, context: Context, event: dict, **kwargs):
                     xcom_sidecar_output = self.extract_xcom(pod=pod)
                     return xcom_sidecar_output
         finally:
-            pod = self.pod_manager.await_pod_completion(pod)
+            pod = self.pod_manager.await_pod_completion(pod, self.base_container_name)
             if pod is not None:
                 self.post_complete_action(
                     pod=pod,

From 5cfb4011acc2120e9fbe981faa7695da91ece6b5 Mon Sep 17 00:00:00 2001
From: Owen Leung <owen.leung2@gmail.com>
Date: Thu, 10 Aug 2023 15:39:31 +0800
Subject: [PATCH 22/22] Fix failing helm test

---
 .../cncf/kubernetes/operators/pod.py          | 54 +++++++++++++++++--
 .../cncf/kubernetes/utils/pod_manager.py      | 14 +++--
 2 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
index a9a80058a0130..c0c481b80e2f8 100644
--- a/airflow/providers/cncf/kubernetes/operators/pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -22,6 +22,7 @@
 import logging
 import re
 import secrets
+import shlex
 import string
 import warnings
 from collections.abc import Container
@@ -29,7 +30,8 @@
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Iterable, Sequence
 
-from kubernetes.client import CoreV1Api, models as k8s
+from kubernetes.client import CoreV1Api, V1Pod, models as k8s
+from kubernetes.stream import stream
 from slugify import slugify
 from urllib3.exceptions import HTTPError
 
@@ -601,7 +603,11 @@ def execute_sync(self, context: Context):
             if self.do_xcom_push:
                 self.pod_manager.await_xcom_sidecar_container_start(pod=self.pod)
                 result = self.extract_xcom(pod=self.pod)
-            self.remote_pod = self.pod_manager.await_pod_completion(self.pod, self.base_container_name)
+
+            istio_enabled = is_istio_enabled(self.client, self.pod)
+            self.remote_pod = self.pod_manager.await_pod_completion(
+                istio_enabled, self.pod, self.base_container_name
+            )
         finally:
             self.cleanup(
                 pod=self.pod or self.pod_request_obj,
@@ -664,9 +670,12 @@ def execute_complete(self, context: Context, event: dict, **kwargs):
                     xcom_sidecar_output = self.extract_xcom(pod=pod)
                     return xcom_sidecar_output
         finally:
-            pod = self.pod_manager.await_pod_completion(pod, self.base_container_name)
+            istio_enabled = is_istio_enabled(self.client, pod)
+
+            pod = self.pod_manager.await_pod_completion(istio_enabled, pod, self.base_container_name)
             if pod is not None:
                 self.post_complete_action(
+                    istio_enabled=istio_enabled,
                     pod=pod,
                     remote_pod=pod,
                 )
@@ -750,6 +759,7 @@ def _read_pod_events(self, pod, *, reraise=True):
                 self.log.error("Pod Event: %s - %s", event.reason, event.message)
 
     def process_pod_deletion(self, pod: k8s.V1Pod, *, reraise=True):
+        istio_enabled = is_istio_enabled(self.client, pod)
         with _optionally_suppress(reraise=reraise):
             if pod is not None:
                 should_delete_pod = (
@@ -763,9 +773,30 @@ def process_pod_deletion(self, pod: k8s.V1Pod, *, reraise=True):
                         and self.pod_manager.container_is_succeeded(pod, self.base_container_name)
                     )
                 )
-                if should_delete_pod:
+                if should_delete_pod and not istio_enabled:
                     self.log.info("Deleting pod: %s", pod.metadata.name)
                     self.pod_manager.delete_pod(pod)
+                elif should_delete_pod and istio_enabled:
+                    command = "/bin/sh -c curl -fsI -X POST http://localhost:15020/quitquitquit && exit 0"
+                    command_to_container = shlex.split(command)
+                    self.log.info("Deleting istio-proxy sidecar inside %s: ", pod.metadata.name)
+                    try:
+                        resp = stream(
+                            self.client.connect_get_namespaced_pod_exec,
+                            name=pod.metadata.name,
+                            namespace=pod.metadata.namespace,
+                            container="istio-proxy",
+                            command=command_to_container,
+                            stderr=True,
+                            stdin=True,
+                            stdout=True,
+                            tty=False,
+                            _preload_content=False,
+                        )
+                        resp.close()
+                    except Exception as e:
+                        self.log.error("Error while deleting istio-proxy sidecar: %s", e)
+                        raise e
                 else:
                     self.log.info("Skipping deleting pod: %s", pod.metadata.name)
 
@@ -958,3 +989,18 @@ def __exit__(self, exctype, excinst, exctb):
             return True
         else:
             return True
+
+
+def is_istio_enabled(client: CoreV1Api, pod: V1Pod) -> bool:
+    """Checks if istio is enabled for the namespace of the pod by inspecting the namespace labels."""
+    if not pod:
+        return False
+
+    namespace_labels = client.read_namespace(name=pod.metadata.namespace).metadata.labels
+
+    istio_enabled = False
+    for key, value in namespace_labels.items():
+        if key == "istio-injection" and value == "enabled":
+            istio_enabled = True
+
+    return istio_enabled
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
index 14c09bf6684d6..03d6b9f3969b3 100644
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -519,11 +519,13 @@ def await_container_completion(self, pod: V1Pod, container_name: str) -> None:
             self.log.info("Waiting for container '%s' state to be completed", container_name)
             time.sleep(1)
 
-    def await_pod_completion(self, pod: V1Pod) -> V1Pod:
+    def await_pod_completion(self, istio_enabled: bool, pod: V1Pod, container_name: str) -> V1Pod:
         """
         Monitors a pod and returns the final state.
 
+        :param istio_enabled: whether istio is enabled in the namespace,
         :param pod: pod spec that will be monitored
+        :param container_name: name of the base container
         :return: tuple[State, str | None]
         """
         while True:
@@ -531,8 +533,10 @@ def await_pod_completion(self, pod: V1Pod) -> V1Pod:
             if remote_pod.status.phase in PodPhase.terminal_states:
                 break
 
-            if remote_pod.status.phase == PodPhase.RUNNING and self.container_is_completed(
-                remote_pod, "base"
+            if (
+                istio_enabled
+                and remote_pod.status.phase == PodPhase.RUNNING
+                and self.container_is_completed(remote_pod, container_name)
             ):
                 break
             self.log.info("Pod %s has phase %s", pod.metadata.name, remote_pod.status.phase)
@@ -574,12 +578,12 @@ def container_is_terminated(self, pod: V1Pod, container_name: str) -> bool:
         return container_is_terminated(pod=remote_pod, container_name=container_name)
 
     def container_is_completed(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is completed"""
+        """Reads pod and checks if container is completed."""
         remote_pod = self.read_pod(pod)
         return container_is_completed(pod=remote_pod, container_name=container_name)
 
     def container_is_succeeded(self, pod: V1Pod, container_name: str) -> bool:
-        """Reads pod and checks if container is succeeded"""
+        """Reads pod and checks if container is succeeded."""
         remote_pod = self.read_pod(pod)
         return container_is_succeeded(pod=remote_pod, container_name=container_name)
 
