From 1c24555b99771c0615b47e60da0755056e38e940 Mon Sep 17 00:00:00 2001 From: John Horan Date: Fri, 21 Feb 2025 10:02:41 +0000 Subject: [PATCH 01/15] test print --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 1 + 1 file changed, 1 insertion(+) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 9c387ec3d3254..7142792df8ba2 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -635,6 +635,7 @@ def extract_xcom(self, pod: k8s.V1Pod) -> dict[Any, Any] | None: def execute(self, context: Context): """Based on the deferrable parameter runs the pod asynchronously or synchronously.""" + self.log.info("Hello World") self.name = self._set_name(self.name) if not self.deferrable: return self.execute_sync(context) From 2e1ae848ac995196391243f9310cc396e24319e5 Mon Sep 17 00:00:00 2001 From: John Horan Date: Fri, 21 Feb 2025 10:16:46 +0000 Subject: [PATCH 02/15] uv test --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 7142792df8ba2..aa5978320ad57 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -635,7 +635,7 @@ def extract_xcom(self, pod: k8s.V1Pod) -> dict[Any, Any] | None: def execute(self, context: Context): """Based on the deferrable parameter runs the pod asynchronously or synchronously.""" - self.log.info("Hello World") + self.log.info("Hello World!") self.name = self._set_name(self.name) if not self.deferrable: return self.execute_sync(context) From 3fcb8c992ee198a7421d606f311bedef596ca71d Mon Sep 17 00:00:00 2001 From: John Horan Date: Mon, 24 Feb 2025 13:43:15 +0000 Subject: [PATCH 03/15] add defer callbacks --- .../cncf/kubernetes/operators/pod.py | 65 +++++++++++++------ 1 file changed, 44 insertions(+), 21 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index aa5978320ad57..ded730a05f546 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -798,6 +798,14 @@ def _refresh_cached_properties(self): def execute_async(self, context: Context) -> None: self.pod_request_obj = self.build_pod_request_obj(context) + for callback in self.callbacks: + callback.on_pod_manifest_created( + pod_request=self.pod_request_obj, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) self.pod = self.get_or_create_pod( # must set `self.pod` for `on_kill` pod_request_obj=self.pod_request_obj, context=context, @@ -887,20 +895,37 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: follow = self.logging_interval is None last_log_time = event.get("last_log_time") - if event["status"] in ("error", "failed", "timeout"): - event_message = event.get("message", "No message provided") - self.log.error( - "Trigger emitted an %s event, failing the task: %s", event["status"], event_message - ) - # fetch some logs when pod is failed + if event["status"] in ("error", "failed", "timeout", "success"): if self.get_logs: self._write_logs(self.pod, follow=follow, since_time=last_log_time) - if self.do_xcom_push: - _ = self.extract_xcom(pod=self.pod) - - message = event.get("stack_trace", event["message"]) - raise AirflowException(message) + if self.callbacks: + pod = self.find_pod(self.pod.metadata.namespace, context=context) + for callback in self.callbacks: + callback.on_pod_completion( + pod=pod, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) + for callback in self.callbacks: + callback.on_pod_teardown( + pod=pod, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) + + xcom_sidecar_output = self.extract_xcom(pod=self.pod) if self.do_xcom_push else None + + if event["status"] == "success": + return xcom_sidecar_output + else: + self.log.error("Trigger emitted an %s event, failing the task: %s", event["status"], event["message"]) + message = event.get("stack_trace", event["message"]) + raise AirflowException(message) if event["status"] == "running": if self.get_logs: @@ -916,20 +941,18 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: self.invoke_defer_method(pod_log_status.last_log_time) else: self.invoke_defer_method() - - elif event["status"] == "success": - # fetch some logs when pod is executed successfully - if self.get_logs: - self._write_logs(self.pod, follow=follow, since_time=last_log_time) - - if self.do_xcom_push: - xcom_sidecar_output = self.extract_xcom(pod=self.pod) - return xcom_sidecar_output - return except TaskDeferred: raise finally: self._clean(event, context) + for callback in self.callbacks: + callback.on_pod_cleanup( + pod=self.pod or self.pod_request_obj, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) def _clean(self, event: dict[str, Any], context: Context) -> None: if event["status"] == "running": From 316528228e00a9e3641bef8a17e01e3de76a7a23 Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 25 Feb 2025 10:55:28 +0000 Subject: [PATCH 04/15] cleanup --- .../cncf/kubernetes/operators/pod.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index ded730a05f546..075398a70096e 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -635,7 +635,6 @@ def extract_xcom(self, pod: k8s.V1Pod) -> dict[Any, Any] | None: def execute(self, context: Context): """Based on the deferrable parameter runs the pod asynchronously or synchronously.""" - self.log.info("Hello World!") self.name = self._set_name(self.name) if not self.deferrable: return self.execute_sync(context) @@ -920,13 +919,13 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: xcom_sidecar_output = self.extract_xcom(pod=self.pod) if self.do_xcom_push else None - if event["status"] == "success": - return xcom_sidecar_output - else: + if event["status"] != "success": self.log.error("Trigger emitted an %s event, failing the task: %s", event["status"], event["message"]) message = event.get("stack_trace", event["message"]) raise AirflowException(message) + return xcom_sidecar_output + if event["status"] == "running": if self.get_logs: self.log.info("Resuming logs read from time %r", last_log_time) @@ -945,14 +944,15 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: raise finally: self._clean(event, context) - for callback in self.callbacks: - callback.on_pod_cleanup( - pod=self.pod or self.pod_request_obj, - client=self.client, - mode=ExecutionMode.SYNC, - context=context, - operator=self, - ) + if self.pod and event["status"] in ("error", "failed", "timeout", "success"): + for callback in self.callbacks: + callback.on_pod_cleanup( + pod=self.pod or self.pod_request_obj, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) def _clean(self, event: dict[str, Any], context: Context) -> None: if event["status"] == "running": From e336e19b1a230e4283db88ae21215d2983765fa1 Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 25 Feb 2025 12:34:57 +0000 Subject: [PATCH 05/15] remove redundant call --- .../cncf/kubernetes/operators/pod.py | 28 +++---------------- 1 file changed, 4 insertions(+), 24 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 075398a70096e..c407d52482db8 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -718,20 +718,7 @@ def execute_sync(self, context: Context): ) finally: pod_to_clean = self.pod or self.pod_request_obj - self.cleanup( - pod=pod_to_clean, - remote_pod=self.remote_pod, - xcom_result=result, - context=context, - ) - for callback in self.callbacks: - callback.on_pod_cleanup( - pod=pod_to_clean, - client=self.client, - mode=ExecutionMode.SYNC, - context=context, - operator=self, - ) + self.post_complete_action(pod=pod_to_clean, remote_pod=self.remote_pod, context=context, result=result) if self.do_xcom_push: return result @@ -944,15 +931,6 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: raise finally: self._clean(event, context) - if self.pod and event["status"] in ("error", "failed", "timeout", "success"): - for callback in self.callbacks: - callback.on_pod_cleanup( - pod=self.pod or self.pod_request_obj, - client=self.client, - mode=ExecutionMode.SYNC, - context=context, - operator=self, - ) def _clean(self, event: dict[str, Any], context: Context) -> None: if event["status"] == "running": @@ -1007,11 +985,13 @@ def _write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime e if not isinstance(e, ApiException) else e.reason, ) - def post_complete_action(self, *, pod, remote_pod, context: Context, **kwargs) -> None: + def post_complete_action(self, *, pod, remote_pod, context: Context, result: dict, **kwargs) -> None: """Actions that must be done after operator finishes logic of the deferrable_execution.""" self.cleanup( pod=pod, remote_pod=remote_pod, + xcom_result=result, + context=context, ) for callback in self.callbacks: callback.on_pod_cleanup( From 3cb67a02620456d770f353e46702af04b57ba589 Mon Sep 17 00:00:00 2001 From: John Horan Date: Wed, 26 Feb 2025 14:42:01 +0000 Subject: [PATCH 06/15] move test to async --- .../cncf/kubernetes/operators/test_pod.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index 37af22925d1b0..57256b8d51298 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -2541,6 +2541,64 @@ def test_trigger_error(self, find_pod, cleanup, mock_write_log): ) + @patch(HOOK_CLASS, new=MagicMock) + def test_execute_async_callbacks(self): + from airflow.providers.cncf.kubernetes.callbacks import ExecutionMode + from unit.cncf.kubernetes.test_callbacks import ( + MockKubernetesPodOperatorCallback, + MockWrapper, + ) + + MockWrapper.reset() + mock_callbacks = MockWrapper.mock_callbacks + remote_pod_mock = MagicMock() + remote_pod_mock.status.phase = "Succeeded" + self.await_pod_mock.return_value = remote_pod_mock + + k = KubernetesPodOperator( + namespace="default", + image="ubuntu:16.04", + cmds=["bash", "-cx"], + arguments=["echo 10"], + labels={"foo": "bar"}, + name="test", + task_id="task", + do_xcom_push=False, + callbacks=MockKubernetesPodOperatorCallback, + ) + context = create_context(k) + + k.trigger_reentry( + context=context, + event={ + "status": "success", + "message": TEST_SUCCESS_MESSAGE, + "name": TEST_NAME, + "namespace": TEST_NAMESPACE, + }, + ) + + # check on_operator_resuming callback + mock_callbacks.on_pod_cleanup.assert_called_once() + assert mock_callbacks.on_pod_cleanup.call_args.kwargs == { + "client": k.client, + "mode": ExecutionMode.SYNC, + "pod": remote_pod_mock, + "operator": k, + "context": context, + } + + # check on_pod_cleanup callback + mock_callbacks.on_pod_cleanup.assert_called_once() + assert mock_callbacks.on_pod_cleanup.call_args.kwargs == { + "client": k.client, + "mode": ExecutionMode.SYNC, + "pod": remote_pod_mock, + "operator": k, + "context": context, + } + + @pytest.mark.parametrize("do_xcom_push", [True, False]) @patch(KUB_OP_PATH.format("extract_xcom")) @patch(KUB_OP_PATH.format("post_complete_action")) From 871a54104c45777eb04f3d9c1bd16d52060dab71 Mon Sep 17 00:00:00 2001 From: John Horan Date: Wed, 26 Feb 2025 15:10:15 +0000 Subject: [PATCH 07/15] add callbacks to test --- .../cncf/kubernetes/operators/pod.py | 34 +++++++------- .../cncf/kubernetes/operators/test_pod.py | 46 ++++++++++++++----- 2 files changed, 51 insertions(+), 29 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index c407d52482db8..94e3f5af378e2 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -885,24 +885,22 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: if self.get_logs: self._write_logs(self.pod, follow=follow, since_time=last_log_time) - if self.callbacks: - pod = self.find_pod(self.pod.metadata.namespace, context=context) - for callback in self.callbacks: - callback.on_pod_completion( - pod=pod, - client=self.client, - mode=ExecutionMode.SYNC, - context=context, - operator=self, - ) - for callback in self.callbacks: - callback.on_pod_teardown( - pod=pod, - client=self.client, - mode=ExecutionMode.SYNC, - context=context, - operator=self, - ) + for callback in self.callbacks: + callback.on_pod_completion( + pod=self.pod, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) + for callback in self.callbacks: + callback.on_pod_teardown( + pod=self.pod, + client=self.client, + mode=ExecutionMode.SYNC, + context=context, + operator=self, + ) xcom_sidecar_output = self.extract_xcom(pod=self.pod) if self.do_xcom_push else None diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index 57256b8d51298..2ae5bd87f25e8 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -2541,8 +2541,8 @@ def test_trigger_error(self, find_pod, cleanup, mock_write_log): ) - @patch(HOOK_CLASS, new=MagicMock) - def test_execute_async_callbacks(self): + @patch(f"{HOOK_CLASS}.get_pod") + def test_execute_async_callbacks(self, get_pod): from airflow.providers.cncf.kubernetes.callbacks import ExecutionMode from unit.cncf.kubernetes.test_callbacks import ( MockKubernetesPodOperatorCallback, @@ -2568,24 +2568,26 @@ def test_execute_async_callbacks(self): ) context = create_context(k) + callback_event = { + "status": "success", + "message": TEST_SUCCESS_MESSAGE, + "name": TEST_NAME, + "namespace": TEST_NAMESPACE, + } k.trigger_reentry( context=context, - event={ - "status": "success", - "message": TEST_SUCCESS_MESSAGE, - "name": TEST_NAME, - "namespace": TEST_NAMESPACE, - }, + event=callback_event ) # check on_operator_resuming callback - mock_callbacks.on_pod_cleanup.assert_called_once() - assert mock_callbacks.on_pod_cleanup.call_args.kwargs == { + mock_callbacks.on_operator_resuming.assert_called_once() + assert mock_callbacks.on_operator_resuming.call_args.kwargs == { "client": k.client, "mode": ExecutionMode.SYNC, - "pod": remote_pod_mock, + "pod": get_pod.return_value, "operator": k, "context": context, + "event": callback_event, } # check on_pod_cleanup callback @@ -2598,6 +2600,28 @@ def test_execute_async_callbacks(self): "context": context, } + # check on_pod_completion callback + mock_callbacks.on_pod_completion.assert_called_once() + assert mock_callbacks.on_pod_completion.call_args.kwargs == { + "client": k.client, + "mode": ExecutionMode.SYNC, + "pod": get_pod.return_value, + "operator": k, + "context": context, + } + + # check on_pod_teardown callback + mock_callbacks.on_pod_teardown.assert_called_once() + assert mock_callbacks.on_pod_teardown.call_args.kwargs == { + "client": k.client, + "mode": ExecutionMode.SYNC, + "pod": get_pod.return_value, + "operator": k, + "context": context, + } + + + @pytest.mark.parametrize("do_xcom_push", [True, False]) @patch(KUB_OP_PATH.format("extract_xcom")) From b996c5a94af7fdbacb534baa8f4d8d4cee4a486e Mon Sep 17 00:00:00 2001 From: John Horan Date: Wed, 26 Feb 2025 15:15:54 +0000 Subject: [PATCH 08/15] cleanup --- .../tests/unit/cncf/kubernetes/operators/test_pod.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index 2ae5bd87f25e8..e23936c340cb2 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -2540,9 +2540,8 @@ def test_trigger_error(self, find_pod, cleanup, mock_write_log): }, ) - - @patch(f"{HOOK_CLASS}.get_pod") - def test_execute_async_callbacks(self, get_pod): + @patch(HOOK_CLASS) + def test_execute_async_callbacks(self, mocked_hook): from airflow.providers.cncf.kubernetes.callbacks import ExecutionMode from unit.cncf.kubernetes.test_callbacks import ( MockKubernetesPodOperatorCallback, @@ -2554,6 +2553,7 @@ def test_execute_async_callbacks(self, get_pod): remote_pod_mock = MagicMock() remote_pod_mock.status.phase = "Succeeded" self.await_pod_mock.return_value = remote_pod_mock + mocked_hook.return_value.get_pod.return_value = remote_pod_mock k = KubernetesPodOperator( namespace="default", @@ -2584,7 +2584,7 @@ def test_execute_async_callbacks(self, get_pod): assert mock_callbacks.on_operator_resuming.call_args.kwargs == { "client": k.client, "mode": ExecutionMode.SYNC, - "pod": get_pod.return_value, + "pod": remote_pod_mock, "operator": k, "context": context, "event": callback_event, @@ -2605,7 +2605,7 @@ def test_execute_async_callbacks(self, get_pod): assert mock_callbacks.on_pod_completion.call_args.kwargs == { "client": k.client, "mode": ExecutionMode.SYNC, - "pod": get_pod.return_value, + "pod": remote_pod_mock, "operator": k, "context": context, } @@ -2615,7 +2615,7 @@ def test_execute_async_callbacks(self, get_pod): assert mock_callbacks.on_pod_teardown.call_args.kwargs == { "client": k.client, "mode": ExecutionMode.SYNC, - "pod": get_pod.return_value, + "pod": remote_pod_mock, "operator": k, "context": context, } From 07285e94be8a65f384e86d3c49b486a9f8eaf14e Mon Sep 17 00:00:00 2001 From: John Horan Date: Wed, 26 Feb 2025 16:23:06 +0000 Subject: [PATCH 09/15] format --- .../tests/unit/cncf/kubernetes/operators/test_pod.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index e23936c340cb2..c39e9364a8d9d 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -2574,10 +2574,7 @@ def test_execute_async_callbacks(self, mocked_hook): "name": TEST_NAME, "namespace": TEST_NAMESPACE, } - k.trigger_reentry( - context=context, - event=callback_event - ) + k.trigger_reentry(context=context, event=callback_event) # check on_operator_resuming callback mock_callbacks.on_operator_resuming.assert_called_once() @@ -2621,8 +2618,6 @@ def test_execute_async_callbacks(self, mocked_hook): } - - @pytest.mark.parametrize("do_xcom_push", [True, False]) @patch(KUB_OP_PATH.format("extract_xcom")) @patch(KUB_OP_PATH.format("post_complete_action")) From 86e5a998cfe608296c8e696305a9d389215377a6 Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 8 Apr 2025 22:46:03 +0100 Subject: [PATCH 10/15] Update pod.py --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 94e3f5af378e2..cac3c79eed428 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -905,7 +905,9 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: xcom_sidecar_output = self.extract_xcom(pod=self.pod) if self.do_xcom_push else None if event["status"] != "success": - self.log.error("Trigger emitted an %s event, failing the task: %s", event["status"], event["message"]) + self.log.error( + "Trigger emitted an %s event, failing the task: %s", event["status"], event["message"] + ) message = event.get("stack_trace", event["message"]) raise AirflowException(message) From a2bb12f1ad2ee3b4875ea14bdb2fcbcf1da3cc83 Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 8 Apr 2025 23:10:22 +0100 Subject: [PATCH 11/15] space From f17abee80d541e397bcc5b1e0a45c18e79c7eb4e Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 8 Apr 2025 23:18:49 +0100 Subject: [PATCH 12/15] ruff fix --- .../kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py | 1 + 1 file changed, 1 insertion(+) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index c39e9364a8d9d..c4a194f772f7e 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -2543,6 +2543,7 @@ def test_trigger_error(self, find_pod, cleanup, mock_write_log): @patch(HOOK_CLASS) def test_execute_async_callbacks(self, mocked_hook): from airflow.providers.cncf.kubernetes.callbacks import ExecutionMode + from unit.cncf.kubernetes.test_callbacks import ( MockKubernetesPodOperatorCallback, MockWrapper, From 03fcd9f173126df03c81f3185362615c22cf55b4 Mon Sep 17 00:00:00 2001 From: John Horan Date: Mon, 16 Jun 2025 13:43:08 +0100 Subject: [PATCH 13/15] merge recent changes --- .../providers/cncf/kubernetes/operators/pod.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index cac3c79eed428..56834ec43ca2c 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -718,7 +718,9 @@ def execute_sync(self, context: Context): ) finally: pod_to_clean = self.pod or self.pod_request_obj - self.post_complete_action(pod=pod_to_clean, remote_pod=self.remote_pod, context=context, result=result) + self.post_complete_action( + pod=pod_to_clean, remote_pod=self.remote_pod, context=context, result=result + ) if self.do_xcom_push: return result @@ -858,6 +860,7 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: grab the latest logs and defer back to the trigger again. """ self.pod = None + xcom_sidecar_output = None try: pod_name = event["name"] pod_namespace = event["namespace"] @@ -930,9 +933,9 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: except TaskDeferred: raise finally: - self._clean(event, context) + self._clean(event, context, result=xcom_sidecar_output) - def _clean(self, event: dict[str, Any], context: Context) -> None: + def _clean(self, event: dict[str, Any], result: dict, context: Context) -> None: if event["status"] == "running": return istio_enabled = self.is_istio_enabled(self.pod) @@ -956,6 +959,7 @@ def _clean(self, event: dict[str, Any], context: Context) -> None: pod=self.pod, remote_pod=self.pod, context=context, + result=result, ) def _write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime | None = None) -> None: @@ -985,7 +989,9 @@ def _write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime e if not isinstance(e, ApiException) else e.reason, ) - def post_complete_action(self, *, pod, remote_pod, context: Context, result: dict, **kwargs) -> None: + def post_complete_action( + self, *, pod: k8s.V1Pod, remote_pod: k8s.V1Pod, context: Context, result: dict, **kwargs + ) -> None: """Actions that must be done after operator finishes logic of the deferrable_execution.""" self.cleanup( pod=pod, From c34305663ff8d920b25b223a3f1d02694cf82dd0 Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 22 Jul 2025 16:00:21 +0100 Subject: [PATCH 14/15] fix pass args --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 987a075fd986b..f27d5f052196c 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -940,7 +940,7 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: except TaskDeferred: raise finally: - self._clean(event, context, result=xcom_sidecar_output) + self._clean(event=event, context=context, result=xcom_sidecar_output) def _clean(self, event: dict[str, Any], result: dict, context: Context) -> None: if event["status"] == "running": From c816f8d4603d8b9f7570d51f6e4147f7c35ea485 Mon Sep 17 00:00:00 2001 From: John Horan Date: Tue, 22 Jul 2025 16:25:42 +0100 Subject: [PATCH 15/15] fix type --- .../src/airflow/providers/cncf/kubernetes/operators/pod.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index f27d5f052196c..c5ec5fc54d114 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -942,7 +942,7 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any: finally: self._clean(event=event, context=context, result=xcom_sidecar_output) - def _clean(self, event: dict[str, Any], result: dict, context: Context) -> None: + def _clean(self, event: dict[str, Any], result: dict | None, context: Context) -> None: if event["status"] == "running": return istio_enabled = self.is_istio_enabled(self.pod) @@ -997,7 +997,7 @@ def _write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime ) def post_complete_action( - self, *, pod: k8s.V1Pod, remote_pod: k8s.V1Pod, context: Context, result: dict, **kwargs + self, *, pod: k8s.V1Pod, remote_pod: k8s.V1Pod, context: Context, result: dict | None, **kwargs ) -> None: """Actions that must be done after operator finishes logic of the deferrable_execution.""" self.cleanup(