add service_file support to GKEPodAsyncHook (#37081)

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
apache · Feb 5, 2024 · 2372e21 · 2372e21
1 parent 2e95a2a
commit 2372e21
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 43 deletions.
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,