From 02726aaff78a38efcf0a2ec363997ef6872e7b31 Mon Sep 17 00:00:00 2001 From: AlejandroMorgante <62363051+AlejandroMorgante@users.noreply.github.com> Date: Thu, 6 Aug 2026 12:14:03 -0300 Subject: [PATCH 1/4] Add operator to execute commands in existing Kubernetes Pods Reusing a running Pod avoids the startup latency incurred when each task creates a new Pod. --- .../cncf/kubernetes/docs/kubernetes_rbac.rst | 21 ++ providers/cncf/kubernetes/docs/operators.rst | 37 +++ providers/cncf/kubernetes/provider.yaml | 1 + .../providers/cncf/kubernetes/exceptions.py | 4 + .../cncf/kubernetes/get_provider_info.py | 1 + .../cncf/kubernetes/operators/pod_exec.py | 247 ++++++++++++++ .../kubernetes/operators/test_pod_exec.py | 313 ++++++++++++++++++ 7 files changed, 624 insertions(+) create mode 100644 providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod_exec.py create mode 100644 providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py diff --git a/providers/cncf/kubernetes/docs/kubernetes_rbac.rst b/providers/cncf/kubernetes/docs/kubernetes_rbac.rst index fbe64cf418a46..971c7c0677016 100644 --- a/providers/cncf/kubernetes/docs/kubernetes_rbac.rst +++ b/providers/cncf/kubernetes/docs/kubernetes_rbac.rst @@ -89,6 +89,27 @@ deployment commonly needs these permissions: retrieving XCom from the sidecar container. ``events`` access is used to read Kubernetes events for diagnostics. +Existing Pod exec permissions +----------------------------- + +``KubernetesPodExecOperator`` executes a command in an existing Pod without managing its lifecycle. +When the Pod name and namespace are provided directly, it needs only these permissions: + +.. code-block:: yaml + + apiVersion: rbac.authorization.k8s.io/v1 + kind: Role + metadata: + name: airflow-pod-exec + namespace: airflow + rules: + - apiGroups: [""] + resources: ["pods"] + verbs: ["get"] + - apiGroups: [""] + resources: ["pods/exec"] + verbs: ["get"] + Job launch permissions ---------------------- diff --git a/providers/cncf/kubernetes/docs/operators.rst b/providers/cncf/kubernetes/docs/operators.rst index 05da87f11dd4e..2acf8f162ae1f 100644 --- a/providers/cncf/kubernetes/docs/operators.rst +++ b/providers/cncf/kubernetes/docs/operators.rst @@ -19,6 +19,43 @@ .. contents:: Table of Contents :depth: 2 +.. _howto/operator:KubernetesPodExecOperator: + +KubernetesPodExecOperator +========================= + +The :class:`~airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator` +executes a command in a running container of an existing Kubernetes Pod. It does not create, +restart, or delete the target Pod. + +.. code-block:: python + + from airflow.providers.cncf.kubernetes.operators.pod_exec import KubernetesPodExecOperator + + run_command = KubernetesPodExecOperator( + task_id="run_command", + pod_name="existing-worker", + namespace="default", + container_name="worker", + command=["python", "-m", "worker.run_job"], + kubernetes_conn_id="kubernetes_default", + ) + +Commands are executed directly rather than through a shell. Include a shell explicitly when using +pipes, redirects, variable expansion, or other shell features. +Standard output and standard error are streamed to the task log. Set ``do_xcom_push=True`` to +also return standard output through XCom. + +The target Pod and container must already be running. When ``container_name`` is omitted, the +operator uses the ``kubectl.kubernetes.io/default-container`` annotation when present, or the +first container otherwise. API-visible static Pods are supported through their mirror Pod name; +components that are not exposed by the Kubernetes API cannot be targeted. The Kubernetes connection +requires ``get`` access to ``pods`` and ``pods/exec``; see :doc:`kubernetes_rbac`. + +If the task or its worker stops while the command is running, Airflow closes the exec connection +but does not modify the target Pod. Kubernetes cannot always determine whether a command completed +before a connection failure, so configure task retries only when the command is safe to repeat. + .. _howto/operator:kubernetespodoperator: KubernetesPodOperator diff --git a/providers/cncf/kubernetes/provider.yaml b/providers/cncf/kubernetes/provider.yaml index cfecb10cba9b7..e637506efa558 100644 --- a/providers/cncf/kubernetes/provider.yaml +++ b/providers/cncf/kubernetes/provider.yaml @@ -151,6 +151,7 @@ operators: - airflow.providers.cncf.kubernetes.operators.custom_object_launcher - airflow.providers.cncf.kubernetes.operators.kueue - airflow.providers.cncf.kubernetes.operators.pod + - airflow.providers.cncf.kubernetes.operators.pod_exec - airflow.providers.cncf.kubernetes.operators.spark_kubernetes - airflow.providers.cncf.kubernetes.operators.resource - airflow.providers.cncf.kubernetes.operators.job diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/exceptions.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/exceptions.py index 522769a4534e8..a809b7c49e894 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/exceptions.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/exceptions.py @@ -39,3 +39,7 @@ class KubernetesApiError(AirflowException): class KubernetesApiPermissionError(AirflowException): """Raised when an error is encountered while trying access Kubernetes API.""" + + +class PodExecException(AirflowException): + """Raised when a command cannot be executed successfully in a Kubernetes pod.""" diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py index 8a2daec956f4e..86430daa9e24e 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py @@ -48,6 +48,7 @@ def get_provider_info(): "airflow.providers.cncf.kubernetes.operators.custom_object_launcher", "airflow.providers.cncf.kubernetes.operators.kueue", "airflow.providers.cncf.kubernetes.operators.pod", + "airflow.providers.cncf.kubernetes.operators.pod_exec", "airflow.providers.cncf.kubernetes.operators.spark_kubernetes", "airflow.providers.cncf.kubernetes.operators.resource", "airflow.providers.cncf.kubernetes.operators.job", diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod_exec.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod_exec.py new file mode 100644 index 0000000000000..70affb0b2efe2 --- /dev/null +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod_exec.py @@ -0,0 +1,247 @@ +# 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. +"""Execute commands in existing Kubernetes pods.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING + +from kubernetes.client.rest import ApiException +from kubernetes.stream import stream as kubernetes_stream + +from airflow.providers.cncf.kubernetes.exceptions import PodExecException +from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook +from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_1_PLUS + +if AIRFLOW_V_3_1_PLUS: + from airflow.sdk import BaseOperator +else: + from airflow.models import BaseOperator + +if TYPE_CHECKING: + from kubernetes.client import CoreV1Api, V1Pod + from kubernetes.stream.ws_client import WSClient + + from airflow.sdk import Context + +__all__ = ["KubernetesPodExecOperator"] + + +class KubernetesPodExecOperator(BaseOperator): + """ + Execute a command in a running container of an existing Kubernetes pod. + + The operator does not create, restart, or delete the target pod. Commands are executed directly, + without a shell; include a shell explicitly in ``command`` when shell features are required. + + :param pod_name: Name of the existing Kubernetes pod. (templated) + :param command: Command and arguments to execute in the container. (templated) + :param namespace: Namespace containing the pod. Defaults to the namespace configured in the + Kubernetes connection, then ``default``. (templated) + :param container_name: Name of the container in which to execute the command. When omitted, the + ``kubectl.kubernetes.io/default-container`` annotation or the first container is used. (templated) + :param kubernetes_conn_id: The :ref:`Kubernetes connection ` to use. + (templated) + :param in_cluster: Use in-cluster Kubernetes configuration. + :param cluster_context: Context to use from the kubeconfig. (templated) + :param config_file: Path to the kubeconfig file. (templated) + :param do_xcom_push: Return standard output for XCom when ``True``. Defaults to ``False``. + """ + + template_fields: Sequence[str] = ( + "pod_name", + "command", + "namespace", + "container_name", + "kubernetes_conn_id", + "cluster_context", + "config_file", + ) + template_fields_renderers = {"command": "py"} + + def __init__( + self, + *, + pod_name: str, + command: Sequence[str], + namespace: str | None = None, + container_name: str | None = None, + kubernetes_conn_id: str | None = KubernetesHook.default_conn_name, + in_cluster: bool | None = None, + cluster_context: str | None = None, + config_file: str | None = None, + do_xcom_push: bool = False, + **kwargs, + ) -> None: + super().__init__(do_xcom_push=do_xcom_push, **kwargs) + self.pod_name = pod_name + self.command = command + self.namespace = namespace + self.container_name = container_name + self.kubernetes_conn_id = kubernetes_conn_id + self.in_cluster = in_cluster + self.cluster_context = cluster_context + self.config_file = config_file + self._exec_client: WSClient | None = None + + @cached_property + def hook(self) -> KubernetesHook: + return KubernetesHook( + conn_id=self.kubernetes_conn_id, + in_cluster=self.in_cluster, + config_file=self.config_file, + cluster_context=self.cluster_context, + ) + + @cached_property + def client(self) -> CoreV1Api: + return self.hook.core_v1_client + + def _resolve_namespace(self) -> str: + return self.namespace or self.hook.get_namespace() or KubernetesHook.DEFAULT_NAMESPACE + + def _validate_command(self) -> list[str]: + if isinstance(self.command, str) or not isinstance(self.command, Sequence): + raise TypeError("`command` must be a sequence of strings, not a single string") + if not self.command: + raise ValueError("`command` must contain at least one element") + if not all(isinstance(argument, str) for argument in self.command): + raise TypeError("Every element of `command` must be a string") + return list(self.command) + + def _resolve_container_name(self, pod: V1Pod) -> str: + containers = pod.spec.containers if pod.spec and pod.spec.containers else [] + container_names = [container.name for container in containers] + if not container_names: + raise PodExecException(f"Pod {self.pod_name!r} does not define any containers") + + if self.container_name: + if self.container_name not in container_names: + raise PodExecException( + f"Container {self.container_name!r} does not exist in pod {self.pod_name!r}" + ) + return self.container_name + + annotations = pod.metadata.annotations if pod.metadata and pod.metadata.annotations else {} + default_container = annotations.get("kubectl.kubernetes.io/default-container") + if isinstance(default_container, str) and default_container in container_names: + return default_container + return container_names[0] + + def _validate_container_is_running(self, pod: V1Pod, container_name: str) -> None: + if not pod.status or pod.status.phase != "Running": + phase = pod.status.phase if pod.status else None + raise PodExecException( + f"Cannot execute a command in pod {self.pod_name!r} while it is in phase {phase!r}" + ) + + statuses = pod.status.container_statuses or [] + container_status = next((status for status in statuses if status.name == container_name), None) + if ( + container_status is None + or container_status.state is None + or container_status.state.running is None + ): + raise PodExecException(f"Container {container_name!r} in pod {self.pod_name!r} is not running") + + def _log_output(self, output: str, *, stream_name: str) -> None: + log_method = self.log.warning if stream_name == "stderr" else self.log.info + for line in output.splitlines(): + log_method("[%s] %s", stream_name, line) + + def _consume_output(self, exec_client: WSClient) -> str: + stdout_chunks: list[str] = [] + while exec_client.is_open(): + exec_client.update(timeout=1) + while exec_client.peek_stdout(): + output = exec_client.read_stdout() + self._log_output(output, stream_name="stdout") + if self.do_xcom_push: + stdout_chunks.append(output) + while exec_client.peek_stderr(): + self._log_output(exec_client.read_stderr(), stream_name="stderr") + return "".join(stdout_chunks) + + def _close_exec_client(self) -> None: + exec_client = self._exec_client + self._exec_client = None + if exec_client is None: + return + try: + exec_client.close() + except Exception: + self.log.exception("Failed to close Kubernetes exec connection") + + def execute(self, context: Context) -> str | None: + command = self._validate_command() + namespace = self._resolve_namespace() + if not self.pod_name: + raise ValueError("`pod_name` must not be empty") + + try: + pod = self.hook.get_pod(name=self.pod_name, namespace=namespace) + except ApiException as error: + raise PodExecException( + f"Unable to read pod {namespace}/{self.pod_name}: {error.reason or error}" + ) from error + + container_name = self._resolve_container_name(pod) + self._validate_container_is_running(pod, container_name) + self.log.info( + "Executing command in container %s of pod %s/%s", container_name, namespace, self.pod_name + ) + + try: + exec_client = kubernetes_stream( + self.client.connect_get_namespaced_pod_exec, + name=self.pod_name, + namespace=namespace, + container=container_name, + command=command, + stdin=False, + stdout=True, + stderr=True, + tty=False, + _preload_content=False, + ) + self._exec_client = exec_client + output = self._consume_output(exec_client) + return_code = exec_client.returncode + except ApiException as error: + raise PodExecException( + f"Unable to execute command in pod {namespace}/{self.pod_name}: {error.reason or error}" + ) from error + finally: + self._close_exec_client() + + if return_code is None: + raise PodExecException( + f"Command in container {container_name!r} of pod {namespace}/{self.pod_name} ended " + "without reporting an exit code" + ) + if return_code != 0: + raise PodExecException( + f"Command in container {container_name!r} of pod {namespace}/{self.pod_name} " + f"failed with exit code {return_code}" + ) + return output if self.do_xcom_push else None + + def on_kill(self) -> None: + """Close the active Kubernetes exec connection without modifying the target pod.""" + self._close_exec_client() diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py new file mode 100644 index 0000000000000..972341e14c99f --- /dev/null +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py @@ -0,0 +1,313 @@ +# 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 + +from types import SimpleNamespace +from unittest import mock + +import pytest +from kubernetes.client.rest import ApiException +from kubernetes.stream.ws_client import WSClient + +from airflow.providers.cncf.kubernetes.exceptions import PodExecException +from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook +from airflow.providers.cncf.kubernetes.operators.pod_exec import KubernetesPodExecOperator + +MODULE = "airflow.providers.cncf.kubernetes.operators.pod_exec" + + +def create_pod( + *, + container_names: tuple[str, ...] = ("main",), + annotations: dict[str, str] | None = None, + phase: str = "Running", + container_statuses=None, + with_spec: bool = True, + with_status: bool = True, +): + if container_statuses is None: + container_statuses = [ + SimpleNamespace(name=name, state=SimpleNamespace(running=object())) for name in container_names + ] + spec = SimpleNamespace(containers=[SimpleNamespace(name=name) for name in container_names]) + status = SimpleNamespace(phase=phase, container_statuses=container_statuses) + return SimpleNamespace( + metadata=SimpleNamespace(annotations=annotations), + spec=spec if with_spec else None, + status=status if with_status else None, + ) + + +def create_exec_client(*, stdout=(), stderr=(), returncode=0): + exec_client = mock.MagicMock(spec=WSClient) + exec_client.is_open.side_effect = [True, False] + exec_client.peek_stdout.side_effect = [*stdout, ""] + exec_client.read_stdout.side_effect = stdout + exec_client.peek_stderr.side_effect = [*stderr, ""] + exec_client.read_stderr.side_effect = stderr + exec_client.returncode = returncode + return exec_client + + +def create_operator(*, pod=None, hook_namespace=None, **kwargs): + operator = KubernetesPodExecOperator( + task_id="exec", + pod_name="existing-pod", + command=["echo", "hello"], + **kwargs, + ) + hook = mock.MagicMock(spec=KubernetesHook) + hook.get_namespace.return_value = hook_namespace + hook.get_pod.return_value = pod or create_pod() + operator.__dict__["hook"] = hook + return operator, hook + + +class TestKubernetesPodExecOperator: + def test_template_fields(self): + assert set(KubernetesPodExecOperator.template_fields) == { + "pod_name", + "command", + "namespace", + "container_name", + "kubernetes_conn_id", + "cluster_context", + "config_file", + } + + @mock.patch(f"{MODULE}.KubernetesHook", autospec=True) + def test_hook_configuration(self, kubernetes_hook_mock): + operator = KubernetesPodExecOperator( + task_id="exec", + pod_name="existing-pod", + command=["date"], + kubernetes_conn_id="kubernetes_test", + in_cluster=True, + config_file="/tmp/kubeconfig", + cluster_context="test-context", + ) + + assert operator.hook is kubernetes_hook_mock.return_value + kubernetes_hook_mock.assert_called_once_with( + conn_id="kubernetes_test", + in_cluster=True, + config_file="/tmp/kubeconfig", + cluster_context="test-context", + ) + + @pytest.mark.parametrize( + ("do_xcom_push", "expected_result"), + [(False, None), (True, "hello\nworld\n")], + ) + @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True) + def test_execute(self, kubernetes_stream_mock, do_xcom_push, expected_result, caplog): + operator, hook = create_operator( + namespace="test-namespace", + container_name="main", + do_xcom_push=do_xcom_push, + ) + exec_client = create_exec_client(stdout=("hello\n", "world\n"), stderr=("warning\n",), returncode=0) + kubernetes_stream_mock.return_value = exec_client + + result = operator.execute(context={}) + + assert result == expected_result + hook.get_pod.assert_called_once_with(name="existing-pod", namespace="test-namespace") + kubernetes_stream_mock.assert_called_once_with( + hook.core_v1_client.connect_get_namespaced_pod_exec, + name="existing-pod", + namespace="test-namespace", + container="main", + command=["echo", "hello"], + stdin=False, + stdout=True, + stderr=True, + tty=False, + _preload_content=False, + ) + exec_client.update.assert_called_once_with(timeout=1) + assert exec_client.read_stdout.call_count == 2 + exec_client.read_stderr.assert_called_once_with() + exec_client.close.assert_called_once_with() + assert operator._exec_client is None + assert "[stdout] hello" in caplog + assert "[stdout] world" in caplog + assert "[stderr] warning" in caplog + + @pytest.mark.parametrize( + ("namespace", "hook_namespace", "expected_namespace"), + [ + ("task-namespace", "connection-namespace", "task-namespace"), + (None, "connection-namespace", "connection-namespace"), + (None, None, KubernetesHook.DEFAULT_NAMESPACE), + ], + ) + def test_resolve_namespace(self, namespace, hook_namespace, expected_namespace): + operator, _ = create_operator(namespace=namespace, hook_namespace=hook_namespace) + + assert operator._resolve_namespace() == expected_namespace + + @pytest.mark.parametrize( + ("container_name", "annotations", "expected_container"), + [ + ("secondary", {"kubectl.kubernetes.io/default-container": "main"}, "secondary"), + (None, {"kubectl.kubernetes.io/default-container": "secondary"}, "secondary"), + (None, {"kubectl.kubernetes.io/default-container": "missing"}, "main"), + (None, None, "main"), + ], + ) + def test_resolve_container_name(self, container_name, annotations, expected_container): + pod = create_pod(container_names=("main", "secondary"), annotations=annotations) + operator, _ = create_operator(pod=pod, container_name=container_name) + + assert operator._resolve_container_name(pod) == expected_container + + @pytest.mark.parametrize("with_spec", [False, True]) + def test_rejects_pod_without_containers(self, with_spec): + pod = create_pod(container_names=(), with_spec=with_spec) + operator, _ = create_operator(pod=pod) + + with pytest.raises(PodExecException, match="does not define any containers"): + operator._resolve_container_name(pod) + + def test_rejects_unknown_container(self): + pod = create_pod() + operator, _ = create_operator(pod=pod, container_name="missing") + + with pytest.raises(PodExecException, match="does not exist"): + operator._resolve_container_name(pod) + + @pytest.mark.parametrize( + ("pod", "expected_message"), + [ + (create_pod(with_status=False), "phase None"), + (create_pod(phase="Pending"), "phase 'Pending'"), + (create_pod(container_statuses=[]), "is not running"), + ( + create_pod(container_statuses=[SimpleNamespace(name="main", state=None)]), + "is not running", + ), + ( + create_pod( + container_statuses=[SimpleNamespace(name="main", state=SimpleNamespace(running=None))] + ), + "is not running", + ), + ], + ) + def test_rejects_unavailable_target(self, pod, expected_message): + operator, _ = create_operator(pod=pod) + + with pytest.raises(PodExecException, match=expected_message): + operator._validate_container_is_running(pod, "main") + + @pytest.mark.parametrize( + ("command", "error", "expected_message"), + [ + ("echo hello", TypeError, "sequence of strings"), + (42, TypeError, "sequence of strings"), + ([], ValueError, "at least one element"), + (["echo", 42], TypeError, "Every element"), + ], + ) + def test_rejects_invalid_command(self, command, error, expected_message): + operator, _ = create_operator() + operator.command = command + + with pytest.raises(error, match=expected_message): + operator.execute(context={}) + + def test_rejects_empty_pod_name(self): + operator, _ = create_operator() + operator.pod_name = "" + + with pytest.raises(ValueError, match="must not be empty"): + operator.execute(context={}) + + @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True) + def test_wraps_pod_read_error(self, kubernetes_stream_mock): + operator, hook = create_operator(namespace="test-namespace") + hook.get_pod.side_effect = ApiException(status=404, reason="Not Found") + + with pytest.raises(PodExecException, match="Unable to read pod.*Not Found"): + operator.execute(context={}) + + kubernetes_stream_mock.assert_not_called() + + @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True) + def test_wraps_exec_api_error(self, kubernetes_stream_mock): + operator, _ = create_operator(namespace="test-namespace") + kubernetes_stream_mock.side_effect = ApiException(status=403, reason="Forbidden") + + with pytest.raises(PodExecException, match="Unable to execute command.*Forbidden"): + operator.execute(context={}) + + assert operator._exec_client is None + + @pytest.mark.parametrize( + ("returncode", "expected_message"), + [(None, "without reporting an exit code"), (17, "failed with exit code 17")], + ) + @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True) + def test_rejects_unsuccessful_command(self, kubernetes_stream_mock, returncode, expected_message): + operator, _ = create_operator(namespace="test-namespace") + exec_client = create_exec_client(returncode=returncode) + kubernetes_stream_mock.return_value = exec_client + + with pytest.raises(PodExecException, match=expected_message): + operator.execute(context={}) + + exec_client.close.assert_called_once_with() + + def test_on_kill_closes_active_connection(self): + operator, _ = create_operator() + exec_client = mock.MagicMock(spec=WSClient) + operator._exec_client = exec_client + + operator.on_kill() + + exec_client.close.assert_called_once_with() + assert operator._exec_client is None + + @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True) + def test_on_kill_while_consuming_output(self, kubernetes_stream_mock): + operator, _ = create_operator(namespace="test-namespace") + exec_client = create_exec_client() + exec_client.is_open.side_effect = lambda: (operator.on_kill(), False)[1] + kubernetes_stream_mock.return_value = exec_client + + assert operator.execute(context={}) is None + + exec_client.close.assert_called_once_with() + assert operator._exec_client is None + + def test_on_kill_without_active_connection(self): + operator, _ = create_operator() + + operator.on_kill() + + assert operator._exec_client is None + + def test_close_error_does_not_mask_task_shutdown(self): + operator, _ = create_operator() + exec_client = mock.MagicMock(spec=WSClient) + exec_client.close.side_effect = RuntimeError("connection already closed") + operator._exec_client = exec_client + + operator.on_kill() + + assert operator._exec_client is None From 6c5c9b32319be0af14900d83533a8c69ec943c94 Mon Sep 17 00:00:00 2001 From: AlejandroMorgante <62363051+AlejandroMorgante@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:57:45 -0300 Subject: [PATCH 2/4] Add system test for existing Kubernetes Pod execution --- providers/cncf/kubernetes/docs/operators.rst | 17 +-- .../kubernetes/example_kubernetes_pod_exec.py | 134 ++++++++++++++++++ 2 files changed, 139 insertions(+), 12 deletions(-) create mode 100644 providers/cncf/kubernetes/tests/system/cncf/kubernetes/example_kubernetes_pod_exec.py diff --git a/providers/cncf/kubernetes/docs/operators.rst b/providers/cncf/kubernetes/docs/operators.rst index 2acf8f162ae1f..c86cb142ad151 100644 --- a/providers/cncf/kubernetes/docs/operators.rst +++ b/providers/cncf/kubernetes/docs/operators.rst @@ -28,18 +28,11 @@ The :class:`~airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodE executes a command in a running container of an existing Kubernetes Pod. It does not create, restart, or delete the target Pod. -.. code-block:: python - - from airflow.providers.cncf.kubernetes.operators.pod_exec import KubernetesPodExecOperator - - run_command = KubernetesPodExecOperator( - task_id="run_command", - pod_name="existing-worker", - namespace="default", - container_name="worker", - command=["python", "-m", "worker.run_job"], - kubernetes_conn_id="kubernetes_default", - ) +.. exampleinclude:: /../tests/system/cncf/kubernetes/example_kubernetes_pod_exec.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_k8s_pod_exec] + :end-before: [END howto_operator_k8s_pod_exec] Commands are executed directly rather than through a shell. Include a shell explicitly when using pipes, redirects, variable expansion, or other shell features. diff --git a/providers/cncf/kubernetes/tests/system/cncf/kubernetes/example_kubernetes_pod_exec.py b/providers/cncf/kubernetes/tests/system/cncf/kubernetes/example_kubernetes_pod_exec.py new file mode 100644 index 0000000000000..bfb616a6f8c94 --- /dev/null +++ b/providers/cncf/kubernetes/tests/system/cncf/kubernetes/example_kubernetes_pod_exec.py @@ -0,0 +1,134 @@ +# 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. +"""Example Dag for executing a command in an existing Kubernetes Pod.""" + +from __future__ import annotations + +import os +import time +from datetime import datetime + +from kubernetes.client.rest import ApiException + +from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook +from airflow.providers.cncf.kubernetes.operators.pod_exec import KubernetesPodExecOperator +from airflow.providers.cncf.kubernetes.operators.resource import ( + KubernetesCreateResourceOperator, + KubernetesDeleteResourceOperator, +) +from airflow.sdk import DAG, TriggerRule, task + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default").lower().replace("_", "-") +DAG_ID = "example_kubernetes_pod_exec_operator" +NAMESPACE = "default" +POD_NAME = f"airflow-pod-exec-{ENV_ID}" +CONTAINER_NAME = "worker" +EXPECTED_OUTPUT = "command executed in existing pod" + +pod_conf = f""" +apiVersion: v1 +kind: Pod +metadata: + name: {POD_NAME} + namespace: {NAMESPACE} +spec: + restartPolicy: Never + containers: + - name: {CONTAINER_NAME} + image: busybox:1.38.0 + command: ["sleep", "3600"] +""" + + +@task +def wait_for_running_pod() -> None: + hook = KubernetesHook() + deadline = time.monotonic() + 120 + last_phase = None + + while time.monotonic() < deadline: + try: + pod = hook.get_pod(name=POD_NAME, namespace=NAMESPACE) + except ApiException as error: + if error.status != 404: + raise + else: + last_phase = pod.status.phase if pod.status else None + container_statuses = pod.status.container_statuses if pod.status else None + container_status = next( + (status for status in container_statuses or [] if status.name == CONTAINER_NAME), None + ) + if ( + last_phase == "Running" + and container_status + and container_status.state + and container_status.state.running + ): + return + time.sleep(2) + + raise TimeoutError(f"Pod {NAMESPACE}/{POD_NAME} did not start; last phase was {last_phase!r}") + + +@task +def verify_output(output: str) -> None: + if output != EXPECTED_OUTPUT: + raise ValueError(f"Unexpected command output: {output!r}") + + +with DAG( + dag_id=DAG_ID, + schedule=None, + start_date=datetime(2021, 1, 1), + catchup=False, + tags=["example", "kubernetes"], +) as dag: + create_pod = KubernetesCreateResourceOperator( + task_id="create_pod", + yaml_conf=pod_conf, + ) + + pod_is_running = wait_for_running_pod() + + # [START howto_operator_k8s_pod_exec] + run_command = KubernetesPodExecOperator( + task_id="run_command", + pod_name=POD_NAME, + namespace=NAMESPACE, + container_name=CONTAINER_NAME, + command=["sh", "-c", f"printf '{EXPECTED_OUTPUT}'"], + do_xcom_push=True, + ) + # [END howto_operator_k8s_pod_exec] + + output_is_valid = verify_output(run_command.output) + + delete_pod = KubernetesDeleteResourceOperator( + task_id="delete_pod", + yaml_conf=pod_conf, + trigger_rule=TriggerRule.ALL_DONE, + ) + + create_pod >> pod_is_running >> run_command >> output_is_valid >> delete_pod + + from tests_common.test_utils.watcher import watcher + + list(dag.tasks) >> watcher() + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +test_run = get_test_run(dag) From 81b488a1244e931bbfc408d5bf38b31841f96a9d Mon Sep 17 00:00:00 2001 From: AlejandroMorgante <62363051+AlejandroMorgante@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:46:06 -0300 Subject: [PATCH 3/4] Trigger CI checks From 52b1cc6040a5199e9aa4614e9a75cc43ab87f472 Mon Sep 17 00:00:00 2001 From: AlejandroMorgante <62363051+AlejandroMorgante@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:58:52 -0300 Subject: [PATCH 4/4] Keep Kubernetes Pod exec tests compatible with Airflow 2 --- .../unit/cncf/kubernetes/operators/test_pod_exec.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py index 972341e14c99f..2dc813b88f7e6 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py @@ -113,8 +113,9 @@ def test_hook_configuration(self, kubernetes_hook_mock): ("do_xcom_push", "expected_result"), [(False, None), (True, "hello\nworld\n")], ) + @mock.patch(f"{MODULE}.KubernetesPodExecOperator.log", spec=["info", "warning"]) @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True) - def test_execute(self, kubernetes_stream_mock, do_xcom_push, expected_result, caplog): + def test_execute(self, kubernetes_stream_mock, log_mock, do_xcom_push, expected_result): operator, hook = create_operator( namespace="test-namespace", container_name="main", @@ -144,9 +145,13 @@ def test_execute(self, kubernetes_stream_mock, do_xcom_push, expected_result, ca exec_client.read_stderr.assert_called_once_with() exec_client.close.assert_called_once_with() assert operator._exec_client is None - assert "[stdout] hello" in caplog - assert "[stdout] world" in caplog - assert "[stderr] warning" in caplog + log_mock.info.assert_has_calls( + [ + mock.call("[%s] %s", "stdout", "hello"), + mock.call("[%s] %s", "stdout", "world"), + ] + ) + log_mock.warning.assert_called_once_with("[%s] %s", "stderr", "warning") @pytest.mark.parametrize( ("namespace", "hook_namespace", "expected_namespace"),