diff --git a/providers/src/airflow/providers/amazon/aws/hooks/eks.py b/providers/src/airflow/providers/amazon/aws/hooks/eks.py index 2b31f5afb66d3..421e66b606a26 100644 --- a/providers/src/airflow/providers/amazon/aws/hooks/eks.py +++ b/providers/src/airflow/providers/amazon/aws/hooks/eks.py @@ -20,6 +20,7 @@ import base64 import json +import os import sys import tempfile from collections.abc import Generator @@ -32,6 +33,7 @@ from botocore.signers import RequestSigner from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.providers.amazon.aws.hooks.sts import StsHook from airflow.utils import yaml from airflow.utils.json import AirflowJsonEncoder @@ -612,9 +614,14 @@ def generate_config_file( def fetch_access_token_for_cluster(self, eks_cluster_name: str) -> str: session = self.get_session() service_id = self.conn.meta.service_model.service_id - sts_url = ( - f"https://sts.{session.region_name}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15" - ) + # This env variable is required so that we get a regionalized endpoint for STS in regions that + # otherwise default to global endpoints. The mechanism below to generate the token is very picky that + # the endpoint is regional. + os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional" + try: + sts_url = f"{StsHook(region_name=session.region_name).conn_client_meta.endpoint_url}/?Action=GetCallerIdentity&Version=2011-06-15" + finally: + del os.environ["AWS_STS_REGIONAL_ENDPOINTS"] signer = RequestSigner( service_id=service_id, diff --git a/providers/tests/amazon/aws/hooks/test_eks.py b/providers/tests/amazon/aws/hooks/test_eks.py index cf3638d30d37d..a5b1471098480 100644 --- a/providers/tests/amazon/aws/hooks/test_eks.py +++ b/providers/tests/amazon/aws/hooks/test_eks.py @@ -1283,10 +1283,12 @@ def test_generate_config_file(self, mock_conn, aws_conn_id, region_name, expecte } @mock.patch("airflow.providers.amazon.aws.hooks.eks.RequestSigner") + @mock.patch("airflow.providers.amazon.aws.hooks.eks.StsHook") @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn") @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.get_session") - def test_fetch_access_token_for_cluster(self, mock_get_session, mock_conn, mock_signer): + def test_fetch_access_token_for_cluster(self, mock_get_session, mock_conn, mock_sts_hook, mock_signer): mock_signer.return_value.generate_presigned_url.return_value = "http://example.com" + mock_sts_hook.return_value.conn_client_meta.endpoint_url = "https://sts.us-east-1.amazonaws.com" mock_get_session.return_value.region_name = "us-east-1" hook = EksHook() token = hook.fetch_access_token_for_cluster(eks_cluster_name="test-cluster")