This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dbaf9c63199 Add KubernetesPodExecOperator for existing Kubernetes Pods
(#71244)
dbaf9c63199 is described below
commit dbaf9c631998ccb6487bf0102d2c63fae9684a83
Author: Alejandro Morgante <[email protected]>
AuthorDate: Wed Sep 2 11:04:20 2026 -0300
Add KubernetesPodExecOperator for existing Kubernetes Pods (#71244)
---
providers/cncf/kubernetes/docs/kubernetes_rbac.rst | 21 ++
providers/cncf/kubernetes/docs/operators.rst | 32 ++
providers/cncf/kubernetes/provider.yaml | 1 +
.../providers/cncf/kubernetes/get_provider_info.py | 1 +
.../cncf/kubernetes/operators/pod_exec.py | 297 ++++++++++++++++
.../cncf/kubernetes/example_kubernetes_pod_exec.py | 119 +++++++
.../cncf/kubernetes/operators/test_pod_exec.py | 381 +++++++++++++++++++++
7 files changed, 852 insertions(+)
diff --git a/providers/cncf/kubernetes/docs/kubernetes_rbac.rst
b/providers/cncf/kubernetes/docs/kubernetes_rbac.rst
index fbe64cf418a..971c7c06770 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 a7d1c242b12..f5e19c27a7a 100644
--- a/providers/cncf/kubernetes/docs/operators.rst
+++ b/providers/cncf/kubernetes/docs/operators.rst
@@ -416,6 +416,38 @@ For further information, look at:
* `Kubernetes Documentation <https://kubernetes.io/docs/home/>`__
* `Pull an Image from a Private Registry
<https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/>`__
+.. _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.
+
+.. 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.
+Standard output and standard error are streamed to the task log. Set
``do_xcom_push=True`` to also
+return standard output through XCom. Captured output is limited to 49,344
UTF-8 bytes by default;
+use ``max_xcom_output_size`` to configure a different finite limit. The task
fails instead of
+returning truncated output when the limit is exceeded.
+
+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.
+
SparkKubernetesOperator
==========================
The
:class:`~airflow.providers.cncf.kubernetes.operators.spark_kubernetes.SparkKubernetesOperator`
allows
diff --git a/providers/cncf/kubernetes/provider.yaml
b/providers/cncf/kubernetes/provider.yaml
index 04346ad7e06..24940c63ba1 100644
--- a/providers/cncf/kubernetes/provider.yaml
+++ b/providers/cncf/kubernetes/provider.yaml
@@ -152,6 +152,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/get_provider_info.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
index 50861688550..6291fdb2bc4 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 00000000000..510f2d56190
--- /dev/null
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod_exec.py
@@ -0,0 +1,297 @@
+# 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.hooks.kubernetes import KubernetesHook
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
+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"]
+
+_DEFAULT_MAX_XCOM_OUTPUT_SIZE = 49_344
+
+
+def _extract_complete_lines(buffer: str, chunk: str) -> tuple[list[str], str]:
+ """Extract newline-delimited lines while retaining the incomplete
remainder."""
+ *lines, remainder = f"{buffer}{chunk}".split("\n")
+ return lines, remainder
+
+
+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
<howto/connection:kubernetes>` 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 through XCom when ``True``.
Defaults to ``False``.
+ :param max_xcom_output_size: Maximum UTF-8 byte size retained for XCom.
Defaults to 49,344 bytes.
+ """
+
+ 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,
+ max_xcom_output_size: int = _DEFAULT_MAX_XCOM_OUTPUT_SIZE,
+ **kwargs,
+ ) -> None:
+ super().__init__(do_xcom_push=do_xcom_push, **kwargs)
+ if (
+ isinstance(max_xcom_output_size, bool)
+ or not isinstance(max_xcom_output_size, int)
+ or max_xcom_output_size <= 0
+ ):
+ raise ValueError("`max_xcom_output_size` must be a positive
integer")
+ 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.max_xcom_output_size = max_xcom_output_size
+ self._exec_client: WSClient | None = None
+ self._exec_target: tuple[str, str, str] | 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 RuntimeError(f"Pod {self.pod_name!r} does not define any
containers")
+
+ if self.container_name:
+ if self.container_name not in container_names:
+ raise ValueError(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 != PodPhase.RUNNING:
+ phase = pod.status.phase if pod.status else None
+ raise RuntimeError(
+ 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 RuntimeError(
+ f"Container {container_name!r} in pod {self.pod_name!r} is not
running or "
+ "container status cannot be retrieved"
+ )
+
+ def _log_lines(self, lines: Sequence[str], *, stream_name: str) -> None:
+ for line in lines:
+ self.log.info("[%s] %s", stream_name, line.removesuffix("\r"))
+
+ def _consume_output(self, exec_client: WSClient) -> str:
+ stdout_chunks: list[str] = []
+ stdout_size = 0
+ stdout_buffer = ""
+ stderr_buffer = ""
+ try:
+ while exec_client.is_open():
+ exec_client.update(timeout=1)
+ while exec_client.peek_stdout():
+ output = exec_client.read_stdout()
+ stdout_lines, stdout_buffer =
_extract_complete_lines(stdout_buffer, output)
+ self._log_lines(stdout_lines, stream_name="stdout")
+ if self.do_xcom_push:
+ stdout_size += len(output.encode("utf-8"))
+ if stdout_size > self.max_xcom_output_size:
+ raise RuntimeError(
+ "Standard output exceeded the configured XCom
limit of "
+ f"{self.max_xcom_output_size} bytes"
+ )
+ stdout_chunks.append(output)
+ while exec_client.peek_stderr():
+ stderr_lines, stderr_buffer = _extract_complete_lines(
+ stderr_buffer, exec_client.read_stderr()
+ )
+ self._log_lines(stderr_lines, stream_name="stderr")
+ finally:
+ if stdout_buffer:
+ self._log_lines((stdout_buffer,), stream_name="stdout")
+ if stderr_buffer:
+ self._log_lines((stderr_buffer,), stream_name="stderr")
+ return "".join(stdout_chunks)
+
+ def _close_exec_client(self) -> None:
+ exec_client = self._exec_client
+ exec_target = self._exec_target
+ self._exec_client = None
+ self._exec_target = None
+ if exec_client is None:
+ return
+ try:
+ exec_client.close()
+ except Exception:
+ namespace, pod_name, container_name = exec_target or (
+ self.namespace,
+ self.pod_name,
+ self.container_name,
+ )
+ self.log.exception(
+ "Failed to close Kubernetes exec connection for container %s
in pod %s/%s",
+ container_name,
+ namespace,
+ pod_name,
+ )
+
+ 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 RuntimeError(
+ 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
+ )
+
+ self._exec_target = (namespace, self.pod_name, container_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 RuntimeError(
+ 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 RuntimeError(
+ f"Command in container {container_name!r} of pod
{namespace}/{self.pod_name} ended "
+ "without reporting an exit code"
+ )
+ if return_code != 0:
+ raise RuntimeError(
+ 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/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 00000000000..aebb2fe838c
--- /dev/null
+++
b/providers/cncf/kubernetes/tests/system/cncf/kubernetes/example_kubernetes_pod_exec.py
@@ -0,0 +1,119 @@
+# 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
+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.providers.cncf.kubernetes.utils.container import
container_is_running
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
+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"]
+"""
+
+
[email protected](poke_interval=2, timeout=120, mode="reschedule")
+def wait_for_running_pod() -> bool:
+ hook = KubernetesHook()
+ try:
+ pod = hook.get_pod(name=POD_NAME, namespace=NAMESPACE)
+ except ApiException as error:
+ if error.status == 404:
+ return False
+ raise
+ return bool(
+ pod.status and pod.status.phase == PodPhase.RUNNING and
container_is_running(pod, CONTAINER_NAME)
+ )
+
+
+@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)
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 00000000000..afe3eda6b6c
--- /dev/null
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod_exec.py
@@ -0,0 +1,381 @@
+# 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.hooks.kubernetes import KubernetesHook
+from airflow.providers.cncf.kubernetes.operators.pod_exec import
KubernetesPodExecOperator
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
+
+MODULE = "airflow.providers.cncf.kubernetes.operators.pod_exec"
+
+
+def create_pod(
+ *,
+ container_names: tuple[str, ...] = ("main",),
+ annotations: dict[str, str] | None = None,
+ phase: str = PodPhase.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}.KubernetesPodExecOperator.log", spec=["info",
"warning"])
+ @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True)
+ def test_execute(self, kubernetes_stream_mock, log_mock, do_xcom_push,
expected_result):
+ operator, hook = create_operator(
+ namespace="test-namespace",
+ container_name="main",
+ do_xcom_push=do_xcom_push,
+ )
+ exec_client = create_exec_client(
+ stdout=("hello", "\nworld", "\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 == 3
+ assert exec_client.read_stderr.call_count == 2
+ exec_client.close.assert_called_once_with()
+ assert operator._exec_client is None
+ assert operator._exec_target is None
+ log_mock.info.assert_has_calls(
+ [
+ mock.call("[%s] %s", "stdout", "hello"),
+ mock.call("[%s] %s", "stdout", "world"),
+ mock.call("[%s] %s", "stderr", "warning"),
+ ]
+ )
+ log_mock.warning.assert_not_called()
+
+ @pytest.mark.parametrize("max_xcom_output_size", [0, -1, True, 1.5])
+ def test_rejects_invalid_max_xcom_output_size(self, max_xcom_output_size):
+ with pytest.raises(ValueError, match="must be a positive integer"):
+ create_operator(max_xcom_output_size=max_xcom_output_size)
+
+ @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True)
+ def test_rejects_xcom_output_over_limit(self, kubernetes_stream_mock):
+ operator, _ = create_operator(namespace="test-namespace",
do_xcom_push=True, max_xcom_output_size=3)
+ exec_client = create_exec_client(stdout=("é", "é"))
+ kubernetes_stream_mock.return_value = exec_client
+
+ with pytest.raises(RuntimeError, match="XCom limit of 3 bytes"):
+ operator.execute(context={})
+
+ exec_client.close.assert_called_once_with()
+
+ @mock.patch(f"{MODULE}.KubernetesPodExecOperator.log", spec=["info"])
+ @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True)
+ def test_execute_flushes_incomplete_log_lines(self,
kubernetes_stream_mock, log_mock):
+ operator, _ = create_operator(namespace="test-namespace",
container_name="main")
+ kubernetes_stream_mock.return_value = create_exec_client(
+ stdout=("partial ", "stdout"), stderr=("partial ", "stderr")
+ )
+
+ operator.execute(context={})
+
+ log_mock.info.assert_has_calls(
+ [
+ mock.call("[%s] %s", "stdout", "partial stdout"),
+ mock.call("[%s] %s", "stderr", "partial stderr"),
+ ]
+ )
+
+ @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(RuntimeError, 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(ValueError, 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 or container
status cannot be retrieved"),
+ (
+ create_pod(container_statuses=[SimpleNamespace(name="main",
state=None)]),
+ "is not running or container status cannot be retrieved",
+ ),
+ (
+ create_pod(
+ container_statuses=[SimpleNamespace(name="main",
state=SimpleNamespace(running=None))]
+ ),
+ "is not running or container status cannot be retrieved",
+ ),
+ ],
+ )
+ def test_rejects_unavailable_target(self, pod, expected_message):
+ operator, _ = create_operator(pod=pod)
+
+ with pytest.raises(RuntimeError, 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(RuntimeError, 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(RuntimeError, 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(RuntimeError, 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
+ assert operator._exec_target 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
+ assert operator._exec_target is None
+
+ @mock.patch(f"{MODULE}.kubernetes_stream", autospec=True)
+ def test_execute_closes_connection_when_output_consumption_fails(self,
kubernetes_stream_mock):
+ operator, _ = create_operator(namespace="test-namespace")
+ exec_client = create_exec_client()
+ exec_client.update.side_effect = RuntimeError("WebSocket update
failed")
+ kubernetes_stream_mock.return_value = exec_client
+
+ with pytest.raises(RuntimeError, match="WebSocket update failed"):
+ operator.execute(context={})
+
+ exec_client.close.assert_called_once_with()
+ assert operator._exec_client is None
+ assert operator._exec_target is None
+
+ def test_on_kill_without_active_connection(self):
+ operator, _ = create_operator()
+
+ operator.on_kill()
+
+ assert operator._exec_client is None
+ assert operator._exec_target is None
+
+ @mock.patch(f"{MODULE}.KubernetesPodExecOperator.log", spec=["exception"])
+ def test_close_error_does_not_mask_task_shutdown(self, log_mock):
+ operator, _ = create_operator()
+ exec_client = mock.MagicMock(spec=WSClient)
+ exec_client.close.side_effect = RuntimeError("connection already
closed")
+ operator._exec_client = exec_client
+ operator._exec_target = ("test-namespace", "existing-pod", "main")
+
+ operator.on_kill()
+
+ assert operator._exec_client is None
+ assert operator._exec_target is None
+ log_mock.exception.assert_called_once_with(
+ "Failed to close Kubernetes exec connection for container %s in
pod %s/%s",
+ "main",
+ "test-namespace",
+ "existing-pod",
+ )
