feat(kubernetes): check state before deferring KubernetesPodOperatorAsync

Author: Lee-W

astronomer/providers/cncf/kubernetes/operators/kubernetes_pod.py

@@ -5,6 +5,10 @@
 from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import (
     KubernetesPodOperator,
 )
+from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import (
+    ContainerState,
+    KubernetesPodTrigger,
+)
 from kubernetes.client import models as k8s
 from pendulum import DateTime
 
@@ -50,6 +54,21 @@ def raise_for_trigger_status(event: Dict[str, Any]) -> None:
             else:
                 raise AirflowException(description)
 
+    def define_container_state(self) -> ContainerState:
+        """Define the state of container"""
+        pod_containers = self.pod.status.container_statuses
+        if not pod_containers:
+            return ContainerState.UNDEFINED
+        container = [c for c in pod_containers if c.name == self.BASE_CONTAINER_NAME][0]
+        for state in (ContainerState.RUNNING, ContainerState.WAITING, ContainerState.TERMINATED):
+            state_obj = getattr(container.state, state)
+            if state_obj:
+                if state != ContainerState.TERMINATED:
+                    return state
+                else:
+                    return ContainerState.TERMINATED if state_obj.exit_code == 0 else ContainerState.FAILED
+        return ContainerState.UNDEFINED
+
     def defer(self, last_log_time: Optional[DateTime] = None, **kwargs: Any) -> None:
         """Defers to ``WaitContainerTrigger`` optionally with last log time."""
         if kwargs:
@@ -76,10 +95,30 @@ def defer(self, last_log_time: Optional[DateTime] = None, **kwargs: Any) -> None
             method_name=self.trigger_reentry.__name__,
         )
 
-    def execute(self, context: Context) -> None:  # noqa: D102
+    def execute(self, context: Context) -> Any:  # noqa: D102
         self.pod_request_obj = self.build_pod_request_obj(context)
         self.pod: k8s.V1Pod = self.get_or_create_pod(self.pod_request_obj, context)
-        self.defer()
+        container_state = self.define_container_state()
+        pod_status = self.pod.status.phase
+        if KubernetesPodTrigger.should_wait(pod_status, container_state):
+            self.defer()
+            return
+
+        if container_state == ContainerState.TERMINATED:
+            event = {
+                "name": self.pod.metadata.name,
+                "namespace": self.pod.metadata.namespace,
+                "status": "success",
+                "message": "All containers inside pod have started successfully.",
+            }
+        else:
+            event = {
+                "name": self.pod.metadata.name,
+                "namespace": self.pod.metadata.namespace,
+                "status": "failed",
+                "message": self.pod.status.message,
+            }
+        return self.trigger_reentry(context=context, event=event)
 
     def execute_complete(self, context: Context, event: Dict[str, Any]) -> Any:  # type: ignore[override]
         """Deprecated; replaced by trigger_reentry."""

tests/cncf/kubernetes/operators/test_kubernetes_pod.py

@@ -3,7 +3,9 @@
 
 import pytest
 from airflow.exceptions import TaskDeferred
+from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import ContainerState
 from airflow.providers.cncf.kubernetes.utils.pod_manager import PodLoggingStatus
+from kubernetes.client import models as k8s
 
 from astronomer.providers.cncf.kubernetes.operators.kubernetes_pod import (
     KubernetesPodOperatorAsync,
@@ -17,6 +19,24 @@
 KUBE_POD_MOD = "astronomer.providers.cncf.kubernetes.operators.kubernetes_pod"
 
 
+def _build_mock_pod(state: k8s.V1ContainerState) -> k8s.V1Pod:
+    return k8s.V1Pod(
+        metadata=k8s.V1ObjectMeta(name="base", namespace="default"),
+        status=k8s.V1PodStatus(
+            container_statuses=[
+                k8s.V1ContainerStatus(
+                    name="base",
+                    image="alpine",
+                    image_id="1",
+                    ready=True,
+                    restart_count=1,
+                    state=state,
+                )
+            ]
+        ),
+    )
+
+
 class TestKubernetesPodOperatorAsync:
     def test_raise_for_trigger_status_pending_timeout(self):
         """Assert trigger raise exception in case of timeout"""
@@ -152,12 +172,79 @@ def test_defer_with_kwargs(self):
         with pytest.raises(ValueError):
             op.defer(kwargs={"timeout": 10})
 
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.trigger_reentry")
+    @mock.patch(
+        f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.define_container_state",
+        return_value=ContainerState.FAILED,
+    )
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.build_pod_request_obj")
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.get_or_create_pod")
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.defer")
+    def test_execute_failed_before_defer(
+        self,
+        mock_defer,
+        mock_get_or_create_pod,
+        mock_build_pod_request_obj,
+        mock_define_container_state,
+        mock_trigger_reentry,
+    ):
+        mock_get_or_create_pod.return_value = _build_mock_pod(
+            k8s.V1ContainerState({"running": k8s.V1ContainerStateTerminated(exit_code=1), "waiting": None})
+        )
+        mock_build_pod_request_obj.return_value = {}
+        mock_defer.return_value = {}
+        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
+
+        op.execute(context=create_context(op))
+        assert mock_trigger_reentry.called
+        assert not mock_defer.called
+
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.trigger_reentry")
+    @mock.patch(
+        f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.define_container_state",
+        return_value=ContainerState.TERMINATED,
+    )
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.build_pod_request_obj")
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.get_or_create_pod")
+    @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.defer")
+    def test_execute_succeeded_before_defer(
+        self,
+        mock_defer,
+        mock_get_or_create_pod,
+        mock_build_pod_request_obj,
+        mock_define_container_state,
+        mock_trigger_reentry,
+    ):
+        mock_get_or_create_pod.return_value = _build_mock_pod(
+            k8s.V1ContainerState({"running": k8s.V1ContainerStateTerminated(exit_code=0), "waiting": None})
+        )
+        mock_build_pod_request_obj.return_value = {}
+        mock_defer.return_value = {}
+        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
+        assert op.execute(context=create_context(op))
+        assert mock_trigger_reentry.called
+        assert not mock_defer.called
+
+    @mock.patch(
+        "astronomer.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperatorAsync.define_container_state",
+        return_value=ContainerState.RUNNING,
+    )
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.build_pod_request_obj")
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.get_or_create_pod")
     @mock.patch(f"{KUBE_POD_MOD}.KubernetesPodOperatorAsync.defer")
-    def test_execute(self, mock_defer, mock_get_or_create_pod, mock_build_pod_request_obj):
+    def test_execute(
+        self,
+        mock_defer,
+        mock_get_or_create_pod,
+        mock_build_pod_request_obj,
+        mock_define_container_state,
+    ):
         """Assert that execute succeeded"""
-        mock_get_or_create_pod.return_value = {}
+        mock_get_or_create_pod.return_value = _build_mock_pod(
+            k8s.V1ContainerState(
+                {"running": k8s.V1ContainerStateRunning(), "terminated": None, "waiting": None}
+            )
+        )
         mock_build_pod_request_obj.return_value = {}
         mock_defer.return_value = {}
         op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
@@ -169,3 +256,40 @@ def test_execute_complete(self, mock_trigger_reentry):
         mock_trigger_reentry.return_value = {}
         op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
         assert op.execute_complete(context=create_context(op), event={}) is None
+
+    @pytest.mark.parametrize(
+        "container_state, expected_state",
+        [
+            (
+                {"running": k8s.V1ContainerStateRunning(), "terminated": None, "waiting": None},
+                ContainerState.RUNNING,
+            ),
+            (
+                {"running": None, "terminated": k8s.V1ContainerStateTerminated(exit_code=0), "waiting": None},
+                ContainerState.TERMINATED,
+            ),
+            (
+                {"running": None, "terminated": None, "waiting": k8s.V1ContainerStateWaiting()},
+                ContainerState.WAITING,
+            ),
+        ],
+    )
+    def test_define_container_state_should_execute_successfully(self, container_state, expected_state):
+        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
+        op.pod = _build_mock_pod(k8s.V1ContainerState(**container_state))
+        assert expected_state == op.define_container_state()
+
+    @pytest.mark.parametrize(
+        "pod",
+        (
+            _build_mock_pod(k8s.V1ContainerState(running=None, terminated=None, waiting=None)),
+            k8s.V1Pod(
+                metadata=k8s.V1ObjectMeta(name="base", namespace="default"),
+                status=k8s.V1PodStatus(container_statuses=[]),
+            ),
+        ),
+    )
+    def test_define_container_state_with_undefined_state(self, pod):
+        op = KubernetesPodOperatorAsync(task_id="test_task", name="test-pod", get_logs=True)
+        op.pod = pod
+        assert op.define_container_state() == ContainerState.UNDEFINED