SameerMesiah97 commented on code in PR #72542:
URL: https://github.com/apache/airflow/pull/72542#discussion_r3961746994
##########
providers/amazon/docs/operators/eks.rst:
##########
@@ -205,6 +205,26 @@ 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`.
+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`
Review Comment:
I think you should add blank lines between each paragraph so the formatting
is consistent with the rest of the file.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_eks.py:
##########
@@ -1293,3 +1295,123 @@ 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()
+
+
+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 is None
+ 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
+
+ @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"],
+ )
+
Review Comment:
It wouldn't hurt to add the below assertions to ensure that the inherited
fields are excluded during templating:
```
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" not in operator.template_fields
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
