Skip to content

Commit 28126c1

Browse files
MrGeorgeOwlMrGeorgeOwl
authored
Add defer mode to GKECreateClusterOperator and GKEDeleteClusterOperator (#28406)
* Add defer mode to GKECreateClusterOperator and GKEDeleteClusterOperator
1 parent 47edfe9 commit 28126c1

8 files changed

Lines changed: 767 additions & 51 deletions

File tree

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

Lines changed: 83 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -30,19 +30,23 @@
3030
import warnings
3131
from typing import Sequence
3232

33-
from google.api_core.exceptions import AlreadyExists, NotFound
33+
from google.api_core.exceptions import NotFound
3434
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
3535
from google.api_core.retry import Retry
3636

3737
# not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable
3838
from google.cloud import container_v1, exceptions # type: ignore[attr-defined]
39-
from google.cloud.container_v1 import ClusterManagerClient
39+
from google.cloud.container_v1 import ClusterManagerAsyncClient, ClusterManagerClient
4040
from google.cloud.container_v1.types import Cluster, Operation
4141

4242
from airflow import version
4343
from airflow.exceptions import AirflowException
4444
from airflow.providers.google.common.consts import CLIENT_INFO
45-
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
45+
from airflow.providers.google.common.hooks.base_google import (
46+
PROVIDE_PROJECT_ID,
47+
GoogleBaseAsyncHook,
48+
GoogleBaseHook,
49+
)
4650

4751
OPERATIONAL_POLL_INTERVAL = 15
4852

@@ -156,9 +160,10 @@ def delete_cluster(
156160
self,
157161
name: str,
158162
project_id: str = PROVIDE_PROJECT_ID,
163+
wait_to_complete: bool = True,
159164
retry: Retry | _MethodDefault = DEFAULT,
160165
timeout: float | None = None,
161-
) -> str | None:
166+
) -> Operation | None:
162167
"""
163168
Deletes the cluster, including the Kubernetes endpoint and all
164169
worker nodes. Firewalls and routes that were configured during
@@ -169,6 +174,8 @@ def delete_cluster(
169174
170175
:param name: The name of the cluster to delete
171176
:param project_id: Google Cloud project ID
177+
:param wait_to_complete: A boolean value which makes method to sleep while
178+
operation of deletion is not finished.
172179
:param retry: Retry object used to determine when/if to retry requests.
173180
If None is specified, requests will not be retried.
174181
:param timeout: The amount of time, in seconds, to wait for the request to
@@ -179,26 +186,28 @@ def delete_cluster(
179186
self.log.info("Deleting (project_id=%s, location=%s, cluster_id=%s)", project_id, self.location, name)
180187

181188
try:
182-
resource = self.get_cluster_manager_client().delete_cluster(
189+
operation = self.get_cluster_manager_client().delete_cluster(
183190
name=f"projects/{project_id}/locations/{self.location}/clusters/{name}",
184191
retry=retry,
185192
timeout=timeout,
186193
)
187-
resource = self.wait_for_operation(resource, project_id)
194+
if wait_to_complete:
195+
operation = self.wait_for_operation(operation, project_id)
188196
# Returns server-defined url for the resource
189-
return resource.self_link
197+
return operation
190198
except NotFound as error:
191199
self.log.info("Assuming Success: %s", error.message)
192200
return None
193201

194202
@GoogleBaseHook.fallback_to_default_project_id
195203
def create_cluster(
196204
self,
197-
cluster: dict | Cluster | None,
205+
cluster: dict | Cluster,
198206
project_id: str = PROVIDE_PROJECT_ID,
207+
wait_to_complete: bool = True,
199208
retry: Retry | _MethodDefault = DEFAULT,
200209
timeout: float | None = None,
201-
) -> str:
210+
) -> Operation | Cluster:
202211
"""
203212
Creates a cluster, consisting of the specified number and type of Google Compute
204213
Engine instances.
@@ -207,6 +216,8 @@ def create_cluster(
207216
be of the same form as the protobuf message
208217
:class:`google.cloud.container_v1.types.Cluster`
209218
:param project_id: Google Cloud project ID
219+
:param wait_to_complete: A boolean value which makes method to sleep while
220+
operation of creation is not finished.
210221
:param retry: A retry object (``google.api_core.retry.Retry``) used to
211222
retry requests.
212223
If None is specified, requests will not be retried.
@@ -231,19 +242,17 @@ def create_cluster(
231242
self.location,
232243
cluster.name, # type: ignore
233244
)
234-
try:
235-
resource = self.get_cluster_manager_client().create_cluster(
236-
parent=f"projects/{project_id}/locations/{self.location}",
237-
cluster=cluster, # type: ignore
238-
retry=retry,
239-
timeout=timeout,
240-
)
241-
resource = self.wait_for_operation(resource, project_id)
245+
operation = self.get_cluster_manager_client().create_cluster(
246+
parent=f"projects/{project_id}/locations/{self.location}",
247+
cluster=cluster, # type: ignore
248+
retry=retry,
249+
timeout=timeout,
250+
)
242251

243-
return resource.target_link
244-
except AlreadyExists as error:
245-
self.log.info("Assuming Success: %s", error.message)
246-
return self.get_cluster(name=cluster.name, project_id=project_id) # type: ignore
252+
if wait_to_complete:
253+
operation = self.wait_for_operation(operation, project_id)
254+
255+
return operation
247256

248257
@GoogleBaseHook.fallback_to_default_project_id
249258
def get_cluster(
@@ -272,12 +281,58 @@ def get_cluster(
272281
name,
273282
)
274283

275-
return (
276-
self.get_cluster_manager_client()
277-
.get_cluster(
278-
name=f"projects/{project_id}/locations/{self.location}/clusters/{name}",
279-
retry=retry,
280-
timeout=timeout,
284+
return self.get_cluster_manager_client().get_cluster(
285+
name=f"projects/{project_id}/locations/{self.location}/clusters/{name}",
286+
retry=retry,
287+
timeout=timeout,
288+
)
289+
290+
291+
class AsyncGKEHook(GoogleBaseAsyncHook):
292+
"""Hook implemented with usage of asynchronous client of GKE."""
293+
294+
sync_hook_class = GKEHook
295+
296+
def __init__(
297+
self,
298+
gcp_conn_id: str = "google_cloud_default",
299+
delegate_to: str | None = None,
300+
location: str | None = None,
301+
impersonation_chain: str | Sequence[str] | None = None,
302+
) -> None:
303+
super().__init__(
304+
gcp_conn_id=gcp_conn_id,
305+
delegate_to=delegate_to,
306+
impersonation_chain=impersonation_chain,
307+
)
308+
self._client: ClusterManagerAsyncClient | None = None
309+
self.location = location
310+
311+
async def _get_client(self) -> ClusterManagerAsyncClient:
312+
if self._client is None:
313+
self._client = ClusterManagerAsyncClient(
314+
credentials=(await self.get_sync_hook()).get_credentials(),
315+
client_info=CLIENT_INFO,
281316
)
282-
.self_link
317+
return self._client
318+
319+
@GoogleBaseHook.fallback_to_default_project_id
320+
async def get_operation(
321+
self,
322+
operation_name: str,
323+
project_id: str = PROVIDE_PROJECT_ID,
324+
) -> Operation:
325+
"""
326+
Fetches the operation from Google Cloud.
327+
328+
:param operation_name: Name of operation to fetch.
329+
:param project_id: Google Cloud project ID.
330+
:return: The new, updated operation from Google Cloud.
331+
"""
332+
project_id = project_id or (await self.get_sync_hook()).project_id
333+
334+
operation_path = f"projects/{project_id}/locations/{self.location}/operations/{operation_name}"
335+
client = await self._get_client()
336+
return await client.get_operation(
337+
name=operation_path,
283338
)

0 commit comments

Comments
 (0)