add service_file support to GKEPodAsyncHook (#37081)

Lee-W · web-flow · commit 2372e21d9dd4 · 2024-02-05T11:50:13.000+05:30
Currently, GKEPodAsyncHook does not support service_file. Thus, passing credentials through " Keyfile Path " or " Keyfile JSON " will be ignored. This PR intends to fix this issue. As the default value of service_file in Token is None (https://github.com/talkiq/gcloud-aio/blob/8c8e1b39ec2e40b42212c270acb98c039267fbc5/auth/gcloud/aio/auth/token.py#L157) and the return value of service_file_as_context when both key file path and key file json are not provided is None. This change won't affect existing behavior
diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -508,33 +508,37 @@ async def get_pod(self, name: str, namespace: str) -> V1Pod:
         :param name: Name of the pod.
         :param namespace: Name of the pod's namespace.
         """
-        async with Token(scopes=self.scopes) as token:
-            async with self.get_conn(token) as connection:
-                v1_api = async_client.CoreV1Api(connection)
-                pod: V1Pod = await v1_api.read_namespaced_pod(
-                    name=name,
-                    namespace=namespace,
-                )
-            return pod
+        async with self.service_file_as_context() as service_file:  # type: ignore[attr-defined]
+            async with Token(scopes=self.scopes, service_file=service_file) as token:
+                async with self.get_conn(token) as connection:
+                    v1_api = async_client.CoreV1Api(connection)
+                    pod: V1Pod = await v1_api.read_namespaced_pod(
+                        name=name,
+                        namespace=namespace,
+                    )
+                return pod
 
     async def delete_pod(self, name: str, namespace: str):
         """Delete a pod.
 
         :param name: Name of the pod.
         :param namespace: Name of the pod's namespace.
         """
-        async with Token(scopes=self.scopes) as token, self.get_conn(token) as connection:
-            try:
-                v1_api = async_client.CoreV1Api(connection)
-                await v1_api.delete_namespaced_pod(
-                    name=name,
-                    namespace=namespace,
-                    body=client.V1DeleteOptions(),
-                )
-            except async_client.ApiException as e:
-                # If the pod is already deleted
-                if e.status != 404:
-                    raise
+        async with self.service_file_as_context() as service_file:  # type: ignore[attr-defined]
+            async with Token(scopes=self.scopes, service_file=service_file) as token, self.get_conn(
+                token
+            ) as connection:
+                try:
+                    v1_api = async_client.CoreV1Api(connection)
+                    await v1_api.delete_namespaced_pod(
+                        name=name,
+                        namespace=namespace,
+                        body=client.V1DeleteOptions(),
+                    )
+                except async_client.ApiException as e:
+                    # If the pod is already deleted
+                    if e.status != 404:
+                        raise
 
     async def read_logs(self, name: str, namespace: str):
         """Read logs inside the pod while starting containers inside.
@@ -547,19 +551,22 @@ async def read_logs(self, name: str, namespace: str):
         :param name: Name of the pod.
         :param namespace: Name of the pod's namespace.
         """
-        async with Token(scopes=self.scopes) as token, self.get_conn(token) as connection:
-            try:
-                v1_api = async_client.CoreV1Api(connection)
-                logs = await v1_api.read_namespaced_pod_log(
-                    name=name,
-                    namespace=namespace,
-                    follow=False,
-                    timestamps=True,
-                )
-                logs = logs.splitlines()
-                for line in logs:
-                    self.log.info("Container logs from %s", line)
-                return logs
-            except HTTPError:
-                self.log.exception("There was an error reading the kubernetes API.")
-                raise
+        async with self.service_file_as_context() as service_file:  # type: ignore[attr-defined]
+            async with Token(scopes=self.scopes, service_file=service_file) as token, self.get_conn(
+                token
+            ) as connection:
+                try:
+                    v1_api = async_client.CoreV1Api(connection)
+                    logs = await v1_api.read_namespaced_pod_log(
+                        name=name,
+                        namespace=namespace,
+                        follow=False,
+                        timestamps=True,
+                    )
+                    logs = logs.splitlines()
+                    for line in logs:
+                        self.log.info("Container logs from %s", line)
+                    return logs
+                except HTTPError:
+                    self.log.exception("There was an error reading the kubernetes API.")
+                    raise
diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py
@@ -317,29 +317,46 @@ def async_hook(self):
         )
 
     @pytest.mark.asyncio
-    @mock.patch(GKE_STRING.format("Token"), mock.MagicMock())
+    @pytest.mark.parametrize("mock_service_file", ("/tmp/service_file.json", None))
+    @mock.patch(GKE_STRING.format("Token"))
     @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn"))
     @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod"))
-    async def test_get_pod(self, read_namespace_pod_mock, get_conn_mock, async_hook):
+    async def test_get_pod(
+        self, read_namespace_pod_mock, get_conn_mock, mock_token, async_hook, mock_service_file
+    ):
+        async_hook.service_file_as_context = mock.MagicMock()
+        async_hook.service_file_as_context.return_value.__aenter__.return_value = mock_service_file
+
         self.make_mock_awaitable(read_namespace_pod_mock)
 
         await async_hook.get_pod(name=POD_NAME, namespace=POD_NAMESPACE)
-
+        mock_token.assert_called_with(
+            scopes=["https://www.googleapis.com/auth/cloud-platform"], service_file=mock_service_file
+        )
         get_conn_mock.assert_called_once()
         read_namespace_pod_mock.assert_called_with(
             name=POD_NAME,
             namespace=POD_NAMESPACE,
         )
 
     @pytest.mark.asyncio
-    @mock.patch(GKE_STRING.format("Token"), mock.MagicMock())
+    @pytest.mark.parametrize("mock_service_file", ("/tmp/service_file.json", None))
+    @mock.patch(GKE_STRING.format("Token"))
     @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn"))
     @mock.patch(GKE_STRING.format("async_client.CoreV1Api.delete_namespaced_pod"))
-    async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook):
+    async def test_delete_pod(
+        self, delete_namespaced_pod, get_conn_mock, mock_token, async_hook, mock_service_file
+    ):
+        async_hook.service_file_as_context = mock.MagicMock()
+        async_hook.service_file_as_context.return_value.__aenter__.return_value = mock_service_file
+
         self.make_mock_awaitable(delete_namespaced_pod)
 
         await async_hook.delete_pod(name=POD_NAME, namespace=POD_NAMESPACE)
 
+        mock_token.assert_called_with(
+            scopes=["https://www.googleapis.com/auth/cloud-platform"], service_file=mock_service_file
+        )
         get_conn_mock.assert_called_once()
         delete_namespaced_pod.assert_called_with(
             name=POD_NAME,
@@ -348,14 +365,23 @@ async def test_delete_pod(self, delete_namespaced_pod, get_conn_mock, async_hook
         )
 
     @pytest.mark.asyncio
-    @mock.patch(GKE_STRING.format("Token"), mock.MagicMock())
+    @pytest.mark.parametrize("mock_service_file", ("/tmp/service_file.json", None))
+    @mock.patch(GKE_STRING.format("Token"))
     @mock.patch(GKE_STRING.format("GKEPodAsyncHook.get_conn"))
     @mock.patch(GKE_STRING.format("async_client.CoreV1Api.read_namespaced_pod_log"))
-    async def test_read_logs(self, read_namespaced_pod_log, get_conn_mock, async_hook, caplog):
+    async def test_read_logs(
+        self, read_namespaced_pod_log, get_conn_mock, mock_token, async_hook, mock_service_file, caplog
+    ):
+        async_hook.service_file_as_context = mock.MagicMock()
+        async_hook.service_file_as_context.return_value.__aenter__.return_value = mock_service_file
+
         self.make_mock_awaitable(read_namespaced_pod_log, result="Test string #1\nTest string #2\n")
 
         await async_hook.read_logs(name=POD_NAME, namespace=POD_NAMESPACE)
 
+        mock_token.assert_called_with(
+            scopes=["https://www.googleapis.com/auth/cloud-platform"], service_file=mock_service_file
+        )
         get_conn_mock.assert_called_once()
         read_namespaced_pod_log.assert_called_with(
             name=POD_NAME,
