This is an automated email from the ASF dual-hosted git repository.
o-nikolas 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 9a02819b4e4 Fix AwsEksExecutor dying with a 401 after its credentials
rotate (#76332)
9a02819b4e4 is described below
commit 9a02819b4e4ab7cfd711f116c2f5273c796a8c90
Author: Sean Ghaeli <[email protected]>
AuthorDate: Wed Nov 11 22:54:24 2026 -0700
Fix AwsEksExecutor dying with a 401 after its credentials rotate (#76332)
The client signed every EKS token with the credentials it resolved when it
was
built. With static keys that are rotated outside the process, such as a
credentials file rewritten by a sidecar, every token signed after those keys
expired was rejected, and the unhandled 401 killed the executor. The ECS,
Batch
and Lambda executors already recover from expired credentials by rebuilding
their client.
---
providers/amazon/docs/executors/eks-executor.rst | 5 +
.../amazon/aws/executors/eks/_client_factory.py | 94 +++++++--
.../aws/executors/eks/test__client_factory.py | 232 ++++++++++++++++++++-
3 files changed, 311 insertions(+), 20 deletions(-)
diff --git a/providers/amazon/docs/executors/eks-executor.rst
b/providers/amazon/docs/executors/eks-executor.rst
index d996f9e38a1..5db2aa163b8 100644
--- a/providers/amazon/docs/executors/eks-executor.rst
+++ b/providers/amazon/docs/executors/eks-executor.rst
@@ -68,6 +68,11 @@ roughly fifteen minutes, so the executor registers a refresh
hook that the Kuber
calls on every authenticated request. The hook reuses the current token and
gets a new one
a minute before it expires.
+If something outside Airflow rotates the executor's AWS keys, such as a
sidecar that rewrites
+the credentials file or a secrets backend that stores them in the connection,
the executor picks
+up the new keys without a restart. When the old keys stop working, the
scheduler may log a
+warning that the cluster rejected the token, and the request is retried once
with the new keys.
+
Before it accepts any tasks, the executor checks that the cluster is
``ACTIVE`` or
``UPDATING``, since EKS keeps the Kubernetes API available during an update
(see `Update
existing cluster to new Kubernetes version
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
index 14846b26812..f3a7ca49eeb 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py
@@ -22,15 +22,20 @@
from __future__ import annotations
import atexit
+import logging
import os
import tempfile
import threading
import time
from base64 import b64decode
+from collections.abc import Callable
from contextlib import suppress
-from typing import TYPE_CHECKING
+from contextvars import ContextVar
+from http import HTTPStatus
import kubernetes
+from kubernetes import client
+from kubernetes_asyncio import client as async_client
from packaging.version import Version
from airflow.providers.amazon.aws.executors.eks.utils import (
@@ -47,9 +52,7 @@ from airflow.providers.amazon.aws.utils.eks_get_token import (
)
from airflow.providers.common.compat.sdk import conf
-if TYPE_CHECKING:
- from kubernetes import client
- from kubernetes_asyncio import client as async_client
+log = logging.getLogger(__name__)
# UPDATING still serves the Kubernetes API, so every other status is refused.
_USABLE_CLUSTER_STATUSES = ("ACTIVE", "UPDATING")
@@ -69,30 +72,68 @@ class _EksExecutorHook(EksHook):
return "AwsEksExecutor"
+_TOKEN_REJECTED = "The EKS cluster rejected the token (401); retrying once
with a new token."
+
+# The token the current request was sent with. Each thread and asyncio task
has its own.
+_sent_token: ContextVar[str | None] = ContextVar("_sent_token", default=None)
+
+
+# A 401 means the cluster rejected the token before acting on the request, so
the request, even a
+# create, is safe to send again. The new token is signed with the credentials
the connection
+# resolves to now, which may have rotated since the rejected token was signed.
+class _EksApiClient(client.ApiClient):
+ def __init__(self, configuration: client.Configuration, expire_token:
Callable[[], None]):
+ super().__init__(configuration=configuration)
+ self._expire_token = expire_token
+
+ def call_api(self, *args, **kwargs):
+ try:
+ return super().call_api(*args, **kwargs)
+ except client.ApiException as e:
+ if e.status != HTTPStatus.UNAUTHORIZED:
+ raise
+ log.warning(_TOKEN_REJECTED)
+ self._expire_token()
+ return super().call_api(*args, **kwargs)
+
+
+class _EksAsyncApiClient(async_client.ApiClient):
+ def __init__(self, configuration: async_client.Configuration,
expire_token: Callable[[], None]):
+ super().__init__(configuration=configuration)
+ self._expire_token = expire_token
+
+ async def call_api(self, *args, **kwargs):
+ try:
+ return await super().call_api(*args, **kwargs)
+ except async_client.ApiException as e:
+ if e.status != HTTPStatus.UNAUTHORIZED:
+ raise
+ log.warning(_TOKEN_REJECTED)
+ self._expire_token()
+ return await super().call_api(*args, **kwargs)
+
+
def _get_eks_kube_client(team_name: str | None = None) -> client.CoreV1Api:
"""Build a Kubernetes client for the team's EKS cluster, in memory and
without a kubeconfig."""
- from kubernetes import client
-
configuration = client.Configuration()
- _configure_eks_auth(configuration, _SYNC_TOKEN_KEY, team_name)
- return client.CoreV1Api(client.ApiClient(configuration=configuration))
+ expire_token = _configure_eks_auth(configuration, _SYNC_TOKEN_KEY,
team_name)
+ return client.CoreV1Api(_EksApiClient(configuration, expire_token))
def _get_eks_async_kube_client(team_name: str | None = None) ->
async_client.CoreV1Api:
"""Build the asynchronous Kubernetes client used when
``async_pod_creation`` is enabled."""
- from kubernetes_asyncio import client as async_client
-
configuration = async_client.Configuration()
# kubernetes_asyncio has always used "BearerToken".
- _configure_eks_auth(configuration, "BearerToken", team_name)
- return
async_client.CoreV1Api(async_client.ApiClient(configuration=configuration))
+ expire_token = _configure_eks_auth(configuration, "BearerToken", team_name)
+ return async_client.CoreV1Api(_EksAsyncApiClient(configuration,
expire_token))
def _configure_eks_auth(
configuration: client.Configuration | async_client.Configuration,
token_key: str,
team_name: str | None,
-) -> None:
+) -> Callable[[], None]:
+ """Set up EKS token auth, and return a function that makes the next
request get a new token."""
# The factories run where there is no executor, such as the pod watcher
process, so they look
# up the team's config themselves, the same way ExecutorConf does. Airflow
releases without
# multi-team don't accept team_name.
@@ -134,18 +175,33 @@ def _configure_eks_auth(
# makes that happen now, so the first request already has a token.
token: str | None = None
refresh_at = 0.0
+ token_lock = threading.Lock()
def refresh_api_key(config: client.Configuration |
async_client.Configuration) -> None:
nonlocal token, refresh_at
- if token is None or time.monotonic() >= refresh_at:
- token = fetch_access_token_for_cluster(
- cluster_name, sts_url, region_name=session.region_name,
session=session
- )
- refresh_at = time.monotonic() + TOKEN_EXPIRATION_MINUTES * 60
- config.api_key[token_key] = token
+ with token_lock:
+ if token is None or time.monotonic() >= refresh_at:
+ # Sign with the credentials the connection resolves to now:
only a new session picks up
+ # static keys rotated outside the process, such as in a
rewritten credentials file.
+ token = fetch_access_token_for_cluster(
+ cluster_name, sts_url, region_name=session.region_name,
session=eks_hook.get_session()
+ )
+ refresh_at = time.monotonic() + TOKEN_EXPIRATION_MINUTES * 60
+ config.api_key[token_key] = token
+ _sent_token.set(token)
+
+ def expire_token() -> None:
+ nonlocal token, eks_hook
+ with token_lock:
+ # Requests rejected with the same token, such as concurrent pod
creations, share one new token.
+ if token == _sent_token.get():
+ # The hook caches the connection, so read it again in case the
keys stored in it rotated.
+ eks_hook = _EksExecutorHook(aws_conn_id=conn_id,
region_name=region_name)
+ token = None
configuration.refresh_api_key_hook = refresh_api_key
refresh_api_key(configuration)
+ return expire_token
# The CA file written for each cluster CA, and the process that wrote it.
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
b/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
index 454206e0061..4c0f10b23e0 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/eks/test__client_factory.py
@@ -16,16 +16,18 @@
# under the License.
from __future__ import annotations
+import asyncio
import inspect
import os
import signal
import tempfile
import threading
import time
-from base64 import b64encode
+from base64 import b64encode, urlsafe_b64decode
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest import mock
+from urllib.parse import parse_qs, urlsplit
import boto3
import pytest
@@ -275,6 +277,234 @@ async def
test_async_client_refreshes_bearer_token(mock_time, _, mock_aws):
await core_v1.api_client.close()
+def _signing_key_id(configuration: client.Configuration) -> str:
+ """Return the access key ID that signed the client's current EKS token."""
+ token =
configuration.auth_settings()["BearerToken"]["value"].removeprefix("Bearer
k8s-aws-v1.")
+ url = urlsafe_b64decode(token + "=" * (-len(token) % 4)).decode()
+ return parse_qs(urlsplit(url).query)["X-Amz-Credential"][0].split("/")[0]
+
+
[email protected](_client_factory, "time", autospec=True)
[email protected](_client_factory, "StsHook", autospec=True)
[email protected]("botocore.client.BaseClient._make_api_call", autospec=True)
+def test_new_token_is_signed_with_rotated_credentials(
+ mock_api_call, mock_sts_hook, mock_time, tmp_path, monkeypatch
+):
+ """Static keys in a credentials file that is rewritten while the executor
runs."""
+ credentials_file = tmp_path / "credentials"
+
+ def write_keys(access_key_id):
+ credentials_file.write_text(
+ f"[rotating]\naws_access_key_id =
{access_key_id}\naws_secret_access_key = secret\n"
+ )
+
+ for var in ("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",
"AWS_SESSION_TOKEN"):
+ monkeypatch.delenv(var, raising=False)
+ monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(credentials_file))
+ monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "config"))
+ monkeypatch.setenv("AWS_PROFILE", "rotating")
+ monkeypatch.setenv("AIRFLOW__AWS_EKS_EXECUTOR__CLUSTER_NAME", CLUSTER_NAME)
+ monkeypatch.setenv("AIRFLOW__AWS_EKS_EXECUTOR__REGION_NAME", "us-east-1")
+ monkeypatch.setenv("AIRFLOW_CONN_AWS_DEFAULT", "aws://")
+ mock_api_call.return_value = {
+ "cluster": {
+ "status": "ACTIVE",
+ "endpoint": CLUSTER_ENDPOINT,
+ "certificateAuthority": {"data": b64encode(CA_PEM).decode()},
+ }
+ }
+ mock_sts_hook.return_value.conn_client_meta.endpoint_url =
"https://sts.us-east-1.amazonaws.com";
+ mock_time.monotonic.return_value = 1000.0
+ write_keys("AKIAFIRSTKEY")
+ configuration = _get_eks_kube_client().api_client.configuration
+ assert _signing_key_id(configuration) == "AKIAFIRSTKEY"
+
+ write_keys("AKIASECONDKEY")
+ mock_time.monotonic.return_value = 1000.0 + TOKEN_REFRESH_SECONDS
+
+ assert _signing_key_id(configuration) == "AKIASECONDKEY"
+
+
+def _unauthorized(api_exception_cls):
+ error = api_exception_cls(status=401, reason="Unauthorized")
+ # The client decodes the body of a failed response before raising it.
+ error.body = b'{"kind":"Status","code":401,"reason":"Unauthorized"}'
+ return error
+
+
[email protected](
+ ("factory", "api_exception_cls"),
+ [(_get_eks_kube_client, client.ApiException), (_get_eks_async_kube_client,
async_client.ApiException)],
+ ids=["sync", "async"],
+)
[email protected](
+ ("unauthorized_responses", "succeeds"),
+ [pytest.param(1, True, id="retried-with-new-token"), pytest.param(2,
False, id="retried-only-once")],
+)
[email protected]
+# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
[email protected](_client_factory, "_write_cluster_ca_file",
return_value=None)
+async def test_rejected_token_is_replaced_and_request_retried_once(
+ _, mock_aws, factory, api_exception_cls, unauthorized_responses, succeeds
+):
+ core_v1 = factory()
+ mock_aws["fetch_token"].side_effect = ["k8s-aws-v1.token-2",
"k8s-aws-v1.token-3"]
+ sent_tokens = []
+ responses: list = [_unauthorized(api_exception_cls) for _ in
range(unauthorized_responses)]
+ responses.append(mock.sentinel.pods)
+
+ def respond(method, url, headers=None, **kwargs):
+ sent_tokens.append(headers["authorization"])
+ response = responses.pop(0)
+ if isinstance(response, Exception):
+ raise response
+ return response
+
+ with mock.patch.object(core_v1.api_client.rest_client, "request",
autospec=True, side_effect=respond):
+ try:
+ result = core_v1.list_namespaced_pod("default",
_preload_content=False)
+ if inspect.isawaitable(result):
+ result = await result
+ except api_exception_cls as e:
+ result = e.status
+
+ assert sent_tokens == ["Bearer k8s-aws-v1.token-1", "Bearer
k8s-aws-v1.token-2"]
+ assert result == (mock.sentinel.pods if succeeds else 401)
+ if inspect.iscoroutinefunction(core_v1.api_client.close):
+ await core_v1.api_client.close()
+
+
[email protected](
+ ("factory", "api_exception_cls"),
+ [(_get_eks_kube_client, client.ApiException), (_get_eks_async_kube_client,
async_client.ApiException)],
+ ids=["sync", "async"],
+)
[email protected]("status", [403, 500])
[email protected]
+# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
[email protected](_client_factory, "_write_cluster_ca_file",
return_value=None)
+async def test_other_errors_are_raised_without_a_new_token_or_retry(
+ _, mock_aws, factory, api_exception_cls, status
+):
+ core_v1 = factory()
+ sent_tokens = []
+
+ def respond(method, url, headers=None, **kwargs):
+ sent_tokens.append(headers["authorization"])
+ error = api_exception_cls(status=status, reason="Error")
+ error.body = b'{"kind":"Status"}'
+ raise error
+
+ with mock.patch.object(core_v1.api_client.rest_client, "request",
autospec=True, side_effect=respond):
+ try:
+ result = core_v1.list_namespaced_pod("default",
_preload_content=False)
+ if inspect.isawaitable(result):
+ result = await result
+ except api_exception_cls as e:
+ result = e
+
+ assert result.status == status
+ assert sent_tokens == ["Bearer k8s-aws-v1.token-1"]
+ mock_aws["fetch_token"].assert_called_once()
+ if inspect.iscoroutinefunction(core_v1.api_client.close):
+ await core_v1.api_client.close()
+
+
[email protected](
+ ("factory", "api_exception_cls"),
+ [(_get_eks_kube_client, client.ApiException), (_get_eks_async_kube_client,
async_client.ApiException)],
+ ids=["sync", "async"],
+)
[email protected]
+# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
[email protected](_client_factory, "_write_cluster_ca_file",
return_value=None)
+async def test_requests_rejected_with_the_same_token_share_one_new_token(
+ _, mock_aws, factory, api_exception_cls
+):
+ requests = 2
+ rejected = "Bearer k8s-aws-v1.token-1"
+ core_v1 = factory()
+ mock_aws["fetch_token"].side_effect = ["k8s-aws-v1.token-2",
"k8s-aws-v1.token-3"]
+ sent_tokens = []
+
+ def list_pods(_=None):
+ return core_v1.list_namespaced_pod("default", _preload_content=False)
+
+ # Both requests are sent with the same token, and the second is rejected
after the first's retry.
+ if factory is _get_eks_kube_client:
+ both_sent, replaced = threading.Barrier(requests), threading.Event()
+
+ def respond(method, url, headers=None, **kwargs):
+ sent_tokens.append(headers["authorization"])
+ if headers["authorization"] != rejected:
+ replaced.set()
+ return mock.sentinel.pods
+ if both_sent.wait():
+ replaced.wait()
+ raise _unauthorized(api_exception_cls)
+
+ with (
+ mock.patch.object(core_v1.api_client.rest_client, "request",
autospec=True, side_effect=respond),
+ ThreadPoolExecutor(max_workers=requests) as pool,
+ ):
+ results = list(pool.map(list_pods, range(requests)))
+ else:
+ all_sent, replaced = asyncio.Event(), asyncio.Event()
+
+ async def respond_async(method, url, headers=None, **kwargs):
+ sent_tokens.append(headers["authorization"])
+ if headers["authorization"] != rejected:
+ replaced.set()
+ return mock.sentinel.pods
+ if len(sent_tokens) < requests:
+ await all_sent.wait()
+ else:
+ all_sent.set()
+ await replaced.wait()
+ raise _unauthorized(api_exception_cls)
+
+ with mock.patch.object(
+ core_v1.api_client.rest_client, "request", autospec=True,
side_effect=respond_async
+ ):
+ results = await asyncio.gather(*(list_pods() for _ in
range(requests)))
+ await core_v1.api_client.close()
+
+ assert results == [mock.sentinel.pods] * requests
+ assert sent_tokens == [rejected] * requests + ["Bearer
k8s-aws-v1.token-2"] * requests
+ assert mock_aws["fetch_token"].call_count == len(set(sent_tokens))
+
+
[email protected](
+ ("factory", "api_exception_cls"),
+ [(_get_eks_kube_client, client.ApiException), (_get_eks_async_kube_client,
async_client.ApiException)],
+ ids=["sync", "async"],
+)
[email protected]
+# kubernetes_asyncio loads the CA eagerly and the fake PEM is not loadable.
[email protected](_client_factory, "_write_cluster_ca_file",
return_value=None)
[email protected](_client_factory, "time", autospec=True)
+async def test_retry_after_401_reads_the_connection_again(mock_time, _,
mock_aws, factory, api_exception_cls):
+ """Keys stored in the Airflow connection, such as in a secrets backend,
may have rotated."""
+ mock_time.monotonic.return_value = 1000.0
+ core_v1 = factory()
+ rebuilt_hook = mock_aws["eks_hook_cls"].return_value = mock.MagicMock()
+ responses = [_unauthorized(api_exception_cls), mock.sentinel.pods,
mock.sentinel.pods]
+
+ with mock.patch.object(core_v1.api_client.rest_client, "request",
autospec=True, side_effect=responses):
+ # The second request comes after a timed refresh, which keeps the hook
rebuilt on the retry.
+ for elapsed in (0, TOKEN_REFRESH_SECONDS):
+ mock_time.monotonic.return_value = 1000.0 + elapsed
+ result = core_v1.list_namespaced_pod("default",
_preload_content=False)
+ if inspect.isawaitable(result):
+ await result
+
+ first, rebuilt = (hook.get_session.return_value for hook in
(mock_aws["eks_hook"], rebuilt_hook))
+ sessions = [call.kwargs["session"] for call in
mock_aws["fetch_token"].call_args_list]
+ assert sessions == [first, rebuilt, rebuilt]
+ if inspect.iscoroutinefunction(core_v1.api_client.close):
+ await core_v1.api_client.close()
+
+
class TestWriteClusterCaFile:
CA_DATA = b64encode(CA_PEM).decode()
