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 8aa1cf26ea3 Let each team's KubernetesExecutor connect to its own
cluster (#75251)
8aa1cf26ea3 is described below
commit 8aa1cf26ea349a9aec6da774286d14ac4f1c06b8
Author: Sean Ghaeli <[email protected]>
AuthorDate: Fri Oct 23 11:51:30 2026 -0700
Let each team's KubernetesExecutor connect to its own cluster (#75251)
In multi-team mode, a team's KubernetesExecutor read its pod settings from
the team's config, but connected to the cluster set in the global
[kubernetes_executor] section, so every team ran its pods on the same
cluster. A client_factory could be set per team, but it was called without
arguments, so it couldn't tell which team it was building a client for.
A KubernetesExecutor subclass that builds its own client, such as the
upcoming AwsEksExecutor, could only install its factory through
[kubernetes_executor] client_factory. A plain KubernetesExecutor in the
same scope then picked it up too and sent its pods to the subclass's
cluster.
airflow kubernetes cleanup-pods also only connected to the global cluster.
It now takes --team, like generate-dag-yaml, and connects the way the
team's executor does, with --executor to pick one when the team has more
than one. It refuses a namespace that another team's KubernetesExecutor, or
a global one, also uses, since it would delete their pods too.
closes: #74844
---
providers/cncf/kubernetes/docs/changelog.rst | 6 +
.../cncf/kubernetes/docs/kubernetes_executor.rst | 44 +++++
providers/cncf/kubernetes/provider.yaml | 16 +-
.../providers/cncf/kubernetes/cli/definition.py | 35 +++-
.../cncf/kubernetes/cli/kubernetes_command.py | 79 +++++++-
.../kubernetes/executors/kubernetes_executor.py | 39 ++--
.../executors/kubernetes_executor_utils.py | 19 +-
.../providers/cncf/kubernetes/get_provider_info.py | 4 +-
.../providers/cncf/kubernetes/kube_client.py | 89 ++++++---
.../cncf/kubernetes/cli/test_kubernetes_command.py | 219 +++++++++++++++++++--
.../executors/test_kubernetes_executor.py | 119 +++++++++--
.../tests/unit/cncf/kubernetes/test_client.py | 79 ++++++++
.../tests/unit/cncf/kubernetes/test_kube_client.py | 124 +++++++++++-
13 files changed, 782 insertions(+), 90 deletions(-)
diff --git a/providers/cncf/kubernetes/docs/changelog.rst
b/providers/cncf/kubernetes/docs/changelog.rst
index 082be601055..829cc47ba9e 100644
--- a/providers/cncf/kubernetes/docs/changelog.rst
+++ b/providers/cncf/kubernetes/docs/changelog.rst
@@ -27,6 +27,12 @@
Changelog
---------
+.. note::
+ In multi-team mode, a team's ``KubernetesExecutor`` now reads
``in_cluster``, ``config_file``,
+ ``cluster_context``, ``verify_ssl``, ``ssl_ca_cert`` and
``api_client_retry_configuration`` from the
+ team's ``[kubernetes_executor]`` config. Like all team config, these don't
fall back to the
+ un-prefixed section. If your teams connected using the un-prefixed values,
set them for each team.
+
10.25.0
.......
diff --git a/providers/cncf/kubernetes/docs/kubernetes_executor.rst
b/providers/cncf/kubernetes/docs/kubernetes_executor.rst
index 06e1c632798..7e5e5a43650 100644
--- a/providers/cncf/kubernetes/docs/kubernetes_executor.rst
+++ b/providers/cncf/kubernetes/docs/kubernetes_executor.rst
@@ -206,6 +206,50 @@ To get task logs out of the workers, you can:
If you don't enable logging persistence, and if you have not enabled
remote logging, logs will be lost after the worker pods shut down.
+Connecting each team to its own cluster
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+In a :doc:`multi-team <apache-airflow:core-concepts/multi-team>` deployment, a
team's KubernetesExecutor
+reads the settings it connects with from the team's ``[kubernetes_executor]``
config, like every other
+setting: ``in_cluster``, ``config_file``, ``cluster_context``, ``verify_ssl``,
``ssl_ca_cert``,
+``api_client_retry_configuration``, ``client_factory`` and
``async_client_factory``. So each team can run
+its pods on a different cluster. Team config doesn't fall back to the
un-prefixed section, so a team that
+connects from outside the cluster has to set ``in_cluster = False`` itself.
``enable_tcp_keepalive`` and
+the ``tcp_keep_*`` settings are only read from the un-prefixed section, and
apply to every team's client.
+
+.. code-block:: ini
+
+ [core]
+ multi_team = True
+ executor =
LocalExecutor;team_a=KubernetesExecutor;team_b=KubernetesExecutor
+
+ [team_a=kubernetes_executor]
+ in_cluster = False
+ config_file = /opt/airflow/kube/clusters.yaml
+ cluster_context = cluster-a
+ namespace = team-a
+
+ [team_b=kubernetes_executor]
+ in_cluster = False
+ config_file = /opt/airflow/kube/clusters.yaml
+ cluster_context = cluster-b
+ namespace = team-b
+
+A ``client_factory`` or ``async_client_factory`` that takes a ``team_name``
argument is called with the
+team it builds the client for (``None`` for an executor with no team). One
that takes no arguments is
+still called without any.
+
+``airflow kubernetes cleanup-pods --team <team>`` connects the way the team's
default executor does, or
+the one named with ``--executor``, and cleans up the pods in the team's
namespace unless ``--namespace`` is
+given. It refuses a namespace that another team's KubernetesExecutor, or a
global one, also uses, since
+it would delete their pods too, unless you pass ``--allow-shared-namespace``
because they run on another
+cluster.
+
+A subclass of KubernetesExecutor that builds its own clients sets the
``client_factory`` and
+``async_client_factory`` class attributes to the import paths of its
factories, which are called with
+``team_name``. They take precedence over the config settings of the same name,
so a plain
+KubernetesExecutor running alongside the subclass keeps its own client.
+
Comparison with CeleryExecutor
------------------------------
diff --git a/providers/cncf/kubernetes/provider.yaml
b/providers/cncf/kubernetes/provider.yaml
index 32af80a1b3b..3d8f4aa3191 100644
--- a/providers/cncf/kubernetes/provider.yaml
+++ b/providers/cncf/kubernetes/provider.yaml
@@ -440,9 +440,11 @@ config:
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``,
+ Import path of a callable returning the
``kubernetes.client.CoreV1Api`` the executor
+ should use, for deployments that mint their own credentials. If the
callable takes a
+ ``team_name`` argument, it is called with the executor's team
(``None`` for an executor
+ with no team); otherwise it is called with no arguments. 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
@@ -454,10 +456,10 @@ config:
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.
+ Import path of a callable returning the
``kubernetes_asyncio.client.CoreV1Api`` used for
+ concurrent pod creation, called the same way as ``client_factory``.
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"
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/definition.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/definition.py
index cbcbfac9e45..35cb25e49c0 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/definition.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/definition.py
@@ -29,7 +29,6 @@ from airflow.cli.cli_config import (
positive_int,
)
from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.providers.common.compat.sdk import conf
if TYPE_CHECKING:
import argparse
@@ -54,8 +53,10 @@ else:
# CLI Args
ARG_NAMESPACE = Arg(
("--namespace",),
- default=conf.get("kubernetes_executor", "namespace"),
- help="Kubernetes Namespace. Default value is `[kubernetes] namespace` in
configuration.",
+ help=(
+ "Kubernetes Namespace. Default value is `[kubernetes_executor]
namespace` in configuration, "
+ "the team's one with --team."
+ ),
)
ARG_MIN_PENDING_MINUTES = Arg(
@@ -84,6 +85,24 @@ ARG_TEAM = Arg(
help="Team name for multi-team configuration. When set, team-specific
config overrides are applied.",
)
+ARG_EXECUTOR = Arg(
+ ("--executor",),
+ default=None,
+ help=(
+ "With --team, the team's executor whose pods to clean up, by alias,
module path or class name. "
+ "Default is the team's default executor."
+ ),
+)
+
+ARG_ALLOW_SHARED_NAMESPACE = Arg(
+ ("--allow-shared-namespace",),
+ action="store_true",
+ help=(
+ "With --team, clean up even if another team's or a global
KubernetesExecutor uses the same "
+ "namespace, such as on another cluster."
+ ),
+)
+
# CLI Commands
KUBERNETES_COMMANDS = (
ActionCommand(
@@ -94,7 +113,15 @@ KUBERNETES_COMMANDS = (
"in evicted/failed/succeeded/pending states"
),
func=lazy_load_command("airflow.providers.cncf.kubernetes.cli.kubernetes_command.cleanup_pods"),
- args=(ARG_NAMESPACE, ARG_MIN_PENDING_MINUTES,
ARG_MIN_COMPLETED_MINUTES, ARG_VERBOSE),
+ args=(
+ ARG_NAMESPACE,
+ ARG_MIN_PENDING_MINUTES,
+ ARG_MIN_COMPLETED_MINUTES,
+ ARG_TEAM,
+ ARG_EXECUTOR,
+ ARG_ALLOW_SHARED_NAMESPACE,
+ ARG_VERBOSE,
+ ),
),
ActionCommand(
name="generate-dag-yaml",
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/kubernetes_command.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/kubernetes_command.py
index 2d2490134dd..b4e2c1c325b 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/kubernetes_command.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/cli/kubernetes_command.py
@@ -26,10 +26,16 @@ from kubernetes import client
from kubernetes.client.api_client import ApiClient
from kubernetes.client.rest import ApiException
+from airflow.exceptions import UnknownExecutorException
+from airflow.executors.executor_loader import ExecutorLoader
from airflow.models import DagModel, DagRun, TaskInstance
from airflow.providers.cncf.kubernetes import pod_generator
-from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import
KubeConfig
-from airflow.providers.cncf.kubernetes.kube_client import get_kube_client
+from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import
KubeConfig, KubernetesExecutor
+from airflow.providers.cncf.kubernetes.kube_client import (
+ _get_executor_kube_client,
+ _team_kwargs,
+ get_kube_client,
+)
from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import
create_unique_id
from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator,
generate_pod_command_args
from airflow.providers.cncf.kubernetes.version_compat import (
@@ -37,6 +43,7 @@ from airflow.providers.cncf.kubernetes.version_compat import (
AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_2_PLUS,
)
+from airflow.providers.common.compat.sdk import conf
from airflow.utils import cli as cli_utils, yaml
from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
from airflow.utils.types import DagRunType
@@ -148,7 +155,16 @@ def _get_pod_completion_time(pod):
@providers_configuration_loaded
def cleanup_pods(args):
"""Clean up k8s pods in evicted/failed/succeeded/pending states."""
- namespace = args.namespace
+ team_name = args.team
+ if team_name and not AIRFLOW_V_3_2_PLUS:
+ raise SystemExit("cleanup-pods --team requires Airflow 3.2+.")
+ if args.executor and not team_name:
+ raise SystemExit("cleanup-pods --executor requires --team.")
+ namespace = args.namespace or conf.get("kubernetes_executor", "namespace",
**_team_kwargs(team_name))
+ if team_name:
+ executor_cls = _get_team_executor_cls(team_name, args.executor)
+ if not args.allow_shared_namespace:
+ _check_namespace_not_shared(team_name, namespace)
min_pending_minutes = args.min_pending_minutes
# protect newly created pods from deletion
@@ -178,7 +194,12 @@ def cleanup_pods(args):
pod_restart_policy_never = "never"
if args.verbose:
print("Loading Kubernetes configuration")
- kube_client = get_kube_client()
+ if team_name:
+ kube_client = _get_executor_kube_client(
+ team_name=team_name,
executor_client_factory=executor_cls.client_factory
+ )
+ else:
+ kube_client = get_kube_client()
if args.verbose:
print(f"Listing pods in namespace {namespace}")
airflow_pod_labels = [
@@ -219,7 +240,7 @@ def cleanup_pods(args):
f'restart policy "{pod_restart_policy}"'
)
try:
- _delete_pod(pod.metadata.name, namespace)
+ _delete_pod(pod.metadata.name, namespace, kube_client)
except ApiException as e:
print(f"Can't remove POD: {e}", file=sys.stderr)
else:
@@ -231,13 +252,57 @@ def cleanup_pods(args):
list_kwargs["_continue"] = continue_token
-def _delete_pod(name, namespace):
+def _get_team_executor_cls(team_name: str, executor: str | None) ->
type[KubernetesExecutor]:
+ """Return the class of the team's executor named ``executor``, or of its
default one."""
+ # Only the team's own executors: a global one doesn't connect with the
team's config.
+ names = [n for n in
ExecutorLoader.get_executor_names(validate_teams=False) if n.team_name ==
team_name]
+ if not names:
+ raise SystemExit(f"Team {team_name!r} has no executors configured.")
+ try:
+ if hasattr(ExecutorLoader, "resolve_executor_name"):
+ name = ExecutorLoader.resolve_executor_name(executor, team_name,
executor_names=names)
+ elif executor: # Airflow < 3.4
+ name = ExecutorLoader.lookup_executor_name_by_str(
+ executor, team_name=team_name, validate_teams=False
+ )
+ else:
+ name = names[0]
+ except UnknownExecutorException as e:
+ raise SystemExit(str(e))
+ executor_cls, _ = ExecutorLoader.import_executor_cls(name)
+ if not issubclass(executor_cls, KubernetesExecutor):
+ raise SystemExit(
+ f"Executor {name.module_path} of team {team_name!r} isn't a
KubernetesExecutor. "
+ "Pass --executor to pick one that is."
+ )
+ return executor_cls
+
+
+def _check_namespace_not_shared(team_name: str, namespace: str) -> None:
+ """Refuse a namespace another scope's KubernetesExecutor also uses, since
their pods can't be told apart."""
+ # Only scopes that run a KubernetesExecutor make pods; the others still
have a default namespace.
+ scopes = {
+ n.team_name
+ for n in ExecutorLoader.get_executor_names(validate_teams=False)
+ if n.team_name != team_name
+ and issubclass(ExecutorLoader.import_executor_cls(n)[0],
KubernetesExecutor)
+ }
+ for other in sorted(scopes, key=lambda t: t or ""):
+ if conf.get("kubernetes_executor", "namespace", **_team_kwargs(other))
== namespace:
+ owner = f"a KubernetesExecutor of team {other!r}" if other else "a
global KubernetesExecutor"
+ raise SystemExit(
+ f"Namespace {namespace!r} of team {team_name!r} is also used
by {owner}, so cleanup-pods "
+ "would delete its pods too. Pass --allow-shared-namespace if
that executor runs on "
+ "another cluster."
+ )
+
+
+def _delete_pod(name, namespace, kube_client):
"""
Delete a namespaced pod.
Helper Function for cleanup_pods.
"""
- kube_client = get_kube_client()
delete_options = client.V1DeleteOptions()
print(f'Deleting POD "{name}" from "{namespace}" namespace')
api_response = kube_client.delete_namespaced_pod(name=name,
namespace=namespace, body=delete_options)
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 0534f10799a..68f61b03bfa 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
@@ -115,6 +115,11 @@ class KubernetesExecutor(BaseExecutor):
RUNNING_POD_LOG_LINES = 100
supports_ad_hoc_ti_run: bool = True
+ # Import paths of the factories building a subclass's own sync and async
clients, called with
+ # ``team_name``. They take precedence over [kubernetes_executor]
client_factory and
+ # async_client_factory, which stay free for a KubernetesExecutor running
alongside.
+ client_factory: str | None = None
+ async_client_factory: str | None = None
supports_multi_team: bool = True
supports_task_instance_uuid = hasattr(BaseExecutor, "get_task_key")
@@ -277,16 +282,24 @@ class KubernetesExecutor(BaseExecutor):
)
from airflow.providers.cncf.kubernetes.kube_client import
_get_executor_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)
+ if self.kube_config.async_pod_creation:
+ if self.client_factory and not self.async_client_factory:
+ raise AirflowConfigException(
+ f"{type(self).__name__} sets client_factory but not
async_client_factory, "
+ "which is required when async_pod_creation is enabled."
+ )
+ if (
+ not self.client_factory
+ 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,
executor_client_factory=self.client_factory
+ )
self.kube_scheduler = AirflowKubernetesScheduler(
kube_config=self.kube_config,
result_queue=self.result_queue,
@@ -294,6 +307,8 @@ class KubernetesExecutor(BaseExecutor):
scheduler_job_id=self.scheduler_job_id,
team_name=self.team_name,
supports_task_instance_uuid=self.supports_task_instance_uuid,
+ client_factory=self.client_factory,
+ async_client_factory=self.async_client_factory,
)
def _coordinator_extra(self, queue: str | None) -> dict[str, Any] | None:
@@ -1045,7 +1060,9 @@ class KubernetesExecutor(BaseExecutor):
from airflow.providers.cncf.kubernetes.kube_client import
_get_executor_kube_client
from airflow.providers.cncf.kubernetes.pod_generator import
PodGenerator
- client = _get_executor_kube_client(team_name=self.team_name)
+ client = _get_executor_kube_client(
+ team_name=self.team_name,
executor_client_factory=self.client_factory
+ )
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 29dede7d10f..39997af0006 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
@@ -95,6 +95,7 @@ class KubernetesJobWatcher(multiprocessing.Process,
LoggingMixin):
kube_config: Configuration,
supports_task_instance_uuid: bool = False,
team_name: str | None = None,
+ client_factory: str | None = None,
):
super().__init__()
self.namespace = namespace
@@ -103,16 +104,19 @@ 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.
+ # Strings survive the spawn start method, so the watcher can build its
team's client in its
+ # own process.
self.team_name = team_name
+ self.client_factory = client_factory
def run(self) -> None:
"""Perform watching."""
if TYPE_CHECKING:
assert self.scheduler_job_id
- kube_client: client.CoreV1Api =
_get_executor_kube_client(team_name=self.team_name)
+ kube_client: client.CoreV1Api = _get_executor_kube_client(
+ team_name=self.team_name,
executor_client_factory=self.client_factory
+ )
while True:
try:
self.resource_version = self._run(
@@ -507,6 +511,8 @@ class AirflowKubernetesScheduler(LoggingMixin):
scheduler_job_id: str,
team_name: str | None = None,
supports_task_instance_uuid: bool = False,
+ client_factory: str | None = None,
+ async_client_factory: str | None = None,
):
super().__init__()
self.log.debug("Creating Kubernetes executor")
@@ -521,6 +527,8 @@ class AirflowKubernetesScheduler(LoggingMixin):
self.scheduler_job_id = scheduler_job_id
# Set before the watchers are built, since they are told which team
they are watching for.
self.team_name = team_name
+ self.client_factory = client_factory
+ self.async_client_factory = async_client_factory
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
@@ -573,6 +581,7 @@ class AirflowKubernetesScheduler(LoggingMixin):
kube_config=self.kube_config,
supports_task_instance_uuid=self.supports_task_instance_uuid,
team_name=self.team_name,
+ client_factory=self.client_factory,
)
watcher.start()
return watcher
@@ -720,7 +729,9 @@ 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_executor_async_kube_client(team_name=self.team_name)
+ self._async_pod_client = await _get_executor_async_kube_client(
+ team_name=self.team_name,
executor_client_factory=self.async_client_factory
+ )
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 7e1c8d2cd9e..47e5e5cf536 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
@@ -273,14 +273,14 @@ def get_provider_info():
"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, [...]
+ "description": "Import path of a callable returning
the ``kubernetes.client.CoreV1Api`` the executor\nshould use, for deployments
that mint their own credentials. If the callable takes a\n``team_name``
argument, it is called with the executor's team (``None`` for an executor\nwith
no team); otherwise it is called with no arguments. When set it replaces
the\ndefault client construction entirely, so ``in_cluster``,
``cluster_context``,\n``config_file``, ``verify_ssl [...]
"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",
+ "description": "Import path of a callable returning
the ``kubernetes_asyncio.client.CoreV1Api`` used for\nconcurrent pod creation,
called the same way as ``client_factory``. Required when\n``client_factory`` is
set and ``async_pod_creation`` is enabled, so that both clients\ncarry the same
credentials; the executor refuses to start without it.\n",
"version_added": "10.24.0",
"type": "string",
"example": "my_company.kubernetes.build_async_client",
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 9b6acbfdb75..1c14a8b2247 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
@@ -18,15 +18,18 @@
from __future__ import annotations
+import inspect
import logging
import socket
import warnings
+from collections.abc import Callable
from typing import Any
import urllib3.util
from urllib3.connection import HTTPConnection
from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import conf
log = logging.getLogger(__name__)
@@ -188,6 +191,7 @@ def get_kube_client(
in_cluster: bool | None = None,
cluster_context: str | None = None,
config_file: str | None = None,
+ team_name: str | None = None,
) -> client.CoreV1Api:
"""
Retrieve Kubernetes client.
@@ -195,22 +199,31 @@ def get_kube_client(
:param in_cluster: whether we are in cluster
:param cluster_context: context of the cluster
:param config_file: configuration file
+ :param team_name: team whose ``[kubernetes_executor]`` connection settings
to use. Like all team
+ config, they don't fall back to the un-prefixed section.
:return: kubernetes client
"""
+ team_kwargs = _team_kwargs(team_name)
if in_cluster is None:
- in_cluster = conf.getboolean("kubernetes_executor", "in_cluster")
+ in_cluster = conf.getboolean("kubernetes_executor", "in_cluster",
**team_kwargs)
if not has_kubernetes:
raise _import_err
configuration = _get_default_configuration()
+ # Keepalive stays a global setting, so it is not read from the team's
config.
if conf.getboolean("kubernetes_executor", "enable_tcp_keepalive"):
configure_tcp_keepalive(configuration)
api_client_retry_configuration = conf.getjson(
- "kubernetes_executor", "api_client_retry_configuration", fallback={}
+ "kubernetes_executor", "api_client_retry_configuration", fallback={},
**team_kwargs
)
- if not conf.getboolean("kubernetes_executor", "verify_ssl"):
+ verify_ssl = conf.getboolean("kubernetes_executor", "verify_ssl",
**team_kwargs)
+ if team_name:
+ # The default configuration is shared by the whole process, so a
team's client sets its
+ # own value rather than changing the default or inheriting another
executor's change.
+ configuration.verify_ssl = verify_ssl
+ elif not verify_ssl:
_disable_verify_ssl()
if isinstance(api_client_retry_configuration, dict):
@@ -222,17 +235,17 @@ def get_kube_client(
config.load_incluster_config(client_configuration=configuration)
else:
if cluster_context is None:
- cluster_context = conf.get("kubernetes_executor",
"cluster_context", fallback=None)
+ cluster_context = conf.get("kubernetes_executor",
"cluster_context", fallback=None, **team_kwargs)
if config_file is None:
- config_file = conf.get("kubernetes_executor", "config_file",
fallback=None)
+ config_file = conf.get("kubernetes_executor", "config_file",
fallback=None, **team_kwargs)
config.load_kube_config(
config_file=config_file, context=cluster_context,
client_configuration=configuration
)
- if not conf.getboolean("kubernetes_executor", "verify_ssl"):
+ if not verify_ssl:
configuration.verify_ssl = False
- ssl_ca_cert = conf.get("kubernetes_executor", "ssl_ca_cert")
+ ssl_ca_cert = conf.get("kubernetes_executor", "ssl_ca_cert", **team_kwargs)
if ssl_ca_cert:
configuration.ssl_ca_cert = ssl_ca_cert
@@ -244,6 +257,7 @@ async def get_async_kube_client(
in_cluster: bool | None = None,
cluster_context: str | None = None,
config_file: str | None = None,
+ team_name: str | None = None,
) -> async_client.CoreV1Api:
"""
Retrieve an asynchronous Kubernetes client.
@@ -256,29 +270,31 @@ async def get_async_kube_client(
:param in_cluster: whether we are in cluster
:param cluster_context: context of the cluster
:param config_file: configuration file
+ :param team_name: team whose connection settings to use, as in
:func:`get_kube_client`
:return: asynchronous kubernetes client
"""
if not has_kubernetes:
raise _import_err
+ team_kwargs = _team_kwargs(team_name)
if in_cluster is None:
- in_cluster = conf.getboolean("kubernetes_executor", "in_cluster")
+ in_cluster = conf.getboolean("kubernetes_executor", "in_cluster",
**team_kwargs)
configuration = async_client.Configuration()
- if not conf.getboolean("kubernetes_executor", "verify_ssl"):
+ if not conf.getboolean("kubernetes_executor", "verify_ssl", **team_kwargs):
configuration.verify_ssl = False
if in_cluster:
async_config.load_incluster_config(client_configuration=configuration)
else:
if cluster_context is None:
- cluster_context = conf.get("kubernetes_executor",
"cluster_context", fallback=None)
+ cluster_context = conf.get("kubernetes_executor",
"cluster_context", fallback=None, **team_kwargs)
if config_file is None:
- config_file = conf.get("kubernetes_executor", "config_file",
fallback=None)
+ config_file = conf.get("kubernetes_executor", "config_file",
fallback=None, **team_kwargs)
await async_config.load_kube_config(
config_file=config_file, context=cluster_context,
client_configuration=configuration
)
- ssl_ca_cert = conf.get("kubernetes_executor", "ssl_ca_cert")
+ ssl_ca_cert = conf.get("kubernetes_executor", "ssl_ca_cert", **team_kwargs)
if ssl_ca_cert:
configuration.ssl_ca_cert = ssl_ca_cert
@@ -290,39 +306,60 @@ def _team_kwargs(team_name: str | None) -> dict[str, str]:
return {"team_name": team_name} if team_name else {}
-def _get_executor_kube_client(team_name: str | None = None) ->
client.CoreV1Api:
+def _call_client_factory(factory: Callable[..., Any], team_name: str | None)
-> Any:
+ """Pass ``team_name`` only to a factory that takes it, since factories
used to take no arguments."""
+ parameters = inspect.signature(factory).parameters.values()
+ if any(p.name == "team_name" or p.kind is p.VAR_KEYWORD for p in
parameters):
+ return factory(team_name=team_name)
+ return factory()
+
+
+def _get_executor_kube_client(
+ team_name: str | None = None, executor_client_factory: 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.
+ Uses the executor's own factory when it has one, then the
``client_factory`` setting, 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
+ own cluster. Like all team config, the team's settings don't fall back
to the un-prefixed
+ section. A factory that takes ``team_name`` is told the team.
+ :param executor_client_factory: import path of the executor's own factory,
see
+ ``KubernetesExecutor.client_factory``
: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.
+ # Factories are import paths, since the pod watcher runs in its own
process and resolves its
+ # factory there itself.
+ if executor_client_factory:
+ return _call_client_factory(import_string(executor_client_factory),
team_name)
if client_factory := conf.getimport(
"kubernetes_executor", "client_factory", fallback=None,
**_team_kwargs(team_name)
):
- return client_factory()
- return get_kube_client()
+ return _call_client_factory(client_factory, team_name)
+ return get_kube_client(team_name=team_name)
-async def _get_executor_async_kube_client(team_name: str | None = None) ->
async_client.CoreV1Api:
+async def _get_executor_async_kube_client(
+ team_name: str | None = None, executor_client_factory: 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.
+ Uses the executor's own factory when it has one, then the
``async_client_factory`` setting,
+ and :func:`get_async_kube_client` otherwise.
- :param team_name: team the executor is running for, resolved the same way
as in
+ :param team_name: team the executor is running for, used the same way as in
:func:`_get_executor_kube_client`
+ :param executor_client_factory: import path of the executor's own factory,
see
+ ``KubernetesExecutor.async_client_factory``
:return: asynchronous kubernetes client
"""
+ if executor_client_factory:
+ return _call_client_factory(import_string(executor_client_factory),
team_name)
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()
+ return _call_client_factory(async_client_factory, team_name)
+ return await get_async_kube_client(team_name=team_name)
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/cli/test_kubernetes_command.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/cli/test_kubernetes_command.py
index 29eadcc55a2..ce2291a7a46 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/cli/test_kubernetes_command.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/cli/test_kubernetes_command.py
@@ -31,9 +31,10 @@ from kubernetes.client import models as k8s
from airflow.cli import cli_parser
from airflow.executors import executor_loader
from airflow.providers.cncf.kubernetes.cli import kubernetes_command
+from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import
KubernetesExecutor
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
pytestmark = pytest.mark.db_test
@@ -118,6 +119,24 @@ class TestGenerateDagYamlCommand:
assert (output_path / path.name).stat().st_size > 0
+class _ExecutorWithOwnClient(KubernetesExecutor):
+ client_factory = "my_company.build_client"
+
+
+TEAM_CONF = {
+ ("core", "multi_team"): "True",
+ ("core", "executor"): (
+
f"KubernetesExecutor;team_a={__name__}._ExecutorWithOwnClient,KubernetesExecutor;"
+ "team_b=LocalExecutor,KubernetesExecutor"
+ ),
+ ("kubernetes_executor", "namespace"): "global-namespace",
+}
+TEAM_NAMESPACES = {
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__NAMESPACE": "team-a-namespace",
+ "AIRFLOW__TEAM_B___KUBERNETES_EXECUTOR__NAMESPACE": "team-b-namespace",
+}
+
+
class TestCleanUpPodsCommand:
label_selector = "dag_id,task_id,try_number,airflow_version"
@@ -128,12 +147,12 @@ class TestCleanUpPodsCommand:
importlib.reload(cli_parser)
cls.parser = cli_parser.get_parser()
- @mock.patch("kubernetes.client.CoreV1Api.delete_namespaced_pod")
-
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.config.load_incluster_config")
- def test_delete_pod(self, load_incluster_config, delete_namespaced_pod):
- kubernetes_command._delete_pod("dummy", "awesome-namespace")
- delete_namespaced_pod.assert_called_with(body=mock.ANY, name="dummy",
namespace="awesome-namespace")
- load_incluster_config.assert_called_once()
+ def test_delete_pod(self):
+ kube_client = MagicMock()
+ kubernetes_command._delete_pod("dummy", "awesome-namespace",
kube_client)
+ kube_client.delete_namespaced_pod.assert_called_with(
+ body=mock.ANY, name="dummy", namespace="awesome-namespace"
+ )
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod")
@mock.patch("kubernetes.client.CoreV1Api.list_namespaced_pod")
@@ -176,7 +195,7 @@ class TestCleanUpPodsCommand:
list_namespaced_pod.assert_called_once_with(
namespace="awesome-namespace", limit=500,
label_selector=self.label_selector
)
- delete_pod.assert_called_with("dummy", "awesome-namespace")
+ delete_pod.assert_called_with("dummy", "awesome-namespace", mock.ANY)
load_incluster_config.assert_called_once()
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod")
@@ -226,7 +245,7 @@ class TestCleanUpPodsCommand:
list_namespaced_pod.assert_called_once_with(
namespace="awesome-namespace", limit=500,
label_selector=self.label_selector
)
- delete_pod.assert_called_with("dummy3", "awesome-namespace")
+ delete_pod.assert_called_with("dummy3", "awesome-namespace", mock.ANY)
load_incluster_config.assert_called_once()
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod")
@@ -249,7 +268,7 @@ class TestCleanUpPodsCommand:
list_namespaced_pod.assert_called_once_with(
namespace="awesome-namespace", limit=500,
label_selector=self.label_selector
)
- delete_pod.assert_called_with("dummy4", "awesome-namespace")
+ delete_pod.assert_called_with("dummy4", "awesome-namespace", mock.ANY)
load_incluster_config.assert_called_once()
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod")
@@ -271,7 +290,7 @@ class TestCleanUpPodsCommand:
list_namespaced_pod.assert_called_once_with(
namespace="awesome-namespace", limit=500,
label_selector=self.label_selector
)
- delete_pod.assert_called_with("dummy5", "awesome-namespace")
+ delete_pod.assert_called_with("dummy5", "awesome-namespace", mock.ANY)
load_incluster_config.assert_called_once()
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod")
@@ -325,9 +344,183 @@ class TestCleanUpPodsCommand:
),
]
list_namespaced_pod.assert_has_calls(calls)
- delete_pod.assert_called_with("dummy", "awesome-namespace")
+ delete_pod.assert_called_with("dummy", "awesome-namespace", mock.ANY)
load_incluster_config.assert_called_once()
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.usefixtures("clean_executor_loader")
+ @pytest.mark.parametrize("has_resolve_executor_name", [True, False],
ids=["resolve", "pre-3.4"])
+ @pytest.mark.parametrize(
+ ("cli_args", "team_client_kwargs", "namespace"),
+ [
+ pytest.param([], None, "global-namespace", id="global"),
+ pytest.param(
+ ["--team", "team_a"],
+ {"team_name": "team_a", "executor_client_factory":
"my_company.build_client"},
+ "team-a-namespace",
+ id="team-default-executor",
+ ),
+ pytest.param(
+ ["--team", "team_a", "--executor", "KubernetesExecutor"],
+ {"team_name": "team_a", "executor_client_factory": None},
+ "team-a-namespace",
+ id="team-named-executor",
+ ),
+ pytest.param(
+ ["--team", "team_a", "--namespace", "other-namespace"],
+ {"team_name": "team_a", "executor_client_factory":
"my_company.build_client"},
+ "other-namespace",
+ id="namespace-flag",
+ ),
+ pytest.param(
+ ["--team", "team_a", "--namespace", "global-namespace",
"--allow-shared-namespace"],
+ {"team_name": "team_a", "executor_client_factory":
"my_company.build_client"},
+ "global-namespace",
+ id="allow-shared-namespace",
+ ),
+ ],
+ )
+ @conf_vars(TEAM_CONF)
+ @mock.patch.dict(os.environ, TEAM_NAMESPACES)
+
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command._delete_pod",
autospec=True)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.cli.kubernetes_command._get_executor_kube_client",
autospec=True
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.cli.kubernetes_command.get_kube_client",
autospec=True)
+ def test_cleanup_pods_uses_the_teams_executor_client(
+ self,
+ get_kube_client,
+ get_executor_kube_client,
+ delete_pod,
+ monkeypatch,
+ cli_args,
+ team_client_kwargs,
+ namespace,
+ has_resolve_executor_name,
+ ):
+ if not has_resolve_executor_name:
+ monkeypatch.delattr(kubernetes_command.ExecutorLoader,
"resolve_executor_name", raising=False)
+ kube_client = (get_executor_kube_client if team_client_kwargs else
get_kube_client).return_value
+ pod = MagicMock()
+ pod.metadata.name = "dummy"
+ pod.metadata.creation_timestamp = parse("2021-12-20T08:01:07Z")
+ pod.status.phase = "Succeeded"
+ pod.status.reason = None
+ kube_client.list_namespaced_pod.return_value = MagicMock(
+ items=[pod], metadata=MagicMock(_continue=None)
+ )
+ kubernetes_command.cleanup_pods(self.parser.parse_args(["kubernetes",
"cleanup-pods", *cli_args]))
+ if team_client_kwargs:
+
get_executor_kube_client.assert_called_once_with(**team_client_kwargs)
+ get_kube_client.assert_not_called()
+ else:
+ get_kube_client.assert_called_once_with()
+ get_executor_kube_client.assert_not_called()
+ kube_client.list_namespaced_pod.assert_called_once_with(
+ namespace=namespace, limit=500, label_selector=self.label_selector
+ )
+ delete_pod.assert_called_once_with("dummy", namespace, kube_client)
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.usefixtures("clean_executor_loader")
+ @pytest.mark.parametrize("has_resolve_executor_name", [True, False],
ids=["resolve", "pre-3.4"])
+ @pytest.mark.parametrize(
+ ("cli_args", "error"),
+ [
+ pytest.param(["--executor", "KubernetesExecutor"], "--executor
requires --team", id="no-team"),
+ pytest.param(["--team", "team_c"], "Team 'team_c' has no
executors", id="unknown-team"),
+ pytest.param(
+ ["--team", "team_a", "--executor", "CeleryExecutor"],
+ "Unknown executor",
+ id="unknown-executor",
+ ),
+ pytest.param(
+ ["--team", "team_b"],
+ "LocalExecutor of team 'team_b' isn't a KubernetesExecutor",
+ id="not-k8s",
+ ),
+ pytest.param(
+ ["--team", "team_a", "--namespace", "global-namespace"],
+ "also used by a global KubernetesExecutor",
+ id="shared-with-global",
+ ),
+ pytest.param(
+ ["--team", "team_a", "--namespace", "team-b-namespace"],
+ "also used by a KubernetesExecutor of team 'team_b'",
+ id="shared-with-team",
+ ),
+ ],
+ )
+ @conf_vars(TEAM_CONF)
+ @mock.patch.dict(os.environ, TEAM_NAMESPACES)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.cli.kubernetes_command._get_executor_kube_client",
+ autospec=True,
+ side_effect=AssertionError("connected"),
+ )
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.cli.kubernetes_command.get_kube_client",
+ autospec=True,
+ side_effect=AssertionError("connected"),
+ )
+ def test_cleanup_pods_team_refuses(
+ self,
+ get_kube_client,
+ get_executor_kube_client,
+ monkeypatch,
+ cli_args,
+ error,
+ has_resolve_executor_name,
+ ):
+ if not has_resolve_executor_name:
+ monkeypatch.delattr(kubernetes_command.ExecutorLoader,
"resolve_executor_name", raising=False)
+ with pytest.raises(SystemExit, match=error):
+
kubernetes_command.cleanup_pods(self.parser.parse_args(["kubernetes",
"cleanup-pods", *cli_args]))
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.usefixtures("clean_executor_loader")
+ @pytest.mark.parametrize(
+ "extra_conf",
+ [
+ pytest.param({("core", "executor"):
"LocalExecutor;team_a=KubernetesExecutor"}, id="global"),
+ pytest.param(
+ {
+ ("core", "executor"):
"KubernetesExecutor;team_a=KubernetesExecutor;team_b=LocalExecutor",
+ ("kubernetes_executor", "namespace"): "global-namespace",
+ },
+ id="team",
+ ),
+ ],
+ )
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.cli.kubernetes_command._get_executor_kube_client",
autospec=True
+ )
+ def test_cleanup_pods_team_ignores_scopes_without_kubernetes_executor(
+ self, get_executor_kube_client, extra_conf
+ ):
+ # The other scope has no namespace set, so it has team_a's built-in
default, but makes no pods.
+ kube_client = get_executor_kube_client.return_value
+ kube_client.list_namespaced_pod.return_value = MagicMock(items=[],
metadata=MagicMock(_continue=None))
+ with conf_vars({("core", "multi_team"): "True", **extra_conf}):
+ kubernetes_command.cleanup_pods(
+ self.parser.parse_args(["kubernetes", "cleanup-pods",
"--team", "team_a"])
+ )
+ kube_client.list_namespaced_pod.assert_called_once_with(
+ namespace="default", limit=500, label_selector=self.label_selector
+ )
+
+ @mock.patch.object(kubernetes_command, "AIRFLOW_V_3_2_PLUS", False)
+ @mock.patch(
+
"airflow.providers.cncf.kubernetes.cli.kubernetes_command.get_kube_client",
+ autospec=True,
+ side_effect=AssertionError("connected"),
+ )
+ def test_cleanup_pods_team_requires_airflow_3_2(self, get_kube_client):
+ with pytest.raises(SystemExit, match=r"--team requires Airflow
3\.2\+"):
+ kubernetes_command.cleanup_pods(
+ self.parser.parse_args(["kubernetes", "cleanup-pods",
"--team", "team_a"])
+ )
+
@pytest.mark.parametrize(
("pod_kwargs", "min_completed_minutes", "expect_deleted"),
[
@@ -405,7 +598,7 @@ class TestCleanUpPodsCommand:
)
)
if expect_deleted:
- delete_pod.assert_called_once_with("run-o1sxc2on",
"awesome-namespace")
+ delete_pod.assert_called_once_with("run-o1sxc2on",
"awesome-namespace", mock.ANY)
else:
delete_pod.assert_not_called()
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 26d0aa254d6..568747c038e 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
@@ -4108,35 +4108,71 @@ class TestKubernetesExecutorMultiTeam:
executor.end()
+class _ExecutorWithOwnClients(KubernetesExecutor):
+ client_factory = "my_company.build_client"
+ async_client_factory = "my_company.build_async_client"
+
+
+EXECUTOR_CLIENT_FACTORIES = pytest.mark.parametrize(
+ ("executor_class", "client_factory", "async_client_factory"),
+ [
+ pytest.param(KubernetesExecutor, None, None, id="kubernetes-executor"),
+ pytest.param(
+ _ExecutorWithOwnClients,
+ "my_company.build_client",
+ "my_company.build_async_client",
+ id="subclass-with-own-clients",
+ ),
+ ],
+)
+
+
@pytest.mark.skipif(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.
+ there while every other test stays green. Each site also has to pass its
team and the
+ executor's own factory, or the client would be built for the wrong team or
by the wrong
+ factory.
"""
+ @EXECUTOR_CLIENT_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.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()
+ def test_start_uses_executor_client(
+ self,
+ mock_watcher,
+ mock_client,
+ mock_get_executor_kube_client,
+ executor_class,
+ client_factory,
+ async_client_factory,
+ ):
+ executor = executor_class()
executor.team_name = "team_a"
executor.job_id = 1
try:
executor.start()
+ assert executor.kube_scheduler.async_client_factory ==
async_client_factory
finally:
executor.end()
-
mock_get_executor_kube_client.assert_called_once_with(team_name="team_a")
+ mock_get_executor_kube_client.assert_called_once_with(
+ team_name="team_a", executor_client_factory=client_factory
+ )
+ assert mock_watcher.call_args.kwargs["client_factory"] ==
client_factory
+ @EXECUTOR_CLIENT_FACTORIES
@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):
+ def test_get_streaming_task_log_uses_executor_client(
+ self, mock_get_executor_kube_client, executor_class, client_factory,
async_client_factory
+ ):
ti = mock.MagicMock(
dag_id="dag",
task_id="task",
@@ -4146,12 +4182,14 @@ class TestClientFactoryCallSites:
hostname="",
executor_config={},
)
- executor = KubernetesExecutor()
+ executor = executor_class()
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_get_executor_kube_client.assert_called_once_with(
+ team_name="team_a", executor_client_factory=client_factory
+ )
@mock.patch.object(KubernetesJobWatcher, "_run",
side_effect=RuntimeError("stop"), autospec=True)
@mock.patch(
@@ -4166,24 +4204,30 @@ class TestClientFactoryCallSites:
scheduler_job_id="1",
kube_config=mock.MagicMock(),
team_name="team_a",
+ client_factory="my_company.build_client",
)
with pytest.raises(RuntimeError, match="stop"):
watcher.run()
-
mock_get_executor_kube_client.assert_called_once_with(team_name="team_a")
+ mock_get_executor_kube_client.assert_called_once_with(
+ team_name="team_a",
executor_client_factory="my_company.build_client"
+ )
@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")
+ """The watcher runs in its own process, so it only learns its team and
factory by being told."""
+ scheduler = mock.Mock(
+ team_name="team_a", scheduler_job_id="1",
client_factory="my_company.build_client"
+ )
AirflowKubernetesScheduler._make_kube_watcher(scheduler, "ns")
assert mock_watcher.call_args.kwargs["team_name"] == "team_a"
+ assert mock_watcher.call_args.kwargs["client_factory"] ==
"my_company.build_client"
@pytest.mark.asyncio
@mock.patch(
@@ -4191,12 +4235,19 @@ class TestClientFactoryCallSites:
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 = mock.Mock(
+ pod_creation_max_concurrency=1,
+ _async_pod_client=None,
+ team_name="team_a",
+ async_client_factory="my_company.build_async_client",
+ )
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")
+ mock_get_executor_async_kube_client.assert_awaited_once_with(
+ team_name="team_a",
executor_client_factory="my_company.build_async_client"
+ )
@pytest.mark.skipif(AirflowKubernetesScheduler is None, reason="kubernetes
python package is not installed")
@@ -4265,6 +4316,48 @@ class TestAsyncClientFactoryValidation:
executor.start()
mock_get_executor_kube_client.assert_not_called()
+ @pytest.mark.parametrize(
+ ("async_client_factory", "expect_raise"),
+ [
+ pytest.param(None, True, id="subclass-without-async-factory"),
+ pytest.param("my_company.build_async_client", False,
id="subclass-with-both-factories"),
+ ],
+ )
+ @conf_vars(
+ {
+ ("kubernetes_executor", "async_pod_creation"): "True",
+ ("kubernetes_executor", "client_factory"):
"my_company.build_config_client",
+ }
+ )
+
@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_checks_the_subclass_factories_not_the_config(
+ self, mock_manager, mock_scheduler, mock_get_executor_kube_client,
async_client_factory, expect_raise
+ ):
+ """A subclass's own factories replace the config ones, so its pair is
what has to be complete."""
+ executor_class = type(
+ "ExecutorWithOwnClients",
+ (KubernetesExecutor,),
+ {"client_factory": "my_company.build_client",
"async_client_factory": async_client_factory},
+ )
+ executor = executor_class()
+ executor.job_id = 1
+ if expect_raise:
+ with pytest.raises(
+ AirflowConfigException, match="sets client_factory but not
async_client_factory"
+ ):
+ executor.start()
+ else:
+ executor.start()
+ mock_get_executor_kube_client.assert_called_once()
+
@pytest.fixture
def kubernetes_scheduler_factory(mocker):
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_client.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_client.py
index 19909341135..c8c1aab827d 100644
--- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_client.py
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_client.py
@@ -33,6 +33,30 @@ from airflow.providers.cncf.kubernetes.kube_client import (
)
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_PLUS
+
+GLOBAL_CONNECTION_SETTINGS = {
+ ("kubernetes_executor", "in_cluster"): "False",
+ ("kubernetes_executor", "config_file"): "/global/kubeconfig",
+ ("kubernetes_executor", "cluster_context"): "global-context",
+ ("kubernetes_executor", "ssl_ca_cert"): "/global/ca.crt",
+ ("kubernetes_executor", "api_client_retry_configuration"): '{"total": 1}',
+}
+TEAM_A_CONNECTION_SETTINGS = {
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__IN_CLUSTER": "False",
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CONFIG_FILE": "/team-a/kubeconfig",
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CLUSTER_CONTEXT": "team-a-context",
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__SSL_CA_CERT": "/team-a/ca.crt",
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__API_CLIENT_RETRY_CONFIGURATION":
'{"total": 7}',
+}
+
+
[email protected]
+def fresh_default_configuration():
+ default = Configuration.get_default_copy()
+ Configuration.set_default(Configuration())
+ yield
+ Configuration.set_default(default)
@pytest.fixture
@@ -209,3 +233,58 @@ class TestClient:
client_configuration =
mock_in_cluster_loader().load_and_set.call_args.args[0]
assert client_configuration.retries.total == 3
assert client_configuration.retries.backoff_factor == 0.5
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @conf_vars(GLOBAL_CONNECTION_SETTINGS)
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.config")
+ def test_team_uses_its_own_connection_settings(self, config, monkeypatch):
+ for key, value in TEAM_A_CONNECTION_SETTINGS.items():
+ monkeypatch.setenv(key, value)
+
+ client = get_kube_client(team_name="team_a")
+
+ config.load_incluster_config.assert_not_called()
+ config.load_kube_config.assert_called_once_with(
+ config_file="/team-a/kubeconfig", context="team-a-context",
client_configuration=mock.ANY
+ )
+ assert client.api_client.configuration.ssl_ca_cert == "/team-a/ca.crt"
+ assert client.api_client.configuration.retries.total == 7
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @conf_vars(GLOBAL_CONNECTION_SETTINGS)
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.config")
+ def test_team_does_not_fall_back_to_the_global_connection_settings(self,
config):
+ client = get_kube_client(team_name="team_b")
+
+ config.load_incluster_config.assert_called_once()
+ config.load_kube_config.assert_not_called()
+ assert client.api_client.configuration.ssl_ca_cert is None
+ assert client.api_client.configuration.retries.total != 1
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @conf_vars({("kubernetes_executor", "verify_ssl"): "False"})
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.config")
+ def test_team_verify_ssl_does_not_leak_between_clients(
+ self, config, monkeypatch, fresh_default_configuration
+ ):
+ """The process-wide default configuration must not carry one
executor's verify_ssl to another."""
+
monkeypatch.setenv("AIRFLOW__TEAM_B___KUBERNETES_EXECUTOR__VERIFY_SSL", "False")
+
+ team_b_client = get_kube_client(in_cluster=True, team_name="team_b")
+ assert Configuration.get_default_copy().verify_ssl
+ global_client = get_kube_client(in_cluster=True)
+ team_a_client = get_kube_client(in_cluster=True, team_name="team_a")
+
+ assert not team_b_client.api_client.configuration.verify_ssl
+ assert not global_client.api_client.configuration.verify_ssl
+ assert team_a_client.api_client.configuration.verify_ssl
+
+ @pytest.mark.platform("linux")
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @conf_vars({("kubernetes_executor", "enable_tcp_keepalive"): "True"})
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.config",
autospec=True)
+ def test_team_client_gets_tcp_keepalive(self, mock_config):
+ client = get_kube_client(in_cluster=False, team_name="team_a")
+
+ pool_options =
client.api_client.rest_client.pool_manager.connection_pool_kw["socket_options"]
+ assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in pool_options
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 de2d61718bc..c195c943b58 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
@@ -43,13 +43,21 @@ def build_other_client():
return OTHER_MARKER
+def build_team_client(team_name):
+ return FACTORY_MARKER, team_name
+
+
+def build_client_from_kwargs(**kwargs):
+ return FACTORY_MARKER, kwargs
+
+
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()
+ mock_get_kube_client.assert_called_once_with(team_name=None)
@conf_vars({("kubernetes_executor", "client_factory"):
f"{__name__}.build_fake_client"})
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client",
autospec=True)
@@ -80,7 +88,35 @@ class TestGetExecutorKubeClient:
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()
+ mock_get_kube_client.assert_called_once_with(team_name="team_b")
+
+ @pytest.mark.parametrize(
+ ("factory", "team_name", "expected"),
+ [
+ pytest.param("build_team_client", None, (FACTORY_MARKER, None),
id="no-team"),
+ pytest.param("build_team_client", "team_a", (FACTORY_MARKER,
"team_a"), id="team"),
+ pytest.param(
+ "build_client_from_kwargs", "team_a", (FACTORY_MARKER,
{"team_name": "team_a"}), id="kwargs"
+ ),
+ ],
+ )
+ def test_factory_is_told_the_team(self, factory, team_name, expected,
monkeypatch):
+ if team_name and not AIRFLOW_V_3_2_PLUS:
+ pytest.skip("team config requires Airflow 3.2+")
+ team_prefix = f"{team_name.upper()}___" if team_name else ""
+ monkeypatch.setenv(
+ f"AIRFLOW__{team_prefix}KUBERNETES_EXECUTOR__CLIENT_FACTORY",
f"{__name__}.{factory}"
+ )
+
+ assert _get_executor_kube_client(team_name=team_name) == expected
+
+ @conf_vars({("kubernetes_executor", "client_factory"):
f"{__name__}.build_other_client"})
+ def test_executor_factory_takes_precedence_over_the_config(self):
+ client = _get_executor_kube_client(
+ team_name="team_a",
executor_client_factory=f"{__name__}.build_team_client"
+ )
+
+ assert client == (FACTORY_MARKER, "team_a")
@conf_vars({("kubernetes_executor", "client_factory"):
"no.such.module.build_client"})
def test_unimportable_factory_raises(self):
@@ -95,7 +131,15 @@ class TestGetExecutorAsyncKubeClient:
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()
+ mock_get_async_kube_client.assert_awaited_once_with(team_name=None)
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_async_kube_client",
autospec=True)
+ async def test_team_without_factory_gets_its_own_default_client(self,
mock_get_async_kube_client):
+ await _get_executor_async_kube_client(team_name="team_b")
+
+ mock_get_async_kube_client.assert_awaited_once_with(team_name="team_b")
@pytest.mark.asyncio
@mock.patch("airflow.providers.cncf.kubernetes.kube_client.get_async_kube_client",
autospec=True)
@@ -118,6 +162,24 @@ class TestGetExecutorAsyncKubeClient:
assert client == FACTORY_MARKER
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.asyncio
+ async def test_async_factory_is_told_the_team(self, monkeypatch):
+ monkeypatch.setenv(
+ "AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__ASYNC_CLIENT_FACTORY",
f"{__name__}.build_team_client"
+ )
+
+ assert await _get_executor_async_kube_client(team_name="team_a") ==
(FACTORY_MARKER, "team_a")
+
+ @pytest.mark.asyncio
+ async def
test_executor_async_factory_takes_precedence_over_the_config(self):
+ with conf_vars({("kubernetes_executor", "async_client_factory"):
f"{__name__}.build_other_client"}):
+ client = await _get_executor_async_kube_client(
+ team_name="team_a",
executor_client_factory=f"{__name__}.build_team_client"
+ )
+
+ assert client == (FACTORY_MARKER, "team_a")
+
class TestGetAsyncKubeClient:
@pytest.mark.asyncio
@@ -131,3 +193,59 @@ class TestGetAsyncKubeClient:
assert isinstance(api.api_client, _TimeoutAsyncK8sApiClient)
mock_load_incluster.assert_called_once()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="team config requires
Airflow 3.2+")
+ @pytest.mark.parametrize(
+ ("team_name", "kube_config_kwargs", "ssl_ca_cert", "verify_ssl"),
+ [
+ pytest.param(
+ "team_a",
+ {"config_file": "/team-a/kubeconfig", "context":
"team-a-context"},
+ "/team-a/ca.crt",
+ False,
+ id="team-settings",
+ ),
+ pytest.param("team_b", None, None, True,
id="no-fallback-to-global"),
+ ],
+ )
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.cncf.kubernetes.kube_client._TimeoutAsyncK8sApiClient",
autospec=True)
+ @mock.patch("airflow.providers.cncf.kubernetes.kube_client.async_config")
+ async def test_team_uses_its_own_connection_settings(
+ self,
+ mock_async_config,
+ mock_api_client,
+ team_name,
+ kube_config_kwargs,
+ ssl_ca_cert,
+ verify_ssl,
+ monkeypatch,
+ ):
+ mock_async_config.load_kube_config = mock.AsyncMock()
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__IN_CLUSTER", "False")
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CONFIG_FILE",
"/team-a/kubeconfig")
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__CLUSTER_CONTEXT",
"team-a-context")
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__SSL_CA_CERT",
"/team-a/ca.crt")
+
monkeypatch.setenv("AIRFLOW__TEAM_A___KUBERNETES_EXECUTOR__VERIFY_SSL", "False")
+
+ global_settings = {
+ ("kubernetes_executor", "in_cluster"): "False",
+ ("kubernetes_executor", "config_file"): "/global/kubeconfig",
+ ("kubernetes_executor", "cluster_context"): "global-context",
+ ("kubernetes_executor", "ssl_ca_cert"): "/global/ca.crt",
+ ("kubernetes_executor", "verify_ssl"): "False",
+ }
+ with conf_vars(global_settings):
+ await get_async_kube_client(team_name=team_name)
+
+ if kube_config_kwargs:
+ mock_async_config.load_incluster_config.assert_not_called()
+ mock_async_config.load_kube_config.assert_awaited_once_with(
+ **kube_config_kwargs, client_configuration=mock.ANY
+ )
+ else:
+ mock_async_config.load_incluster_config.assert_called_once()
+ mock_async_config.load_kube_config.assert_not_awaited()
+ configuration = mock_api_client.call_args.args[0]
+ assert configuration.ssl_ca_cert == ssl_ca_cert
+ assert configuration.verify_ssl is verify_ssl
