Enable '_enable_tcp_keepalive' functionality for GKEPodHook (#36999)

airflow/providers/google/cloud/hooks/kubernetes_engine.py

@@ -48,6 +48,7 @@
 
 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 @@ def __init__(
         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 @@ def __init__(
         )
         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 @@ def get_xcom_sidecar_container_resources(self):
     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):

tests/providers/google/cloud/hooks/test_kubernetes_engine.py

@@ -450,3 +450,34 @@ def _get_config(self):
 
     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)