From 06cd6d8378ce4f6c0fd6c2e2b9855838eec8380c Mon Sep 17 00:00:00 2001 From: AutomationDev85 Date: Thu, 8 Jan 2026 09:35:19 +0100 Subject: [PATCH 1/5] Kubernetes API calls use timeouts --- .../cncf/kubernetes/hooks/kubernetes.py | 30 +++++++++++++------ .../kubernetes/kubernetes_helper_functions.py | 8 +++++ .../cncf/kubernetes/utils/pod_manager.py | 14 ++++++--- .../cncf/kubernetes/hooks/test_kubernetes.py | 9 +++++- .../test_kubernetes_helper_functions.py | 21 +++++++++++++ .../cncf/kubernetes/utils/test_pod_manager.py | 13 ++++++-- 6 files changed, 79 insertions(+), 16 deletions(-) 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 47ef247fac6c2..6b3b08afc2973 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 @@ -36,7 +36,7 @@ from airflow.models import Connection from airflow.providers.cncf.kubernetes.exceptions import KubernetesApiError, KubernetesApiPermissionError from airflow.providers.cncf.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive -from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import generic_api_retry +from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import generic_api_retry, with_timeout from airflow.providers.cncf.kubernetes.utils.container import ( container_is_completed, container_is_running, @@ -409,6 +409,7 @@ def get_custom_object( namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, plural=plural, name=name, + **with_timeout(), ) return response @@ -432,7 +433,7 @@ def delete_custom_object( namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, plural=plural, name=name, - **kwargs, + **with_timeout(kwargs), ) def get_namespace(self) -> str | None: @@ -473,6 +474,7 @@ def get_pod_log_stream( name=pod_name, container=container, namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, + **with_timeout(), ), ) @@ -494,6 +496,7 @@ def get_pod_logs( container=container, _preload_content=False, namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, + **with_timeout(), ) def get_pod(self, name: str, namespace: str) -> V1Pod: @@ -501,6 +504,7 @@ def get_pod(self, name: str, namespace: str) -> V1Pod: return self.core_v1_client.read_namespaced_pod( name=name, namespace=namespace, + **with_timeout(), ) def get_namespaced_pod_list( @@ -522,7 +526,7 @@ def get_namespaced_pod_list( watch=watch, label_selector=label_selector, _preload_content=False, - **kwargs, + **with_timeout(kwargs), ) def get_deployment_status( @@ -558,7 +562,7 @@ def create_job( self.log.debug("Job Creation Request: \n%s", json_job) try: resp = self.batch_v1_client.create_namespaced_job( - body=sanitized_job, namespace=job.metadata.namespace, **kwargs + body=sanitized_job, namespace=job.metadata.namespace, **with_timeout(kwargs) ) self.log.debug("Job Creation Response: %s", resp) except Exception as e: @@ -577,7 +581,9 @@ def get_job(self, job_name: str, namespace: str) -> V1Job: :param namespace: Namespace of the Job. :return: Job object """ - return self.batch_v1_client.read_namespaced_job(name=job_name, namespace=namespace, pretty=True) + return self.batch_v1_client.read_namespaced_job( + name=job_name, namespace=namespace, pretty=True, **with_timeout() + ) @generic_api_retry def get_job_status(self, job_name: str, namespace: str) -> V1Job: @@ -589,7 +595,7 @@ def get_job_status(self, job_name: str, namespace: str) -> V1Job: :return: Job object """ return self.batch_v1_client.read_namespaced_job_status( - name=job_name, namespace=namespace, pretty=True + name=job_name, namespace=namespace, pretty=True, **with_timeout() ) def wait_until_job_complete(self, job_name: str, namespace: str, job_poll_interval: float = 10) -> V1Job: @@ -616,7 +622,7 @@ def list_jobs_all_namespaces(self) -> V1JobList: :return: V1JobList object """ - return self.batch_v1_client.list_job_for_all_namespaces(pretty=True) + return self.batch_v1_client.list_job_for_all_namespaces(pretty=True, **with_timeout()) @generic_api_retry def list_jobs_from_namespace(self, namespace: str) -> V1JobList: @@ -626,7 +632,7 @@ def list_jobs_from_namespace(self, namespace: str) -> V1JobList: :param namespace: Namespace of the Job. :return: V1JobList object """ - return self.batch_v1_client.list_namespaced_job(namespace=namespace, pretty=True) + return self.batch_v1_client.list_namespaced_job(namespace=namespace, pretty=True, **with_timeout()) def is_job_complete(self, job: V1Job) -> bool: """ @@ -687,6 +693,7 @@ def patch_namespaced_job(self, job_name: str, namespace: str, body: object) -> V name=job_name, namespace=namespace, body=body, + **with_timeout(), ) def apply_from_yaml_file( @@ -897,6 +904,7 @@ async def get_pod(self, name: str, namespace: str) -> V1Pod: pod: V1Pod = await v1_api.read_namespaced_pod( name=name, namespace=namespace, + **with_timeout(), ) return pod except HTTPError as e: @@ -916,7 +924,7 @@ async def delete_pod(self, name: str, namespace: str): try: v1_api = async_client.CoreV1Api(connection) await v1_api.delete_namespaced_pod( - name=name, namespace=namespace, body=client.V1DeleteOptions() + name=name, namespace=namespace, body=client.V1DeleteOptions(), **with_timeout() ) except async_client.ApiException as e: # If the pod is already deleted @@ -950,6 +958,7 @@ async def read_logs( follow=False, timestamps=True, since_seconds=since_seconds, + **with_timeout(), ) logs = logs.splitlines() return logs @@ -975,6 +984,7 @@ async def get_pod_events( namespace=namespace, resource_version=resource_version, resource_version_match="NotOlderThan" if resource_version else None, + **with_timeout(), ) return events except HTTPError as e: @@ -1015,6 +1025,7 @@ async def watch_pod_events( field_selector=f"involvedObject.name={name}", resource_version=resource_version, timeout_seconds=timeout_seconds, + **with_timeout(), ): event: CoreV1Event = event_watched.get("object") yield event @@ -1069,6 +1080,7 @@ async def get_job_status(self, name: str, namespace: str) -> V1Job: job: V1Job = await v1_api.read_namespaced_job_status( name=name, namespace=namespace, + **with_timeout(), ) return job diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py index 3646ff8b75c31..055fd35eb9686 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py @@ -52,6 +52,7 @@ class KubernetesApiException(AirflowException): """When communication with kubernetes API fails.""" +API_TIMEOUT = 60 # allow 1 min of timeout for kubernetes api calls API_RETRIES = conf.getint("workers", "api_retries", fallback=5) API_RETRY_WAIT_MIN = conf.getfloat("workers", "api_retry_wait_min", fallback=1) API_RETRY_WAIT_MAX = conf.getfloat("workers", "api_retry_wait_max", fallback=15) @@ -100,6 +101,13 @@ def generic_api_retry(func): )(func) +def with_timeout(kwargs: dict | None = None) -> dict: + """Attach a total request timeout unless explicitly provided.""" + out = dict(kwargs or {}) + out.setdefault("_request_timeout", API_TIMEOUT) + return out + + def rand_str(num): """ Generate random lowercase alphanumeric string of length num. diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py index c9b48158d960a..4f2cf4f94e9b5 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py @@ -42,6 +42,7 @@ KubernetesApiException, PodLaunchFailedException, generic_api_retry, + with_timeout, ) from airflow.providers.cncf.kubernetes.utils.container import ( container_is_completed, @@ -331,7 +332,7 @@ def run_pod_async(self, pod: V1Pod, **kwargs) -> V1Pod: self.log.debug("Pod Creation Request: \n%s", json_pod) try: resp = self._client.create_namespaced_pod( - body=sanitized_pod, namespace=pod.metadata.namespace, **kwargs + body=sanitized_pod, namespace=pod.metadata.namespace, **with_timeout(kwargs) ) self.log.debug("Pod Creation Response: %s", resp) except Exception as e: @@ -346,7 +347,7 @@ def delete_pod(self, pod: V1Pod) -> None: """Delete POD.""" try: self._client.delete_namespaced_pod( - pod.metadata.name, pod.metadata.namespace, body=client.V1DeleteOptions() + pod.metadata.name, pod.metadata.namespace, body=client.V1DeleteOptions(), **with_timeout() ) except ApiException as e: # If the pod is already deleted @@ -750,7 +751,7 @@ def read_pod_logs( follow=follow, timestamps=timestamps, _preload_content=False, - **additional_kwargs, + **with_timeout(additional_kwargs), ) except HTTPError: self.log.exception("There was an error reading the kubernetes API.") @@ -800,6 +801,7 @@ def read_pod_events(self, pod: V1Pod, resource_version: str | None = None) -> Co field_selector=f"involvedObject.name={pod.metadata.name}", resource_version=resource_version, resource_version_match="NotOlderThan" if resource_version else None, + **with_timeout(), ) except HTTPError as e: raise KubernetesApiException(f"There was an error reading the kubernetes API: {e}") @@ -808,7 +810,9 @@ def read_pod_events(self, pod: V1Pod, resource_version: str | None = None) -> Co def read_pod(self, pod: V1Pod) -> V1Pod: """Read POD information.""" try: - return self._client.read_namespaced_pod(pod.metadata.name, pod.metadata.namespace) + return self._client.read_namespaced_pod( + pod.metadata.name, pod.metadata.namespace, **with_timeout() + ) except HTTPError as e: raise KubernetesApiException(f"There was an error reading the kubernetes API: {e}") @@ -873,6 +877,7 @@ def extract_xcom_json(self, pod: V1Pod) -> str: stderr=True, tty=False, _preload_content=False, + **with_timeout(), ) ) as client: self.log.info("Running command... %s", command) @@ -905,6 +910,7 @@ def extract_xcom_kill(self, pod: V1Pod): stderr=True, tty=False, _preload_content=False, + **with_timeout(), ) ) as resp: self._exec_pod_command(resp, "kill -2 $(pgrep -u $(id -u) -f 'sh')") 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 ea0a64938c69c..6be1755bcf67c 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 @@ -36,6 +36,7 @@ from airflow.models import Connection from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook +from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import with_timeout from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException from tests_common.test_utils.db import clear_test_connections @@ -475,6 +476,7 @@ def test_delete_custom_object( name="name", namespace="namespace", _preload_content="_preload_content", + **with_timeout(), ) @patch("kubernetes.config.kube_config.KubeConfigLoader") @@ -487,7 +489,7 @@ def test_get_job_status(self, mock_client, mock_kube_config_merger, mock_kube_co job_actual = hook.get_job_status(job_name=JOB_NAME, namespace=NAMESPACE) mock_client.read_namespaced_job_status.assert_called_once_with( - name=JOB_NAME, namespace=NAMESPACE, pretty=True + name=JOB_NAME, namespace=NAMESPACE, pretty=True, **with_timeout() ) assert job_actual == job_expected @@ -1038,6 +1040,7 @@ async def test_async_get_pod_events_with_resource_version( namespace=NAMESPACE, resource_version="12345", resource_version_match="NotOlderThan", + **with_timeout(), ) assert result == mock_events @@ -1067,6 +1070,7 @@ async def test_async_get_pod_events_without_resource_version( namespace=NAMESPACE, resource_version=None, resource_version_match=None, + **with_timeout(), ) assert result == mock_events @@ -1186,6 +1190,7 @@ async def test_async_watch_pod_events_polling_fallback( namespace=NAMESPACE, resource_version="12345", resource_version_match="NotOlderThan", + **with_timeout(), ) mock_sleep.assert_called_once_with(10) @@ -1242,6 +1247,7 @@ async def test_get_pod(self, lib_method, kube_config_loader): lib_method.assert_called_with( name=POD_NAME, namespace=NAMESPACE, + **with_timeout(), ) @pytest.mark.asyncio @@ -1288,6 +1294,7 @@ async def test_read_logs(self, lib_method, kube_config_loader): follow=False, timestamps=True, since_seconds=10, + **with_timeout(), ) assert len(logs) == 1 assert "2023-01-11 Some string logs..." in logs diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py index 3622a5c34f48d..ce6711bfb730b 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py @@ -26,10 +26,12 @@ from urllib3.exceptions import HTTPError from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( + API_TIMEOUT, KubernetesApiException, WaitRetryAfterOrExponential, _should_retry_api, create_unique_id, + with_timeout, ) pod_name_regex = r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$" @@ -42,6 +44,25 @@ def __init__(self, exception=None): self.outcome.exception = mock.Mock(return_value=exception) +def test_with_timeout(): + # default value + timeout_kwargs = with_timeout() + assert "_request_timeout" in timeout_kwargs + assert timeout_kwargs["_request_timeout"] == API_TIMEOUT + + # check override + timeout_kwargs = with_timeout({"_request_timeout": 5}) + assert "_request_timeout" in timeout_kwargs + assert timeout_kwargs["_request_timeout"] == 5 + + # check additional params + timeout_kwargs = with_timeout({"some_param": 123}) + assert "some_param" in timeout_kwargs + assert timeout_kwargs["some_param"] == 123 + assert "_request_timeout" in timeout_kwargs + assert timeout_kwargs["_request_timeout"] == API_TIMEOUT + + def test_should_retry_api(): exc = HTTPError() assert _should_retry_api(exc) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py index 062266540004b..e60831745b093 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py @@ -30,6 +30,7 @@ from urllib3.exceptions import HTTPError as BaseHTTPError from airflow.providers.cncf.kubernetes.exceptions import KubernetesApiError +from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import with_timeout from airflow.providers.cncf.kubernetes.utils.pod_manager import ( AsyncPodManager, PodLogsConsumer, @@ -178,6 +179,7 @@ def test_read_pod_logs_retries_successfully(self): timestamps=False, name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, + **with_timeout(), ), mock.call( _preload_content=False, @@ -186,6 +188,7 @@ def test_read_pod_logs_retries_successfully(self): timestamps=False, name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, + **with_timeout(), ), ] ) @@ -234,6 +237,7 @@ def test_read_pod_logs_successfully_with_tail_lines(self): name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, tail_lines=100, + **with_timeout(), ), ] ) @@ -254,6 +258,7 @@ def test_read_pod_logs_successfully_with_since_seconds(self): name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, since_seconds=2, + **with_timeout(), ), ] ) @@ -313,6 +318,7 @@ def test_read_pod_events_with_resource_version(self): field_selector="involvedObject.name=test-pod", resource_version="12345", resource_version_match="NotOlderThan", + **with_timeout(), ) def test_read_pod_events_without_resource_version(self): @@ -331,6 +337,7 @@ def test_read_pod_events_without_resource_version(self): field_selector="involvedObject.name=test-pod", resource_version=None, resource_version_match=None, + **with_timeout(), ) @pytest.mark.asyncio @@ -436,12 +443,14 @@ def test_read_pod_events_retries_successfully(self): field_selector=f"involvedObject.name={mock.sentinel.metadata.name}", resource_version=None, resource_version_match=None, + **with_timeout(), ), mock.call( namespace=mock.sentinel.metadata.namespace, field_selector=f"involvedObject.name={mock.sentinel.metadata.name}", resource_version=None, resource_version_match=None, + **with_timeout(), ), ] ) @@ -474,8 +483,8 @@ def test_read_pod_retries_successfully(self): assert mock.sentinel.pod_info == pod_info self.mock_kube_client.read_namespaced_pod.assert_has_calls( [ - mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace), - mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace), + mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace, **with_timeout()), + mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace, **with_timeout()), ] ) From 3bd53dec2400b93f985ab37b3bce584610436f15 Mon Sep 17 00:00:00 2001 From: AutomationDev85 Date: Fri, 9 Jan 2026 11:02:09 +0100 Subject: [PATCH 2/5] Add wrapper Api clients to add default timeout --- .../cncf/kubernetes/hooks/kubernetes.py | 58 +++++++-------- .../kubernetes/kubernetes_helper_functions.py | 25 +++++-- .../cncf/kubernetes/utils/pod_manager.py | 14 ++-- .../cncf/kubernetes/hooks/test_kubernetes.py | 9 +-- .../test_kubernetes_helper_functions.py | 71 +++++++++++++------ .../cncf/kubernetes/utils/test_pod_manager.py | 13 +--- 6 files changed, 101 insertions(+), 89 deletions(-) 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 6b3b08afc2973..be842c9eca411 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 @@ -36,7 +36,11 @@ from airflow.models import Connection from airflow.providers.cncf.kubernetes.exceptions import KubernetesApiError, KubernetesApiPermissionError from airflow.providers.cncf.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive -from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import generic_api_retry, with_timeout +from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( + TimeoutAsyncK8sApiClient, + TimeoutK8sApiClient, + generic_api_retry, +) from airflow.providers.cncf.kubernetes.utils.container import ( container_is_completed, container_is_running, @@ -239,7 +243,7 @@ def _get_field(self, field_name): prefixed_name = f"extra__kubernetes__{field_name}" return self.conn_extras.get(prefixed_name) or None - def get_conn(self) -> client.ApiClient: + def get_conn(self) -> TimeoutK8sApiClient: """Return kubernetes api session for use with requests.""" in_cluster = self._coalesce_param(self.in_cluster, self._get_field("in_cluster")) cluster_context = self._coalesce_param(self.cluster_context, self._get_field("cluster_context")) @@ -272,7 +276,7 @@ def get_conn(self) -> client.ApiClient: self.log.debug("loading kube_config from: in_cluster configuration") self._is_in_cluster = True config.load_incluster_config() - return client.ApiClient() + return TimeoutK8sApiClient() if kubeconfig_path is not None: self.log.debug("loading kube_config from: %s", kubeconfig_path) @@ -282,7 +286,7 @@ def get_conn(self) -> client.ApiClient: client_configuration=self.client_configuration, context=cluster_context, ) - return client.ApiClient() + return TimeoutK8sApiClient() if kubeconfig is not None: with tempfile.NamedTemporaryFile() as temp_config: @@ -297,7 +301,7 @@ def get_conn(self) -> client.ApiClient: client_configuration=self.client_configuration, context=cluster_context, ) - return client.ApiClient() + return TimeoutK8sApiClient() if self.config_dict: self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("config dictionary")) @@ -307,11 +311,11 @@ def get_conn(self) -> client.ApiClient: client_configuration=self.client_configuration, context=cluster_context, ) - return client.ApiClient() + return TimeoutK8sApiClient() return self._get_default_client(cluster_context=cluster_context) - def _get_default_client(self, *, cluster_context: str | None = None) -> client.ApiClient: + def _get_default_client(self, *, cluster_context: str | None = None) -> TimeoutK8sApiClient: # if we get here, then no configuration has been supplied # we should try in_cluster since that's most likely # but failing that just load assuming a kubeconfig file @@ -326,7 +330,7 @@ def _get_default_client(self, *, cluster_context: str | None = None) -> client.A client_configuration=self.client_configuration, context=cluster_context, ) - return client.ApiClient() + return TimeoutK8sApiClient() @property def is_in_cluster(self) -> bool: @@ -409,7 +413,6 @@ def get_custom_object( namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, plural=plural, name=name, - **with_timeout(), ) return response @@ -433,7 +436,7 @@ def delete_custom_object( namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, plural=plural, name=name, - **with_timeout(kwargs), + **kwargs, ) def get_namespace(self) -> str | None: @@ -474,7 +477,6 @@ def get_pod_log_stream( name=pod_name, container=container, namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, - **with_timeout(), ), ) @@ -496,7 +498,6 @@ def get_pod_logs( container=container, _preload_content=False, namespace=namespace or self.get_namespace() or self.DEFAULT_NAMESPACE, - **with_timeout(), ) def get_pod(self, name: str, namespace: str) -> V1Pod: @@ -504,7 +505,6 @@ def get_pod(self, name: str, namespace: str) -> V1Pod: return self.core_v1_client.read_namespaced_pod( name=name, namespace=namespace, - **with_timeout(), ) def get_namespaced_pod_list( @@ -526,7 +526,7 @@ def get_namespaced_pod_list( watch=watch, label_selector=label_selector, _preload_content=False, - **with_timeout(kwargs), + **kwargs, ) def get_deployment_status( @@ -562,7 +562,7 @@ def create_job( self.log.debug("Job Creation Request: \n%s", json_job) try: resp = self.batch_v1_client.create_namespaced_job( - body=sanitized_job, namespace=job.metadata.namespace, **with_timeout(kwargs) + body=sanitized_job, namespace=job.metadata.namespace, **kwargs ) self.log.debug("Job Creation Response: %s", resp) except Exception as e: @@ -581,9 +581,7 @@ def get_job(self, job_name: str, namespace: str) -> V1Job: :param namespace: Namespace of the Job. :return: Job object """ - return self.batch_v1_client.read_namespaced_job( - name=job_name, namespace=namespace, pretty=True, **with_timeout() - ) + return self.batch_v1_client.read_namespaced_job(name=job_name, namespace=namespace, pretty=True) @generic_api_retry def get_job_status(self, job_name: str, namespace: str) -> V1Job: @@ -595,7 +593,7 @@ def get_job_status(self, job_name: str, namespace: str) -> V1Job: :return: Job object """ return self.batch_v1_client.read_namespaced_job_status( - name=job_name, namespace=namespace, pretty=True, **with_timeout() + name=job_name, namespace=namespace, pretty=True ) def wait_until_job_complete(self, job_name: str, namespace: str, job_poll_interval: float = 10) -> V1Job: @@ -622,7 +620,7 @@ def list_jobs_all_namespaces(self) -> V1JobList: :return: V1JobList object """ - return self.batch_v1_client.list_job_for_all_namespaces(pretty=True, **with_timeout()) + return self.batch_v1_client.list_job_for_all_namespaces(pretty=True) @generic_api_retry def list_jobs_from_namespace(self, namespace: str) -> V1JobList: @@ -632,7 +630,7 @@ def list_jobs_from_namespace(self, namespace: str) -> V1JobList: :param namespace: Namespace of the Job. :return: V1JobList object """ - return self.batch_v1_client.list_namespaced_job(namespace=namespace, pretty=True, **with_timeout()) + return self.batch_v1_client.list_namespaced_job(namespace=namespace, pretty=True) def is_job_complete(self, job: V1Job) -> bool: """ @@ -693,7 +691,6 @@ def patch_namespaced_job(self, job_name: str, namespace: str, body: object) -> V name=job_name, namespace=namespace, body=body, - **with_timeout(), ) def apply_from_yaml_file( @@ -810,7 +807,7 @@ async def api_client_from_kubeconfig_file(_kubeconfig_path: str | None): client_configuration=self.client_configuration, context=cluster_context, ) - return async_client.ApiClient() + return TimeoutAsyncK8sApiClient() if num_selected_configuration > 1: raise AirflowException( @@ -823,13 +820,13 @@ async def api_client_from_kubeconfig_file(_kubeconfig_path: str | None): self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("within a pod")) self._is_in_cluster = True async_config.load_incluster_config() - return async_client.ApiClient() + return TimeoutAsyncK8sApiClient() if self.config_dict: self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("config dictionary")) self._is_in_cluster = False await async_config.load_kube_config_from_dict(self.config_dict, context=cluster_context) - return async_client.ApiClient() + return TimeoutAsyncK8sApiClient() if kubeconfig_path is not None: self.log.debug("loading kube_config from: %s", kubeconfig_path) @@ -881,10 +878,10 @@ async def _get_field(self, field_name): return extras.get(prefixed_name) @contextlib.asynccontextmanager - async def get_conn(self) -> async_client.ApiClient: + async def get_conn(self) -> TimeoutAsyncK8sApiClient: kube_client = None try: - kube_client = await self._load_config() or async_client.ApiClient() + kube_client = await self._load_config() or TimeoutAsyncK8sApiClient() yield kube_client finally: if kube_client is not None: @@ -904,7 +901,6 @@ async def get_pod(self, name: str, namespace: str) -> V1Pod: pod: V1Pod = await v1_api.read_namespaced_pod( name=name, namespace=namespace, - **with_timeout(), ) return pod except HTTPError as e: @@ -924,7 +920,7 @@ async def delete_pod(self, name: str, namespace: str): try: v1_api = async_client.CoreV1Api(connection) await v1_api.delete_namespaced_pod( - name=name, namespace=namespace, body=client.V1DeleteOptions(), **with_timeout() + name=name, namespace=namespace, body=client.V1DeleteOptions() ) except async_client.ApiException as e: # If the pod is already deleted @@ -958,7 +954,6 @@ async def read_logs( follow=False, timestamps=True, since_seconds=since_seconds, - **with_timeout(), ) logs = logs.splitlines() return logs @@ -984,7 +979,6 @@ async def get_pod_events( namespace=namespace, resource_version=resource_version, resource_version_match="NotOlderThan" if resource_version else None, - **with_timeout(), ) return events except HTTPError as e: @@ -1025,7 +1019,6 @@ async def watch_pod_events( field_selector=f"involvedObject.name={name}", resource_version=resource_version, timeout_seconds=timeout_seconds, - **with_timeout(), ): event: CoreV1Event = event_watched.get("object") yield event @@ -1080,7 +1073,6 @@ async def get_job_status(self, name: str, namespace: str) -> V1Job: job: V1Job = await v1_api.read_namespaced_job_status( name=name, namespace=namespace, - **with_timeout(), ) return job diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py index 055fd35eb9686..621354d026858 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py @@ -24,7 +24,9 @@ import pendulum import tenacity +from kubernetes import client as sync_k8s_client from kubernetes.client.rest import ApiException as SyncApiException +from kubernetes_asyncio import client as async_k8s_client from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException from slugify import slugify from sqlalchemy import select @@ -62,6 +64,22 @@ class KubernetesApiException(AirflowException): TRANSIENT_STATUS_CODES = {409, 429, 500, 502, 503, 504} +class TimeoutK8sApiClient(sync_k8s_client.ApiClient): + """Wrapper around kubernetes sync ApiClient to set default timeout.""" + + def call_api(self, *args, **kwargs): + kwargs.setdefault("_request_timeout", API_TIMEOUT) + return super().call_api(*args, **kwargs) + + +class TimeoutAsyncK8sApiClient(async_k8s_client.ApiClient): + """Wrapper around kubernetes async ApiClient to set default timeout.""" + + async def call_api(self, *args, **kwargs): + kwargs.setdefault("_request_timeout", API_TIMEOUT) + return await super().call_api(*args, **kwargs) + + def _should_retry_api(exc: BaseException) -> bool: """Retry on selected ApiException status codes, plus plain HTTP/timeout errors.""" if isinstance(exc, (SyncApiException, AsyncApiException)): @@ -101,13 +119,6 @@ def generic_api_retry(func): )(func) -def with_timeout(kwargs: dict | None = None) -> dict: - """Attach a total request timeout unless explicitly provided.""" - out = dict(kwargs or {}) - out.setdefault("_request_timeout", API_TIMEOUT) - return out - - def rand_str(num): """ Generate random lowercase alphanumeric string of length num. diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py index 4f2cf4f94e9b5..c9b48158d960a 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py @@ -42,7 +42,6 @@ KubernetesApiException, PodLaunchFailedException, generic_api_retry, - with_timeout, ) from airflow.providers.cncf.kubernetes.utils.container import ( container_is_completed, @@ -332,7 +331,7 @@ def run_pod_async(self, pod: V1Pod, **kwargs) -> V1Pod: self.log.debug("Pod Creation Request: \n%s", json_pod) try: resp = self._client.create_namespaced_pod( - body=sanitized_pod, namespace=pod.metadata.namespace, **with_timeout(kwargs) + body=sanitized_pod, namespace=pod.metadata.namespace, **kwargs ) self.log.debug("Pod Creation Response: %s", resp) except Exception as e: @@ -347,7 +346,7 @@ def delete_pod(self, pod: V1Pod) -> None: """Delete POD.""" try: self._client.delete_namespaced_pod( - pod.metadata.name, pod.metadata.namespace, body=client.V1DeleteOptions(), **with_timeout() + pod.metadata.name, pod.metadata.namespace, body=client.V1DeleteOptions() ) except ApiException as e: # If the pod is already deleted @@ -751,7 +750,7 @@ def read_pod_logs( follow=follow, timestamps=timestamps, _preload_content=False, - **with_timeout(additional_kwargs), + **additional_kwargs, ) except HTTPError: self.log.exception("There was an error reading the kubernetes API.") @@ -801,7 +800,6 @@ def read_pod_events(self, pod: V1Pod, resource_version: str | None = None) -> Co field_selector=f"involvedObject.name={pod.metadata.name}", resource_version=resource_version, resource_version_match="NotOlderThan" if resource_version else None, - **with_timeout(), ) except HTTPError as e: raise KubernetesApiException(f"There was an error reading the kubernetes API: {e}") @@ -810,9 +808,7 @@ def read_pod_events(self, pod: V1Pod, resource_version: str | None = None) -> Co def read_pod(self, pod: V1Pod) -> V1Pod: """Read POD information.""" try: - return self._client.read_namespaced_pod( - pod.metadata.name, pod.metadata.namespace, **with_timeout() - ) + return self._client.read_namespaced_pod(pod.metadata.name, pod.metadata.namespace) except HTTPError as e: raise KubernetesApiException(f"There was an error reading the kubernetes API: {e}") @@ -877,7 +873,6 @@ def extract_xcom_json(self, pod: V1Pod) -> str: stderr=True, tty=False, _preload_content=False, - **with_timeout(), ) ) as client: self.log.info("Running command... %s", command) @@ -910,7 +905,6 @@ def extract_xcom_kill(self, pod: V1Pod): stderr=True, tty=False, _preload_content=False, - **with_timeout(), ) ) as resp: self._exec_pod_command(resp, "kill -2 $(pgrep -u $(id -u) -f 'sh')") 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 6be1755bcf67c..ea0a64938c69c 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 @@ -36,7 +36,6 @@ from airflow.models import Connection from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook -from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import with_timeout from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException from tests_common.test_utils.db import clear_test_connections @@ -476,7 +475,6 @@ def test_delete_custom_object( name="name", namespace="namespace", _preload_content="_preload_content", - **with_timeout(), ) @patch("kubernetes.config.kube_config.KubeConfigLoader") @@ -489,7 +487,7 @@ def test_get_job_status(self, mock_client, mock_kube_config_merger, mock_kube_co job_actual = hook.get_job_status(job_name=JOB_NAME, namespace=NAMESPACE) mock_client.read_namespaced_job_status.assert_called_once_with( - name=JOB_NAME, namespace=NAMESPACE, pretty=True, **with_timeout() + name=JOB_NAME, namespace=NAMESPACE, pretty=True ) assert job_actual == job_expected @@ -1040,7 +1038,6 @@ async def test_async_get_pod_events_with_resource_version( namespace=NAMESPACE, resource_version="12345", resource_version_match="NotOlderThan", - **with_timeout(), ) assert result == mock_events @@ -1070,7 +1067,6 @@ async def test_async_get_pod_events_without_resource_version( namespace=NAMESPACE, resource_version=None, resource_version_match=None, - **with_timeout(), ) assert result == mock_events @@ -1190,7 +1186,6 @@ async def test_async_watch_pod_events_polling_fallback( namespace=NAMESPACE, resource_version="12345", resource_version_match="NotOlderThan", - **with_timeout(), ) mock_sleep.assert_called_once_with(10) @@ -1247,7 +1242,6 @@ async def test_get_pod(self, lib_method, kube_config_loader): lib_method.assert_called_with( name=POD_NAME, namespace=NAMESPACE, - **with_timeout(), ) @pytest.mark.asyncio @@ -1294,7 +1288,6 @@ async def test_read_logs(self, lib_method, kube_config_loader): follow=False, timestamps=True, since_seconds=10, - **with_timeout(), ) assert len(logs) == 1 assert "2023-01-11 Some string logs..." in logs diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py index ce6711bfb730b..0b3dd195a6cd4 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py @@ -28,15 +28,65 @@ from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( API_TIMEOUT, KubernetesApiException, + TimeoutAsyncK8sApiClient, + TimeoutK8sApiClient, WaitRetryAfterOrExponential, _should_retry_api, create_unique_id, - with_timeout, ) pod_name_regex = r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$" +class TestTimeoutK8sApiClient: + @pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + [ + pytest.param({}, API_TIMEOUT, id="default-timeout"), + pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), + ], + ) + def test_call_api_timeout_inject(self, kwargs, expected_timeout): + with mock.patch("kubernetes.client.ApiClient.call_api") as mocked_call_api: + mocked_call_api.return_value = "ok" + cli = TimeoutK8sApiClient() + + out = cli.call_api("arg1", kwargs_arg1="fake", **kwargs) + + mocked_call_api.assert_called_once() + call_args, call_kwargs = mocked_call_api.call_args + assert call_args[0] == "arg1" + assert call_kwargs["kwargs_arg1"] == "fake" + assert call_kwargs["_request_timeout"] == expected_timeout + assert out == "ok" + + +class TestTimeoutAsyncK8sApiClient: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + [ + pytest.param({}, API_TIMEOUT, id="default-timeout"), + pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), + ], + ) + async def test_call_api_timeout_inject(self, kwargs, expected_timeout): + with mock.patch( + "kubernetes_asyncio.client.ApiClient.call_api", new_callable=mock.AsyncMock + ) as mocked_call_api: + mocked_call_api.return_value = "ok" + cli = TimeoutAsyncK8sApiClient() + + out = await cli.call_api("arg1", kwargs_arg1="fake", **kwargs) + + mocked_call_api.assert_called_once() + call_args, call_kwargs = mocked_call_api.call_args + assert call_args[0] == "arg1" + assert call_kwargs["kwargs_arg1"] == "fake" + assert call_kwargs["_request_timeout"] == expected_timeout + assert out == "ok" + + class DummyRetryState: def __init__(self, exception=None): self.outcome = mock.Mock() if exception is not None else None @@ -44,25 +94,6 @@ def __init__(self, exception=None): self.outcome.exception = mock.Mock(return_value=exception) -def test_with_timeout(): - # default value - timeout_kwargs = with_timeout() - assert "_request_timeout" in timeout_kwargs - assert timeout_kwargs["_request_timeout"] == API_TIMEOUT - - # check override - timeout_kwargs = with_timeout({"_request_timeout": 5}) - assert "_request_timeout" in timeout_kwargs - assert timeout_kwargs["_request_timeout"] == 5 - - # check additional params - timeout_kwargs = with_timeout({"some_param": 123}) - assert "some_param" in timeout_kwargs - assert timeout_kwargs["some_param"] == 123 - assert "_request_timeout" in timeout_kwargs - assert timeout_kwargs["_request_timeout"] == API_TIMEOUT - - def test_should_retry_api(): exc = HTTPError() assert _should_retry_api(exc) diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py index e60831745b093..062266540004b 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/utils/test_pod_manager.py @@ -30,7 +30,6 @@ from urllib3.exceptions import HTTPError as BaseHTTPError from airflow.providers.cncf.kubernetes.exceptions import KubernetesApiError -from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import with_timeout from airflow.providers.cncf.kubernetes.utils.pod_manager import ( AsyncPodManager, PodLogsConsumer, @@ -179,7 +178,6 @@ def test_read_pod_logs_retries_successfully(self): timestamps=False, name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, - **with_timeout(), ), mock.call( _preload_content=False, @@ -188,7 +186,6 @@ def test_read_pod_logs_retries_successfully(self): timestamps=False, name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, - **with_timeout(), ), ] ) @@ -237,7 +234,6 @@ def test_read_pod_logs_successfully_with_tail_lines(self): name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, tail_lines=100, - **with_timeout(), ), ] ) @@ -258,7 +254,6 @@ def test_read_pod_logs_successfully_with_since_seconds(self): name=mock.sentinel.metadata.name, namespace=mock.sentinel.metadata.namespace, since_seconds=2, - **with_timeout(), ), ] ) @@ -318,7 +313,6 @@ def test_read_pod_events_with_resource_version(self): field_selector="involvedObject.name=test-pod", resource_version="12345", resource_version_match="NotOlderThan", - **with_timeout(), ) def test_read_pod_events_without_resource_version(self): @@ -337,7 +331,6 @@ def test_read_pod_events_without_resource_version(self): field_selector="involvedObject.name=test-pod", resource_version=None, resource_version_match=None, - **with_timeout(), ) @pytest.mark.asyncio @@ -443,14 +436,12 @@ def test_read_pod_events_retries_successfully(self): field_selector=f"involvedObject.name={mock.sentinel.metadata.name}", resource_version=None, resource_version_match=None, - **with_timeout(), ), mock.call( namespace=mock.sentinel.metadata.namespace, field_selector=f"involvedObject.name={mock.sentinel.metadata.name}", resource_version=None, resource_version_match=None, - **with_timeout(), ), ] ) @@ -483,8 +474,8 @@ def test_read_pod_retries_successfully(self): assert mock.sentinel.pod_info == pod_info self.mock_kube_client.read_namespaced_pod.assert_has_calls( [ - mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace, **with_timeout()), - mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace, **with_timeout()), + mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace), + mock.call(mock.sentinel.metadata.name, mock.sentinel.metadata.namespace), ] ) From 25ec50c3f937f9903fdb5251e16e839dcfdcb53d Mon Sep 17 00:00:00 2001 From: AutomationDev85 Date: Fri, 9 Jan 2026 15:01:15 +0100 Subject: [PATCH 3/5] fix mypy and docu issues --- .../cncf/kubernetes/hooks/kubernetes.py | 45 ++++++++++----- .../kubernetes/kubernetes_helper_functions.py | 18 ------ .../cncf/kubernetes/hooks/test_kubernetes.py | 57 ++++++++++++++++++- .../test_kubernetes_helper_functions.py | 52 ----------------- .../google/cloud/hooks/kubernetes_engine.py | 2 +- 5 files changed, 87 insertions(+), 87 deletions(-) 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 be842c9eca411..450f8f1bc63be 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 @@ -37,8 +37,7 @@ from airflow.providers.cncf.kubernetes.exceptions import KubernetesApiError, KubernetesApiPermissionError from airflow.providers.cncf.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( - TimeoutAsyncK8sApiClient, - TimeoutK8sApiClient, + API_TIMEOUT, generic_api_retry, ) from airflow.providers.cncf.kubernetes.utils.container import ( @@ -49,7 +48,7 @@ from airflow.utils import yaml if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Generator + from collections.abc import AsyncGenerator, AsyncIterator, Generator from kubernetes.client import V1JobList from kubernetes.client.models import CoreV1Event, CoreV1EventList, V1Job, V1Pod @@ -72,6 +71,22 @@ def _load_body_to_dict(body: str) -> dict: return body_dict +class _TimeoutK8sApiClient(client.ApiClient): + """Wrapper around kubernetes sync ApiClient to set default timeout.""" + + def call_api(self, *args, **kwargs): + kwargs.setdefault("_request_timeout", API_TIMEOUT) + return super().call_api(*args, **kwargs) + + +class _TimeoutAsyncK8sApiClient(async_client.ApiClient): + """Wrapper around kubernetes async ApiClient to set default timeout.""" + + async def call_api(self, *args, **kwargs): + kwargs.setdefault("_request_timeout", API_TIMEOUT) + return await super().call_api(*args, **kwargs) + + class PodOperatorHookProtocol(Protocol): """ Protocol to define methods relied upon by KubernetesPodOperator. @@ -243,7 +258,7 @@ def _get_field(self, field_name): prefixed_name = f"extra__kubernetes__{field_name}" return self.conn_extras.get(prefixed_name) or None - def get_conn(self) -> TimeoutK8sApiClient: + def get_conn(self) -> _TimeoutK8sApiClient: """Return kubernetes api session for use with requests.""" in_cluster = self._coalesce_param(self.in_cluster, self._get_field("in_cluster")) cluster_context = self._coalesce_param(self.cluster_context, self._get_field("cluster_context")) @@ -276,7 +291,7 @@ def get_conn(self) -> TimeoutK8sApiClient: self.log.debug("loading kube_config from: in_cluster configuration") self._is_in_cluster = True config.load_incluster_config() - return TimeoutK8sApiClient() + return _TimeoutK8sApiClient() if kubeconfig_path is not None: self.log.debug("loading kube_config from: %s", kubeconfig_path) @@ -286,7 +301,7 @@ def get_conn(self) -> TimeoutK8sApiClient: client_configuration=self.client_configuration, context=cluster_context, ) - return TimeoutK8sApiClient() + return _TimeoutK8sApiClient() if kubeconfig is not None: with tempfile.NamedTemporaryFile() as temp_config: @@ -301,7 +316,7 @@ def get_conn(self) -> TimeoutK8sApiClient: client_configuration=self.client_configuration, context=cluster_context, ) - return TimeoutK8sApiClient() + return _TimeoutK8sApiClient() if self.config_dict: self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("config dictionary")) @@ -311,11 +326,11 @@ def get_conn(self) -> TimeoutK8sApiClient: client_configuration=self.client_configuration, context=cluster_context, ) - return TimeoutK8sApiClient() + return _TimeoutK8sApiClient() return self._get_default_client(cluster_context=cluster_context) - def _get_default_client(self, *, cluster_context: str | None = None) -> TimeoutK8sApiClient: + def _get_default_client(self, *, cluster_context: str | None = None) -> _TimeoutK8sApiClient: # if we get here, then no configuration has been supplied # we should try in_cluster since that's most likely # but failing that just load assuming a kubeconfig file @@ -330,7 +345,7 @@ def _get_default_client(self, *, cluster_context: str | None = None) -> TimeoutK client_configuration=self.client_configuration, context=cluster_context, ) - return TimeoutK8sApiClient() + return _TimeoutK8sApiClient() @property def is_in_cluster(self) -> bool: @@ -807,7 +822,7 @@ async def api_client_from_kubeconfig_file(_kubeconfig_path: str | None): client_configuration=self.client_configuration, context=cluster_context, ) - return TimeoutAsyncK8sApiClient() + return _TimeoutAsyncK8sApiClient() if num_selected_configuration > 1: raise AirflowException( @@ -820,13 +835,13 @@ async def api_client_from_kubeconfig_file(_kubeconfig_path: str | None): self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("within a pod")) self._is_in_cluster = True async_config.load_incluster_config() - return TimeoutAsyncK8sApiClient() + return _TimeoutAsyncK8sApiClient() if self.config_dict: self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("config dictionary")) self._is_in_cluster = False await async_config.load_kube_config_from_dict(self.config_dict, context=cluster_context) - return TimeoutAsyncK8sApiClient() + return _TimeoutAsyncK8sApiClient() if kubeconfig_path is not None: self.log.debug("loading kube_config from: %s", kubeconfig_path) @@ -878,10 +893,10 @@ async def _get_field(self, field_name): return extras.get(prefixed_name) @contextlib.asynccontextmanager - async def get_conn(self) -> TimeoutAsyncK8sApiClient: + async def get_conn(self) -> AsyncIterator[_TimeoutAsyncK8sApiClient]: # type: ignore[override] kube_client = None try: - kube_client = await self._load_config() or TimeoutAsyncK8sApiClient() + kube_client = await self._load_config() or _TimeoutAsyncK8sApiClient() yield kube_client finally: if kube_client is not None: diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py index 621354d026858..451a0ba9cadcf 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py @@ -24,9 +24,7 @@ import pendulum import tenacity -from kubernetes import client as sync_k8s_client from kubernetes.client.rest import ApiException as SyncApiException -from kubernetes_asyncio import client as async_k8s_client from kubernetes_asyncio.client.exceptions import ApiException as AsyncApiException from slugify import slugify from sqlalchemy import select @@ -64,22 +62,6 @@ class KubernetesApiException(AirflowException): TRANSIENT_STATUS_CODES = {409, 429, 500, 502, 503, 504} -class TimeoutK8sApiClient(sync_k8s_client.ApiClient): - """Wrapper around kubernetes sync ApiClient to set default timeout.""" - - def call_api(self, *args, **kwargs): - kwargs.setdefault("_request_timeout", API_TIMEOUT) - return super().call_api(*args, **kwargs) - - -class TimeoutAsyncK8sApiClient(async_k8s_client.ApiClient): - """Wrapper around kubernetes async ApiClient to set default timeout.""" - - async def call_api(self, *args, **kwargs): - kwargs.setdefault("_request_timeout", API_TIMEOUT) - return await super().call_api(*args, **kwargs) - - def _should_retry_api(exc: BaseException) -> bool: """Retry on selected ApiException status codes, plus plain HTTP/timeout errors.""" if isinstance(exc, (SyncApiException, AsyncApiException)): 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 ea0a64938c69c..465d8aacb3532 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 @@ -35,7 +35,13 @@ from kubernetes_asyncio import client as async_client from airflow.models import Connection -from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook +from airflow.providers.cncf.kubernetes.hooks.kubernetes import ( + AsyncKubernetesHook, + KubernetesHook, + _TimeoutAsyncK8sApiClient, + _TimeoutK8sApiClient, +) +from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import API_TIMEOUT from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException from tests_common.test_utils.db import clear_test_connections @@ -78,6 +84,55 @@ class DeprecationRemovalRequired(AirflowException): ... DEFAULT_CONN_ID = "kubernetes_default" +class TestTimeoutK8sApiClient: + @pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + [ + pytest.param({}, API_TIMEOUT, id="default-timeout"), + pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), + ], + ) + def test_call_api_timeout_inject(self, kwargs, expected_timeout): + with mock.patch("kubernetes.client.ApiClient.call_api") as mocked_call_api: + mocked_call_api.return_value = "ok" + cli = _TimeoutK8sApiClient() + + out = cli.call_api("arg1", kwargs_arg1="fake", **kwargs) + + mocked_call_api.assert_called_once() + call_args, call_kwargs = mocked_call_api.call_args + assert call_args[0] == "arg1" + assert call_kwargs["kwargs_arg1"] == "fake" + assert call_kwargs["_request_timeout"] == expected_timeout + assert out == "ok" + + +class TestTimeoutAsyncK8sApiClient: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + [ + pytest.param({}, API_TIMEOUT, id="default-timeout"), + pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), + ], + ) + async def test_call_api_timeout_inject(self, kwargs, expected_timeout): + with mock.patch( + "kubernetes_asyncio.client.ApiClient.call_api", new_callable=mock.AsyncMock + ) as mocked_call_api: + mocked_call_api.return_value = "ok" + cli = _TimeoutAsyncK8sApiClient() + + out = await cli.call_api("arg1", kwargs_arg1="fake", **kwargs) + + mocked_call_api.assert_called_once() + call_args, call_kwargs = mocked_call_api.call_args + assert call_args[0] == "arg1" + assert call_kwargs["kwargs_arg1"] == "fake" + assert call_kwargs["_request_timeout"] == expected_timeout + assert out == "ok" + + @pytest.fixture def remove_default_conn(monkeypatch): original_env_var = os.environ.get(f"AIRFLOW_CONN_{DEFAULT_CONN_ID.upper()}") diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py index 0b3dd195a6cd4..3622a5c34f48d 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_kubernetes_helper_functions.py @@ -26,10 +26,7 @@ from urllib3.exceptions import HTTPError from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( - API_TIMEOUT, KubernetesApiException, - TimeoutAsyncK8sApiClient, - TimeoutK8sApiClient, WaitRetryAfterOrExponential, _should_retry_api, create_unique_id, @@ -38,55 +35,6 @@ pod_name_regex = r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$" -class TestTimeoutK8sApiClient: - @pytest.mark.parametrize( - ("kwargs", "expected_timeout"), - [ - pytest.param({}, API_TIMEOUT, id="default-timeout"), - pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), - ], - ) - def test_call_api_timeout_inject(self, kwargs, expected_timeout): - with mock.patch("kubernetes.client.ApiClient.call_api") as mocked_call_api: - mocked_call_api.return_value = "ok" - cli = TimeoutK8sApiClient() - - out = cli.call_api("arg1", kwargs_arg1="fake", **kwargs) - - mocked_call_api.assert_called_once() - call_args, call_kwargs = mocked_call_api.call_args - assert call_args[0] == "arg1" - assert call_kwargs["kwargs_arg1"] == "fake" - assert call_kwargs["_request_timeout"] == expected_timeout - assert out == "ok" - - -class TestTimeoutAsyncK8sApiClient: - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("kwargs", "expected_timeout"), - [ - pytest.param({}, API_TIMEOUT, id="default-timeout"), - pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), - ], - ) - async def test_call_api_timeout_inject(self, kwargs, expected_timeout): - with mock.patch( - "kubernetes_asyncio.client.ApiClient.call_api", new_callable=mock.AsyncMock - ) as mocked_call_api: - mocked_call_api.return_value = "ok" - cli = TimeoutAsyncK8sApiClient() - - out = await cli.call_api("arg1", kwargs_arg1="fake", **kwargs) - - mocked_call_api.assert_called_once() - call_args, call_kwargs = mocked_call_api.call_args - assert call_args[0] == "arg1" - assert call_kwargs["kwargs_arg1"] == "fake" - assert call_kwargs["_request_timeout"] == expected_timeout - assert out == "ok" - - class DummyRetryState: def __init__(self, exception=None): self.outcome = mock.Mock() if exception is not None else None diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 865b5b67c1545..68788ddf05abd 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -498,7 +498,7 @@ def __init__( ) @contextlib.asynccontextmanager - async def get_conn(self) -> async_client.ApiClient: + async def get_conn(self) -> async_client.ApiClient: # type: ignore[override] kube_client = None try: kube_client = await self._load_config() From 718a59d7d7eb0424c6598afb8a9313ac10a0a607 Mon Sep 17 00:00:00 2001 From: AutomationDev85 Date: Mon, 12 Jan 2026 11:56:56 +0100 Subject: [PATCH 4/5] Fixed typing and rework server side and client timeout --- .../cncf/kubernetes/hooks/kubernetes.py | 22 ++++++++++---- .../kubernetes/kubernetes_helper_functions.py | 1 + .../cncf/kubernetes/hooks/test_kubernetes.py | 29 +++++++++++++++++-- .../google/cloud/hooks/kubernetes_engine.py | 2 +- 4 files changed, 44 insertions(+), 10 deletions(-) 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 450f8f1bc63be..6cab89856bce5 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 @@ -38,6 +38,7 @@ from airflow.providers.cncf.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( API_TIMEOUT, + API_TIMEOUT_OFFSET_SERVER_SIDE, generic_api_retry, ) from airflow.providers.cncf.kubernetes.utils.container import ( @@ -48,7 +49,7 @@ from airflow.utils import yaml if TYPE_CHECKING: - from collections.abc import AsyncGenerator, AsyncIterator, Generator + from collections.abc import AsyncGenerator, Generator from kubernetes.client import V1JobList from kubernetes.client.models import CoreV1Event, CoreV1EventList, V1Job, V1Pod @@ -71,11 +72,19 @@ def _load_body_to_dict(body: str) -> dict: return body_dict +def _get_request_timeout(timeout_seconds: int | None) -> float: + """Get the client-side request timeout.""" + if timeout_seconds is not None and timeout_seconds > API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE: + return timeout_seconds + API_TIMEOUT_OFFSET_SERVER_SIDE + return API_TIMEOUT + + class _TimeoutK8sApiClient(client.ApiClient): """Wrapper around kubernetes sync ApiClient to set default timeout.""" def call_api(self, *args, **kwargs): - kwargs.setdefault("_request_timeout", API_TIMEOUT) + timeout_seconds = kwargs.get("timeout_seconds") # get server-side timeout + kwargs.setdefault("_request_timeout", _get_request_timeout(timeout_seconds)) # client-side timeout return super().call_api(*args, **kwargs) @@ -83,7 +92,8 @@ class _TimeoutAsyncK8sApiClient(async_client.ApiClient): """Wrapper around kubernetes async ApiClient to set default timeout.""" async def call_api(self, *args, **kwargs): - kwargs.setdefault("_request_timeout", API_TIMEOUT) + timeout_seconds = kwargs.get("timeout_seconds") # server-side timeout + kwargs.setdefault("_request_timeout", _get_request_timeout(timeout_seconds)) # client-side timeout return await super().call_api(*args, **kwargs) @@ -258,7 +268,7 @@ def _get_field(self, field_name): prefixed_name = f"extra__kubernetes__{field_name}" return self.conn_extras.get(prefixed_name) or None - def get_conn(self) -> _TimeoutK8sApiClient: + def get_conn(self) -> client.ApiClient: """Return kubernetes api session for use with requests.""" in_cluster = self._coalesce_param(self.in_cluster, self._get_field("in_cluster")) cluster_context = self._coalesce_param(self.cluster_context, self._get_field("cluster_context")) @@ -330,7 +340,7 @@ def get_conn(self) -> _TimeoutK8sApiClient: return self._get_default_client(cluster_context=cluster_context) - def _get_default_client(self, *, cluster_context: str | None = None) -> _TimeoutK8sApiClient: + def _get_default_client(self, *, cluster_context: str | None = None) -> client.ApiClient: # if we get here, then no configuration has been supplied # we should try in_cluster since that's most likely # but failing that just load assuming a kubeconfig file @@ -893,7 +903,7 @@ async def _get_field(self, field_name): return extras.get(prefixed_name) @contextlib.asynccontextmanager - async def get_conn(self) -> AsyncIterator[_TimeoutAsyncK8sApiClient]: # type: ignore[override] + async def get_conn(self) -> async_client.ApiClient: kube_client = None try: kube_client = await self._load_config() or _TimeoutAsyncK8sApiClient() diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py index 451a0ba9cadcf..c2334b1c9a2e2 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py @@ -53,6 +53,7 @@ class KubernetesApiException(AirflowException): API_TIMEOUT = 60 # allow 1 min of timeout for kubernetes api calls +API_TIMEOUT_OFFSET_SERVER_SIDE = 5 # offset to the server side timeout for the client side timeout API_RETRIES = conf.getint("workers", "api_retries", fallback=5) API_RETRY_WAIT_MIN = conf.getfloat("workers", "api_retry_wait_min", fallback=1) API_RETRY_WAIT_MAX = conf.getfloat("workers", "api_retry_wait_max", fallback=15) 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 465d8aacb3532..d1fac78704989 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 @@ -41,7 +41,10 @@ _TimeoutAsyncK8sApiClient, _TimeoutK8sApiClient, ) -from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import API_TIMEOUT +from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import ( + API_TIMEOUT, + API_TIMEOUT_OFFSET_SERVER_SIDE, +) from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException from tests_common.test_utils.db import clear_test_connections @@ -89,7 +92,17 @@ class TestTimeoutK8sApiClient: ("kwargs", "expected_timeout"), [ pytest.param({}, API_TIMEOUT, id="default-timeout"), - pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), + pytest.param({"timeout_seconds": 5678, "_request_timeout": 1234}, 1234, id="explicit-timeout"), + pytest.param( + {"timeout_seconds": API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE}, + API_TIMEOUT, + id="server-side-timeout-limit", + ), + pytest.param( + {"timeout_seconds": API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE + 1}, + API_TIMEOUT + 1, + id="server-side-timeout-above-limit", + ), ], ) def test_call_api_timeout_inject(self, kwargs, expected_timeout): @@ -113,7 +126,17 @@ class TestTimeoutAsyncK8sApiClient: ("kwargs", "expected_timeout"), [ pytest.param({}, API_TIMEOUT, id="default-timeout"), - pytest.param({"_request_timeout": 1234}, 1234, id="explicit-timeout"), + pytest.param({"timeout_seconds": 5678, "_request_timeout": 1234}, 1234, id="explicit-timeout"), + pytest.param( + {"timeout_seconds": API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE}, + API_TIMEOUT, + id="server-side-timeout-limit", + ), + pytest.param( + {"timeout_seconds": API_TIMEOUT - API_TIMEOUT_OFFSET_SERVER_SIDE + 1}, + API_TIMEOUT + 1, + id="server-side-timeout-above-limit", + ), ], ) async def test_call_api_timeout_inject(self, kwargs, expected_timeout): diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 68788ddf05abd..865b5b67c1545 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -498,7 +498,7 @@ def __init__( ) @contextlib.asynccontextmanager - async def get_conn(self) -> async_client.ApiClient: # type: ignore[override] + async def get_conn(self) -> async_client.ApiClient: kube_client = None try: kube_client = await self._load_config() From dc1e4835fdc008c3d30dc15828cd4d579ded5cd7 Mon Sep 17 00:00:00 2001 From: AutomationDev85 Date: Tue, 13 Jan 2026 08:43:27 +0100 Subject: [PATCH 5/5] Extend function note --- .../src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 6cab89856bce5..30499d2c47790 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 @@ -1025,7 +1025,7 @@ async def watch_pod_events( :param name: Pod name to watch events for :param namespace: Kubernetes namespace :param resource_version: Only return events not older than this resource version - :param timeout_seconds: Timeout in seconds for the watch stream + :param timeout_seconds: Timeout in seconds for the watch stream. A small additional buffer may be applied internally. """ if self._event_polling_fallback: async for event_polled in self.watch_pod_events_polling_fallback(