Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions providers/amazon/src/airflow/providers/amazon/aws/hooks/eks.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,32 @@ class NodegroupStates(Enum):
echo $json_string
"""

# Cache directory used by the AWS CLI when running exec-based credential
# plugins such as ``aws eks get-token``. Older versions of the AWS CLI
# create this directory with ``os.makedirs()`` without ``exist_ok=True``,
# causing a ``FileExistsError`` race condition when parallel tasks on the
# same Celery worker invoke the CLI simultaneously.
# See: https://github.com/apache/airflow/issues/60943
_AWS_CLI_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".aws", "cli", "cache")


def _ensure_eks_cache_dirs() -> None:
"""
Pre-create the AWS CLI cache directory to avoid a race condition.

When multiple tasks run concurrently on the same worker and each
invokes ``aws eks get-token`` via exec-based kubeconfig auth, the
AWS CLI may try to create ``~/.aws/cli/cache`` simultaneously.
Older botocore versions do not pass ``exist_ok=True``, so a
``FileExistsError`` is raised when two processes race.

Pre-creating the directory here means it already exists by the
time the CLI runs.

.. seealso:: https://github.com/apache/airflow/issues/60943
"""
os.makedirs(_AWS_CLI_CACHE_DIR, exist_ok=True)


class EksHook(AwsBaseHook):
"""
Expand Down Expand Up @@ -604,6 +630,10 @@ def generate_config_file(
:param eks_cluster_name: The name of the cluster to generate kubeconfig file for.
:param pod_namespace: The namespace to run within kubernetes.
"""
# Pre-create the AWS CLI cache directory so that concurrent
# exec-based auth calls on the same worker do not race.
_ensure_eks_cache_dirs()

args = ""
if self.region_name is not None:
args = args + f" --region-name {self.region_name}"
Expand Down
46 changes: 45 additions & 1 deletion providers/amazon/tests/unit/amazon/aws/hooks/test_eks.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@
NODEGROUP_NOT_FOUND_MSG,
)

from airflow.providers.amazon.aws.hooks.eks import EksHook
from airflow.providers.amazon.aws.hooks.eks import EksHook, _ensure_eks_cache_dirs

from unit.amazon.aws.utils.eks_test_constants import (
DEFAULT_CONN_ID,
Expand Down Expand Up @@ -1274,6 +1274,50 @@ def test_generate_config_file(self, mock_conn, aws_conn_id, region_name, expecte
assert expected_region_args in command_arg


EKS_HOOK_MODULE = "airflow.providers.amazon.aws.hooks.eks"


class TestEnsureEksCacheDirs:
"""Tests for _ensure_eks_cache_dirs."""

def test_creates_cache_dir(self, tmp_path):
"""Verify the AWS CLI cache directory is created when it does not exist."""
cache_dir = tmp_path / ".aws" / "cli" / "cache"

with mock.patch(f"{EKS_HOOK_MODULE}._AWS_CLI_CACHE_DIR", str(cache_dir)):
_ensure_eks_cache_dirs()

assert cache_dir.is_dir()

def test_no_error_when_dir_already_exists(self, tmp_path):
"""Verify no error is raised when the cache directory already exists."""
cache_dir = tmp_path / ".aws" / "cli" / "cache"
cache_dir.mkdir(parents=True)

with mock.patch(f"{EKS_HOOK_MODULE}._AWS_CLI_CACHE_DIR", str(cache_dir)):
_ensure_eks_cache_dirs()

assert cache_dir.is_dir()

@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn")
def test_generate_config_file_calls_ensure_cache_dirs(self, mock_conn):
"""Verify that generate_config_file() pre-creates the cache directory."""
mock_conn.describe_cluster.return_value = {
"cluster": {"certificateAuthority": {"data": "test-cert"}, "endpoint": "test-endpoint"}
}
hook = EksHook(aws_conn_id=None, region_name=None)
hook.get_connection = lambda _: None

with mock.patch(f"{EKS_HOOK_MODULE}._ensure_eks_cache_dirs") as mock_ensure:
with hook.generate_config_file(
eks_cluster_name="test-cluster",
pod_namespace="default",
credentials_file="/tmp/test_creds.aws_creds",
):
pass
mock_ensure.assert_called_once()


# Helper methods for repeated assert combinations.
def assert_all_arn_values_are_valid(expected_arn_values, pattern, arn_under_test) -> None:
"""
Expand Down
Loading