This is an automated email from the ASF dual-hosted git repository.
shahar1 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 35e8183358a Add EKS operator for commands in existing Pods (#72542)
35e8183358a is described below
commit 35e8183358a82e9ad53a3b380910bd7465548310
Author: Alejandro Morgante <[email protected]>
AuthorDate: Wed Oct 21 16:23:27 2026 -0300
Add EKS operator for commands in existing Pods (#72542)
Allow Dags to execute commands in externally managed EKS Pods without
making Airflow responsible for their lifecycle.
---
.../tests/unit/always/test_project_structure.py | 2 +
providers/amazon/docs/operators/eks.rst | 30 +++
.../airflow/providers/amazon/aws/operators/eks.py | 121 ++++++++++-
.../amazon/aws/example_eks_with_nodegroups.py | 89 ++++++++-
.../tests/unit/amazon/aws/operators/test_eks.py | 222 ++++++++++++++++++++-
5 files changed, 458 insertions(+), 6 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py
b/airflow-core/tests/unit/always/test_project_structure.py
index a9d12c05600..88a17745b89 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -511,6 +511,8 @@ class
TestAmazonProviderProjectStructure(ExampleCoverageTest):
BASE_CLASSES = {
"airflow.providers.amazon.aws.operators.base_aws.AwsBaseOperator",
+ # Compatibility fallback for EksPodExecOperator, not a standalone
Amazon operator.
+ "airflow.providers.amazon.aws.operators.eks.KubernetesPodExecOperator",
"airflow.providers.amazon.aws.operators.glue_crawler._GlueCrawlerBaseOperator",
"airflow.providers.amazon.aws.operators.rds.RdsBaseOperator",
"airflow.providers.amazon.aws.operators.sagemaker.SageMakerBaseOperator",
diff --git a/providers/amazon/docs/operators/eks.rst
b/providers/amazon/docs/operators/eks.rst
index 76c0f5ea4ed..3d8c99fb4d6 100644
--- a/providers/amazon/docs/operators/eks.rst
+++ b/providers/amazon/docs/operators/eks.rst
@@ -205,6 +205,36 @@ Note: An Amazon EKS Cluster with underlying compute
infrastructure is required.
:start-after: [START howto_operator_eks_pod_operator]
:end-before: [END howto_operator_eks_pod_operator]
+.. _howto/operator:EksPodExecOperator:
+
+Execute a command in an existing Pod on Amazon EKS
+==================================================
+
+To execute a command in a running container without managing the Pod
lifecycle, use
+:class:`~airflow.providers.amazon.aws.operators.eks.EksPodExecOperator`.
+
+This operator requires ``apache-airflow-providers-cncf-kubernetes>=10.22.0``.
+Existing EKS operators remain available with older supported versions of the
Kubernetes provider.
+
+As with ``EksPodOperator``, ``kubernetes_conn_id`` defaults to
``kubernetes_default`` and can be
+set to another Kubernetes connection. If the default connection contains
``kube_config`` or
+``cluster_context``, use a separate connection without those settings, since
EKS generates its own kubeconfig.
+
+The Pod must already exist and be running. The operator streams command
output, waits for the exit code,
+and does not create, restart, or delete the Pod.
+
+The AWS identity must have permission to call ``eks:DescribeCluster`` and be
authorized to access the
+EKS cluster. Kubernetes RBAC must allow ``get`` on ``pods`` and both ``get``
and ``create`` on ``pods/exec``.
+
+See
:class:`~airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator`
+for command, output, XCom, and retry behavior.
+
+.. exampleinclude::
/../../amazon/tests/system/amazon/aws/example_eks_with_nodegroups.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_eks_pod_exec]
+ :end-before: [END howto_operator_eks_pod_exec]
+
Sensors
-------
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
index bd12e2d3075..4840c1a38a0 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
@@ -47,10 +47,28 @@ from airflow.providers.amazon.aws.utils import (
build_resource_in_use_retry_args,
validate_execute_complete_event,
)
-from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.amazon.aws.utils.mixins import AwsHookParams,
aws_template_fields
from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
-from airflow.providers.common.compat.sdk import AirflowException, conf
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowOptionalProviderFeatureException,
+ BaseOperator,
+ conf,
+)
+
+try:
+ from airflow.providers.cncf.kubernetes.operators.pod_exec import
KubernetesPodExecOperator
+except ImportError:
+
+ class KubernetesPodExecOperator(BaseOperator): # type: ignore[no-redef]
+ """Keep existing EKS operators importable with older Kubernetes
providers."""
+
+ def __init__(self, **kwargs):
+ raise AirflowOptionalProviderFeatureException(
+ "EksPodExecOperator requires
apache-airflow-providers-cncf-kubernetes>=10.22.0."
+ )
+
try:
from airflow.providers.cncf.kubernetes.operators.pod import
KubernetesPodOperator
@@ -1348,3 +1366,102 @@ class EksPodOperator(KubernetesPodOperator):
self.log.exception("Failed to refresh AWS credentials.")
raise
super()._refresh_cached_properties()
+
+
+class EksPodExecOperator(KubernetesPodExecOperator):
+ """
+ Execute a command in a running container of an existing Pod on Amazon EKS.
+
+ The operator authenticates with Amazon EKS and delegates command execution
to
+
:class:`~airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator`.
+ It does not create, restart, or delete the target Pod.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:EksPodExecOperator`
+
+ :param cluster_name: The name of the Amazon EKS Cluster containing the
Pod. (templated)
+ :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 ``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.
+ Defaults to ``None``. (templated)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
(templated)
+ Defaults to ``aws_default``. If this is ``None`` or empty, the default
boto3 credential
+ strategy is used without an Airflow connection lookup.
+ :param region_name: AWS region containing the Amazon EKS Cluster.
(templated)
+ Defaults to ``None``, which uses the region from the AWS connection
when available and
+ otherwise falls back to the default boto3 region strategy.
+ :param verify: Whether to verify SSL certificates, or the path to a CA
bundle. Defaults to
+ ``None``, which uses the value from the AWS connection when available.
(templated)
+ :param botocore_config: Configuration dictionary for the botocore client.
Defaults to ``None``,
+ which uses ``config_kwargs`` from the AWS connection when available.
+ :param kubernetes_conn_id: Kubernetes connection used for additional
client configuration.
+ Defaults to ``kubernetes_default``, as with ``EksPodOperator``.
(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] = aws_template_fields(
+ "cluster_name",
+ *(
+ field
+ for field in KubernetesPodExecOperator.template_fields
+ if field not in {"cluster_context", "config_file"}
+ ),
+ )
+
+ def __init__(
+ self,
+ *,
+ cluster_name: str,
+ pod_name: str,
+ command: Sequence[str],
+ namespace: str = DEFAULT_NAMESPACE_NAME,
+ container_name: str | None = None,
+ aws_conn_id: str | None = DEFAULT_CONN_ID,
+ region_name: str | None = None,
+ verify: bool | str | None = None,
+ botocore_config: dict | None = None,
+ **kwargs,
+ ) -> None:
+ for parameter in ("in_cluster", "config_file", "cluster_context"):
+ if parameter in kwargs:
+ raise ValueError(f"The {parameter} is not an allowed parameter
for the EksPodExecOperator.")
+ hook_params = AwsHookParams.from_constructor(
+ aws_conn_id, region_name, verify, botocore_config,
additional_params=kwargs
+ )
+ super().__init__(
+ pod_name=pod_name,
+ command=command,
+ namespace=namespace,
+ container_name=container_name,
+ in_cluster=False,
+ cluster_context=None,
+ config_file=None,
+ **kwargs,
+ )
+ self.cluster_name = cluster_name
+ self.aws_conn_id = hook_params.aws_conn_id
+ self.region_name = hook_params.region_name
+ self.verify = hook_params.verify
+ self.botocore_config = hook_params.botocore_config
+
+ def execute(self, context: Context) -> str | None:
+ eks_hook = EksHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+ with eks_hook.generate_config_file_with_credentials(
+ eks_cluster_name=self.cluster_name, pod_namespace=self.namespace
+ ) as (config_file, _):
+ self.config_file = config_file
+ try:
+ return super().execute(context)
+ finally:
+ self.config_file = None
+ self.__dict__.pop("client", None)
+ self.__dict__.pop("hook", None)
diff --git
a/providers/amazon/tests/system/amazon/aws/example_eks_with_nodegroups.py
b/providers/amazon/tests/system/amazon/aws/example_eks_with_nodegroups.py
index 018d49c0831..8ed41186500 100644
--- a/providers/amazon/tests/system/amazon/aws/example_eks_with_nodegroups.py
+++ b/providers/amazon/tests/system/amazon/aws/example_eks_with_nodegroups.py
@@ -16,19 +16,26 @@
# under the License.
from __future__ import annotations
+import asyncio
+from contextlib import contextmanager
from datetime import datetime
+from typing import TYPE_CHECKING
import boto3
+from kubernetes.client import V1Container, V1ObjectMeta, V1Pod, V1PodSpec
-from airflow.providers.amazon.aws.hooks.eks import ClusterStates,
NodegroupStates
+from airflow.providers.amazon.aws.hooks.eks import ClusterStates, EksHook,
NodegroupStates
from airflow.providers.amazon.aws.operators.eks import (
EksCreateClusterOperator,
EksCreateNodegroupOperator,
EksDeleteClusterOperator,
EksDeleteNodegroupOperator,
+ EksPodExecOperator,
EksPodOperator,
)
from airflow.providers.amazon.aws.sensors.eks import EksClusterStateSensor,
EksNodegroupStateSensor
+from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodManager
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
@@ -49,7 +56,13 @@ except ImportError:
from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
from system.amazon.aws.utils.k8s import get_describe_pod_operator
+if TYPE_CHECKING:
+ from collections.abc import Generator
+
+ from kubernetes.client import CoreV1Api
+
DAG_ID = "example_eks_with_nodegroups"
+EXPECTED_EXEC_OUTPUT = "command executed in existing EKS pod"
# Externally fetched variables:
ROLE_ARN_KEY = "ROLE_ARN"
@@ -76,6 +89,43 @@ def delete_launch_template(template_name: str):
boto3.client("ec2").delete_launch_template(LaunchTemplateName=template_name)
+@contextmanager
+def get_eks_kubernetes_client(cluster_name: str) -> Generator[CoreV1Api, None,
None]:
+ eks_hook = EksHook()
+ with eks_hook.generate_config_file_with_credentials(
+ eks_cluster_name=cluster_name, pod_namespace="default"
+ ) as (config_file, _):
+ yield KubernetesHook(kubernetes_conn_id=None,
config_file=config_file).core_v1_client
+
+
+@task
+def create_exec_pod(cluster_name: str, pod_name: str) -> None:
+ pod = V1Pod(
+ metadata=V1ObjectMeta(name=pod_name, namespace="default"),
+ spec=V1PodSpec(
+ containers=[V1Container(name="main", image="busybox:1.38.0",
command=["sleep", "3600"])],
+ restart_policy="Never",
+ ),
+ )
+ with get_eks_kubernetes_client(cluster_name) as kube_client:
+ pod_manager = PodManager(kube_client=kube_client)
+ created_pod = pod_manager.create_pod(pod)
+ asyncio.run(pod_manager.await_pod_start(created_pod))
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_exec_pod(cluster_name: str, pod_name: str) -> None:
+ pod = V1Pod(metadata=V1ObjectMeta(name=pod_name, namespace="default"))
+ with get_eks_kubernetes_client(cluster_name) as kube_client:
+ PodManager(kube_client=kube_client).delete_pod(pod)
+
+
+@task
+def verify_exec_output(output: str) -> None:
+ if output != EXPECTED_EXEC_OUTPUT:
+ raise ValueError(f"Unexpected command output: {output!r}")
+
+
with DAG(
dag_id=DAG_ID,
schedule="@once",
@@ -88,6 +138,7 @@ with DAG(
cluster_name = f"{env_id}-cluster"
nodegroup_name = f"{env_id}-nodegroup"
launch_template_name = f"{env_id}-launch-template"
+ exec_pod_name = f"{env_id}-exec-pod"
# [START howto_operator_eks_create_cluster]
# Create an Amazon EKS Cluster control plane without attaching compute
service.
@@ -151,6 +202,22 @@ with DAG(
# it is cleaned anyway with the cluster later on.
start_pod.is_delete_operator_pod = False
+ create_exec_pod_task = create_exec_pod(cluster_name, exec_pod_name)
+
+ # [START howto_operator_eks_pod_exec]
+ run_command = EksPodExecOperator(
+ task_id="run_command_in_existing_pod",
+ cluster_name=cluster_name,
+ pod_name=exec_pod_name,
+ command=["sh", "-c", f"printf '{EXPECTED_EXEC_OUTPUT}'"],
+ do_xcom_push=True,
+ )
+ # [END howto_operator_eks_pod_exec]
+
+ exec_output_is_valid = verify_exec_output(run_command.output)
+
+ delete_exec_pod_task = delete_exec_pod(cluster_name, exec_pod_name)
+
describe_pod = get_describe_pod_operator(
cluster_name, pod_name="{{ ti.xcom_pull(key='pod_name',
task_ids='run_pod') }}"
)
@@ -211,15 +278,31 @@ with DAG(
# TEST SETUP
test_context,
create_launch_template(launch_template_name),
- # TEST BODY
create_cluster,
await_create_cluster,
create_nodegroup,
await_create_nodegroup,
+ )
+ chain(
+ # TEST BODY: EksPodOperator
+ await_create_nodegroup,
start_pod,
- # TEST TEARDOWN
describe_pod,
await_nodegroup_stable,
+ )
+ chain(
+ # TEST BODY: EksPodExecOperator
+ await_create_nodegroup,
+ create_exec_pod_task,
+ run_command,
+ exec_output_is_valid,
+ # TEST TEARDOWN
+ delete_exec_pod_task,
+ await_nodegroup_stable,
+ )
+ chain(
+ # TEST TEARDOWN
+ await_nodegroup_stable,
delete_nodegroup, # part of the test AND teardown
await_delete_nodegroup,
await_cluster_stable,
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
index 8d100e6391f..418878c777c 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import runpy
from contextlib import AbstractContextManager, nullcontext
from typing import Any, TypedDict
from unittest import mock
@@ -27,6 +28,7 @@ from botocore.waiter import Waiter
from airflow.exceptions import AirflowException,
AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.eks import ClusterStates, EksHook
+from airflow.providers.amazon.aws.operators import eks
from airflow.providers.amazon.aws.operators.eks import (
EksCreateClusterOperator,
EksCreateFargateProfileOperator,
@@ -34,6 +36,7 @@ from airflow.providers.amazon.aws.operators.eks import (
EksDeleteClusterOperator,
EksDeleteFargateProfileOperator,
EksDeleteNodegroupOperator,
+ EksPodExecOperator,
EksPodOperator,
)
from airflow.providers.amazon.aws.triggers.eks import (
@@ -43,7 +46,7 @@ from airflow.providers.amazon.aws.triggers.eks import (
EksPodTrigger,
)
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
-from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, TaskDeferred
from unit.amazon.aws.utils.eks_test_constants import (
NODEROLE_ARN,
@@ -1463,3 +1466,220 @@ class TestEksPodOperator:
with pytest.raises(RuntimeError, match="Pod must be created with
metadata before deferring"):
op.invoke_defer_method()
+
+
[email protected]("sys.modules",
{"airflow.providers.cncf.kubernetes.operators.pod_exec": None})
+def test_eks_operators_remain_available_without_pod_exec():
+ operators = runpy.run_path(eks.__file__)
+
+ operators["EksPodOperator"](task_id="existing_operator",
cluster_name=CLUSTER_NAME)
+
+
[email protected]("sys.modules",
{"airflow.providers.cncf.kubernetes.operators.pod_exec": None})
+def test_eks_pod_exec_requires_available_kubernetes_pod_exec():
+ operators = runpy.run_path(eks.__file__)
+
+ with pytest.raises(
+ AirflowOptionalProviderFeatureException,
+ match=r"EksPodExecOperator requires
apache-airflow-providers-cncf-kubernetes>=10\.22\.0",
+ ):
+ operators["EksPodExecOperator"](
+ task_id="exec_operator", cluster_name=CLUSTER_NAME,
pod_name="existing-pod", command=["true"]
+ )
+
+
+class TestEksPodExecOperator:
+ @pytest.mark.parametrize("parameter", ["in_cluster", "config_file",
"cluster_context"])
+ @pytest.mark.parametrize("value", [None, "custom-value"])
+ def test_rejects_managed_kubernetes_parameters(self, parameter, value):
+ with pytest.raises(
+ ValueError, match=f"The {parameter} is not an allowed parameter
for the EksPodExecOperator"
+ ):
+ EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ command=["echo", "hello"],
+ **{parameter: value},
+ )
+
+ @staticmethod
+ def configure_eks_auth(eks_hook_mock):
+ eks_hook = eks_hook_mock.return_value
+ auth_context = mock.MagicMock(spec=AbstractContextManager)
+ auth_context.__enter__.return_value = (
+ "/tmp/eks-kubeconfig",
+ "/tmp/aws-credentials",
+ )
+ eks_hook.generate_config_file_with_credentials.return_value =
auth_context
+ return eks_hook
+
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator.execute",
+ autospec=True,
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.eks.EksHook",
autospec=True)
+ def test_execute_uses_shared_auth_context(self, eks_hook_mock,
pod_exec_execute_mock):
+ eks_hook = self.configure_eks_auth(eks_hook_mock)
+
+ def execute_with_generated_config(operator, context):
+ assert operator.config_file == "/tmp/eks-kubeconfig"
+
eks_hook.generate_config_file_with_credentials.return_value.__exit__.assert_not_called()
+ return "command output"
+
+ pod_exec_execute_mock.side_effect = execute_with_generated_config
+ operator = EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ namespace="workloads",
+ container_name="worker",
+ command=["dbt", "run"],
+ aws_conn_id="aws_test",
+ region_name="us-east-2",
+ verify=False,
+ botocore_config={"retries": {"max_attempts": 5}},
+ do_xcom_push=True,
+ max_xcom_output_size=1024,
+ )
+
+ result = operator.execute({})
+
+ assert result == "command output"
+ eks_hook_mock.assert_called_once_with(
+ aws_conn_id="aws_test",
+ region_name="us-east-2",
+ verify=False,
+ config={"retries": {"max_attempts": 5}},
+ )
+ eks_hook.generate_config_file_with_credentials.assert_called_once_with(
+ eks_cluster_name=CLUSTER_NAME,
+ pod_namespace="workloads",
+ )
+
eks_hook.generate_config_file_with_credentials.return_value.__exit__.assert_called_once_with(
+ None, None, None
+ )
+ pod_exec_execute_mock.assert_called_once_with(operator, {})
+ assert operator.config_file is None
+
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator.execute",
+ autospec=True,
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.eks.EksHook",
autospec=True)
+ def test_execute_clears_config_file_on_failure(self, eks_hook_mock,
pod_exec_execute_mock):
+ eks_hook = self.configure_eks_auth(eks_hook_mock)
+ error = RuntimeError("command failed")
+ pod_exec_execute_mock.side_effect = error
+ operator = EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ command=["false"],
+ )
+
+ with pytest.raises(RuntimeError, match="command failed") as raised:
+ operator.execute({})
+
+ assert raised.value is error
+
eks_hook.generate_config_file_with_credentials.return_value.__exit__.assert_called_once_with(
+ RuntimeError, error, mock.ANY
+ )
+ assert operator.config_file is None
+
+ @pytest.mark.parametrize("first_execution_fails", [False, True])
+
@mock.patch("airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesHook",
autospec=True)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator.execute",
+ autospec=True,
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.eks.EksHook",
autospec=True)
+ def test_execute_recreates_client_with_new_config(
+ self, eks_hook_mock, pod_exec_execute_mock, kubernetes_hook_mock,
first_execution_fails
+ ):
+ eks_hook = self.configure_eks_auth(eks_hook_mock)
+
eks_hook.generate_config_file_with_credentials.return_value.__enter__.side_effect
= [
+ ("/tmp/first-kubeconfig", "/tmp/first-credentials"),
+ ("/tmp/second-kubeconfig", "/tmp/second-credentials"),
+ ]
+ clients = []
+
+ def execute_with_client(operator, context):
+ clients.append(operator.client)
+ if first_execution_fails and len(clients) == 1:
+ raise RuntimeError("command failed")
+ return "command output"
+
+ pod_exec_execute_mock.side_effect = execute_with_client
+ operator = EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ command=["true"],
+ )
+
+ if first_execution_fails:
+ with pytest.raises(RuntimeError, match="command failed"):
+ operator.execute({})
+ else:
+ operator.execute({})
+ assert operator.execute({}) == "command output"
+
+ assert kubernetes_hook_mock.call_args_list == [
+ mock.call(
+ conn_id="kubernetes_default", in_cluster=False,
config_file=config_file, cluster_context=None
+ )
+ for config_file in ("/tmp/first-kubeconfig",
"/tmp/second-kubeconfig")
+ ]
+
+ @pytest.mark.parametrize("kubernetes_conn_id", ["eks_kubernetes", None])
+ def test_kubernetes_connection_can_be_configured(self, kubernetes_conn_id):
+ operator = EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ command=["true"],
+ kubernetes_conn_id=kubernetes_conn_id,
+ )
+
+ assert operator.kubernetes_conn_id == kubernetes_conn_id
+ assert operator.hook.conn_id == (kubernetes_conn_id or
"kubernetes_default")
+
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.operators.pod_exec.KubernetesPodExecOperator.execute",
+ autospec=True,
+ )
+ @mock.patch("airflow.providers.amazon.aws.operators.eks.EksHook",
autospec=True)
+ def test_execute_propagates_auth_context_failure(self, eks_hook_mock,
pod_exec_execute_mock):
+ eks_hook = self.configure_eks_auth(eks_hook_mock)
+ error = RuntimeError("Unable to retrieve AWS credentials")
+
eks_hook.generate_config_file_with_credentials.return_value.__enter__.side_effect
= error
+ operator = EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ command=["true"],
+ )
+
+ with pytest.raises(RuntimeError, match="Unable to retrieve AWS
credentials") as raised:
+ operator.execute({})
+
+ assert raised.value is error
+ pod_exec_execute_mock.assert_not_called()
+
+ def test_template_fields(self):
+ operator = EksPodExecOperator(
+ task_id="run_command",
+ cluster_name=CLUSTER_NAME,
+ pod_name="existing-pod",
+ command=["dbt", "run"],
+ )
+
+ validate_template_fields(operator)
+ assert "cluster_name" in operator.template_fields
+ assert "pod_name" in operator.template_fields
+ assert "command" in operator.template_fields
+
+ assert "cluster_context" not in operator.template_fields
+ assert "config_file" not in operator.template_fields
+ assert "kubernetes_conn_id" in operator.template_fields
