This is an automated email from the ASF dual-hosted git repository.
shahar1 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 7fc063d5964 Reduce SSL setup overhead in async Kubernetes triggers
(#73399)
7fc063d5964 is described below
commit 7fc063d5964dd47a47efcfaf9b18299369b8194e
Author: Xavi Vega <[email protected]>
AuthorDate: Sun Nov 8 14:48:30 2026 +0100
Reduce SSL setup overhead in async Kubernetes triggers (#73399)
---
providers/cncf/kubernetes/docs/changelog.rst | 7 +
.../providers/cncf/kubernetes/hooks/kubernetes.py | 52 ++--
.../providers/cncf/kubernetes/triggers/job.py | 5 +
.../providers/cncf/kubernetes/triggers/pod.py | 80 +++---
.../unit/cncf/kubernetes/hooks/test_kubernetes.py | 308 ++++++++++++---------
.../unit/cncf/kubernetes/triggers/test_job.py | 8 +
.../unit/cncf/kubernetes/triggers/test_pod.py | 43 ++-
7 files changed, 321 insertions(+), 182 deletions(-)
diff --git a/providers/cncf/kubernetes/docs/changelog.rst
b/providers/cncf/kubernetes/docs/changelog.rst
index df8c1612f2c..a3c0c19545a 100644
--- a/providers/cncf/kubernetes/docs/changelog.rst
+++ b/providers/cncf/kubernetes/docs/changelog.rst
@@ -30,6 +30,13 @@ Changelog
11.0.0
......
+.. note::
+
+ ``AsyncKubernetesHook`` now reuses its API client for non-exec
authentication.
+ Direct hook users must call ``await hook.close()`` when finished to
release the
+ connection pool and SSL context. Pod and job triggers close the hook
during cleanup.
+ TLS certificate material is retained until the client is closed and
recreated.
+
.. note::
``@task.kubernetes_cmd`` now renders only the command returned by the
decorated function inside
``execute()``. Its other template fields are rendered once before
``execute()`` and are no
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
index 1216fd74429..07f82372cf0 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py
@@ -828,7 +828,13 @@ def _split_log_bytes(raw_bytes: bytes) -> list[str]:
class AsyncKubernetesHook(KubernetesHook):
- """Hook to use Kubernetes SDK asynchronously."""
+ """
+ Hook to use Kubernetes SDK asynchronously.
+
+ Non-exec authentication reuses an API client across calls. Call ``await
hook.close()``
+ when finished to release its connection pool and SSL context. Pod and job
triggers
+ do this in their cleanup methods.
+ """
def __init__(
self,
@@ -843,12 +849,13 @@ class AsyncKubernetesHook(KubernetesHook):
self.config_dict = config_dict
self._extras: dict | None = connection_extras
self._event_polling_fallback = False
- self._config_loaded = False
+ self._client_cacheable = False
# Override the parent's sync-typed client_configuration with the async
type.
self.client_configuration: async_client.Configuration | None =
client_configuration
# Cached result of exec-auth detection. None means not yet detected.
# This is to optimise and not calling _uses_exec_auth repeatedly on
every _load_config() call.
self._is_exec_auth: bool | None = None
+ self._cached_kube_client: async_client.ApiClient | None = None
def _uses_exec_auth(self, kubeconfig_data: dict, context: str | None =
None) -> bool:
"""
@@ -920,16 +927,12 @@ class AsyncKubernetesHook(KubernetesHook):
"""
Load Kubernetes configuration.
- For static auth (token, certificate), configuration is loaded once per
hook instance
- and cached. For exec-based auth (EKS, GKE), the config is reloaded on
every call so
- that short-lived tokens are always refreshed.
+ Reload credentials on every access while retaining the client for
non-exec auth.
+ Exec-based auth continues to use a new client for each access.
"""
if self.client_configuration is None:
self.client_configuration = async_client.Configuration()
- if self._config_loaded:
- return
-
in_cluster = self._coalesce_param(self.in_cluster, await
self._get_field("in_cluster"))
cluster_context = self._coalesce_param(self.cluster_context, await
self._get_field("cluster_context"))
kubeconfig_path = await self._get_field("kube_config_path")
@@ -950,7 +953,7 @@ class AsyncKubernetesHook(KubernetesHook):
self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("within a
pod"))
async_config.load_incluster_config(client_configuration=self.client_configuration)
self._is_in_cluster = True
- self._config_loaded = True
+ self._client_cacheable = True
return
self._is_in_cluster = False
@@ -967,7 +970,7 @@ class AsyncKubernetesHook(KubernetesHook):
self._is_exec_auth = self._uses_exec_auth(self.config_dict,
context=cluster_context)
if not self._is_exec_auth:
- self._config_loaded = True
+ self._client_cacheable = True
return
if kubeconfig_path is not None:
@@ -985,7 +988,7 @@ class AsyncKubernetesHook(KubernetesHook):
)
if not self._is_exec_auth:
- self._config_loaded = True
+ self._client_cacheable = True
return
if kubeconfig is not None:
@@ -1022,7 +1025,7 @@ class AsyncKubernetesHook(KubernetesHook):
)
if not self._is_exec_auth:
- self._config_loaded = True
+ self._client_cacheable = True
return
self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("default
configuration file"))
@@ -1036,7 +1039,7 @@ class AsyncKubernetesHook(KubernetesHook):
self._is_exec_auth = await
self._default_kubeconfig_uses_exec_auth(cluster_context)
if not self._is_exec_auth:
- self._config_loaded = True
+ self._client_cacheable = True
async def get_conn_extras(self) -> dict:
if self._extras is None:
@@ -1067,14 +1070,27 @@ class AsyncKubernetesHook(KubernetesHook):
@contextlib.asynccontextmanager
async def get_conn(self) -> AsyncGenerator[async_client.ApiClient, None]:
- kube_client = None
+ await self._load_config()
+ if self._client_cacheable:
+ # Reuse one client per hook: each construction runs
ssl.create_default_context()
+ # on the event loop and opens a new connection pool. Owners
release it via
+ # close(); triggers do so in cleanup().
+ if self._cached_kube_client is None:
+ self._cached_kube_client =
_TimeoutAsyncK8sApiClient(configuration=self.client_configuration)
+ yield self._cached_kube_client
+ return
+ # Exec-based auth is not cacheable: build and close a client per call.
+ kube_client =
_TimeoutAsyncK8sApiClient(configuration=self.client_configuration)
try:
- await self._load_config()
- kube_client =
_TimeoutAsyncK8sApiClient(configuration=self.client_configuration)
yield kube_client
finally:
- if kube_client is not None:
- await kube_client.close()
+ await kube_client.close()
+
+ async def close(self) -> None:
+ """Release the cached API client, if any. Safe to call multiple
times."""
+ cached, self._cached_kube_client = self._cached_kube_client, None
+ if cached is not None:
+ await cached.close()
@generic_api_retry
async def get_pod(self, name: str, namespace: str) -> V1Pod:
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/job.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/job.py
index 22b60be3a78..6339dc3c48f 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/job.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/job.py
@@ -176,3 +176,8 @@ class KubernetesJobTrigger(BaseTrigger):
cluster_context=self.cluster_context,
)
return PodManager(kube_client=sync_hook.core_v1_client)
+
+ async def cleanup(self) -> None:
+ """Release the hook's cached API client when the trigger exits."""
+ await super().cleanup()
+ await self.hook.close()
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/pod.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/pod.py
index 18a465a0ecf..b27e0ac48fd 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/pod.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/pod.py
@@ -545,46 +545,54 @@ class KubernetesPodTrigger(BaseTrigger):
On Airflow 3.3.0+ pod deletion on user kill is handled in
``on_kill()`` only; this avoids
deleting pods on triggerer restart. On older Airflow versions,
``cleanup()`` still uses
``safe_to_cancel()`` because ``on_kill()`` is not wired for user kills.
- """
- # TODO: Remove this Airflow < 3.3 cleanup branch (early return,
``safe_to_cancel``, and
- # related tests) once the minimum Airflow version supported by this
provider is >= 3.3.
- # In Airflow 3.3+, ``BaseTrigger.on_kill()`` handles user-initiated
kills; keeping the
- # legacy path for backward compatibility with older Airflow versions.
- if AIRFLOW_V_3_3_PLUS:
- return
-
- if self._fired_event:
- self.log.debug("Skipping cleanup since an event has already been
fired.")
- return
-
- if self.on_kill_action == OnKillAction.KEEP_POD:
- self.log.debug("Skipping cleanup since on_kill_action is set to
%r.", self.on_kill_action.value)
- return
+ Always releases the hook's cached API client, since the triggerer
invokes ``cleanup()``
+ after the trigger finishes for any reason (event fired, cancelled, or
killed).
+ """
try:
- safe = await self.safe_to_cancel()
- except Exception:
- self.log.warning(
- "Could not determine task state during cleanup; skipping pod
deletion to be safe.",
- exc_info=True,
- )
- return
+ # TODO: Remove this Airflow < 3.3 cleanup branch (early return,
``safe_to_cancel``, and
+ # related tests) once the minimum Airflow version supported by
this provider is >= 3.3.
+ # In Airflow 3.3+, ``BaseTrigger.on_kill()`` handles
user-initiated kills; keeping the
+ # legacy path for backward compatibility with older Airflow
versions.
+ if AIRFLOW_V_3_3_PLUS:
+ return
+
+ if self._fired_event:
+ self.log.debug("Skipping cleanup since an event has already
been fired.")
+ return
+
+ if self.on_kill_action == OnKillAction.KEEP_POD:
+ self.log.debug(
+ "Skipping cleanup since on_kill_action is set to %r.",
self.on_kill_action.value
+ )
+ return
- if not safe:
- self.log.debug(
- "Skipping cleanup since the task is still in deferred state
(likely a triggerer restart)."
- )
- return
+ try:
+ safe = await self.safe_to_cancel()
+ except Exception:
+ self.log.warning(
+ "Could not determine task state during cleanup; skipping
pod deletion to be safe.",
+ exc_info=True,
+ )
+ return
- self.log.info("Deleting pod %s in namespace %s.", self.pod_name,
self.pod_namespace)
- try:
- await self.hook.delete_pod(
- name=self.pod_name,
- namespace=self.pod_namespace,
- grace_period_seconds=self.termination_grace_period,
- )
- except Exception:
- self.log.exception("Unexpected error while deleting pod %s",
self.pod_name)
+ if not safe:
+ self.log.debug(
+ "Skipping cleanup since the task is still in deferred
state (likely a triggerer restart)."
+ )
+ return
+
+ self.log.info("Deleting pod %s in namespace %s.", self.pod_name,
self.pod_namespace)
+ try:
+ await self.hook.delete_pod(
+ name=self.pod_name,
+ namespace=self.pod_namespace,
+ grace_period_seconds=self.termination_grace_period,
+ )
+ except Exception:
+ self.log.exception("Unexpected error while deleting pod %s",
self.pod_name)
+ finally:
+ await self.hook.close()
def define_container_state(self, pod: V1Pod) -> ContainerState:
if pod.status is None or pod.status.container_statuses is None:
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py
index c8df520d0ee..df74354f83b 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/hooks/test_kubernetes.py
@@ -18,9 +18,11 @@
from __future__ import annotations
import asyncio
+import inspect
import json
import os
import socket
+import ssl
import tempfile
from asyncio import Future
from unittest import mock
@@ -1060,6 +1062,35 @@ class TestAsyncKubernetesHook:
KUBE_API = "kubernetes_asyncio.client.api.core_v1_api.CoreV1Api.{}"
KUBE_BATCH_API = "kubernetes_asyncio.client.api.batch_v1_api.BatchV1Api.{}"
KUBE_ASYNC_HOOK = HOOK_MODULE + ".AsyncKubernetesHook.{}"
+ SSL_CREATE = "kubernetes_asyncio.client.rest.ssl.create_default_context"
+
+ STATIC_AUTH_CONFIG_DICT = {
+ "apiVersion": "v1",
+ "kind": "Config",
+ "clusters": [{"name": "test", "cluster": {"server":
"https://k8s.invalid:6443"}}],
+ "users": [{"name": "test", "user": {"token": "static-token"}}],
+ "contexts": [{"name": "test", "context": {"cluster": "test", "user":
"test"}}],
+ "current-context": "test",
+ }
+
+ EXEC_AUTH_CONFIG_DICT = {
+ "apiVersion": "v1",
+ "kind": "Config",
+ "clusters": [{"name": "test", "cluster": {"server":
"https://k8s.invalid:6443"}}],
+ "users": [
+ {
+ "name": "test",
+ "user": {
+ "exec": {
+ "apiVersion": "client.authentication.k8s.io/v1beta1",
+ "command": "aws",
+ }
+ },
+ }
+ ],
+ "contexts": [{"name": "test", "context": {"cluster": "test", "user":
"test"}}],
+ "current-context": "test",
+ }
@staticmethod
def mock_await_result(return_value):
@@ -1087,6 +1118,17 @@ class TestAsyncKubernetesHook:
yield
clear_test_connections()
+ @pytest_asyncio.fixture
+ async def hook(self):
+ hook = AsyncKubernetesHook(
+ conn_id=None,
+ in_cluster=False,
+ config_file=None,
+ cluster_context=None,
+ )
+ yield hook
+ await hook.close()
+
@pytest.mark.asyncio
@mock.patch(INCLUSTER_CONFIG_LOADER)
@mock.patch(KUBE_LOADER_CONFIG)
@@ -1291,7 +1333,7 @@ class TestAsyncKubernetesHook:
}
hook = AsyncKubernetesHook(conn_id=None, in_cluster=False,
config_dict=exec_config)
await hook._load_config()
- assert hook._config_loaded is False
+ assert hook._client_cacheable is False
@pytest.mark.asyncio
@pytest.mark.parametrize(
@@ -1330,7 +1372,7 @@ class TestAsyncKubernetesHook:
side_effect=lambda field: str(kubeconfig_file) if field ==
"kube_config_path" else None
)
await hook._load_config()
- assert hook._config_loaded is expected_cached
+ assert hook._client_cacheable is expected_cached
@pytest.mark.asyncio
@pytest.mark.parametrize(
@@ -1380,7 +1422,7 @@ class TestAsyncKubernetesHook:
side_effect=lambda field: kubeconfig if field == "kube_config"
else None
)
await hook._load_config()
- assert hook._config_loaded is expected_cached
+ assert hook._client_cacheable is expected_cached
@pytest.mark.asyncio
@pytest.mark.parametrize(
@@ -1438,7 +1480,7 @@ class TestAsyncKubernetesHook:
await hook._load_config()
mock_load_file.assert_awaited_once()
- assert hook._config_loaded is expected_cached
+ assert hook._client_cacheable is expected_cached
@pytest.mark.asyncio
@mock.patch("airflow.providers.cncf.kubernetes.hooks.kubernetes.async_config.load_kube_config")
@@ -1459,7 +1501,7 @@ class TestAsyncKubernetesHook:
with
mock.patch(f"{HOOK_MODULE}.async_config.KUBE_CONFIG_DEFAULT_LOCATION",
default_location):
await hook._load_config()
- assert hook._config_loaded is True
+ assert hook._client_cacheable is True
@pytest.mark.asyncio
@mock.patch("airflow.providers.cncf.kubernetes.hooks.kubernetes.async_config.load_kube_config")
@@ -1481,12 +1523,12 @@ class TestAsyncKubernetesHook:
with
mock.patch(f"{HOOK_MODULE}.async_config.KUBE_CONFIG_DEFAULT_LOCATION",
"~/.kube/config"):
await hook._load_config()
- assert hook._config_loaded is True
+ assert hook._client_cacheable is True
@pytest.mark.asyncio
@mock.patch(KUBE_API.format("list_namespaced_event"))
async def test_async_get_pod_events_with_resource_version(
- self, mock_list_namespaced_event, kube_config_loader
+ self, mock_list_namespaced_event, kube_config_loader, hook
):
"""Test getting pod events with resource_version parameter."""
mock_event = mock.Mock()
@@ -1495,13 +1537,6 @@ class TestAsyncKubernetesHook:
mock_events.items = [mock_event]
mock_list_namespaced_event.return_value =
self.mock_await_result(mock_events)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
result = await hook.get_pod_events(name=POD_NAME, namespace=NAMESPACE,
resource_version="12345")
mock_list_namespaced_event.assert_called_once_with(
@@ -1515,7 +1550,7 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch(KUBE_API.format("list_namespaced_event"))
async def test_async_get_pod_events_without_resource_version(
- self, mock_list_namespaced_event, kube_config_loader
+ self, mock_list_namespaced_event, kube_config_loader, hook
):
"""Test getting pod events without resource_version parameter."""
mock_event = mock.Mock()
@@ -1524,13 +1559,6 @@ class TestAsyncKubernetesHook:
mock_events.items = [mock_event]
mock_list_namespaced_event.return_value =
self.mock_await_result(mock_events)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
result = await hook.get_pod_events(name=POD_NAME, namespace=NAMESPACE)
mock_list_namespaced_event.assert_called_once_with(
@@ -1544,7 +1572,7 @@ class TestAsyncKubernetesHook:
@mock.patch("kubernetes_asyncio.watch.Watch")
@mock.patch(KUBE_API.format("list_namespaced_event"))
async def test_async_watch_pod_events(
- self, mock_list_namespaced_event, mock_watch_class, mock_get_pod,
kube_config_loader
+ self, mock_list_namespaced_event, mock_watch_class, mock_get_pod,
kube_config_loader, hook
):
"""Test watching pod events using Watch API."""
mock_event1 = mock.Mock()
@@ -1564,13 +1592,6 @@ class TestAsyncKubernetesHook:
mock_pod.status.phase = "Running"
mock_get_pod.return_value = mock_pod
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async_event_generator = hook.watch_pod_events(
name=POD_NAME, namespace=NAMESPACE, resource_version="12345",
timeout_seconds=30
@@ -1593,7 +1614,7 @@ class TestAsyncKubernetesHook:
@mock.patch("kubernetes_asyncio.watch.Watch")
@mock.patch(KUBE_API.format("list_namespaced_event"))
async def test_async_watch_pod_events_permission_error_fallback(
- self, mock_list_namespaced_event, mock_watch_class, mock_get_pod,
kube_config_loader
+ self, mock_list_namespaced_event, mock_watch_class, mock_get_pod,
kube_config_loader, hook
):
"""Test fallback to polling when watch permission is denied."""
@@ -1617,13 +1638,6 @@ class TestAsyncKubernetesHook:
mock_pod.status.phase = "Running"
mock_get_pod.return_value = mock_pod
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async for event in hook.watch_pod_events(
name=POD_NAME, namespace=NAMESPACE, resource_version="12345",
timeout_seconds=30
@@ -1639,7 +1653,7 @@ class TestAsyncKubernetesHook:
@mock.patch(KUBE_API.format("list_namespaced_event"))
@mock.patch("asyncio.sleep", new_callable=mock.AsyncMock)
async def test_async_watch_pod_events_polling_fallback(
- self, mock_sleep, mock_list_namespaced_event, kube_config_loader
+ self, mock_sleep, mock_list_namespaced_event, kube_config_loader, hook
):
"""Test polling fallback method."""
mock_event1 = mock.Mock()
@@ -1650,13 +1664,6 @@ class TestAsyncKubernetesHook:
mock_events.items = [mock_event1, mock_event2]
mock_list_namespaced_event.return_value =
self.mock_await_result(mock_events)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async for event in hook.watch_pod_events_polling_fallback(
name=POD_NAME, namespace=NAMESPACE, resource_version="12345",
interval=10
@@ -1678,16 +1685,9 @@ class TestAsyncKubernetesHook:
@mock.patch("kubernetes_asyncio.watch.Watch")
@mock.patch(KUBE_API.format("list_namespaced_event"))
async def test_async_watch_pod_events_uses_fallback_if_already_set(
- self, mock_list_namespaced_event, mock_watch_class, kube_config_loader
+ self, mock_list_namespaced_event, mock_watch_class,
kube_config_loader, hook
):
"""Test that watch uses polling fallback if flag is already set."""
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
hook._event_polling_fallback = True
mock_event = mock.Mock()
@@ -1715,6 +1715,7 @@ class TestAsyncKubernetesHook:
mock_get_pod,
mock_watch_class,
kube_config_loader,
+ hook,
):
"""
The watch should reconnect when the watch stream ends (e.g. timeout)
@@ -1748,13 +1749,6 @@ class TestAsyncKubernetesHook:
mock_watch_class.side_effect = [watch_instance1, watch_instance2]
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async for event in hook.watch_pod_events(
@@ -1775,6 +1769,7 @@ class TestAsyncKubernetesHook:
mock_get_pod,
mock_watch_class,
kube_config_loader,
+ hook,
):
"""
When the Kubernetes API reports resourceVersion too old (410),
@@ -1807,13 +1802,6 @@ class TestAsyncKubernetesHook:
mock_watch_class.side_effect = [watch_instance1, watch_instance2]
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async for event in hook.watch_pod_events(
name=POD_NAME, namespace=NAMESPACE, resource_version="1",
timeout_seconds=1
@@ -1832,6 +1820,7 @@ class TestAsyncKubernetesHook:
mock_get_pod,
mock_watch_class,
kube_config_loader,
+ hook,
):
"""
Verify that watch_pod_events stops cleanly when the pod no longer
exists (404).
@@ -1843,13 +1832,6 @@ class TestAsyncKubernetesHook:
mock_watch = mock.Mock()
mock_watch_class.return_value = mock_watch
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async for event in hook.watch_pod_events(
name=POD_NAME,
@@ -1870,6 +1852,7 @@ class TestAsyncKubernetesHook:
mock_watch_class,
pod_status,
kube_config_loader,
+ hook,
):
"""
Verify that watch_pod_events stops immediately when the pod
@@ -1883,13 +1866,6 @@ class TestAsyncKubernetesHook:
mock_pod.status.phase = pod_status
mock_get_pod.return_value = mock_pod
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
events = []
async for event in hook.watch_pod_events(
name=POD_NAME,
@@ -1902,15 +1878,9 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch(KUBE_API.format("read_namespaced_pod"))
- async def test_get_pod(self, lib_method, kube_config_loader):
+ async def test_get_pod(self, lib_method, kube_config_loader, hook):
lib_method.return_value = self.mock_await_result(None)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
await hook.get_pod(
name=POD_NAME,
namespace=NAMESPACE,
@@ -1924,15 +1894,9 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch(KUBE_API.format("delete_namespaced_pod"))
- async def test_delete_pod(self, lib_method, kube_config_loader):
+ async def test_delete_pod(self, lib_method, kube_config_loader, hook):
lib_method.return_value = self.mock_await_result(None)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
await hook.delete_pod(
name=POD_NAME,
namespace=NAMESPACE,
@@ -1942,18 +1906,11 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch(KUBE_API.format("read_namespaced_pod_log"))
- async def test_read_logs(self, lib_method, kube_config_loader):
+ async def test_read_logs(self, lib_method, kube_config_loader, hook):
mock_raw_resp = mock.AsyncMock()
mock_raw_resp.read = mock.AsyncMock(return_value=b"2023-01-11 Some
string logs...")
lib_method.return_value = self.mock_await_result(mock_raw_resp)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
logs = await hook.read_logs(
name=POD_NAME,
namespace=NAMESPACE,
@@ -1976,7 +1933,7 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch(KUBE_API.format("read_namespaced_pod_log"))
- async def test_read_logs_handles_non_utf8_bytes(self, lib_method,
kube_config_loader):
+ async def test_read_logs_handles_non_utf8_bytes(self, lib_method,
kube_config_loader, hook):
"""Non-UTF-8 bytes in pod logs are replaced instead of raising
UnicodeDecodeError."""
raw_bytes = b"2023-01-11 valid line\n2023-01-11 broken \x80\x81 bytes"
@@ -1984,13 +1941,6 @@ class TestAsyncKubernetesHook:
mock_raw_resp.read = mock.AsyncMock(return_value=raw_bytes)
lib_method.return_value = self.mock_await_result(mock_raw_resp)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
logs = await hook.read_logs(
name=POD_NAME,
namespace=NAMESPACE,
@@ -2007,7 +1957,9 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch("asyncio.to_thread", new_callable=mock.AsyncMock)
@mock.patch(KUBE_API.format("read_namespaced_pod_log"))
- async def test_read_logs_decodes_off_the_event_loop(self, lib_method,
mock_to_thread, kube_config_loader):
+ async def test_read_logs_decodes_off_the_event_loop(
+ self, lib_method, mock_to_thread, kube_config_loader, hook
+ ):
"""The CPU-bound decode/splitlines is offloaded to a worker thread,
not run on the loop."""
raw_bytes = b"2023-01-11 Some string logs..."
mock_raw_resp = mock.AsyncMock()
@@ -2015,13 +1967,6 @@ class TestAsyncKubernetesHook:
lib_method.return_value = self.mock_await_result(mock_raw_resp)
mock_to_thread.return_value = ["decoded line"]
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
-
logs = await hook.read_logs(name=POD_NAME, namespace=NAMESPACE,
container_name=CONTAINER_NAME)
assert logs == ["decoded line"]
@@ -2085,15 +2030,9 @@ class TestAsyncKubernetesHook:
@pytest.mark.asyncio
@mock.patch(KUBE_BATCH_API.format("read_namespaced_job_status"))
- async def test_get_job_status(self, lib_method, kube_config_loader):
+ async def test_get_job_status(self, lib_method, kube_config_loader, hook):
lib_method.return_value = self.mock_await_result(None)
- hook = AsyncKubernetesHook(
- conn_id=None,
- in_cluster=False,
- config_file=None,
- cluster_context=None,
- )
await hook.get_job_status(
name=JOB_NAME,
namespace=NAMESPACE,
@@ -2307,3 +2246,118 @@ class TestAsyncKubernetesHook:
assert hook_a.client_configuration is not hook_b.client_configuration
assert hook_a.client_configuration.host ==
"https://cluster-a.example.com";
assert hook_b.client_configuration.host ==
"https://cluster-b.example.com";
+
+ @pytest.mark.asyncio
+ @mock.patch(SSL_CREATE, autospec=True,
side_effect=ssl.create_default_context)
+ async def test_static_auth_reuses_client_and_ssl_context(self,
mock_create_ssl, hook):
+ hook.config_dict = self.STATIC_AUTH_CONFIG_DICT
+ async with hook.get_conn() as first_client:
+ pass
+
+ for _ in range(2):
+ async with hook.get_conn() as client:
+ assert client is first_client
+
+ assert not first_client.rest_client.pool_manager.closed
+ mock_create_ssl.assert_called_once()
+
+ @staticmethod
+ async def _bearer_auth_headers(kube_client) -> dict[str, str]:
+ headers: dict[str, str] = {}
+ # update_params_for_auth is sync in kubernetes_asyncio<33 and async
from 33 on
+ result = kube_client.update_params_for_auth(headers, [],
["BearerToken"])
+ if inspect.isawaitable(result):
+ await result
+ return headers
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("source", ["file", "dict"])
+ @mock.patch(SSL_CREATE, autospec=True,
side_effect=ssl.create_default_context)
+ async def test_cached_client_reloads_rotated_kubeconfig_token(
+ self, mock_create_ssl, hook, tmp_path, source
+ ):
+ kubeconfig = tmp_path / "config"
+ config_data = json.loads(json.dumps(self.STATIC_AUTH_CONFIG_DICT))
+ if source == "file":
+ kubeconfig.write_text(yaml.safe_dump(config_data))
+ hook._get_field = mock.AsyncMock(
+ side_effect=lambda field: str(kubeconfig) if field ==
"kube_config_path" else None
+ )
+ else:
+ hook.config_dict = config_data
+
+ async with hook.get_conn() as first_client:
+ first_headers = await self._bearer_auth_headers(first_client)
+
+ config_data["users"][0]["user"]["token"] = "rotated-token"
+ if source == "file":
+ kubeconfig.write_text(yaml.safe_dump(config_data))
+
+ async with hook.get_conn() as second_client:
+ second_headers = await self._bearer_auth_headers(second_client)
+
+ assert first_client is second_client
+ assert first_headers["authorization"] == "Bearer static-token"
+ assert second_headers["authorization"] == "Bearer rotated-token"
+ mock_create_ssl.assert_called_once()
+
+ @pytest.mark.asyncio
+ @mock.patch(SSL_CREATE, autospec=True,
side_effect=ssl.create_default_context)
+
@mock.patch("kubernetes_asyncio.config.kube_config.google_auth_credentials",
autospec=True)
+ async def test_cached_client_refreshes_expired_auth_provider_token(
+ self, mock_google_credentials, mock_create_ssl, hook, time_machine
+ ):
+ time_machine.move_to("2026-01-01T00:00:00Z")
+ config_data = json.loads(json.dumps(self.STATIC_AUTH_CONFIG_DICT))
+ provider_config = {"access-token": "initial-token", "expiry":
"2026-01-01T02:00:00Z"}
+ config_data["users"][0]["user"] = {"auth-provider": {"name": "gcp",
"config": provider_config}}
+ hook.config_dict = config_data
+ mock_google_credentials.return_value = mock.Mock(
+ token="refreshed-token", expiry="2026-01-01T04:00:00Z"
+ )
+
+ async with hook.get_conn() as first_client:
+ first_headers = await self._bearer_auth_headers(first_client)
+
+ time_machine.move_to("2026-01-01T03:00:00Z")
+ async with hook.get_conn() as second_client:
+ second_headers = await self._bearer_auth_headers(second_client)
+
+ assert first_client is second_client
+ assert first_headers["authorization"] == "Bearer initial-token"
+ assert second_headers["authorization"] == "Bearer refreshed-token"
+ mock_google_credentials.assert_awaited_once()
+ mock_create_ssl.assert_called_once()
+
+ @pytest.mark.asyncio
+ @mock.patch(SSL_CREATE, autospec=True,
side_effect=ssl.create_default_context)
+ @mock.patch(HOOK_MODULE + ".async_config.load_kube_config_from_dict",
autospec=True)
+ async def test_exec_auth_keeps_per_call_clients(self, mock_load_config,
mock_create_ssl, hook):
+ hook.config_dict = self.EXEC_AUTH_CONFIG_DICT
+ clients = []
+ for _ in range(3):
+ async with hook.get_conn() as client:
+ clients.append(client)
+ assert client.rest_client.pool_manager.closed
+
+ assert len({id(client) for client in clients}) == 3
+ assert mock_create_ssl.call_count == 3
+ assert mock_load_config.await_count == 3
+ assert hook._cached_kube_client is None
+ assert not hook._client_cacheable
+
+ @pytest.mark.asyncio
+ async def test_close_releases_cached_client_and_is_idempotent(self, hook):
+ hook.config_dict = self.STATIC_AUTH_CONFIG_DICT
+ async with hook.get_conn() as cached:
+ pass
+
+ await hook.close()
+ assert hook._cached_kube_client is None
+ assert cached.rest_client.pool_manager.closed
+
+ await hook.close()
+
+ async with hook.get_conn() as reopened:
+ assert reopened is not cached
+ assert not reopened.rest_client.pool_manager.closed
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_job.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_job.py
index 37aeb8c1d9f..9b37fc1c5e4 100644
--- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_job.py
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_job.py
@@ -27,6 +27,7 @@ from airflow.triggers.base import TriggerEvent
TRIGGER_PATH = "airflow.providers.cncf.kubernetes.triggers.job.{}"
TRIGGER_CLASS = TRIGGER_PATH.format("KubernetesJobTrigger")
+HOOK_PATH =
"airflow.providers.cncf.kubernetes.hooks.kubernetes.AsyncKubernetesHook"
JOB_NAME = "test-job-name"
POD_NAME = "test-pod-name"
CONTAINER_NAME = "test-container-name"
@@ -210,3 +211,10 @@ class TestKubernetesJobTrigger:
cluster_context=CLUSTER_CONTEXT,
)
assert hook_actual == hook_expected
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{HOOK_PATH}.close", autospec=True)
+ async def test_cleanup_closes_hook(self, mock_close, trigger):
+ await trigger.cleanup()
+
+ mock_close.assert_awaited_once_with(trigger.hook)
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_pod.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_pod.py
index 2dcbc426da9..7409587a32e 100644
--- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_pod.py
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_pod.py
@@ -32,6 +32,7 @@ from kubernetes.client import models as k8s
from pendulum import DateTime
from sqlalchemy.orm.session import Session
+from airflow.providers.cncf.kubernetes.hooks.kubernetes import
AsyncKubernetesHook
from airflow.providers.cncf.kubernetes.triggers.pod import ContainerState,
KubernetesPodTrigger
from airflow.providers.cncf.kubernetes.utils.pod_manager import PodPhase
from airflow.providers.common.compat.sdk import AirflowException
@@ -890,6 +891,7 @@ class TestKubernetesPodTrigger:
on_kill_action="delete_pod",
on_finish_action="delete_pod",
)
+ mock_hook.close = mock.AsyncMock(spec=AsyncKubernetesHook.close)
with mock.patch(f"{TRIGGER_PATH}.safe_to_cancel",
new_callable=mock.AsyncMock, return_value=False):
await trigger.cleanup()
mock_hook.delete_pod.assert_not_called()
@@ -911,6 +913,7 @@ class TestKubernetesPodTrigger:
on_finish_action="delete_pod",
)
mock_hook.delete_pod = mock.AsyncMock()
+ mock_hook.close = mock.AsyncMock(spec=AsyncKubernetesHook.close)
with mock.patch(f"{TRIGGER_PATH}.safe_to_cancel",
new_callable=mock.AsyncMock, return_value=True):
await trigger.cleanup()
mock_hook.delete_pod.assert_called_once_with(
@@ -926,7 +929,8 @@ class TestKubernetesPodTrigger:
@pytest.mark.asyncio
@mock.patch(f"{TRIGGER_PATH}.hook")
async def test_cleanup_noop_on_airflow_3_3_plus(self, mock_hook):
- """On Airflow 3.3.0+, cleanup does not delete pods (on_kill handles
user kill)."""
+ """On Airflow 3.3.0+, cleanup releases the client and on_kill handles
pod deletion."""
+ mock_hook.close = mock.AsyncMock(spec=AsyncKubernetesHook.close)
trigger = KubernetesPodTrigger(
pod_name=POD_NAME,
pod_namespace=NAMESPACE,
@@ -938,6 +942,7 @@ class TestKubernetesPodTrigger:
)
await trigger.cleanup()
mock_hook.delete_pod.assert_not_called()
+ mock_hook.close.assert_awaited_once()
@pytest.mark.asyncio
@mock.patch(f"{TRIGGER_PATH}.get_task_state", new_callable=mock.AsyncMock)
@@ -1096,6 +1101,7 @@ class TestKubernetesPodTrigger:
on_kill_action="delete_pod",
on_finish_action="delete_pod",
)
+ mock_hook.close = mock.AsyncMock(spec=AsyncKubernetesHook.close)
with mock.patch(
f"{TRIGGER_PATH}.safe_to_cancel",
new_callable=mock.AsyncMock,
@@ -1103,3 +1109,38 @@ class TestKubernetesPodTrigger:
):
await trigger.cleanup()
mock_hook.delete_pod.assert_not_called()
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("fired_event", "on_kill_action", "safe_result", "delete_error"),
+ [
+ pytest.param(True, "delete_pod", True, None, id="event-fired"),
+ pytest.param(False, "keep_pod", True, None, id="keep-pod"),
+ pytest.param(False, "delete_pod", False, None, id="restart"),
+ pytest.param(False, "delete_pod", RuntimeError("API down"), None,
id="state-error"),
+ pytest.param(False, "delete_pod", True, RuntimeError("API down"),
id="delete-error"),
+ pytest.param(False, "delete_pod", True, None, id="pod-deleted"),
+ ],
+ )
+
@mock.patch("airflow.providers.cncf.kubernetes.triggers.pod.AIRFLOW_V_3_3_PLUS",
False)
+ @mock.patch(f"{TRIGGER_PATH}.safe_to_cancel", autospec=True)
+
@mock.patch("airflow.providers.cncf.kubernetes.triggers.pod.AsyncKubernetesHook",
autospec=True)
+ async def test_legacy_cleanup_always_closes_hook(
+ self, mock_hook_class, mock_safe_to_cancel, fired_event,
on_kill_action, safe_result, delete_error
+ ):
+ trigger = KubernetesPodTrigger(
+ pod_name=POD_NAME,
+ pod_namespace=NAMESPACE,
+ base_container_name=BASE_CONTAINER_NAME,
+ trigger_start_time=TRIGGER_START_TIME,
+ schedule_timeout=STARTUP_TIMEOUT_SECS,
+ on_kill_action=on_kill_action,
+ )
+ trigger._fired_event = fired_event
+ mock_safe_to_cancel.side_effect = [safe_result]
+ mock_hook = mock_hook_class.return_value
+ mock_hook.delete_pod.side_effect = delete_error
+
+ await trigger.cleanup()
+
+ mock_hook.close.assert_awaited_once()
