You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2024/01/26 16:33:36 UTC
(airflow) branch main updated: Enable '_enable_tcp_keepalive' functionality for GKEPodHook (#36999)
This is an automated email from the ASF dual-hosted git repository.
potiuk 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 35daa3462b Enable '_enable_tcp_keepalive' functionality for GKEPodHook (#36999)
35daa3462b is described below
commit 35daa3462b128e5972d9e270dc66e6862f934262
Author: Maksim <ma...@google.com>
AuthorDate: Fri Jan 26 17:33:29 2024 +0100
Enable '_enable_tcp_keepalive' functionality for GKEPodHook (#36999)
---
.../google/cloud/hooks/kubernetes_engine.py | 7 +++++
.../google/cloud/hooks/test_kubernetes_engine.py | 31 ++++++++++++++++++++++
2 files changed, 38 insertions(+)
diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
index 30be61dc1f..4fb7534d87 100644
--- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py
+++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -48,6 +48,7 @@ from urllib3.exceptions import HTTPError
from airflow import version
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.cncf.kubernetes.kube_client import _enable_tcp_keepalive
from airflow.providers.cncf.kubernetes.utils.pod_manager import PodOperatorHookProtocol
from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import (
@@ -352,6 +353,7 @@ class GKEPodHook(GoogleBaseHook, PodOperatorHookProtocol):
self,
cluster_url: str,
ssl_ca_cert: str,
+ disable_tcp_keepalive: bool | None = None,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
**kwargs,
@@ -363,6 +365,7 @@ class GKEPodHook(GoogleBaseHook, PodOperatorHookProtocol):
)
self._cluster_url = cluster_url
self._ssl_ca_cert = ssl_ca_cert
+ self.disable_tcp_keepalive = disable_tcp_keepalive
@cached_property
def api_client(self) -> client.ApiClient:
@@ -397,6 +400,10 @@ class GKEPodHook(GoogleBaseHook, PodOperatorHookProtocol):
def get_conn(self) -> client.ApiClient:
configuration = self._get_config()
configuration.refresh_api_key_hook = self._refresh_api_key_hook
+
+ if self.disable_tcp_keepalive is not True:
+ _enable_tcp_keepalive()
+
return client.ApiClient(configuration)
def _refresh_api_key_hook(self, configuration: client.configuration.Configuration):
diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py
index 69189a8ae9..954f3e018e 100644
--- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py
+++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py
@@ -450,3 +450,34 @@ class TestGKEPodHook:
def _get_credentials(self):
return self.credentials
+
+ @pytest.mark.parametrize(
+ "disable_tcp_keepalive, expected",
+ (
+ (True, False),
+ (None, True),
+ (False, True),
+ ),
+ )
+ @mock.patch(GKE_STRING.format("_enable_tcp_keepalive"))
+ def test_disable_tcp_keepalive(
+ self,
+ mock_enable,
+ disable_tcp_keepalive,
+ expected,
+ ):
+ with mock.patch(
+ BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_base_gcp_hook_default_project_id
+ ):
+ gke_hook = GKEPodHook(
+ gcp_conn_id="test",
+ impersonation_chain=IMPERSONATE_CHAIN,
+ ssl_ca_cert=None,
+ cluster_url=None,
+ disable_tcp_keepalive=disable_tcp_keepalive,
+ )
+ gke_hook.get_credentials = self._get_credentials
+
+ api_conn = gke_hook.get_conn()
+ assert mock_enable.called is expected
+ assert isinstance(api_conn, kubernetes.client.api_client.ApiClient)