seanghaeli commented on code in PR #74399: URL: https://github.com/apache/airflow/pull/74399#discussion_r4299994399
########## providers/amazon/src/airflow/providers/amazon/aws/executors/eks/_client_factory.py: ########## @@ -0,0 +1,144 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Kubernetes client factories for the AwsEksExecutor.""" + +# Internal to AwsEksExecutor, which writes these functions' import paths into +# [kubernetes_executor] client_factory and async_client_factory. + +from __future__ import annotations + +import os +import tempfile +import time +from base64 import b64decode +from typing import TYPE_CHECKING + +import kubernetes +from packaging.version import Version + +from airflow.providers.amazon.aws.executors.eks.utils import ( + CONFIG_DEFAULTS, + CONFIG_GROUP_NAME, + AllEksConfigKeys, +) +from airflow.providers.amazon.aws.hooks.eks import EksHook +from airflow.providers.amazon.aws.hooks.sts import StsHook +from airflow.providers.amazon.aws.utils.eks_get_token import ( + TOKEN_EXPIRATION_MINUTES, + fetch_access_token_for_cluster, +) +from airflow.providers.common.compat.sdk import conf + +if TYPE_CHECKING: + from kubernetes import client + from kubernetes_asyncio import client as async_client + +# UPDATING still serves the Kubernetes API, so every other status is refused. +_USABLE_CLUSTER_STATUSES = ("ACTIVE", "UPDATING") + + +def _get_sync_token_key(kubernetes_version: str) -> str: + # kubernetes 36 moved the bearer token from the "authorization" api_key to "BearerToken". + return "BearerToken" if Version(kubernetes_version).major >= 36 else "authorization" + + +_SYNC_TOKEN_KEY = _get_sync_token_key(kubernetes.__version__) + + +def _get_eks_kube_client() -> client.CoreV1Api: + """Build a Kubernetes client for the configured EKS cluster, in memory and without a kubeconfig.""" + from kubernetes import client + + configuration = client.Configuration() + _configure_eks_auth(configuration, _SYNC_TOKEN_KEY) + return client.CoreV1Api(client.ApiClient(configuration=configuration)) + + +def _get_eks_async_kube_client() -> 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") + return async_client.CoreV1Api(async_client.ApiClient(configuration=configuration)) + + +def _configure_eks_auth( + configuration: client.Configuration | async_client.Configuration, token_key: str +) -> None: + cluster_name = conf.get(CONFIG_GROUP_NAME, AllEksConfigKeys.CLUSTER_NAME, fallback=None) Review Comment: I've outlined the issue in this comment from @ferruzzi https://github.com/apache/airflow/pull/74399#discussion_r4288802873 Let me know if you agree with the plan laid out there. -- 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]
