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 3b0388815f0 Centralize EKS kubeconfig and credential setup in EksHook
(#74775)
3b0388815f0 is described below
commit 3b0388815f0d8ec06b5a15e9d61d7b1bd5e54962
Author: Alejandro Morgante <[email protected]>
AuthorDate: Mon Oct 19 02:33:32 2026 -0300
Centralize EKS kubeconfig and credential setup in EksHook (#74775)
EKS execution, trigger reentry and polling repeat the same credential
resolution and temporary-file setup. A shared context keeps their file
lifetimes aligned while preserving access to the credential file for
refresh during long-running tasks.
---------
Co-authored-by: Sean Ghaeli <[email protected]>
---
generated/known_airflow_exceptions.txt | 2 +-
.../src/airflow/providers/amazon/aws/hooks/eks.py | 32 +++++++++
.../airflow/providers/amazon/aws/operators/eks.py | 46 +++---------
.../airflow/providers/amazon/aws/triggers/eks.py | 44 ++++--------
.../amazon/tests/unit/amazon/aws/hooks/test_eks.py | 84 ++++++++++++++++++++++
.../tests/unit/amazon/aws/operators/test_eks.py | 64 +++++++++++++++++
.../tests/unit/amazon/aws/triggers/test_eks.py | 29 ++++++++
7 files changed, 230 insertions(+), 71 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index b26359f86f8..47ebd560a73 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -71,7 +71,7 @@
providers/amazon/src/airflow/providers/amazon/aws/operators/comprehend.py::2
providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py::5
providers/amazon/src/airflow/providers/amazon/aws/operators/ec2.py::1
providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py::7
-providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py::12
+providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py::10
providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py::24
providers/amazon/src/airflow/providers/amazon/aws/operators/eventbridge.py::1
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py::6
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/eks.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/eks.py
index 5a20629ef91..70409e3b62e 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/eks.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/eks.py
@@ -608,6 +608,38 @@ class EksHook(AwsBaseHook):
with contextlib.suppress(OSError):
os.unlink(temp_path)
+ @contextmanager
+ def generate_config_file_with_credentials(
+ self, eks_cluster_name: str, pod_namespace: str | None
+ ) -> Generator[tuple[str, str], None, None]:
+ """
+ Create temporary kubeconfig and credential files using this hook's AWS
session.
+
+ Both files remain available until the context exits and are then
removed, including
+ when the caller raises an exception. The credential file can be
refreshed in place
+ while the context is active without regenerating the kubeconfig.
+
+ :param eks_cluster_name: Name of the EKS cluster.
+ :param pod_namespace: Kubernetes namespace, or ``None`` for the client
default.
+ :return: Tuple of kubeconfig path and credential file path.
+ """
+ credentials = self.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 self._secure_credential_context(
+ frozen_credentials.access_key, frozen_credentials.secret_key,
frozen_credentials.token
+ ) as credentials_file:
+ with self.generate_config_file(
+ eks_cluster_name=eks_cluster_name,
+ pod_namespace=pod_namespace,
+ credentials_file=credentials_file,
+ ) as config_file:
+ yield config_file, credentials_file
+
@contextmanager
def generate_config_file(
self,
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 3eee3434b91..bd12e2d3075 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
@@ -1252,25 +1252,10 @@ class EksPodOperator(KubernetesPodOperator):
aws_conn_id=self.aws_conn_id,
region_name=self.region,
)
- session = eks_hook.get_session()
- credentials_obj = session.get_credentials()
- if credentials_obj is None:
- raise AirflowException(
- "Unable to retrieve AWS credentials. Credentials may have
expired or not been configured. "
- "Please check your AWS connection configuration."
- )
- credentials = credentials_obj.get_frozen_credentials()
- with eks_hook._secure_credential_context(
- credentials.access_key, credentials.secret_key, credentials.token
- ) as credentials_file:
- # Store credentials file path for potential refresh during
long-running tasks
- self._credentials_file_path = credentials_file
- with eks_hook.generate_config_file(
- eks_cluster_name=self.cluster_name,
- pod_namespace=self.namespace,
- credentials_file=credentials_file,
- ) as self.config_file:
- return super().execute(context)
+ with eks_hook.generate_config_file_with_credentials(
+ eks_cluster_name=self.cluster_name, pod_namespace=self.namespace
+ ) as (self.config_file, self._credentials_file_path):
+ return super().execute(context)
def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
eks_hook = EksHook(
@@ -1279,25 +1264,10 @@ class EksPodOperator(KubernetesPodOperator):
)
eks_cluster_name = event["eks_cluster_name"]
pod_namespace = event["namespace"]
- session = eks_hook.get_session()
- credentials_obj = session.get_credentials()
- if credentials_obj is None:
- raise AirflowException(
- "Unable to retrieve AWS credentials. Credentials may have
expired or not been configured. "
- "Please check your AWS connection configuration."
- )
- credentials = credentials_obj.get_frozen_credentials()
- with eks_hook._secure_credential_context(
- credentials.access_key, credentials.secret_key, credentials.token
- ) as credentials_file:
- # Store credentials file path for potential refresh during
long-running tasks
- self._credentials_file_path = credentials_file
- with eks_hook.generate_config_file(
- eks_cluster_name=eks_cluster_name,
- pod_namespace=pod_namespace,
- credentials_file=credentials_file,
- ) as self.config_file:
- return super().trigger_reentry(context, event)
+ with eks_hook.generate_config_file_with_credentials(
+ eks_cluster_name=eks_cluster_name, pod_namespace=pod_namespace
+ ) as (self.config_file, self._credentials_file_path):
+ return super().trigger_reentry(context, event)
def _write_credentials_to_file(
self, credentials_file_path: str, access_key: str, secret_key: str,
session_token: str | None
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py
index 1d142d9dbb6..2f9cce21ae9 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/eks.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import datetime
+from pathlib import Path
from typing import TYPE_CHECKING, Any
import tenacity
@@ -29,6 +30,7 @@ from airflow.providers.amazon.aws.utils.waiter_with_logging
import async_wait
from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger
from airflow.providers.common.compat.sdk import AirflowException
from airflow.triggers.base import TriggerEvent
+from airflow.utils import yaml
from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
@@ -186,45 +188,23 @@ class EksPodTrigger(KubernetesPodTrigger):
async def run(self):
"""Generate fresh kubeconfig, then delegate to the parent trigger."""
- from airflow.utils import yaml
-
eks_hook = EksHook(
aws_conn_id=self._aws_conn_id,
region_name=self.region,
)
- session = eks_hook.get_session()
- credentials_obj = session.get_credentials()
- if credentials_obj is None:
- raise RuntimeError(
- "Unable to retrieve AWS credentials for EKS trigger. "
- "Credentials may have expired or not been configured."
+ with eks_hook.generate_config_file_with_credentials(
+ eks_cluster_name=self.eks_cluster_name,
pod_namespace=self.pod_namespace
+ ) as (config_file_path, _):
+ self.config_dict = yaml.safe_load(
+ Path(config_file_path).read_text() # noqa: ASYNC240
)
- credentials = credentials_obj.get_frozen_credentials()
-
- # Create fresh credential and kubeconfig files. The context managers
- # keep the temp files alive for the entire duration of the trigger.
- with eks_hook._secure_credential_context(
- credentials.access_key, credentials.secret_key, credentials.token
- ) as credentials_file:
- with eks_hook.generate_config_file(
- eks_cluster_name=self.eks_cluster_name,
- pod_namespace=self.pod_namespace,
- credentials_file=credentials_file,
- ) as config_file_path:
- # Reading a small local temp file created by the context
manager above.
- # Blocking I/O is acceptable here as the file is tiny and
local.
- from pathlib import Path
-
- self.config_dict = yaml.safe_load(
- Path(config_file_path).read_text() # noqa: ASYNC240
- )
- # Invalidate any previously cached hook so the new config_dict
- # is picked up when the parent creates the AsyncKubernetesHook.
- self.__dict__.pop("hook", None)
+ # Invalidate any previously cached hook so the new config_dict
+ # is picked up when the parent creates the AsyncKubernetesHook.
+ self.__dict__.pop("hook", None)
- async for event in super().run():
- yield event
+ async for event in super().run():
+ yield event
class EksDeleteClusterTrigger(AwsBaseWaiterTrigger):
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_eks.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_eks.py
index af8b7c0e6ce..26c35383f0a 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_eks.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_eks.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import json
import subprocess
import sys
+from contextlib import AbstractContextManager
from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING
@@ -30,6 +31,8 @@ from urllib.parse import urlsplit
import pytest
import time_machine
import yaml
+from boto3.session import Session
+from botocore.credentials import Credentials
from botocore.exceptions import ClientError
from moto import mock_aws
from moto.core import DEFAULT_ACCOUNT_ID
@@ -1212,6 +1215,87 @@ class TestEksHooks:
class TestEksHook:
python_executable = f"python{sys.version_info[0]}.{sys.version_info[1]}"
+ @mock.patch.object(EksHook, "generate_config_file", autospec=True)
+ @mock.patch.object(EksHook, "get_session", autospec=True)
+ def test_generate_config_file_with_credentials(self, get_session,
generate_config):
+ hook = EksHook(aws_conn_id=None)
+ session = mock.create_autospec(Session, instance=True)
+ session.get_credentials.return_value = Credentials("access-key",
"secret-key", "session-token")
+ get_session.return_value = session
+ config_context = mock.MagicMock(spec=AbstractContextManager)
+ config_context.__enter__.return_value = "/tmp/kubeconfig"
+ generate_config.return_value = config_context
+
+ with hook.generate_config_file_with_credentials("cluster",
"namespace") as (config, credentials):
+ assert config == "/tmp/kubeconfig"
+ assert Path(credentials).read_text() == (
+ "export AWS_ACCESS_KEY_ID='access-key'\n"
+ "export AWS_SECRET_ACCESS_KEY='secret-key'\n"
+ "export AWS_SESSION_TOKEN='session-token'\n"
+ )
+ config_context.__exit__.assert_not_called()
+
+ credentials = generate_config.call_args.kwargs["credentials_file"]
+ generate_config.assert_called_once_with(
+ hook, eks_cluster_name="cluster", pod_namespace="namespace",
credentials_file=credentials
+ )
+ assert not Path(credentials).exists()
+ config_context.__exit__.assert_called_once()
+
+ @mock.patch.object(EksHook, "generate_config_file", autospec=True)
+ @mock.patch.object(EksHook, "get_session", autospec=True)
+ def test_auth_context_cleans_up_after_caller_failure(self, get_session,
generate_config):
+ hook = EksHook(aws_conn_id=None)
+ session = mock.create_autospec(Session, instance=True)
+ session.get_credentials.return_value = Credentials("access-key",
"secret-key", "session-token")
+ get_session.return_value = session
+ config_context = mock.MagicMock(spec=AbstractContextManager)
+ config_context.__enter__.return_value = "/tmp/kubeconfig"
+ generate_config.return_value = config_context
+ error = ValueError("caller failed")
+
+ with pytest.raises(ValueError, match="caller failed") as raised:
+ with hook.generate_config_file_with_credentials("cluster",
"namespace") as (_, credentials):
+ raise error
+
+ assert raised.value is error
+ assert not Path(credentials).exists()
+ config_context.__exit__.assert_called_once_with(ValueError, error,
mock.ANY)
+
+ @mock.patch.object(EksHook, "generate_config_file", autospec=True)
+ @mock.patch.object(EksHook, "get_session", autospec=True)
+ def test_auth_context_cleans_up_after_kubeconfig_failure(self,
get_session, generate_config):
+ hook = EksHook(aws_conn_id=None)
+ session = mock.create_autospec(Session, instance=True)
+ session.get_credentials.return_value = Credentials("access-key",
"secret-key", "session-token")
+ get_session.return_value = session
+ config_context = mock.MagicMock(spec=AbstractContextManager)
+ error = ValueError("kubeconfig failed")
+ config_context.__enter__.side_effect = error
+ generate_config.return_value = config_context
+
+ with pytest.raises(ValueError, match="kubeconfig failed") as raised:
+ with hook.generate_config_file_with_credentials("cluster",
"namespace"):
+ pytest.fail("Kubeconfig failure must prevent entering the auth
context")
+
+ assert raised.value is error
+ credentials = generate_config.call_args.kwargs["credentials_file"]
+ assert not Path(credentials).exists()
+
+ @mock.patch.object(EksHook, "generate_config_file", autospec=True)
+ @mock.patch.object(EksHook, "get_session", autospec=True)
+ def test_generate_config_file_with_credentials_missing_credentials(self,
get_session, generate_config):
+ hook = EksHook(aws_conn_id=None)
+ session = mock.create_autospec(Session, instance=True)
+ session.get_credentials.return_value = None
+ get_session.return_value = session
+
+ with pytest.raises(RuntimeError, match="Unable to retrieve AWS
credentials"):
+ with hook.generate_config_file_with_credentials("cluster",
"namespace"):
+ pytest.fail("Missing credentials must prevent kubeconfig
creation")
+
+ generate_config.assert_not_called()
+
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn")
@pytest.mark.parametrize(
("aws_conn_id", "region_name", "expected_region_args"),
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 ebd756518b9..8d100e6391f 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
+from contextlib import AbstractContextManager, nullcontext
from typing import Any, TypedDict
from unittest import mock
@@ -1062,6 +1063,69 @@ class TestEksDeleteFargateProfileOperator:
class TestEksPodOperator:
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute",
autospec=True
+ )
+ @mock.patch.object(EksHook, "generate_config_file_with_credentials",
autospec=True)
+ @pytest.mark.parametrize("fails", [False, True])
+ def test_execute_uses_shared_auth_context(self, generate_config, execute,
fails):
+ config_context = mock.MagicMock(spec=AbstractContextManager)
+ config_context.__enter__.return_value = ("/tmp/kubeconfig",
"/tmp/credentials")
+ generate_config.return_value = config_context
+ operator = EksPodOperator(
+ task_id="run_pod", cluster_name=CLUSTER_NAME,
namespace="workloads", image="test-image"
+ )
+
+ def run_in_context(op, context):
+ assert op.config_file == "/tmp/kubeconfig"
+ assert op._credentials_file_path == "/tmp/credentials"
+ config_context.__exit__.assert_not_called()
+ if fails:
+ raise RuntimeError("parent failed")
+ return "result"
+
+ execute.side_effect = run_in_context
+ with pytest.raises(RuntimeError, match="parent failed") if fails else
nullcontext():
+ assert operator.execute({}) == "result"
+
+ generate_config.assert_called_once_with(
+ mock.ANY, eks_cluster_name=CLUSTER_NAME, pod_namespace="workloads"
+ )
+ execute.assert_called_once_with(operator, {})
+ config_context.__exit__.assert_called_once()
+
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.trigger_reentry",
autospec=True
+ )
+ @mock.patch.object(EksHook, "generate_config_file_with_credentials",
autospec=True)
+ @pytest.mark.parametrize("fails", [False, True])
+ def test_trigger_reentry_uses_shared_auth_context(self, generate_config,
reentry, fails):
+ config_context = mock.MagicMock(spec=AbstractContextManager)
+ config_context.__enter__.return_value = ("/tmp/kubeconfig",
"/tmp/credentials")
+ generate_config.return_value = config_context
+ operator = EksPodOperator(
+ task_id="run_pod", cluster_name=CLUSTER_NAME,
namespace="workloads", image="test-image"
+ )
+ event = {"eks_cluster_name": "event-cluster", "namespace":
"event-namespace"}
+
+ def run_in_context(op, context, event):
+ assert op.config_file == "/tmp/kubeconfig"
+ assert op._credentials_file_path == "/tmp/credentials"
+ config_context.__exit__.assert_not_called()
+ if fails:
+ raise RuntimeError("parent failed")
+ return "result"
+
+ reentry.side_effect = run_in_context
+ with pytest.raises(RuntimeError, match="parent failed") if fails else
nullcontext():
+ assert operator.trigger_reentry({}, event) == "result"
+
+ generate_config.assert_called_once_with(
+ mock.ANY, eks_cluster_name="event-cluster",
pod_namespace="event-namespace"
+ )
+ reentry.assert_called_once_with(operator, {}, event)
+ config_context.__exit__.assert_called_once()
+
@mock.patch("airflow.providers.cncf.kubernetes.operators.pod.KubernetesPodOperator.execute")
@mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook.generate_config_file")
@mock.patch("airflow.providers.amazon.aws.hooks.eks.EksHook._secure_credential_context")
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_eks.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_eks.py
index 9adb0859a4e..163aaf3963a 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_eks.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_eks.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import datetime
+from contextlib import AbstractContextManager
from unittest.mock import AsyncMock, MagicMock, Mock, call, patch
import pytest
@@ -369,6 +370,34 @@ class
TestEksDeleteClusterTriggerDeleteNodegroupsAndFargateProfiles(TestEksTrigg
class TestEksPodTrigger:
"""Tests for EksPodTrigger."""
+ @pytest.mark.asyncio
+ @patch("pathlib.Path.read_text", autospec=True, return_value="apiVersion:
v1\nkind: Config\n")
+ @patch(
+
"airflow.providers.amazon.aws.hooks.eks.EksHook.generate_config_file_with_credentials",
autospec=True
+ )
+
@patch("airflow.providers.cncf.kubernetes.triggers.pod.KubernetesPodTrigger.run",
autospec=True)
+ async def test_keeps_shared_auth_context_while_polling(self, parent_run,
generate_config, read_text):
+ config_context = MagicMock(spec=AbstractContextManager)
+ config_context.__enter__.return_value = ("/tmp/kubeconfig",
"/tmp/credentials")
+ generate_config.return_value = config_context
+
+ async def events():
+ config_context.__exit__.assert_not_called()
+ yield TriggerEvent({"status": "success"})
+
+ parent_run.return_value = events()
+ trigger = self._create_trigger()
+ events_received = [event async for event in trigger.run()]
+
+ assert events_received == [TriggerEvent({"status": "success"})]
+ assert trigger.config_dict == {"apiVersion": "v1", "kind": "Config"}
+ generate_config.assert_called_once()
+ assert generate_config.call_args.kwargs == {
+ "eks_cluster_name": CLUSTER_NAME,
+ "pod_namespace": "default",
+ }
+ config_context.__exit__.assert_called_once()
+
TRIGGER_START_TIME = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC)
def _create_trigger(self, **overrides):
