phanikumv commented on code in PR #29266:
URL: https://github.com/apache/airflow/pull/29266#discussion_r1143025827
##########
airflow/providers/google/cloud/hooks/kubernetes_engine.py:
##########
@@ -336,3 +343,195 @@ async def get_operation(
return await client.get_operation(
name=operation_path,
)
+
+
+class GKEPodHook(GoogleBaseHook):
+ """Hook for managing Google Kubernetes Engine pod APIs."""
+
+ def __init__(
+ self,
+ cluster_url: str,
+ ssl_ca_cert: str,
+ *args,
+ **kwargs,
+ ):
+ super().__init__(*args, **kwargs)
+ self._cluster_url = cluster_url
+ self._ssl_ca_cert = ssl_ca_cert
+
+ @cached_property
+ def api_client(self) -> client.ApiClient:
+ return self.get_conn()
+
+ @cached_property
+ def core_v1_client(self) -> client.CoreV1Api:
+ return client.CoreV1Api(self.api_client)
+
+ @property
+ def is_in_cluster(self) -> bool:
+ return False
+
+ @staticmethod
+ def get_xcom_sidecar_container_image():
+ """Returns the xcom sidecar image that defined in the connection"""
+ return PodDefaults.SIDECAR_CONTAINER.image
+
+ def get_conn(self) -> client.ApiClient:
+ configuration = self._get_config()
+ return client.ApiClient(configuration)
+
+ def _get_config(self) -> client.configuration.Configuration:
+ configuration = client.Configuration(
+ host=self._cluster_url,
+ api_key_prefix={"authorization": "Bearer"},
+ api_key={"authorization": self._get_token(self.get_credentials())},
+ )
+ configuration.ssl_ca_cert = FileOrData(
+ {
+ "certificate-authority-data": self._ssl_ca_cert,
+ },
+ file_key_name="certificate-authority",
+ ).as_file()
+ return configuration
+
+ @staticmethod
+ def _get_token(creds: google.auth.credentials.Credentials) -> str:
+ if creds.token is None or creds.expired:
+ auth_req = google_requests.Request()
+ creds.refresh(auth_req)
+ return creds.token
+
+ def get_pod(self, name: str, namespace: str) -> V1Pod:
+ """
+ Gets pod's object.
+
+ :param name: Name of the pod.
+ :param namespace: Name of the pod's namespace.
+ """
+ return self.core_v1_client.read_namespaced_pod(
+ name=name,
+ namespace=namespace,
+ )
+
+
+class AsyncGKEPodHook(GoogleBaseAsyncHook):
Review Comment:
```suggestion
class GKEPodAsyncHook(GoogleBaseAsyncHook):
```
just to align with how Async classes are being named
##########
airflow/providers/google/cloud/triggers/kubernetes_engine.py:
##########
@@ -18,14 +18,108 @@
from __future__ import annotations
import asyncio
+from datetime import datetime
from typing import Any, AsyncIterator, Sequence
from google.cloud.container_v1.types import Operation
-from airflow.providers.google.cloud.hooks.kubernetes_engine import AsyncGKEHook
+try:
+ from airflow.providers.cncf.kubernetes.triggers.pod import
KubernetesPodTrigger
+except ImportError:
+ # preserve backward compatibility for older versions of cncf.kubernetes
provider
+ from airflow.providers.cncf.kubernetes.triggers.kubernetes_pod import
KubernetesPodTrigger
+from airflow.providers.google.cloud.hooks.kubernetes_engine import
AsyncGKEHook, AsyncGKEPodHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
+class GKEPodTrigger(KubernetesPodTrigger):
Review Comment:
shouldnt this be `GKEStartPodTrigger` instead of `GKEPodTrigger` - its being
used only by `GKEStartPodOperator`.
##########
airflow/providers/google/cloud/hooks/kubernetes_engine.py:
##########
@@ -336,3 +343,195 @@ async def get_operation(
return await client.get_operation(
name=operation_path,
)
+
+
+class GKEPodHook(GoogleBaseHook):
+ """Hook for managing Google Kubernetes Engine pod APIs."""
+
+ def __init__(
+ self,
+ cluster_url: str,
+ ssl_ca_cert: str,
+ *args,
+ **kwargs,
+ ):
+ super().__init__(*args, **kwargs)
+ self._cluster_url = cluster_url
+ self._ssl_ca_cert = ssl_ca_cert
+
+ @cached_property
+ def api_client(self) -> client.ApiClient:
+ return self.get_conn()
+
+ @cached_property
+ def core_v1_client(self) -> client.CoreV1Api:
+ return client.CoreV1Api(self.api_client)
+
+ @property
+ def is_in_cluster(self) -> bool:
+ return False
+
+ @staticmethod
+ def get_xcom_sidecar_container_image():
+ """Returns the xcom sidecar image that defined in the connection"""
+ return PodDefaults.SIDECAR_CONTAINER.image
+
+ def get_conn(self) -> client.ApiClient:
+ configuration = self._get_config()
+ return client.ApiClient(configuration)
+
+ def _get_config(self) -> client.configuration.Configuration:
+ configuration = client.Configuration(
+ host=self._cluster_url,
+ api_key_prefix={"authorization": "Bearer"},
+ api_key={"authorization": self._get_token(self.get_credentials())},
+ )
+ configuration.ssl_ca_cert = FileOrData(
+ {
+ "certificate-authority-data": self._ssl_ca_cert,
+ },
+ file_key_name="certificate-authority",
+ ).as_file()
+ return configuration
+
+ @staticmethod
+ def _get_token(creds: google.auth.credentials.Credentials) -> str:
+ if creds.token is None or creds.expired:
+ auth_req = google_requests.Request()
+ creds.refresh(auth_req)
+ return creds.token
+
+ def get_pod(self, name: str, namespace: str) -> V1Pod:
+ """
+ Gets pod's object.
+
+ :param name: Name of the pod.
+ :param namespace: Name of the pod's namespace.
+ """
+ return self.core_v1_client.read_namespaced_pod(
+ name=name,
+ namespace=namespace,
+ )
+
+
+class AsyncGKEPodHook(GoogleBaseAsyncHook):
+ """
+ Hook for managing Google Kubernetes Engine pods APIs in asynchronous way.
+
+ :param cluster_url: The URL pointed to the cluster.
+ :param ssl_ca_cert: SSL certificate that is used for authentication to the
pod.
+ """
+
+ sync_hook_class = GKEPodHook
+ scopes = ["https://www.googleapis.com/auth/cloud-platform";]
+
+ def __init__(
+ self,
+ cluster_url: str,
+ ssl_ca_cert: str,
+ **kwargs,
+ ):
+
+ self._cluster_url = cluster_url
+ self._ssl_ca_cert = ssl_ca_cert
+
+ kwargs.update(
+ cluster_url=cluster_url,
+ ssl_ca_cert=ssl_ca_cert,
+ )
+ super().__init__(**kwargs)
+
+ @contextlib.asynccontextmanager
+ async def get_conn(self, token: Token) -> async_client.ApiClient: # type:
ignore[override]
Review Comment:
why did we add `# type: ignore[override]` here
##########
tests/system/providers/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py:
##########
@@ -68,9 +68,11 @@
in_cluster=False,
is_delete_operator_pod=True,
get_logs=True,
+ deferrable=True,
)
- pod_task_xcom = GKEStartPodOperator(
+ # [START howto_operator_gke_start_pod_xcom_async]
+ pod_task_xcom_async = GKEStartPodOperator(
task_id="pod_task_xcom",
Review Comment:
```suggestion
task_id="pod_task_xcom_async",
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
