This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 11cbf4da098 Allow a custom Kubernetes client factory for
KubernetesExecutor (#73014)
11cbf4da098 is described below
commit 11cbf4da098e5a02a7c2dd664131b0ea09ea4a28
Author: Sean Ghaeli <[email protected]>
AuthorDate: Wed Oct 7 10:55:32 2026 -0700
Allow a custom Kubernetes client factory for KubernetesExecutor (#73014)
Deployments that mint their own Kubernetes credentials have no way to supply
their own client today. The executor can only use an in-cluster service
account
token or a kubeconfig, so anything building a client from a cloud
provider's API
has to refresh credentials out of band. That has produced a recurring class
of
stale, racing and malformed token bugs rather than a set of one-off defects.
The new option holds an import path rather than a callable because the pod
watcher runs in its own process. Under the spawn start method that process
shares no objects with the scheduler, so it has to resolve the factory from
configuration itself. Every client the executor builds already funnels
through
get_kube_client, so a single seam there covers the executor, the pod
adoption
path and the watcher alike.
The option is opt-in and unset by default, so in-cluster and kubeconfig
behaviour is unchanged for everyone who does not set it.
---
providers/cncf/kubernetes/provider.yaml | 24 +++
.../kubernetes/executors/kubernetes_executor.py | 19 ++-
.../executors/kubernetes_executor_utils.py | 17 +-
.../providers/cncf/kubernetes/get_provider_info.py | 14 ++
.../providers/cncf/kubernetes/kube_client.py | 43 +++++
.../executors/test_kubernetes_executor.py | 175 ++++++++++++++++++++-
.../tests/unit/cncf/kubernetes/test_kube_client.py | 96 ++++++++++-
7 files changed, 370 insertions(+), 18 deletions(-)
diff --git a/providers/cncf/kubernetes/provider.yaml
b/providers/cncf/kubernetes/provider.yaml
index 4d403b38797..cf046460e29 100644
--- a/providers/cncf/kubernetes/provider.yaml
+++ b/providers/cncf/kubernetes/provider.yaml
@@ -436,6 +436,30 @@ config:
type: string
example: ~
default: ~
+ client_factory:
+ description: |
+ Import path of a zero-argument callable returning the
``kubernetes.client.CoreV1Api``
+ the executor should use, for deployments that mint their own
credentials. When set it
+ replaces the default client construction entirely, so
``in_cluster``, ``cluster_context``,
+ ``config_file``, ``verify_ssl``, ``ssl_ca_cert``,
``enable_tcp_keepalive`` and
+ ``api_client_retry_configuration`` no longer apply. It is resolved
in every process that
+ needs a client, including the pod watcher subprocess, so it must be
importable on the
+ scheduler and the API server. Each process keeps the returned
client, so it has to
+ refresh its own credentials.
+ version_added: 10.24.0
+ type: string
+ example: "my_company.kubernetes.build_client"
+ default: ~
+ async_client_factory:
+ description: |
+ Import path of a zero-argument callable returning the
``kubernetes_asyncio.client.CoreV1Api``
+ used for concurrent pod creation. Required when ``client_factory``
is set and
+ ``async_pod_creation`` is enabled, so that both clients carry the
same credentials;
+ the executor refuses to start without it.
+ version_added: 10.24.0
+ type: string
+ example: "my_company.kubernetes.build_async_client"
+ default: ~
kube_client_request_args:
description: |
Keyword parameters to pass while calling a kubernetes client
core_v1_api methods
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
index 94412c5a1f4..8c2477bfd57 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py
@@ -45,7 +45,7 @@ from kubernetes.client.rest import ApiException
from kubernetes.dynamic import DynamicClient
from sqlalchemy import select
-from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowConfigException,
AirflowProviderDeprecationWarning
from airflow.executors.base_executor import BaseExecutor
from airflow.providers.cncf.kubernetes.exceptions import
PodMutationHookException, PodReconciliationError
from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types
import (
@@ -275,9 +275,18 @@ class KubernetesExecutor(BaseExecutor):
from
airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils import (
AirflowKubernetesScheduler,
)
- from airflow.providers.cncf.kubernetes.kube_client import
get_kube_client
+ from airflow.providers.cncf.kubernetes.kube_client import
_get_executor_kube_client
- self.kube_client = get_kube_client()
+ if (
+ self.kube_config.async_pod_creation
+ and self.conf.get("kubernetes_executor", "client_factory",
fallback=None)
+ and not self.conf.get("kubernetes_executor",
"async_client_factory", fallback=None)
+ ):
+ raise AirflowConfigException(
+ "In the [kubernetes_executor] Airflow config,
async_client_factory is required "
+ "when client_factory is set and async_pod_creation is enabled."
+ )
+ self.kube_client = _get_executor_kube_client(team_name=self.team_name)
self.kube_scheduler = AirflowKubernetesScheduler(
kube_config=self.kube_config,
result_queue=self.result_queue,
@@ -1032,10 +1041,10 @@ class KubernetesExecutor(BaseExecutor):
log_streams: list[RawLogStream] = []
try:
- from airflow.providers.cncf.kubernetes.kube_client import
get_kube_client
+ from airflow.providers.cncf.kubernetes.kube_client import
_get_executor_kube_client
from airflow.providers.cncf.kubernetes.pod_generator import
PodGenerator
- client = get_kube_client()
+ client = _get_executor_kube_client(team_name=self.team_name)
hostname_desc = f" {ti.hostname}" if ti.hostname else ""
messages.append(f"Attempting to fetch logs from pod{hostname_desc}
through kube API")
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
index 576a6fe796f..29dede7d10f 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py
@@ -45,7 +45,10 @@ from
airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types impor
KubernetesWatch,
task_instance_id_from_pod,
)
-from airflow.providers.cncf.kubernetes.kube_client import
get_async_kube_client, get_kube_client
+from airflow.providers.cncf.kubernetes.kube_client import (
+ _get_executor_async_kube_client,
+ _get_executor_kube_client,
+)
from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
annotations_for_logging_task_metadata,
annotations_to_key,
@@ -91,6 +94,7 @@ class KubernetesJobWatcher(multiprocessing.Process,
LoggingMixin):
scheduler_job_id: str,
kube_config: Configuration,
supports_task_instance_uuid: bool = False,
+ team_name: str | None = None,
):
super().__init__()
self.namespace = namespace
@@ -99,13 +103,16 @@ class KubernetesJobWatcher(multiprocessing.Process,
LoggingMixin):
self.resource_version = resource_version
self.kube_config = kube_config
self.supports_task_instance_uuid = supports_task_instance_uuid
+ # A string survives the spawn start method, so the watcher can resolve
its team's factory
+ # from configuration in its own process.
+ self.team_name = team_name
def run(self) -> None:
"""Perform watching."""
if TYPE_CHECKING:
assert self.scheduler_job_id
- kube_client: client.CoreV1Api = get_kube_client()
+ kube_client: client.CoreV1Api =
_get_executor_kube_client(team_name=self.team_name)
while True:
try:
self.resource_version = self._run(
@@ -512,8 +519,9 @@ class AirflowKubernetesScheduler(LoggingMixin):
self._manager = multiprocessing.Manager()
self.watcher_queue = self._manager.Queue()
self.scheduler_job_id = scheduler_job_id
- self.kube_watchers = self._make_kube_watchers()
+ # Set before the watchers are built, since they are told which team
they are watching for.
self.team_name = team_name
+ self.kube_watchers = self._make_kube_watchers()
# Async pod-creation state; populated lazily, only used when
async_pod_creation is enabled.
self._async_loop: asyncio.AbstractEventLoop | None = None
self._async_pod_client: async_client.CoreV1Api | None = None
@@ -564,6 +572,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
scheduler_job_id=self.scheduler_job_id,
kube_config=self.kube_config,
supports_task_instance_uuid=self.supports_task_instance_uuid,
+ team_name=self.team_name,
)
watcher.start()
return watcher
@@ -711,7 +720,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
) -> list[Exception | None]:
"""Issue create_namespaced_pod calls concurrently, bounded by a
semaphore; one result per pod, in order."""
if self._async_pod_client is None:
- self._async_pod_client = await get_async_kube_client()
+ self._async_pod_client = await
_get_executor_async_kube_client(team_name=self.team_name)
api = self._async_pod_client
semaphore = asyncio.Semaphore(self.pod_creation_max_concurrency)
request_kwargs: dict[str, Any] =
self.kube_config.kube_client_request_args or {}
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
index 6291fdb2bc4..7e1c8d2cd9e 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/get_provider_info.py
@@ -272,6 +272,20 @@ def get_provider_info():
"example": None,
"default": None,
},
+ "client_factory": {
+ "description": "Import path of a zero-argument
callable returning the ``kubernetes.client.CoreV1Api``\nthe executor should
use, for deployments that mint their own credentials. When set it\nreplaces the
default client construction entirely, so ``in_cluster``,
``cluster_context``,\n``config_file``, ``verify_ssl``, ``ssl_ca_cert``,
``enable_tcp_keepalive`` and\n``api_client_retry_configuration`` no longer
apply. It is resolved in every process that\nneeds a client, [...]
+ "version_added": "10.24.0",
+ "type": "string",
+ "example": "my_company.kubernetes.build_client",
+ "default": None,
+ },
+ "async_client_factory": {
+ "description": "Import path of a zero-argument
callable returning the ``kubernetes_asyncio.client.CoreV1Api``\nused for
concurrent pod creation. Required when ``client_factory`` is set
and\n``async_pod_creation`` is enabled, so that both clients carry the same
credentials;\nthe executor refuses to start without it.\n",
+ "version_added": "10.24.0",
+ "type": "string",
+ "example": "my_company.kubernetes.build_async_client",
+ "default": None,
+ },
"kube_client_request_args": {
"description": "Keyword parameters to pass while
calling a kubernetes client core_v1_api methods\nfrom Kubernetes Executor
provided as a single line formatted JSON dictionary string.\nList of supported
params are similar for all core_v1_apis, hence a single config\nvariable for
all apis.
See:\nhttps://raw.githubusercontent.com/kubernetes-client/python/41f11a09995efcd0142e25946adc7591431bfb2f/kubernetes/client/api/core_v1_api.py\n";,
"version_added": None,
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
index 08f91abe73e..6c3cad4c428 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py
@@ -263,3 +263,46 @@ async def get_async_kube_client(
configuration.ssl_ca_cert = ssl_ca_cert
return async_client.CoreV1Api(_TimeoutAsyncK8sApiClient(configuration))
+
+
+def _team_kwargs(team_name: str | None) -> dict[str, str]:
+ """Only pass ``team_name`` when set, since older Airflow releases have no
team-aware lookup."""
+ return {"team_name": team_name} if team_name else {}
+
+
+def _get_executor_kube_client(team_name: str | None = None) ->
client.CoreV1Api:
+ """
+ Retrieve the Kubernetes client for the KubernetesExecutor.
+
+ Uses the ``client_factory`` setting when it is set, and
:func:`get_kube_client` otherwise.
+
+ :param team_name: team the executor is running for, so a team can point
its executor at its
+ own cluster; a team that sets no factory gets the default client
rather than the one
+ configured in the un-prefixed section, matching how team config
resolves everywhere else
+ :return: kubernetes client
+ """
+ # Config only stores strings, so the factory is an import path; the pod
watcher runs in its
+ # own process and resolves it there itself.
+ if client_factory := conf.getimport(
+ "kubernetes_executor", "client_factory", fallback=None,
**_team_kwargs(team_name)
+ ):
+ return client_factory()
+ return get_kube_client()
+
+
+async def _get_executor_async_kube_client(team_name: str | None = None) ->
async_client.CoreV1Api:
+ """
+ Retrieve the asynchronous Kubernetes client for the KubernetesExecutor.
+
+ Uses the ``async_client_factory`` setting when it is set, and
:func:`get_async_kube_client`
+ otherwise.
+
+ :param team_name: team the executor is running for, resolved the same way
as in
+ :func:`_get_executor_kube_client`
+ :return: asynchronous kubernetes client
+ """
+ if async_client_factory := conf.getimport(
+ "kubernetes_executor", "async_client_factory", fallback=None,
**_team_kwargs(team_name)
+ ):
+ return async_client_factory()
+ return await get_async_kube_client()
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
index 84e03ce35fb..fc08f0338c7 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_kubernetes_executor.py
@@ -36,6 +36,7 @@ from sqlalchemy.orm import Session
from urllib3 import HTTPConnectionPool, HTTPResponse
from urllib3.exceptions import MaxRetryError, ProtocolError
+from airflow.exceptions import AirflowConfigException
from airflow.executors.base_executor import BaseExecutor
from airflow.jobs.job import Job
from airflow.models.taskinstancekey import TaskInstanceKey
@@ -1289,7 +1290,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1338,7 +1339,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1395,7 +1396,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1447,7 +1448,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1498,7 +1499,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1550,7 +1551,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -1636,7 +1637,7 @@ class TestKubernetesExecutor:
AirflowKubernetesScheduler is None, reason="kubernetes python package
is not installed"
)
@mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_async_kube_client",
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
new_callable=mock.AsyncMock,
)
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher")
@@ -3862,7 +3863,7 @@ class TestKubernetesJobWatcher:
self.watcher._run = mock_underscore_run
with mock.patch(
-
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.get_kube_client"
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_kube_client"
):
with pytest.raises(SystemError, match="sentinel"):
# self.watcher._run() is mocked and return "500" as last
resource_version
@@ -4067,6 +4068,164 @@ class TestKubernetesExecutorMultiTeam:
executor.end()
[email protected](AirflowKubernetesScheduler is None, reason="kubernetes
python package is not installed")
+class TestClientFactoryCallSites:
+ """
+ Pin the call sites that must go through the executor's client helpers.
+
+ If any of them calls ``get_kube_client`` directly instead, the factory
silently stops applying
+ there while every other test stays green. Each site also has to pass its
team, or a team's
+ executor would build clients from the global factory instead of its own.
+ """
+
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client._get_executor_kube_client",
autospec=True)
+
@mock.patch("airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.client")
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher",
+ autospec=True,
+ )
+ def test_start_uses_executor_client(self, mock_watcher, mock_client,
mock_get_executor_kube_client):
+ executor = KubernetesExecutor()
+ executor.team_name = "team_a"
+ executor.job_id = 1
+ try:
+ executor.start()
+ finally:
+ executor.end()
+
+
mock_get_executor_kube_client.assert_called_once_with(team_name="team_a")
+
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client._get_executor_kube_client",
autospec=True)
+ def test_get_streaming_task_log_uses_executor_client(self,
mock_get_executor_kube_client):
+ ti = mock.MagicMock(
+ dag_id="dag",
+ task_id="task",
+ map_index=-1,
+ run_id="run",
+ queued_by_job_id=None,
+ hostname="",
+ executor_config={},
+ )
+ executor = KubernetesExecutor()
+ executor.team_name = "team_a"
+
+ executor.get_streaming_task_log(ti=ti, try_number=1)
+
+
mock_get_executor_kube_client.assert_called_once_with(team_name="team_a")
+
+ @mock.patch.object(KubernetesJobWatcher, "_run",
side_effect=RuntimeError("stop"), autospec=True)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_kube_client",
+ autospec=True,
+ )
+ def test_job_watcher_run_uses_executor_client(self,
mock_get_executor_kube_client, mock_run):
+ watcher = KubernetesJobWatcher(
+ namespace="ns",
+ watcher_queue=mock.MagicMock(),
+ resource_version="0",
+ scheduler_job_id="1",
+ kube_config=mock.MagicMock(),
+ team_name="team_a",
+ )
+
+ with pytest.raises(RuntimeError, match="stop"):
+ watcher.run()
+
+
mock_get_executor_kube_client.assert_called_once_with(team_name="team_a")
+
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.KubernetesJobWatcher",
+ autospec=True,
+ )
+ def test_make_kube_watcher_passes_the_team_to_the_watcher(self,
mock_watcher):
+ """The watcher runs in its own process, so it only learns its team by
being told."""
+ scheduler = mock.Mock(team_name="team_a", scheduler_job_id="1")
+
+ AirflowKubernetesScheduler._make_kube_watcher(scheduler, "ns")
+
+ assert mock_watcher.call_args.kwargs["team_name"] == "team_a"
+
+ @pytest.mark.asyncio
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils._get_executor_async_kube_client",
+ autospec=True,
+ )
+ async def test_create_pods_async_uses_executor_client(self,
mock_get_executor_async_kube_client):
+ scheduler = mock.Mock(pod_creation_max_concurrency=1,
_async_pod_client=None, team_name="team_a")
+ scheduler.kube_config.kube_client_request_args = {}
+
+ await AirflowKubernetesScheduler._create_pods_async(scheduler, [])
+
+
mock_get_executor_async_kube_client.assert_awaited_once_with(team_name="team_a")
+
+
[email protected](AirflowKubernetesScheduler is None, reason="kubernetes
python package is not installed")
+class TestAsyncClientFactoryValidation:
+ @pytest.mark.parametrize(
+ ("async_pod_creation", "async_client_factory", "expect_raise"),
+ [
+ pytest.param("True", "", True, id="async-without-async-factory"),
+ pytest.param("False", "", False, id="sync-only"),
+ pytest.param("True", "my_company.build_async_client", False,
id="async-with-both-factories"),
+ ],
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client._get_executor_kube_client",
autospec=True)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.AirflowKubernetesScheduler",
+ autospec=True,
+ )
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor.multiprocessing.Manager",
+ autospec=True,
+ )
+ def test_start_requires_async_client_factory(
+ self,
+ mock_manager,
+ mock_scheduler,
+ mock_get_executor_kube_client,
+ async_pod_creation,
+ async_client_factory,
+ expect_raise,
+ ):
+ config = {
+ ("kubernetes_executor", "async_pod_creation"): async_pod_creation,
+ ("kubernetes_executor", "client_factory"):
"my_company.build_client",
+ ("kubernetes_executor", "async_client_factory"):
async_client_factory,
+ }
+ with conf_vars(config):
+ executor = KubernetesExecutor()
+ executor.job_id = 1
+ if expect_raise:
+ with pytest.raises(AirflowConfigException,
match="async_client_factory is required"):
+ executor.start()
+ mock_get_executor_kube_client.assert_not_called()
+ else:
+ executor.start()
+ mock_get_executor_kube_client.assert_called_once()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="Multi-team requires
Airflow 3.2+")
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client._get_executor_kube_client",
autospec=True)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils.AirflowKubernetesScheduler",
+ autospec=True,
+ )
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.executors.kubernetes_executor.multiprocessing.Manager",
+ autospec=True,
+ )
+ def test_start_requires_async_client_factory_for_team(
+ self, mock_manager, mock_scheduler, mock_get_executor_kube_client,
monkeypatch
+ ):
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__ASYNC_POD_CREATION",
"True")
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CLIENT_FACTORY",
"my_company.build_client")
+
+ executor = KubernetesExecutor(team_name="team_a")
+ executor.job_id = 1
+ with pytest.raises(AirflowConfigException, match="async_client_factory
is required"):
+ executor.start()
+ mock_get_executor_kube_client.assert_not_called()
+
+
@pytest.fixture
def kubernetes_scheduler_factory(mocker):
mocker.patch("multiprocessing.Manager", autospec=True)
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py
index 53363d49d29..de2d61718bc 100644
--- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kube_client.py
@@ -20,9 +20,103 @@ from unittest import mock
import pytest
-from airflow.providers.cncf.kubernetes.kube_client import
_TimeoutAsyncK8sApiClient, get_async_kube_client
+from airflow.providers.cncf.kubernetes.kube_client import (
+ _get_executor_async_kube_client,
+ _get_executor_kube_client,
+ _TimeoutAsyncK8sApiClient,
+ get_async_kube_client,
+)
+from airflow.providers.common.compat.sdk import AirflowConfigException
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
+
+FACTORY_MARKER = "client-from-factory"
+OTHER_MARKER = "client-from-other-factory"
+
+
+def build_fake_client():
+ return FACTORY_MARKER
+
+
+def build_other_client():
+ return OTHER_MARKER
+
+
+class TestGetExecutorKubeClient:
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client",
autospec=True)
+ def test_without_factory_uses_default_client(self, mock_get_kube_client):
+ client = _get_executor_kube_client()
+
+ assert client is mock_get_kube_client.return_value
+ mock_get_kube_client.assert_called_once_with()
+
+ @conf_vars({("kubernetes_executor", "client_factory"):
f"{__name__}.build_fake_client"})
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client",
autospec=True)
+ def test_factory_replaces_default_client(self, mock_get_kube_client):
+ client = _get_executor_kube_client()
+
+ assert client == FACTORY_MARKER
+ mock_get_kube_client.assert_not_called()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @conf_vars({("kubernetes_executor", "client_factory"):
f"{__name__}.build_other_client"})
+ def test_team_factory_takes_precedence_over_the_global_one(self,
monkeypatch):
+ monkeypatch.setenv(
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CLIENT_FACTORY",
f"{__name__}.build_fake_client"
+ )
+
+ assert _get_executor_kube_client(team_name="team_a") == FACTORY_MARKER
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @conf_vars({("kubernetes_executor", "client_factory"):
f"{__name__}.build_fake_client"})
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client",
autospec=True)
+ def test_team_without_factory_does_not_use_the_global_one(self,
mock_get_kube_client, monkeypatch):
+ """Team config does not inherit the un-prefixed section, so team_b
gets neither factory."""
+ monkeypatch.setenv(
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CLIENT_FACTORY",
f"{__name__}.build_other_client"
+ )
+
+ client = _get_executor_kube_client(team_name="team_b")
+
+ assert client is mock_get_kube_client.return_value
+ mock_get_kube_client.assert_called_once_with()
+
+ @conf_vars({("kubernetes_executor", "client_factory"):
"no.such.module.build_client"})
+ def test_unimportable_factory_raises(self):
+ with pytest.raises(AirflowConfigException):
+ _get_executor_kube_client()
+
+
+class TestGetExecutorAsyncKubeClient:
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_async_kube_client",
autospec=True)
+ async def test_without_factory_uses_default_client(self,
mock_get_async_kube_client):
+ client = await _get_executor_async_kube_client()
+
+ assert client is mock_get_async_kube_client.return_value
+ mock_get_async_kube_client.assert_awaited_once_with()
+
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_async_kube_client",
autospec=True)
+ async def test_async_factory_replaces_default_client(self,
mock_get_async_kube_client):
+ with conf_vars({("kubernetes_executor", "async_client_factory"):
f"{__name__}.build_fake_client"}):
+ client = await _get_executor_async_kube_client()
+
+ assert client == FACTORY_MARKER
+ mock_get_async_kube_client.assert_not_called()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.asyncio
+ async def
test_team_async_factory_takes_precedence_over_the_global_one(self, monkeypatch):
+ monkeypatch.setenv(
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__ASYNC_CLIENT_FACTORY",
+ f"{__name__}.build_fake_client",
+ )
+ with conf_vars({("kubernetes_executor", "async_client_factory"):
f"{__name__}.build_other_client"}):
+ client = await _get_executor_async_kube_client(team_name="team_a")
+
+ assert client == FACTORY_MARKER
class TestGetAsyncKubeClient:
