diff --git a/airflow-core/tests/unit/always/test_project_structure.py b/airflow-core/tests/unit/always/test_project_structure.py index 338279620684b..435c72e4f24f8 100644 --- a/airflow-core/tests/unit/always/test_project_structure.py +++ b/airflow-core/tests/unit/always/test_project_structure.py @@ -498,6 +498,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 76c0f5ea4ed11..e73ee5b7e2eaa 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 ``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 3b54c32eb16ca..1f6529c23130e 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 @@ 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 @@ -1368,3 +1386,113 @@ def _refresh_cached_properties(self) -> None: 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: + 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, + ) + credentials = eks_hook.get_session().get_credentials() + if credentials is None: + raise RuntimeError( + "Unable to retrieve AWS credentials. Credentials may have expired or not been configured. " + "Please check your AWS connection configuration." + ) + frozen_credentials = credentials.get_frozen_credentials() + with eks_hook._secure_credential_context( + frozen_credentials.access_key, + frozen_credentials.secret_key, + frozen_credentials.token, + ) as credentials_file: + with eks_hook.generate_config_file( + eks_cluster_name=self.cluster_name, + pod_namespace=self.namespace, + credentials_file=credentials_file, + ) 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 018d49c083166..6f31b1f9c9141 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 @@ 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,50 @@ 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() + credentials = eks_hook.get_session().get_credentials() + if credentials is None: + raise RuntimeError("Unable to retrieve AWS credentials for the EKS system test.") + frozen_credentials = credentials.get_frozen_credentials() + with eks_hook._secure_credential_context( + frozen_credentials.access_key, + frozen_credentials.secret_key, + frozen_credentials.token, + ) as credentials_file: + with eks_hook.generate_config_file(cluster_name, "default", credentials_file) 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 +145,7 @@ def delete_launch_template(template_name: str): 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 +209,22 @@ def delete_launch_template(template_name: str): # 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 +285,31 @@ def delete_launch_template(template_name: str): # 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 f95be60d51d3d..dffe06768bff4 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,8 @@ # under the License. from __future__ import annotations +import runpy +from types import SimpleNamespace from typing import Any, TypedDict from unittest import mock @@ -26,6 +28,7 @@ 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, @@ -33,6 +36,7 @@ EksDeleteClusterOperator, EksDeleteFargateProfileOperator, EksDeleteNodegroupOperator, + EksPodExecOperator, EksPodOperator, ) from airflow.providers.amazon.aws.triggers.eks import ( @@ -42,7 +46,7 @@ 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, @@ -1300,3 +1304,203 @@ def test_invoke_defer_method_raises_when_pod_is_none(self, mock_convert_config): with pytest.raises(RuntimeError, match="Pod must be created with metadata before deferring"): op.invoke_defer_method() + + +@mock.patch.dict("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__) + + operator = operators["EksPodOperator"](task_id="existing_operator", cluster_name=CLUSTER_NAME) + assert operator.cluster_name == CLUSTER_NAME + 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: + @staticmethod + def configure_eks_auth(eks_hook_mock): + eks_hook = eks_hook_mock.return_value + credentials = eks_hook.get_session.return_value.get_credentials.return_value + credentials.get_frozen_credentials.return_value = SimpleNamespace( + access_key="test_access_key", + secret_key="test_secret_key", + token="test_token", + ) + eks_hook._secure_credential_context.return_value.__enter__.return_value = "/tmp/aws-credentials" + eks_hook.generate_config_file.return_value.__enter__.return_value = "/tmp/eks-kubeconfig" + 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(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" + assert operator.kubernetes_conn_id == "kubernetes_default" + assert operator.in_cluster is False + assert operator.do_xcom_push is True + assert operator.max_xcom_output_size == 1024 + 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.get_session.return_value.get_credentials.assert_called_once_with() + eks_hook._secure_credential_context.assert_called_once_with( + "test_access_key", "test_secret_key", "test_token" + ) + eks_hook.generate_config_file.assert_called_once_with( + eks_cluster_name=CLUSTER_NAME, + pod_namespace="workloads", + credentials_file="/tmp/aws-credentials", + ) + 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): + self.configure_eks_auth(eks_hook_mock) + pod_exec_execute_mock.side_effect = RuntimeError("command failed") + operator = EksPodExecOperator( + task_id="run_command", + cluster_name=CLUSTER_NAME, + pod_name="existing-pod", + command=["false"], + ) + + with pytest.raises(RuntimeError, match="command failed"): + operator.execute({}) + + 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.return_value.__enter__.side_effect = [ + "/tmp/first-kubeconfig", + "/tmp/second-kubeconfig", + ] + 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_rejects_missing_credentials(self, eks_hook_mock, pod_exec_execute_mock): + eks_hook = eks_hook_mock.return_value + eks_hook.get_session.return_value.get_credentials.return_value = None + 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"): + operator.execute({}) + + eks_hook._secure_credential_context.assert_not_called() + eks_hook.generate_config_file.assert_not_called() + 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