VladaZakharova commented on code in PR #29266:
URL: https://github.com/apache/airflow/pull/29266#discussion_r1143059727
##########
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:
This was done to pass pre-commit check, since signature of this method is
not the same as in method in super class BaseHook
--
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]
