apache-airflow-providers-google 11.0.0__py3-none-any.whl → 12.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. airflow/providers/google/__init__.py +3 -3
  2. airflow/providers/google/assets/gcs.py +1 -7
  3. airflow/providers/google/cloud/hooks/alloy_db.py +289 -0
  4. airflow/providers/google/cloud/hooks/cloud_batch.py +13 -5
  5. airflow/providers/google/cloud/hooks/dataproc.py +7 -3
  6. airflow/providers/google/cloud/hooks/dataproc_metastore.py +41 -22
  7. airflow/providers/google/cloud/hooks/kubernetes_engine.py +7 -38
  8. airflow/providers/google/cloud/hooks/translate.py +355 -0
  9. airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +147 -0
  10. airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +10 -0
  11. airflow/providers/google/cloud/links/alloy_db.py +55 -0
  12. airflow/providers/google/cloud/links/translate.py +98 -0
  13. airflow/providers/google/cloud/log/stackdriver_task_handler.py +1 -5
  14. airflow/providers/google/cloud/openlineage/mixins.py +4 -12
  15. airflow/providers/google/cloud/openlineage/utils.py +200 -22
  16. airflow/providers/google/cloud/operators/alloy_db.py +459 -0
  17. airflow/providers/google/cloud/operators/automl.py +55 -44
  18. airflow/providers/google/cloud/operators/bigquery.py +60 -15
  19. airflow/providers/google/cloud/operators/dataproc.py +12 -0
  20. airflow/providers/google/cloud/operators/gcs.py +5 -14
  21. airflow/providers/google/cloud/operators/kubernetes_engine.py +377 -705
  22. airflow/providers/google/cloud/operators/mlengine.py +41 -31
  23. airflow/providers/google/cloud/operators/translate.py +586 -1
  24. airflow/providers/google/cloud/operators/vertex_ai/feature_store.py +163 -0
  25. airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +5 -0
  26. airflow/providers/google/cloud/sensors/dataproc.py +2 -2
  27. airflow/providers/google/cloud/sensors/vertex_ai/__init__.py +16 -0
  28. airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +112 -0
  29. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +6 -11
  30. airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +3 -0
  31. airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +3 -0
  32. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +5 -10
  33. airflow/providers/google/cloud/transfers/gcs_to_gcs.py +3 -15
  34. airflow/providers/google/cloud/transfers/gcs_to_local.py +9 -0
  35. airflow/providers/google/cloud/transfers/local_to_gcs.py +41 -6
  36. airflow/providers/google/cloud/transfers/s3_to_gcs.py +15 -0
  37. airflow/providers/google/get_provider_info.py +30 -18
  38. airflow/providers/google/version_compat.py +36 -0
  39. {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/METADATA +16 -18
  40. {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/RECORD +42 -37
  41. airflow/providers/google/cloud/hooks/datapipeline.py +0 -71
  42. airflow/providers/google/cloud/openlineage/BigQueryErrorRunFacet.json +0 -30
  43. airflow/providers/google/cloud/operators/datapipeline.py +0 -63
  44. {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/WHEEL +0 -0
  45. {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/entry_points.txt +0 -0
@@ -24,17 +24,17 @@ from collections.abc import Sequence
24
24
  from functools import cached_property
25
25
  from typing import TYPE_CHECKING, Any
26
26
 
27
- import requests
28
- import yaml
29
27
  from google.api_core.exceptions import AlreadyExists
30
- from google.cloud.container_v1.types import Cluster
31
28
  from kubernetes.client import V1JobList, models as k8s
32
- from kubernetes.utils.create_from_yaml import FailToCreateError
33
29
  from packaging.version import parse as parse_version
34
30
 
35
31
  from airflow.configuration import conf
36
32
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
37
33
  from airflow.providers.cncf.kubernetes.operators.job import KubernetesJobOperator
34
+ from airflow.providers.cncf.kubernetes.operators.kueue import (
35
+ KubernetesInstallKueueOperator,
36
+ KubernetesStartKueueJobOperator,
37
+ )
38
38
  from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
39
39
  from airflow.providers.cncf.kubernetes.operators.resource import (
40
40
  KubernetesCreateResourceOperator,
@@ -57,6 +57,7 @@ from airflow.providers.google.cloud.triggers.kubernetes_engine import (
57
57
  GKEOperationTrigger,
58
58
  GKEStartPodTrigger,
59
59
  )
60
+ from airflow.providers.google.common.deprecated import deprecated
60
61
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
61
62
  from airflow.providers_manager import ProvidersManager
62
63
  from airflow.utils.timezone import utcnow
@@ -72,7 +73,8 @@ except ImportError:
72
73
  )
73
74
 
74
75
  if TYPE_CHECKING:
75
- from kubernetes.client.models import V1Job, V1Pod
76
+ from google.cloud.container_v1.types import Cluster
77
+ from kubernetes.client.models import V1Job
76
78
  from pendulum import DateTime
77
79
 
78
80
  from airflow.utils.context import Context
@@ -92,17 +94,17 @@ class GKEClusterAuthDetails:
92
94
 
93
95
  def __init__(
94
96
  self,
95
- cluster_name,
96
- project_id,
97
- use_internal_ip,
98
- cluster_hook,
97
+ cluster_name: str,
98
+ project_id: str,
99
+ use_internal_ip: bool,
100
+ cluster_hook: GKEHook,
99
101
  ):
100
102
  self.cluster_name = cluster_name
101
103
  self.project_id = project_id
102
104
  self.use_internal_ip = use_internal_ip
103
105
  self.cluster_hook = cluster_hook
104
- self._cluster_url = None
105
- self._ssl_ca_cert = None
106
+ self._cluster_url: str
107
+ self._ssl_ca_cert: str
106
108
 
107
109
  def fetch_cluster_info(self) -> tuple[str, str]:
108
110
  """Fetch cluster info for connecting to it."""
@@ -119,11 +121,63 @@ class GKEClusterAuthDetails:
119
121
  return self._cluster_url, self._ssl_ca_cert
120
122
 
121
123
 
122
- class GKEDeleteClusterOperator(GoogleCloudBaseOperator):
124
+ class GKEOperatorMixin:
125
+ """Mixin for GKE operators provides proper unified hooks creation."""
126
+
127
+ enable_tcp_keepalive = False
128
+
129
+ template_fields: Sequence[str] = (
130
+ "location",
131
+ "cluster_name",
132
+ "use_internal_ip",
133
+ "project_id",
134
+ "gcp_conn_id",
135
+ "impersonation_chain",
136
+ )
137
+
138
+ @cached_property
139
+ def cluster_hook(self) -> GKEHook:
140
+ return GKEHook(
141
+ gcp_conn_id=self.gcp_conn_id, # type: ignore[attr-defined]
142
+ location=self.location, # type: ignore[attr-defined]
143
+ impersonation_chain=self.impersonation_chain, # type: ignore[attr-defined]
144
+ )
145
+
146
+ @cached_property
147
+ def hook(self) -> GKEKubernetesHook:
148
+ return GKEKubernetesHook(
149
+ gcp_conn_id=self.gcp_conn_id, # type: ignore[attr-defined]
150
+ impersonation_chain=self.impersonation_chain, # type: ignore[attr-defined]
151
+ cluster_url=self.cluster_url,
152
+ ssl_ca_cert=self.ssl_ca_cert,
153
+ enable_tcp_keepalive=self.enable_tcp_keepalive,
154
+ )
155
+
156
+ @cached_property
157
+ def cluster_info(self) -> tuple[str, str]:
158
+ """Fetch cluster info for connecting to it."""
159
+ auth_details = GKEClusterAuthDetails(
160
+ cluster_name=self.cluster_name, # type: ignore[attr-defined]
161
+ project_id=self.project_id, # type: ignore[attr-defined]
162
+ use_internal_ip=self.use_internal_ip, # type: ignore[attr-defined]
163
+ cluster_hook=self.cluster_hook,
164
+ )
165
+ return auth_details.fetch_cluster_info()
166
+
167
+ @property
168
+ def cluster_url(self) -> str:
169
+ return self.cluster_info[0]
170
+
171
+ @property
172
+ def ssl_ca_cert(self) -> str:
173
+ return self.cluster_info[1]
174
+
175
+
176
+ class GKEDeleteClusterOperator(GKEOperatorMixin, GoogleCloudBaseOperator):
123
177
  """
124
178
  Deletes the cluster, including the Kubernetes endpoint and all worker nodes.
125
179
 
126
- To delete a certain cluster, you must specify the ``project_id``, the ``name``
180
+ To delete a certain cluster, you must specify the ``project_id``, the ``cluster_name``
127
181
  of the cluster, the ``location`` that the cluster is in, and the ``task_id``.
128
182
 
129
183
  **Operator Creation**: ::
@@ -132,7 +186,7 @@ class GKEDeleteClusterOperator(GoogleCloudBaseOperator):
132
186
  task_id='cluster_delete',
133
187
  project_id='my-project',
134
188
  location='cluster-location'
135
- name='cluster-name')
189
+ cluster_name='cluster-name')
136
190
 
137
191
  .. seealso::
138
192
  For more detail about deleting clusters have a look at the reference:
@@ -142,71 +196,87 @@ class GKEDeleteClusterOperator(GoogleCloudBaseOperator):
142
196
  For more information on how to use this operator, take a look at the guide:
143
197
  :ref:`howto/operator:GKEDeleteClusterOperator`
144
198
 
145
- :param project_id: The Google Developers Console [project ID or project number]
146
- :param name: The name of the resource to delete, in this case cluster name
147
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
148
- resides.
149
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
150
- :param api_version: The api version to use
199
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
200
+ cluster resides, e.g. 'us-central1-a'
201
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
202
+ :param use_internal_ip: Use the internal IP address as the endpoint.
203
+ :param project_id: The Google Developers Console project id
204
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
205
+ users to specify a service account.
151
206
  :param impersonation_chain: Optional service account to impersonate using short-term
152
- credentials, or chained list of accounts required to get the access_token
207
+ credentials, or list of accounts required to get the access_token
153
208
  of the last account in the list, which will be impersonated in the request.
154
209
  If set as a string, the account must grant the originating account
155
210
  the Service Account Token Creator IAM role.
156
211
  If set as a sequence, the identities from the list must grant
157
212
  Service Account Token Creator IAM role to the directly preceding identity, with first
158
213
  account from the list granting this role to the originating account (templated).
214
+ :param name: (Deprecated) The name of the resource to delete, in this case cluster name
215
+ :param api_version: The api version to use
159
216
  :param deferrable: Run operator in the deferrable mode.
160
217
  :param poll_interval: Interval size which defines how often operation status is checked.
161
218
  """
162
219
 
163
- template_fields: Sequence[str] = (
164
- "project_id",
165
- "gcp_conn_id",
166
- "name",
167
- "location",
168
- "api_version",
169
- "impersonation_chain",
220
+ template_fields: Sequence[str] = tuple(
221
+ {"api_version", "deferrable", "poll_interval"} | set(GKEOperatorMixin.template_fields)
170
222
  )
171
223
 
172
224
  def __init__(
173
225
  self,
174
- *,
175
- name: str,
176
226
  location: str,
227
+ use_internal_ip: bool = False,
177
228
  project_id: str = PROVIDE_PROJECT_ID,
178
229
  gcp_conn_id: str = "google_cloud_default",
179
- api_version: str = "v2",
180
230
  impersonation_chain: str | Sequence[str] | None = None,
231
+ cluster_name: str | None = None,
232
+ name: str | None = None,
233
+ api_version: str = "v2",
181
234
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
182
235
  poll_interval: int = 10,
236
+ *args,
183
237
  **kwargs,
184
238
  ) -> None:
185
- super().__init__(**kwargs)
239
+ super().__init__(*args, **kwargs)
186
240
 
241
+ self.location = location
242
+ self.cluster_name = cluster_name or name
243
+ self.use_internal_ip = use_internal_ip
187
244
  self.project_id = project_id
188
245
  self.gcp_conn_id = gcp_conn_id
189
- self.location = location
190
- self.api_version = api_version
191
- self.name = name
192
246
  self.impersonation_chain = impersonation_chain
247
+ self._name = name
248
+ self.api_version = api_version
193
249
  self.deferrable = deferrable
194
250
  self.poll_interval = poll_interval
195
251
  self._check_input()
196
252
 
197
- self._hook: GKEHook | None = None
253
+ @property
254
+ @deprecated(
255
+ planned_removal_date="May 01, 2025",
256
+ use_instead="cluster_name",
257
+ category=AirflowProviderDeprecationWarning,
258
+ )
259
+ def name(self) -> str | None:
260
+ return self._name
261
+
262
+ @name.setter
263
+ @deprecated(
264
+ planned_removal_date="May 01, 2025",
265
+ use_instead="cluster_name",
266
+ category=AirflowProviderDeprecationWarning,
267
+ )
268
+ def name(self, name: str) -> None:
269
+ self._name = name
198
270
 
199
271
  def _check_input(self) -> None:
200
- if not all([self.project_id, self.name, self.location]):
201
- self.log.error("One of (project_id, name, location) is missing or incorrect")
272
+ if not all([self.project_id, self.cluster_name, self.location]):
273
+ self.log.error("One of (project_id, cluster_name, location) is missing or incorrect")
202
274
  raise AirflowException("Operator has incorrect or missing input.")
203
275
 
204
276
  def execute(self, context: Context) -> str | None:
205
- hook = self._get_hook()
206
-
207
277
  wait_to_complete = not self.deferrable
208
- operation = hook.delete_cluster(
209
- name=self.name,
278
+ operation = self.cluster_hook.delete_cluster(
279
+ name=self.cluster_name,
210
280
  project_id=self.project_id,
211
281
  wait_to_complete=wait_to_complete,
212
282
  )
@@ -236,23 +306,13 @@ class GKEDeleteClusterOperator(GoogleCloudBaseOperator):
236
306
  raise AirflowException(message)
237
307
 
238
308
  self.log.info(message)
239
- operation = self._get_hook().get_operation(
309
+ operation = self.cluster_hook.get_operation(
240
310
  operation_name=event["operation_name"],
241
311
  )
242
312
  return operation.self_link
243
313
 
244
- def _get_hook(self) -> GKEHook:
245
- if self._hook is None:
246
- self._hook = GKEHook(
247
- gcp_conn_id=self.gcp_conn_id,
248
- location=self.location,
249
- impersonation_chain=self.impersonation_chain,
250
- )
251
-
252
- return self._hook
253
314
 
254
-
255
- class GKECreateClusterOperator(GoogleCloudBaseOperator):
315
+ class GKECreateClusterOperator(GKEOperatorMixin, GoogleCloudBaseOperator):
256
316
  """
257
317
  Create a Google Kubernetes Engine Cluster of specified dimensions and wait until the cluster is created.
258
318
 
@@ -284,61 +344,58 @@ class GKECreateClusterOperator(GoogleCloudBaseOperator):
284
344
  For more information on how to use this operator, take a look at the guide:
285
345
  :ref:`howto/operator:GKECreateClusterOperator`
286
346
 
287
- :param project_id: The Google Developers Console [project ID or project number]
288
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
289
- resides.
290
- :param body: The Cluster definition to create, can be protobuf or python dict, if
291
- dict it must match protobuf message Cluster
292
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
293
- :param api_version: The api version to use
347
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
348
+ cluster resides, e.g. 'us-central1-a'
349
+ :param use_internal_ip: Use the internal IP address as the endpoint.
350
+ :param project_id: The Google Developers Console project id
351
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
352
+ users to specify a service account.
294
353
  :param impersonation_chain: Optional service account to impersonate using short-term
295
- credentials, or chained list of accounts required to get the access_token
354
+ credentials, or list of accounts required to get the access_token
296
355
  of the last account in the list, which will be impersonated in the request.
297
356
  If set as a string, the account must grant the originating account
298
357
  the Service Account Token Creator IAM role.
299
358
  If set as a sequence, the identities from the list must grant
300
359
  Service Account Token Creator IAM role to the directly preceding identity, with first
301
360
  account from the list granting this role to the originating account (templated).
361
+ :param body: The Cluster definition to create, can be protobuf or python dict, if
362
+ dict it must match protobuf message Cluster
363
+ :param api_version: The api version to use
302
364
  :param deferrable: Run operator in the deferrable mode.
303
365
  :param poll_interval: Interval size which defines how often operation status is checked.
304
366
  """
305
367
 
306
- template_fields: Sequence[str] = (
307
- "project_id",
308
- "gcp_conn_id",
309
- "location",
310
- "api_version",
311
- "body",
312
- "impersonation_chain",
368
+ template_fields: Sequence[str] = tuple(
369
+ {"body", "api_version", "deferrable", "poll_interval"} | set(GKEOperatorMixin.template_fields)
313
370
  )
314
371
  operator_extra_links = (KubernetesEngineClusterLink(),)
315
372
 
316
373
  def __init__(
317
374
  self,
318
- *,
319
- location: str,
320
375
  body: dict | Cluster,
376
+ location: str,
377
+ use_internal_ip: bool = False,
321
378
  project_id: str = PROVIDE_PROJECT_ID,
322
379
  gcp_conn_id: str = "google_cloud_default",
323
- api_version: str = "v2",
324
380
  impersonation_chain: str | Sequence[str] | None = None,
325
- poll_interval: int = 10,
381
+ api_version: str = "v2",
326
382
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
383
+ poll_interval: int = 10,
384
+ *args,
327
385
  **kwargs,
328
386
  ) -> None:
329
- super().__init__(**kwargs)
330
-
387
+ self.body = body
388
+ self.location = location
389
+ self.use_internal_ip = use_internal_ip
390
+ self.cluster_name = body.get("name") if isinstance(body, dict) else getattr(body, "name", None)
331
391
  self.project_id = project_id
332
392
  self.gcp_conn_id = gcp_conn_id
333
- self.location = location
334
- self.api_version = api_version
335
- self.body = body
336
393
  self.impersonation_chain = impersonation_chain
394
+ self.api_version = api_version
337
395
  self.poll_interval = poll_interval
338
396
  self.deferrable = deferrable
339
397
  self._validate_input()
340
-
341
- self._hook: GKEHook | None = None
398
+ super().__init__(*args, **kwargs)
342
399
 
343
400
  def _validate_input(self) -> None:
344
401
  """Primary validation of the input body."""
@@ -404,36 +461,32 @@ class GKECreateClusterOperator(GoogleCloudBaseOperator):
404
461
  )
405
462
 
406
463
  def execute(self, context: Context) -> str:
407
- hook = self._get_hook()
464
+ KubernetesEngineClusterLink.persist(context=context, task_instance=self, cluster=self.body)
465
+
408
466
  try:
409
- wait_to_complete = not self.deferrable
410
- operation = hook.create_cluster(
467
+ operation = self.cluster_hook.create_cluster(
411
468
  cluster=self.body,
412
469
  project_id=self.project_id,
413
- wait_to_complete=wait_to_complete,
470
+ wait_to_complete=not self.deferrable,
414
471
  )
415
-
416
- KubernetesEngineClusterLink.persist(context=context, task_instance=self, cluster=self.body)
417
-
418
- if self.deferrable:
419
- self.defer(
420
- trigger=GKEOperationTrigger(
421
- operation_name=operation.name,
422
- project_id=self.project_id,
423
- location=self.location,
424
- gcp_conn_id=self.gcp_conn_id,
425
- impersonation_chain=self.impersonation_chain,
426
- poll_interval=self.poll_interval,
427
- ),
428
- method_name="execute_complete",
429
- )
430
-
431
- return operation.target_link
432
-
433
472
  except AlreadyExists as error:
434
473
  self.log.info("Assuming Success: %s", error.message)
435
- name = self.body.name if isinstance(self.body, Cluster) else self.body["name"]
436
- return hook.get_cluster(name=name, project_id=self.project_id).self_link
474
+ return self.cluster_hook.get_cluster(name=self.cluster_name, project_id=self.project_id).self_link
475
+
476
+ if self.deferrable:
477
+ self.defer(
478
+ trigger=GKEOperationTrigger(
479
+ operation_name=operation.name,
480
+ project_id=self.project_id,
481
+ location=self.location,
482
+ gcp_conn_id=self.gcp_conn_id,
483
+ impersonation_chain=self.impersonation_chain,
484
+ poll_interval=self.poll_interval,
485
+ ),
486
+ method_name="execute_complete",
487
+ )
488
+
489
+ return operation.target_link
437
490
 
438
491
  def execute_complete(self, context: Context, event: dict) -> str:
439
492
  status = event["status"]
@@ -444,23 +497,13 @@ class GKECreateClusterOperator(GoogleCloudBaseOperator):
444
497
  raise AirflowException(message)
445
498
 
446
499
  self.log.info(message)
447
- operation = self._get_hook().get_operation(
500
+ operation = self.cluster_hook.get_operation(
448
501
  operation_name=event["operation_name"],
449
502
  )
450
503
  return operation.target_link
451
504
 
452
- def _get_hook(self) -> GKEHook:
453
- if self._hook is None:
454
- self._hook = GKEHook(
455
- gcp_conn_id=self.gcp_conn_id,
456
- location=self.location,
457
- impersonation_chain=self.impersonation_chain,
458
- )
459
-
460
- return self._hook
461
505
 
462
-
463
- class GKEStartKueueInsideClusterOperator(GoogleCloudBaseOperator):
506
+ class GKEStartKueueInsideClusterOperator(GKEOperatorMixin, KubernetesInstallKueueOperator):
464
507
  """
465
508
  Installs Kueue of specific version inside Cluster.
466
509
 
@@ -472,13 +515,15 @@ class GKEStartKueueInsideClusterOperator(GoogleCloudBaseOperator):
472
515
  For more details about Kueue have a look at the reference:
473
516
  https://kueue.sigs.k8s.io/docs/overview/
474
517
 
475
- :param project_id: The Google Developers Console [project ID or project number].
476
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster resides.
477
- :param cluster_name: The Cluster name in which to install Kueue.
478
- :param kueue_version: Version of Kueue to install.
479
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
518
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
519
+ cluster resides, e.g. 'us-central1-a'
520
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
521
+ :param use_internal_ip: Use the internal IP address as the endpoint.
522
+ :param project_id: The Google Developers Console project id
523
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
524
+ users to specify a service account.
480
525
  :param impersonation_chain: Optional service account to impersonate using short-term
481
- credentials, or chained list of accounts required to get the access_token
526
+ credentials, or list of accounts required to get the access_token
482
527
  of the last account in the list, which will be impersonated in the request.
483
528
  If set as a string, the account must grant the originating account
484
529
  the Service Account Token Creator IAM role.
@@ -487,121 +532,44 @@ class GKEStartKueueInsideClusterOperator(GoogleCloudBaseOperator):
487
532
  account from the list granting this role to the originating account (templated).
488
533
  """
489
534
 
490
- template_fields: Sequence[str] = (
491
- "project_id",
492
- "location",
493
- "kueue_version",
494
- "cluster_name",
495
- "gcp_conn_id",
496
- "impersonation_chain",
535
+ enable_tcp_keepalive = True
536
+ template_fields = tuple(
537
+ set(GKEOperatorMixin.template_fields) | set(KubernetesInstallKueueOperator.template_fields)
497
538
  )
498
539
  operator_extra_links = (KubernetesEngineClusterLink(),)
499
540
 
500
541
  def __init__(
501
542
  self,
502
- *,
503
543
  location: str,
504
544
  cluster_name: str,
505
- kueue_version: str,
506
545
  use_internal_ip: bool = False,
507
546
  project_id: str = PROVIDE_PROJECT_ID,
508
547
  gcp_conn_id: str = "google_cloud_default",
509
548
  impersonation_chain: str | Sequence[str] | None = None,
549
+ *args,
510
550
  **kwargs,
511
551
  ) -> None:
512
- super().__init__(**kwargs)
552
+ super().__init__(*args, **kwargs)
513
553
  self.project_id = project_id
514
554
  self.location = location
515
555
  self.cluster_name = cluster_name
516
- self.kueue_version = kueue_version
517
556
  self.gcp_conn_id = gcp_conn_id
518
- self.impersonation_chain = impersonation_chain
519
557
  self.use_internal_ip = use_internal_ip
520
- self._kueue_yaml_url = (
521
- f"https://github.com/kubernetes-sigs/kueue/releases/download/{self.kueue_version}/manifests.yaml"
522
- )
523
-
524
- @cached_property
525
- def cluster_hook(self) -> GKEHook:
526
- return GKEHook(
527
- gcp_conn_id=self.gcp_conn_id,
528
- location=self.location,
529
- impersonation_chain=self.impersonation_chain,
530
- )
531
-
532
- @cached_property
533
- def deployment_hook(self) -> GKEKubernetesHook:
534
- if self._cluster_url is None or self._ssl_ca_cert is None:
535
- raise AttributeError(
536
- "Cluster url and ssl_ca_cert should be defined before using self.deployment_hook method. "
537
- "Try to use self.get_kube_creds method",
538
- )
539
- return GKEKubernetesHook(
540
- gcp_conn_id=self.gcp_conn_id,
541
- impersonation_chain=self.impersonation_chain,
542
- cluster_url=self._cluster_url,
543
- ssl_ca_cert=self._ssl_ca_cert,
544
- )
545
-
546
- @cached_property
547
- def pod_hook(self) -> GKEKubernetesHook:
548
- if self._cluster_url is None or self._ssl_ca_cert is None:
549
- raise AttributeError(
550
- "Cluster url and ssl_ca_cert should be defined before using self.pod_hook method. "
551
- "Try to use self.get_kube_creds method",
552
- )
553
-
554
- return GKEKubernetesHook(
555
- gcp_conn_id=self.gcp_conn_id,
556
- impersonation_chain=self.impersonation_chain,
557
- cluster_url=self._cluster_url,
558
- ssl_ca_cert=self._ssl_ca_cert,
559
- enable_tcp_keepalive=True,
560
- )
561
-
562
- @staticmethod
563
- def _get_yaml_content_from_file(kueue_yaml_url) -> list[dict]:
564
- """Download content of YAML file and separate it into several dictionaries."""
565
- response = requests.get(kueue_yaml_url, allow_redirects=True)
566
- if response.status_code != 200:
567
- raise AirflowException("Was not able to read the yaml file from given URL")
568
-
569
- return list(yaml.safe_load_all(response.text))
558
+ self.impersonation_chain = impersonation_chain
570
559
 
571
560
  def execute(self, context: Context):
572
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
573
- cluster_name=self.cluster_name,
574
- project_id=self.project_id,
575
- use_internal_ip=self.use_internal_ip,
576
- cluster_hook=self.cluster_hook,
577
- ).fetch_cluster_info()
578
-
579
- cluster = self.cluster_hook.get_cluster(
580
- name=self.cluster_name,
581
- project_id=self.project_id,
582
- )
561
+ cluster = self.cluster_hook.get_cluster(name=self.cluster_name, project_id=self.project_id)
583
562
  KubernetesEngineClusterLink.persist(context=context, task_instance=self, cluster=cluster)
584
563
 
585
- yaml_objects = self._get_yaml_content_from_file(kueue_yaml_url=self._kueue_yaml_url)
586
-
587
564
  if self.cluster_hook.check_cluster_autoscaling_ability(cluster=cluster):
588
- try:
589
- self.pod_hook.apply_from_yaml_file(yaml_objects=yaml_objects)
590
-
591
- self.deployment_hook.check_kueue_deployment_running(
592
- name="kueue-controller-manager", namespace="kueue-system"
593
- )
594
-
595
- self.log.info("Kueue installed successfully!")
596
- except FailToCreateError:
597
- self.log.info("Kueue is already enabled for the cluster")
565
+ super().execute(context)
598
566
  else:
599
567
  self.log.info(
600
568
  "Cluster doesn't have ability to autoscale, will not install Kueue inside. Aborting"
601
569
  )
602
570
 
603
571
 
604
- class GKEStartPodOperator(KubernetesPodOperator):
572
+ class GKEStartPodOperator(GKEOperatorMixin, KubernetesPodOperator):
605
573
  """
606
574
  Executes a task in a Kubernetes pod in the specified Google Kubernetes Engine cluster.
607
575
 
@@ -622,8 +590,7 @@ class GKEStartPodOperator(KubernetesPodOperator):
622
590
 
623
591
  :param location: The name of the Google Kubernetes Engine zone or region in which the
624
592
  cluster resides, e.g. 'us-central1-a'
625
- :param cluster_name: The name of the Google Kubernetes Engine cluster the pod
626
- should be spawned in
593
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
627
594
  :param use_internal_ip: Use the internal IP address as the endpoint.
628
595
  :param project_id: The Google Developers Console project id
629
596
  :param gcp_conn_id: The Google cloud connection id to use. This allows for
@@ -636,27 +603,28 @@ class GKEStartPodOperator(KubernetesPodOperator):
636
603
  If set as a sequence, the identities from the list must grant
637
604
  Service Account Token Creator IAM role to the directly preceding identity, with first
638
605
  account from the list granting this role to the originating account (templated).
639
- :param regional: The location param is region name.
640
- :param deferrable: Run operator in the deferrable mode.
606
+ :param regional: (Deprecated) The location param is region name.
641
607
  :param on_finish_action: What to do when the pod reaches its final state, or the execution is interrupted.
642
608
  If "delete_pod", the pod will be deleted regardless its state; if "delete_succeeded_pod",
643
609
  only succeeded pod will be deleted. You can set to "keep_pod" to keep the pod.
644
610
  Current default is `keep_pod`, but this will be changed in the next major release of this provider.
645
- :param is_delete_operator_pod: What to do when the pod reaches its final
611
+ :param is_delete_operator_pod: (Deprecated) What to do when the pod reaches its final
646
612
  state, or the execution is interrupted. If True, delete the
647
613
  pod; if False, leave the pod. Current default is False, but this will be
648
614
  changed in the next major release of this provider.
649
615
  Deprecated - use `on_finish_action` instead.
616
+ :param deferrable: Run operator in the deferrable mode.
650
617
  """
651
618
 
652
619
  template_fields: Sequence[str] = tuple(
653
- {"project_id", "location", "cluster_name"} | set(KubernetesPodOperator.template_fields)
620
+ {"on_finish_action", "deferrable"}
621
+ | (set(KubernetesPodOperator.template_fields) - {"is_delete_operator_pod", "regional"})
622
+ | set(GKEOperatorMixin.template_fields)
654
623
  )
655
624
  operator_extra_links = (KubernetesEnginePodLink(),)
656
625
 
657
626
  def __init__(
658
627
  self,
659
- *,
660
628
  location: str,
661
629
  cluster_name: str,
662
630
  use_internal_ip: bool = False,
@@ -666,141 +634,110 @@ class GKEStartPodOperator(KubernetesPodOperator):
666
634
  regional: bool | None = None,
667
635
  on_finish_action: str | None = None,
668
636
  is_delete_operator_pod: bool | None = None,
637
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
638
+ *args,
669
639
  **kwargs,
670
640
  ) -> None:
671
641
  if is_delete_operator_pod is not None:
672
- warnings.warn(
673
- "`is_delete_operator_pod` parameter is deprecated, please use `on_finish_action`",
674
- AirflowProviderDeprecationWarning,
675
- stacklevel=2,
676
- )
677
642
  kwargs["on_finish_action"] = (
678
643
  OnFinishAction.DELETE_POD if is_delete_operator_pod else OnFinishAction.KEEP_POD
679
644
  )
645
+ elif on_finish_action is not None:
646
+ kwargs["on_finish_action"] = OnFinishAction(on_finish_action)
680
647
  else:
681
- if on_finish_action is not None:
682
- kwargs["on_finish_action"] = OnFinishAction(on_finish_action)
683
- else:
684
- warnings.warn(
685
- f"You have not set parameter `on_finish_action` in class {self.__class__.__name__}. "
686
- "Currently the default for this parameter is `keep_pod` but in a future release"
687
- " the default will be changed to `delete_pod`. To ensure pods are not deleted in"
688
- " the future you will need to set `on_finish_action=keep_pod` explicitly.",
689
- AirflowProviderDeprecationWarning,
690
- stacklevel=2,
691
- )
692
- kwargs["on_finish_action"] = OnFinishAction.KEEP_POD
693
-
694
- if regional is not None:
695
648
  warnings.warn(
696
- f"You have set parameter regional in class {self.__class__.__name__}. "
697
- "In current implementation of the operator the parameter is not used and will "
698
- "be deleted in future.",
649
+ f"You have not set parameter `on_finish_action` in class {self.__class__.__name__}. "
650
+ "Currently the default for this parameter is `keep_pod` but in a future release"
651
+ " the default will be changed to `delete_pod`. To ensure pods are not deleted in"
652
+ " the future you will need to set `on_finish_action=keep_pod` explicitly.",
699
653
  AirflowProviderDeprecationWarning,
700
654
  stacklevel=2,
701
655
  )
656
+ kwargs["on_finish_action"] = OnFinishAction.KEEP_POD
702
657
 
703
- super().__init__(**kwargs)
658
+ super().__init__(*args, **kwargs)
704
659
  self.project_id = project_id
705
660
  self.location = location
706
661
  self.cluster_name = cluster_name
707
662
  self.gcp_conn_id = gcp_conn_id
708
- self.impersonation_chain = impersonation_chain
709
663
  self.use_internal_ip = use_internal_ip
664
+ self.impersonation_chain = impersonation_chain
665
+ self._regional = regional
666
+ if is_delete_operator_pod is not None:
667
+ self.is_delete_operator_pod = is_delete_operator_pod
668
+ self.deferrable = deferrable
710
669
 
711
- self.pod: V1Pod | None = None
712
- self._ssl_ca_cert: str | None = None
713
- self._cluster_url: str | None = None
714
-
715
- if self.gcp_conn_id is None:
716
- raise AirflowException(
717
- "The gcp_conn_id parameter has become required. If you want to use Application Default "
718
- "Credentials (ADC) strategy for authorization, create an empty connection "
719
- "called `google_cloud_default`.",
720
- )
721
670
  # There is no need to manage the kube_config file, as it will be generated automatically.
722
671
  # All Kubernetes parameters (except config_file) are also valid for the GKEStartPodOperator.
723
672
  if self.config_file:
724
673
  raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.")
725
674
 
726
- @cached_property
727
- def cluster_hook(self) -> GKEHook:
728
- return GKEHook(
729
- gcp_conn_id=self.gcp_conn_id,
730
- location=self.location,
731
- impersonation_chain=self.impersonation_chain,
732
- )
733
-
734
- @cached_property
735
- def hook(self) -> GKEKubernetesHook:
736
- if self._cluster_url is None or self._ssl_ca_cert is None:
737
- raise AttributeError(
738
- "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
739
- "Try to use self.get_kube_creds method",
740
- )
741
-
742
- return GKEKubernetesHook(
743
- gcp_conn_id=self.gcp_conn_id,
744
- cluster_url=self._cluster_url,
745
- ssl_ca_cert=self._ssl_ca_cert,
746
- impersonation_chain=self.impersonation_chain,
747
- enable_tcp_keepalive=True,
748
- )
749
-
750
- def execute(self, context: Context):
751
- """Execute process of creating pod and executing provided command inside it."""
752
- self.fetch_cluster_info()
753
- return super().execute(context)
754
-
755
- def fetch_cluster_info(self) -> tuple[str, str | None]:
756
- """Fetch cluster info for connecting to it."""
757
- cluster = self.cluster_hook.get_cluster(
758
- name=self.cluster_name,
759
- project_id=self.project_id,
760
- )
761
-
762
- if not self.use_internal_ip:
763
- self._cluster_url = f"https://{cluster.endpoint}"
764
- else:
765
- self._cluster_url = f"https://{cluster.private_cluster_config.private_endpoint}"
766
- self._ssl_ca_cert = cluster.master_auth.cluster_ca_certificate
767
- return self._cluster_url, self._ssl_ca_cert
675
+ @property
676
+ @deprecated(
677
+ planned_removal_date="May 01, 2025",
678
+ use_instead="on_finish_action",
679
+ category=AirflowProviderDeprecationWarning,
680
+ )
681
+ def is_delete_operator_pod(self) -> bool | None:
682
+ return self._is_delete_operator_pod
683
+
684
+ @is_delete_operator_pod.setter
685
+ @deprecated(
686
+ planned_removal_date="May 01, 2025",
687
+ use_instead="on_finish_action",
688
+ category=AirflowProviderDeprecationWarning,
689
+ )
690
+ def is_delete_operator_pod(self, is_delete_operator_pod) -> None:
691
+ self._is_delete_operator_pod = is_delete_operator_pod
692
+
693
+ @property
694
+ @deprecated(
695
+ planned_removal_date="May 01, 2025",
696
+ reason="The parameter is not in actual use.",
697
+ category=AirflowProviderDeprecationWarning,
698
+ )
699
+ def regional(self) -> bool | None:
700
+ return self._regional
701
+
702
+ @regional.setter
703
+ @deprecated(
704
+ planned_removal_date="May 01, 2025",
705
+ reason="The parameter is not in actual use.",
706
+ category=AirflowProviderDeprecationWarning,
707
+ )
708
+ def regional(self, regional) -> None:
709
+ self._regional = regional
768
710
 
769
711
  def invoke_defer_method(self, last_log_time: DateTime | None = None):
770
712
  """Redefine triggers which are being used in child classes."""
771
713
  trigger_start_time = utcnow()
714
+ on_finish_action = self.on_finish_action
715
+ if type(on_finish_action) is str and self.on_finish_action not in [i.value for i in OnFinishAction]:
716
+ on_finish_action = self.on_finish_action.split(".")[-1].lower() # type: ignore[assignment]
772
717
  self.defer(
773
718
  trigger=GKEStartPodTrigger(
774
719
  pod_name=self.pod.metadata.name, # type: ignore[union-attr]
775
720
  pod_namespace=self.pod.metadata.namespace, # type: ignore[union-attr]
776
721
  trigger_start_time=trigger_start_time,
777
- cluster_url=self._cluster_url, # type: ignore[arg-type]
778
- ssl_ca_cert=self._ssl_ca_cert, # type: ignore[arg-type]
722
+ cluster_url=self.cluster_url,
723
+ ssl_ca_cert=self.ssl_ca_cert,
779
724
  get_logs=self.get_logs,
780
725
  startup_timeout=self.startup_timeout_seconds,
781
726
  cluster_context=self.cluster_context,
782
727
  poll_interval=self.poll_interval,
783
728
  in_cluster=self.in_cluster,
784
729
  base_container_name=self.base_container_name,
785
- on_finish_action=self.on_finish_action,
730
+ on_finish_action=on_finish_action,
786
731
  gcp_conn_id=self.gcp_conn_id,
787
732
  impersonation_chain=self.impersonation_chain,
788
733
  logging_interval=self.logging_interval,
789
734
  last_log_time=last_log_time,
790
735
  ),
791
- method_name="execute_complete",
792
- kwargs={"cluster_url": self._cluster_url, "ssl_ca_cert": self._ssl_ca_cert},
736
+ method_name="trigger_reentry",
793
737
  )
794
738
 
795
- def execute_complete(self, context: Context, event: dict, **kwargs):
796
- # It is required for hook to be initialized
797
- self._cluster_url = kwargs["cluster_url"]
798
- self._ssl_ca_cert = kwargs["ssl_ca_cert"]
799
-
800
- return super().trigger_reentry(context, event)
801
739
 
802
-
803
- class GKEStartJobOperator(KubernetesJobOperator):
740
+ class GKEStartJobOperator(GKEOperatorMixin, KubernetesJobOperator):
804
741
  """
805
742
  Executes a Kubernetes job in the specified Google Kubernetes Engine cluster.
806
743
 
@@ -821,7 +758,7 @@ class GKEStartJobOperator(KubernetesJobOperator):
821
758
 
822
759
  :param location: The name of the Google Kubernetes Engine zone or region in which the
823
760
  cluster resides, e.g. 'us-central1-a'
824
- :param cluster_name: The name of the Google Kubernetes Engine cluster
761
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
825
762
  :param use_internal_ip: Use the internal IP address as the endpoint.
826
763
  :param project_id: The Google Developers Console project id
827
764
  :param gcp_conn_id: The Google cloud connection id to use. This allows for
@@ -834,20 +771,20 @@ class GKEStartJobOperator(KubernetesJobOperator):
834
771
  If set as a sequence, the identities from the list must grant
835
772
  Service Account Token Creator IAM role to the directly preceding identity, with first
836
773
  account from the list granting this role to the originating account (templated).
837
- :param location: The location param is region name.
838
774
  :param deferrable: Run operator in the deferrable mode.
839
775
  :param poll_interval: (Deferrable mode only) polling period in seconds to
840
776
  check for the status of job.
841
777
  """
842
778
 
843
779
  template_fields: Sequence[str] = tuple(
844
- {"project_id", "location", "cluster_name"} | set(KubernetesJobOperator.template_fields)
780
+ {"deferrable", "poll_interval"}
781
+ | set(GKEOperatorMixin.template_fields)
782
+ | set(KubernetesJobOperator.template_fields)
845
783
  )
846
784
  operator_extra_links = (KubernetesEngineJobLink(),)
847
785
 
848
786
  def __init__(
849
787
  self,
850
- *,
851
788
  location: str,
852
789
  cluster_name: str,
853
790
  use_internal_ip: bool = False,
@@ -856,55 +793,24 @@ class GKEStartJobOperator(KubernetesJobOperator):
856
793
  impersonation_chain: str | Sequence[str] | None = None,
857
794
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
858
795
  job_poll_interval: float = 10.0,
796
+ *args,
859
797
  **kwargs,
860
798
  ) -> None:
861
- super().__init__(**kwargs)
799
+ super().__init__(*args, **kwargs)
800
+ self.deferrable = deferrable
801
+ self.job_poll_interval = job_poll_interval
862
802
  self.project_id = project_id
863
803
  self.location = location
864
804
  self.cluster_name = cluster_name
865
805
  self.gcp_conn_id = gcp_conn_id
866
- self.impersonation_chain = impersonation_chain
867
806
  self.use_internal_ip = use_internal_ip
868
- self.deferrable = deferrable
869
- self.job_poll_interval = job_poll_interval
870
-
871
- self.job: V1Job | None = None
872
- self._ssl_ca_cert: str | None = None
873
- self._cluster_url: str | None = None
807
+ self.impersonation_chain = impersonation_chain
874
808
 
875
- if self.gcp_conn_id is None:
876
- raise AirflowException(
877
- "The gcp_conn_id parameter has become required. If you want to use Application Default "
878
- "Credentials (ADC) strategy for authorization, create an empty connection "
879
- "called `google_cloud_default`.",
880
- )
881
809
  # There is no need to manage the kube_config file, as it will be generated automatically.
882
810
  # All Kubernetes parameters (except config_file) are also valid for the GKEStartJobOperator.
883
811
  if self.config_file:
884
812
  raise AirflowException("config_file is not an allowed parameter for the GKEStartJobOperator.")
885
813
 
886
- @cached_property
887
- def cluster_hook(self) -> GKEHook:
888
- return GKEHook(
889
- gcp_conn_id=self.gcp_conn_id,
890
- location=self.location,
891
- impersonation_chain=self.impersonation_chain,
892
- )
893
-
894
- @cached_property
895
- def hook(self) -> GKEKubernetesHook:
896
- if self._cluster_url is None or self._ssl_ca_cert is None:
897
- raise AttributeError(
898
- "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
899
- "Try to use self.get_kube_creds method",
900
- )
901
-
902
- return GKEKubernetesHook(
903
- gcp_conn_id=self.gcp_conn_id,
904
- cluster_url=self._cluster_url,
905
- ssl_ca_cert=self._ssl_ca_cert,
906
- )
907
-
908
814
  def execute(self, context: Context):
909
815
  """Execute process of creating Job."""
910
816
  if self.deferrable:
@@ -918,21 +824,13 @@ class GKEStartJobOperator(KubernetesJobOperator):
918
824
  f"package {kubernetes_provider_name}=={kubernetes_provider_version} which doesn't "
919
825
  f"support this feature. Please upgrade it to version higher than {min_version}."
920
826
  )
921
-
922
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
923
- cluster_name=self.cluster_name,
924
- project_id=self.project_id,
925
- use_internal_ip=self.use_internal_ip,
926
- cluster_hook=self.cluster_hook,
927
- ).fetch_cluster_info()
928
-
929
827
  return super().execute(context)
930
828
 
931
829
  def execute_deferrable(self):
932
830
  self.defer(
933
831
  trigger=GKEJobTrigger(
934
- cluster_url=self._cluster_url,
935
- ssl_ca_cert=self._ssl_ca_cert,
832
+ cluster_url=self.cluster_url,
833
+ ssl_ca_cert=self.ssl_ca_cert,
936
834
  job_name=self.job.metadata.name, # type: ignore[union-attr]
937
835
  job_namespace=self.job.metadata.namespace, # type: ignore[union-attr]
938
836
  pod_name=self.pod.metadata.name, # type: ignore[union-attr]
@@ -945,18 +843,10 @@ class GKEStartJobOperator(KubernetesJobOperator):
945
843
  do_xcom_push=self.do_xcom_push,
946
844
  ),
947
845
  method_name="execute_complete",
948
- kwargs={"cluster_url": self._cluster_url, "ssl_ca_cert": self._ssl_ca_cert},
949
846
  )
950
847
 
951
- def execute_complete(self, context: Context, event: dict, **kwargs):
952
- # It is required for hook to be initialized
953
- self._cluster_url = kwargs["cluster_url"]
954
- self._ssl_ca_cert = kwargs["ssl_ca_cert"]
955
-
956
- return super().execute_complete(context, event)
957
848
 
958
-
959
- class GKEDescribeJobOperator(GoogleCloudBaseOperator):
849
+ class GKEDescribeJobOperator(GKEOperatorMixin, GoogleCloudBaseOperator):
960
850
  """
961
851
  Retrieve information about Job by given name.
962
852
 
@@ -965,15 +855,16 @@ class GKEDescribeJobOperator(GoogleCloudBaseOperator):
965
855
  :ref:`howto/operator:GKEDescribeJobOperator`
966
856
 
967
857
  :param job_name: The name of the Job to delete
968
- :param project_id: The Google Developers Console project id.
969
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
970
- resides.
971
- :param cluster_name: The name of the Google Kubernetes Engine cluster.
972
858
  :param namespace: The name of the Google Kubernetes Engine namespace.
859
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
860
+ cluster resides, e.g. 'us-central1-a'
861
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
973
862
  :param use_internal_ip: Use the internal IP address as the endpoint.
974
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
863
+ :param project_id: The Google Developers Console project id
864
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
865
+ users to specify a service account.
975
866
  :param impersonation_chain: Optional service account to impersonate using short-term
976
- credentials, or chained list of accounts required to get the access_token
867
+ credentials, or list of accounts required to get the access_token
977
868
  of the last account in the list, which will be impersonated in the request.
978
869
  If set as a string, the account must grant the originating account
979
870
  the Service Account Token Creator IAM role.
@@ -982,67 +873,33 @@ class GKEDescribeJobOperator(GoogleCloudBaseOperator):
982
873
  account from the list granting this role to the originating account (templated).
983
874
  """
984
875
 
985
- template_fields: Sequence[str] = (
986
- "project_id",
987
- "gcp_conn_id",
988
- "job_name",
989
- "namespace",
990
- "cluster_name",
991
- "location",
992
- "impersonation_chain",
993
- )
876
+ template_fields: Sequence[str] = tuple({"job_name", "namespace"} | set(GKEOperatorMixin.template_fields))
994
877
  operator_extra_links = (KubernetesEngineJobLink(),)
995
878
 
996
879
  def __init__(
997
880
  self,
998
- *,
999
881
  job_name: str,
1000
- location: str,
1001
882
  namespace: str,
883
+ location: str,
1002
884
  cluster_name: str,
1003
- project_id: str = PROVIDE_PROJECT_ID,
1004
885
  use_internal_ip: bool = False,
886
+ project_id: str = PROVIDE_PROJECT_ID,
1005
887
  gcp_conn_id: str = "google_cloud_default",
1006
888
  impersonation_chain: str | Sequence[str] | None = None,
889
+ *args,
1007
890
  **kwargs,
1008
891
  ) -> None:
1009
- super().__init__(**kwargs)
892
+ super().__init__(*args, **kwargs)
1010
893
 
1011
- self.project_id = project_id
1012
- self.gcp_conn_id = gcp_conn_id
1013
- self.location = location
1014
894
  self.job_name = job_name
1015
895
  self.namespace = namespace
896
+ self.project_id = project_id
897
+ self.location = location
1016
898
  self.cluster_name = cluster_name
899
+ self.gcp_conn_id = gcp_conn_id
1017
900
  self.use_internal_ip = use_internal_ip
1018
901
  self.impersonation_chain = impersonation_chain
1019
-
1020
902
  self.job: V1Job | None = None
1021
- self._ssl_ca_cert: str
1022
- self._cluster_url: str
1023
-
1024
- @cached_property
1025
- def cluster_hook(self) -> GKEHook:
1026
- return GKEHook(
1027
- gcp_conn_id=self.gcp_conn_id,
1028
- location=self.location,
1029
- impersonation_chain=self.impersonation_chain,
1030
- )
1031
-
1032
- @cached_property
1033
- def hook(self) -> GKEKubernetesHook:
1034
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1035
- cluster_name=self.cluster_name,
1036
- project_id=self.project_id,
1037
- use_internal_ip=self.use_internal_ip,
1038
- cluster_hook=self.cluster_hook,
1039
- ).fetch_cluster_info()
1040
-
1041
- return GKEKubernetesHook(
1042
- gcp_conn_id=self.gcp_conn_id,
1043
- cluster_url=self._cluster_url,
1044
- ssl_ca_cert=self._ssl_ca_cert,
1045
- )
1046
903
 
1047
904
  def execute(self, context: Context) -> None:
1048
905
  self.job = self.hook.get_job(job_name=self.job_name, namespace=self.namespace)
@@ -1056,7 +913,7 @@ class GKEDescribeJobOperator(GoogleCloudBaseOperator):
1056
913
  return None
1057
914
 
1058
915
 
1059
- class GKEListJobsOperator(GoogleCloudBaseOperator):
916
+ class GKEListJobsOperator(GKEOperatorMixin, GoogleCloudBaseOperator):
1060
917
  """
1061
918
  Retrieve list of Jobs.
1062
919
 
@@ -1067,83 +924,51 @@ class GKEListJobsOperator(GoogleCloudBaseOperator):
1067
924
  For more information on how to use this operator, take a look at the guide:
1068
925
  :ref:`howto/operator:GKEListJobsOperator`
1069
926
 
1070
- :param project_id: The Google Developers Console project id.
1071
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
1072
- resides.
927
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
928
+ cluster resides, e.g. 'us-central1-a'
1073
929
  :param cluster_name: The name of the Google Kubernetes Engine cluster.
1074
- :param namespace: The name of the Google Kubernetes Engine namespace.
1075
930
  :param use_internal_ip: Use the internal IP address as the endpoint.
1076
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
1077
- :param do_xcom_push: If set to True the result list of Jobs will be pushed to the task result.
931
+ :param project_id: The Google Developers Console project id
932
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
933
+ users to specify a service account.
1078
934
  :param impersonation_chain: Optional service account to impersonate using short-term
1079
- credentials, or chained list of accounts required to get the access_token
935
+ credentials, or list of accounts required to get the access_token
1080
936
  of the last account in the list, which will be impersonated in the request.
1081
937
  If set as a string, the account must grant the originating account
1082
938
  the Service Account Token Creator IAM role.
1083
939
  If set as a sequence, the identities from the list must grant
1084
940
  Service Account Token Creator IAM role to the directly preceding identity, with first
1085
941
  account from the list granting this role to the originating account (templated).
942
+ :param namespace: The name of the Google Kubernetes Engine namespace.
943
+ :param do_xcom_push: If set to True the result list of Jobs will be pushed to the task result.
1086
944
  """
1087
945
 
1088
- template_fields: Sequence[str] = (
1089
- "project_id",
1090
- "gcp_conn_id",
1091
- "namespace",
1092
- "cluster_name",
1093
- "location",
1094
- "impersonation_chain",
1095
- )
946
+ template_fields: Sequence[str] = tuple({"namespace"} | set(GKEOperatorMixin.template_fields))
1096
947
  operator_extra_links = (KubernetesEngineWorkloadsLink(),)
1097
948
 
1098
949
  def __init__(
1099
950
  self,
1100
- *,
1101
951
  location: str,
1102
952
  cluster_name: str,
1103
- namespace: str | None = None,
1104
- project_id: str = PROVIDE_PROJECT_ID,
1105
953
  use_internal_ip: bool = False,
1106
- do_xcom_push: bool = True,
954
+ project_id: str = PROVIDE_PROJECT_ID,
1107
955
  gcp_conn_id: str = "google_cloud_default",
1108
956
  impersonation_chain: str | Sequence[str] | None = None,
957
+ namespace: str | None = None,
958
+ do_xcom_push: bool = True,
959
+ *args,
1109
960
  **kwargs,
1110
961
  ) -> None:
1111
- super().__init__(**kwargs)
962
+ super().__init__(*args, **kwargs)
1112
963
 
1113
964
  self.project_id = project_id
1114
- self.gcp_conn_id = gcp_conn_id
1115
965
  self.location = location
1116
- self.namespace = namespace
1117
966
  self.cluster_name = cluster_name
967
+ self.gcp_conn_id = gcp_conn_id
1118
968
  self.use_internal_ip = use_internal_ip
1119
- self.do_xcom_push = do_xcom_push
1120
969
  self.impersonation_chain = impersonation_chain
1121
-
1122
- self._ssl_ca_cert: str
1123
- self._cluster_url: str
1124
-
1125
- @cached_property
1126
- def cluster_hook(self) -> GKEHook:
1127
- return GKEHook(
1128
- gcp_conn_id=self.gcp_conn_id,
1129
- location=self.location,
1130
- impersonation_chain=self.impersonation_chain,
1131
- )
1132
-
1133
- @cached_property
1134
- def hook(self) -> GKEKubernetesHook:
1135
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1136
- cluster_name=self.cluster_name,
1137
- project_id=self.project_id,
1138
- use_internal_ip=self.use_internal_ip,
1139
- cluster_hook=self.cluster_hook,
1140
- ).fetch_cluster_info()
1141
-
1142
- return GKEKubernetesHook(
1143
- gcp_conn_id=self.gcp_conn_id,
1144
- cluster_url=self._cluster_url,
1145
- ssl_ca_cert=self._ssl_ca_cert,
1146
- )
970
+ self.namespace = namespace
971
+ self.do_xcom_push = do_xcom_push
1147
972
 
1148
973
  def execute(self, context: Context) -> dict:
1149
974
  if self.namespace:
@@ -1159,7 +984,7 @@ class GKEListJobsOperator(GoogleCloudBaseOperator):
1159
984
  return V1JobList.to_dict(jobs)
1160
985
 
1161
986
 
1162
- class GKECreateCustomResourceOperator(KubernetesCreateResourceOperator):
987
+ class GKECreateCustomResourceOperator(GKEOperatorMixin, KubernetesCreateResourceOperator):
1163
988
  """
1164
989
  Create a resource in the specified Google Kubernetes Engine cluster.
1165
990
 
@@ -1192,30 +1017,28 @@ class GKECreateCustomResourceOperator(KubernetesCreateResourceOperator):
1192
1017
  """
1193
1018
 
1194
1019
  template_fields: Sequence[str] = tuple(
1195
- {"project_id", "location", "cluster_name"} | set(KubernetesCreateResourceOperator.template_fields)
1020
+ set(GKEOperatorMixin.template_fields) | set(KubernetesCreateResourceOperator.template_fields)
1196
1021
  )
1197
1022
 
1198
1023
  def __init__(
1199
1024
  self,
1200
- *,
1201
1025
  location: str,
1202
1026
  cluster_name: str,
1203
1027
  use_internal_ip: bool = False,
1204
1028
  project_id: str = PROVIDE_PROJECT_ID,
1205
1029
  gcp_conn_id: str = "google_cloud_default",
1206
1030
  impersonation_chain: str | Sequence[str] | None = None,
1031
+ *args,
1207
1032
  **kwargs,
1208
1033
  ) -> None:
1209
- super().__init__(**kwargs)
1034
+ super().__init__(*args, **kwargs)
1035
+
1210
1036
  self.project_id = project_id
1211
1037
  self.location = location
1212
1038
  self.cluster_name = cluster_name
1213
1039
  self.gcp_conn_id = gcp_conn_id
1214
- self.impersonation_chain = impersonation_chain
1215
1040
  self.use_internal_ip = use_internal_ip
1216
-
1217
- self._ssl_ca_cert: str | None = None
1218
- self._cluster_url: str | None = None
1041
+ self.impersonation_chain = impersonation_chain
1219
1042
 
1220
1043
  if self.gcp_conn_id is None:
1221
1044
  raise AirflowException(
@@ -1224,44 +1047,14 @@ class GKECreateCustomResourceOperator(KubernetesCreateResourceOperator):
1224
1047
  "called `google_cloud_default`.",
1225
1048
  )
1226
1049
  # There is no need to manage the kube_config file, as it will be generated automatically.
1227
- # All Kubernetes parameters (except config_file) are also valid for the GKEStartPodOperator.
1050
+ # All Kubernetes parameters (except config_file) are also valid for the GKECreateCustomResourceOperator.
1228
1051
  if self.config_file:
1229
- raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.")
1230
-
1231
- @cached_property
1232
- def cluster_hook(self) -> GKEHook:
1233
- return GKEHook(
1234
- gcp_conn_id=self.gcp_conn_id,
1235
- location=self.location,
1236
- impersonation_chain=self.impersonation_chain,
1237
- )
1238
-
1239
- @cached_property
1240
- def hook(self) -> GKEKubernetesHook:
1241
- if self._cluster_url is None or self._ssl_ca_cert is None:
1242
- raise AttributeError(
1243
- "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
1244
- "Try to use self.get_kube_creds method",
1052
+ raise AirflowException(
1053
+ "config_file is not an allowed parameter for the GKECreateCustomResourceOperator."
1245
1054
  )
1246
- return GKEKubernetesHook(
1247
- gcp_conn_id=self.gcp_conn_id,
1248
- cluster_url=self._cluster_url,
1249
- ssl_ca_cert=self._ssl_ca_cert,
1250
- impersonation_chain=self.impersonation_chain,
1251
- )
1252
-
1253
- def execute(self, context: Context):
1254
- """Execute process of creating Custom Resource."""
1255
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1256
- cluster_name=self.cluster_name,
1257
- project_id=self.project_id,
1258
- use_internal_ip=self.use_internal_ip,
1259
- cluster_hook=self.cluster_hook,
1260
- ).fetch_cluster_info()
1261
- return super().execute(context)
1262
1055
 
1263
1056
 
1264
- class GKEDeleteCustomResourceOperator(KubernetesDeleteResourceOperator):
1057
+ class GKEDeleteCustomResourceOperator(GKEOperatorMixin, KubernetesDeleteResourceOperator):
1265
1058
  """
1266
1059
  Delete a resource in the specified Google Kubernetes Engine cluster.
1267
1060
 
@@ -1294,30 +1087,28 @@ class GKEDeleteCustomResourceOperator(KubernetesDeleteResourceOperator):
1294
1087
  """
1295
1088
 
1296
1089
  template_fields: Sequence[str] = tuple(
1297
- {"project_id", "location", "cluster_name"} | set(KubernetesDeleteResourceOperator.template_fields)
1090
+ set(GKEOperatorMixin.template_fields) | set(KubernetesDeleteResourceOperator.template_fields)
1298
1091
  )
1299
1092
 
1300
1093
  def __init__(
1301
1094
  self,
1302
- *,
1303
1095
  location: str,
1304
1096
  cluster_name: str,
1305
1097
  use_internal_ip: bool = False,
1306
1098
  project_id: str = PROVIDE_PROJECT_ID,
1307
1099
  gcp_conn_id: str = "google_cloud_default",
1308
1100
  impersonation_chain: str | Sequence[str] | None = None,
1101
+ *args,
1309
1102
  **kwargs,
1310
1103
  ) -> None:
1311
- super().__init__(**kwargs)
1104
+ super().__init__(*args, **kwargs)
1105
+
1312
1106
  self.project_id = project_id
1313
1107
  self.location = location
1314
1108
  self.cluster_name = cluster_name
1315
1109
  self.gcp_conn_id = gcp_conn_id
1316
- self.impersonation_chain = impersonation_chain
1317
1110
  self.use_internal_ip = use_internal_ip
1318
-
1319
- self._ssl_ca_cert: str | None = None
1320
- self._cluster_url: str | None = None
1111
+ self.impersonation_chain = impersonation_chain
1321
1112
 
1322
1113
  if self.gcp_conn_id is None:
1323
1114
  raise AirflowException(
@@ -1326,77 +1117,59 @@ class GKEDeleteCustomResourceOperator(KubernetesDeleteResourceOperator):
1326
1117
  "called `google_cloud_default`.",
1327
1118
  )
1328
1119
  # There is no need to manage the kube_config file, as it will be generated automatically.
1329
- # All Kubernetes parameters (except config_file) are also valid for the GKEStartPodOperator.
1120
+ # All Kubernetes parameters (except config_file) are also valid for the GKEDeleteCustomResourceOperator.
1330
1121
  if self.config_file:
1331
- raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.")
1332
-
1333
- @cached_property
1334
- def cluster_hook(self) -> GKEHook:
1335
- return GKEHook(
1336
- gcp_conn_id=self.gcp_conn_id,
1337
- location=self.location,
1338
- impersonation_chain=self.impersonation_chain,
1339
- )
1340
-
1341
- @cached_property
1342
- def hook(self) -> GKEKubernetesHook:
1343
- if self._cluster_url is None or self._ssl_ca_cert is None:
1344
- raise AttributeError(
1345
- "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
1346
- "Try to use self.get_kube_creds method",
1122
+ raise AirflowException(
1123
+ "config_file is not an allowed parameter for the GKEDeleteCustomResourceOperator."
1347
1124
  )
1348
- return GKEKubernetesHook(
1349
- gcp_conn_id=self.gcp_conn_id,
1350
- cluster_url=self._cluster_url,
1351
- ssl_ca_cert=self._ssl_ca_cert,
1352
- impersonation_chain=self.impersonation_chain,
1353
- )
1354
-
1355
- def execute(self, context: Context):
1356
- """Execute process of deleting Custom Resource."""
1357
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1358
- cluster_name=self.cluster_name,
1359
- project_id=self.project_id,
1360
- use_internal_ip=self.use_internal_ip,
1361
- cluster_hook=self.cluster_hook,
1362
- ).fetch_cluster_info()
1363
- return super().execute(context)
1364
1125
 
1365
1126
 
1366
- class GKEStartKueueJobOperator(GKEStartJobOperator):
1127
+ class GKEStartKueueJobOperator(GKEOperatorMixin, KubernetesStartKueueJobOperator):
1367
1128
  """
1368
1129
  Executes a Kubernetes Job in Kueue in the specified Google Kubernetes Engine cluster.
1369
1130
 
1370
- :param queue_name: The name of the Queue in the cluster
1131
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
1132
+ cluster resides, e.g. 'us-central1-a'
1133
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1134
+ :param use_internal_ip: Use the internal IP address as the endpoint.
1135
+ :param project_id: The Google Developers Console project id
1136
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
1137
+ users to specify a service account.
1138
+ :param impersonation_chain: Optional service account to impersonate using short-term
1139
+ credentials, or list of accounts required to get the access_token
1140
+ of the last account in the list, which will be impersonated in the request.
1141
+ If set as a string, the account must grant the originating account
1142
+ the Service Account Token Creator IAM role.
1143
+ If set as a sequence, the identities from the list must grant
1144
+ Service Account Token Creator IAM role to the directly preceding identity, with first
1145
+ account from the list granting this role to the originating account (templated).
1371
1146
  """
1372
1147
 
1148
+ template_fields = tuple(
1149
+ set(GKEOperatorMixin.template_fields) | set(KubernetesStartKueueJobOperator.template_fields)
1150
+ )
1151
+
1373
1152
  def __init__(
1374
1153
  self,
1375
- *,
1376
- queue_name: str,
1154
+ location: str,
1155
+ cluster_name: str,
1156
+ use_internal_ip: bool = False,
1157
+ project_id: str = PROVIDE_PROJECT_ID,
1158
+ gcp_conn_id: str = "google_cloud_default",
1159
+ impersonation_chain: str | Sequence[str] | None = None,
1160
+ *args,
1377
1161
  **kwargs,
1378
1162
  ) -> None:
1379
- super().__init__(**kwargs)
1380
- self.queue_name = queue_name
1381
-
1382
- if self.suspend is False:
1383
- raise AirflowException(
1384
- "The `suspend` parameter can't be False. If you want to use Kueue for running Job"
1385
- " in a Kubernetes cluster, set the `suspend` parameter to True.",
1386
- )
1387
- elif self.suspend is None:
1388
- warnings.warn(
1389
- f"You have not set parameter `suspend` in class {self.__class__.__name__}. "
1390
- "For running a Job in Kueue the `suspend` parameter should set to True.",
1391
- UserWarning,
1392
- stacklevel=2,
1393
- )
1394
- self.suspend = True
1395
- self.labels.update({"kueue.x-k8s.io/queue-name": queue_name})
1396
- self.annotations.update({"kueue.x-k8s.io/queue-name": queue_name})
1163
+ super().__init__(*args, **kwargs)
1164
+ self.project_id = project_id
1165
+ self.location = location
1166
+ self.cluster_name = cluster_name
1167
+ self.gcp_conn_id = gcp_conn_id
1168
+ self.use_internal_ip = use_internal_ip
1169
+ self.impersonation_chain = impersonation_chain
1397
1170
 
1398
1171
 
1399
- class GKEDeleteJobOperator(KubernetesDeleteJobOperator):
1172
+ class GKEDeleteJobOperator(GKEOperatorMixin, KubernetesDeleteJobOperator):
1400
1173
  """
1401
1174
  Delete a Kubernetes job in the specified Google Kubernetes Engine cluster.
1402
1175
 
@@ -1417,7 +1190,7 @@ class GKEDeleteJobOperator(KubernetesDeleteJobOperator):
1417
1190
 
1418
1191
  :param location: The name of the Google Kubernetes Engine zone or region in which the
1419
1192
  cluster resides, e.g. 'us-central1-a'
1420
- :param cluster_name: The name of the Google Kubernetes Engine cluster
1193
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1421
1194
  :param use_internal_ip: Use the internal IP address as the endpoint.
1422
1195
  :param project_id: The Google Developers Console project id
1423
1196
  :param gcp_conn_id: The Google cloud connection id to use. This allows for
@@ -1433,30 +1206,28 @@ class GKEDeleteJobOperator(KubernetesDeleteJobOperator):
1433
1206
  """
1434
1207
 
1435
1208
  template_fields: Sequence[str] = tuple(
1436
- {"project_id", "location", "cluster_name"} | set(KubernetesDeleteJobOperator.template_fields)
1209
+ set(GKEOperatorMixin.template_fields) | set(KubernetesDeleteJobOperator.template_fields)
1437
1210
  )
1438
1211
 
1439
1212
  def __init__(
1440
1213
  self,
1441
- *,
1442
1214
  location: str,
1443
1215
  cluster_name: str,
1444
1216
  use_internal_ip: bool = False,
1445
1217
  project_id: str = PROVIDE_PROJECT_ID,
1446
1218
  gcp_conn_id: str = "google_cloud_default",
1447
1219
  impersonation_chain: str | Sequence[str] | None = None,
1220
+ *args,
1448
1221
  **kwargs,
1449
1222
  ) -> None:
1450
- super().__init__(**kwargs)
1223
+ super().__init__(*args, **kwargs)
1224
+
1451
1225
  self.project_id = project_id
1452
1226
  self.location = location
1453
1227
  self.cluster_name = cluster_name
1454
1228
  self.gcp_conn_id = gcp_conn_id
1455
- self.impersonation_chain = impersonation_chain
1456
1229
  self.use_internal_ip = use_internal_ip
1457
-
1458
- self._ssl_ca_cert: str | None = None
1459
- self._cluster_url: str | None = None
1230
+ self.impersonation_chain = impersonation_chain
1460
1231
 
1461
1232
  if self.gcp_conn_id is None:
1462
1233
  raise AirflowException(
@@ -1469,41 +1240,8 @@ class GKEDeleteJobOperator(KubernetesDeleteJobOperator):
1469
1240
  if self.config_file:
1470
1241
  raise AirflowException("config_file is not an allowed parameter for the GKEDeleteJobOperator.")
1471
1242
 
1472
- @cached_property
1473
- def cluster_hook(self) -> GKEHook:
1474
- return GKEHook(
1475
- gcp_conn_id=self.gcp_conn_id,
1476
- location=self.location,
1477
- impersonation_chain=self.impersonation_chain,
1478
- )
1479
-
1480
- @cached_property
1481
- def hook(self) -> GKEKubernetesHook:
1482
- if self._cluster_url is None or self._ssl_ca_cert is None:
1483
- raise AttributeError(
1484
- "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
1485
- "Try to use self.get_kube_creds method",
1486
- )
1487
-
1488
- return GKEKubernetesHook(
1489
- gcp_conn_id=self.gcp_conn_id,
1490
- cluster_url=self._cluster_url,
1491
- ssl_ca_cert=self._ssl_ca_cert,
1492
- )
1493
-
1494
- def execute(self, context: Context):
1495
- """Execute process of deleting Job."""
1496
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1497
- cluster_name=self.cluster_name,
1498
- project_id=self.project_id,
1499
- use_internal_ip=self.use_internal_ip,
1500
- cluster_hook=self.cluster_hook,
1501
- ).fetch_cluster_info()
1502
1243
 
1503
- return super().execute(context)
1504
-
1505
-
1506
- class GKESuspendJobOperator(GoogleCloudBaseOperator):
1244
+ class GKESuspendJobOperator(GKEOperatorMixin, GoogleCloudBaseOperator):
1507
1245
  """
1508
1246
  Suspend Job by given name.
1509
1247
 
@@ -1512,15 +1250,16 @@ class GKESuspendJobOperator(GoogleCloudBaseOperator):
1512
1250
  :ref:`howto/operator:GKESuspendJobOperator`
1513
1251
 
1514
1252
  :param name: The name of the Job to suspend
1515
- :param project_id: The Google Developers Console project id.
1516
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
1517
- resides.
1518
- :param cluster_name: The name of the Google Kubernetes Engine cluster.
1519
1253
  :param namespace: The name of the Google Kubernetes Engine namespace.
1254
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
1255
+ cluster resides, e.g. 'us-central1-a'
1256
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1520
1257
  :param use_internal_ip: Use the internal IP address as the endpoint.
1521
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
1258
+ :param project_id: The Google Developers Console project id
1259
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
1260
+ users to specify a service account.
1522
1261
  :param impersonation_chain: Optional service account to impersonate using short-term
1523
- credentials, or chained list of accounts required to get the access_token
1262
+ credentials, or list of accounts required to get the access_token
1524
1263
  of the last account in the list, which will be impersonated in the request.
1525
1264
  If set as a string, the account must grant the originating account
1526
1265
  the Service Account Token Creator IAM role.
@@ -1529,67 +1268,33 @@ class GKESuspendJobOperator(GoogleCloudBaseOperator):
1529
1268
  account from the list granting this role to the originating account (templated).
1530
1269
  """
1531
1270
 
1532
- template_fields: Sequence[str] = (
1533
- "project_id",
1534
- "gcp_conn_id",
1535
- "name",
1536
- "namespace",
1537
- "cluster_name",
1538
- "location",
1539
- "impersonation_chain",
1540
- )
1271
+ template_fields: Sequence[str] = tuple({"name", "namespace"} | set(GKEOperatorMixin.template_fields))
1541
1272
  operator_extra_links = (KubernetesEngineJobLink(),)
1542
1273
 
1543
1274
  def __init__(
1544
1275
  self,
1545
- *,
1546
1276
  name: str,
1547
- location: str,
1548
1277
  namespace: str,
1278
+ location: str,
1549
1279
  cluster_name: str,
1550
- project_id: str = PROVIDE_PROJECT_ID,
1551
1280
  use_internal_ip: bool = False,
1281
+ project_id: str = PROVIDE_PROJECT_ID,
1552
1282
  gcp_conn_id: str = "google_cloud_default",
1553
1283
  impersonation_chain: str | Sequence[str] | None = None,
1284
+ *args,
1554
1285
  **kwargs,
1555
1286
  ) -> None:
1556
- super().__init__(**kwargs)
1287
+ super().__init__(*args, **kwargs)
1557
1288
 
1558
- self.project_id = project_id
1559
- self.gcp_conn_id = gcp_conn_id
1560
- self.location = location
1561
1289
  self.name = name
1562
1290
  self.namespace = namespace
1291
+ self.project_id = project_id
1292
+ self.location = location
1563
1293
  self.cluster_name = cluster_name
1294
+ self.gcp_conn_id = gcp_conn_id
1564
1295
  self.use_internal_ip = use_internal_ip
1565
1296
  self.impersonation_chain = impersonation_chain
1566
-
1567
1297
  self.job: V1Job | None = None
1568
- self._ssl_ca_cert: str
1569
- self._cluster_url: str
1570
-
1571
- @cached_property
1572
- def cluster_hook(self) -> GKEHook:
1573
- return GKEHook(
1574
- gcp_conn_id=self.gcp_conn_id,
1575
- location=self.location,
1576
- impersonation_chain=self.impersonation_chain,
1577
- )
1578
-
1579
- @cached_property
1580
- def hook(self) -> GKEKubernetesHook:
1581
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1582
- cluster_name=self.cluster_name,
1583
- project_id=self.project_id,
1584
- use_internal_ip=self.use_internal_ip,
1585
- cluster_hook=self.cluster_hook,
1586
- ).fetch_cluster_info()
1587
-
1588
- return GKEKubernetesHook(
1589
- gcp_conn_id=self.gcp_conn_id,
1590
- cluster_url=self._cluster_url,
1591
- ssl_ca_cert=self._ssl_ca_cert,
1592
- )
1593
1298
 
1594
1299
  def execute(self, context: Context) -> None:
1595
1300
  self.job = self.hook.patch_namespaced_job(
@@ -1607,7 +1312,7 @@ class GKESuspendJobOperator(GoogleCloudBaseOperator):
1607
1312
  return k8s.V1Job.to_dict(self.job)
1608
1313
 
1609
1314
 
1610
- class GKEResumeJobOperator(GoogleCloudBaseOperator):
1315
+ class GKEResumeJobOperator(GKEOperatorMixin, GoogleCloudBaseOperator):
1611
1316
  """
1612
1317
  Resume Job by given name.
1613
1318
 
@@ -1616,15 +1321,16 @@ class GKEResumeJobOperator(GoogleCloudBaseOperator):
1616
1321
  :ref:`howto/operator:GKEResumeJobOperator`
1617
1322
 
1618
1323
  :param name: The name of the Job to resume
1619
- :param project_id: The Google Developers Console project id.
1620
- :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
1621
- resides.
1622
- :param cluster_name: The name of the Google Kubernetes Engine cluster.
1623
1324
  :param namespace: The name of the Google Kubernetes Engine namespace.
1325
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
1326
+ cluster resides, e.g. 'us-central1-a'
1327
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1624
1328
  :param use_internal_ip: Use the internal IP address as the endpoint.
1625
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
1329
+ :param project_id: The Google Developers Console project id
1330
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
1331
+ users to specify a service account.
1626
1332
  :param impersonation_chain: Optional service account to impersonate using short-term
1627
- credentials, or chained list of accounts required to get the access_token
1333
+ credentials, or list of accounts required to get the access_token
1628
1334
  of the last account in the list, which will be impersonated in the request.
1629
1335
  If set as a string, the account must grant the originating account
1630
1336
  the Service Account Token Creator IAM role.
@@ -1633,67 +1339,33 @@ class GKEResumeJobOperator(GoogleCloudBaseOperator):
1633
1339
  account from the list granting this role to the originating account (templated).
1634
1340
  """
1635
1341
 
1636
- template_fields: Sequence[str] = (
1637
- "project_id",
1638
- "gcp_conn_id",
1639
- "name",
1640
- "namespace",
1641
- "cluster_name",
1642
- "location",
1643
- "impersonation_chain",
1644
- )
1342
+ template_fields: Sequence[str] = tuple({"name", "namespace"} | set(GKEOperatorMixin.template_fields))
1645
1343
  operator_extra_links = (KubernetesEngineJobLink(),)
1646
1344
 
1647
1345
  def __init__(
1648
1346
  self,
1649
- *,
1650
1347
  name: str,
1651
- location: str,
1652
1348
  namespace: str,
1349
+ location: str,
1653
1350
  cluster_name: str,
1654
- project_id: str = PROVIDE_PROJECT_ID,
1655
1351
  use_internal_ip: bool = False,
1352
+ project_id: str = PROVIDE_PROJECT_ID,
1656
1353
  gcp_conn_id: str = "google_cloud_default",
1657
1354
  impersonation_chain: str | Sequence[str] | None = None,
1355
+ *args,
1658
1356
  **kwargs,
1659
1357
  ) -> None:
1660
- super().__init__(**kwargs)
1358
+ super().__init__(*args, **kwargs)
1661
1359
 
1662
- self.project_id = project_id
1663
- self.gcp_conn_id = gcp_conn_id
1664
- self.location = location
1665
1360
  self.name = name
1666
1361
  self.namespace = namespace
1362
+ self.project_id = project_id
1363
+ self.location = location
1667
1364
  self.cluster_name = cluster_name
1365
+ self.gcp_conn_id = gcp_conn_id
1668
1366
  self.use_internal_ip = use_internal_ip
1669
1367
  self.impersonation_chain = impersonation_chain
1670
-
1671
1368
  self.job: V1Job | None = None
1672
- self._ssl_ca_cert: str
1673
- self._cluster_url: str
1674
-
1675
- @cached_property
1676
- def cluster_hook(self) -> GKEHook:
1677
- return GKEHook(
1678
- gcp_conn_id=self.gcp_conn_id,
1679
- location=self.location,
1680
- impersonation_chain=self.impersonation_chain,
1681
- )
1682
-
1683
- @cached_property
1684
- def hook(self) -> GKEKubernetesHook:
1685
- self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1686
- cluster_name=self.cluster_name,
1687
- project_id=self.project_id,
1688
- use_internal_ip=self.use_internal_ip,
1689
- cluster_hook=self.cluster_hook,
1690
- ).fetch_cluster_info()
1691
-
1692
- return GKEKubernetesHook(
1693
- gcp_conn_id=self.gcp_conn_id,
1694
- cluster_url=self._cluster_url,
1695
- ssl_ca_cert=self._ssl_ca_cert,
1696
- )
1697
1369
 
1698
1370
  def execute(self, context: Context) -> None:
1699
1371
  self.job = self.hook.patch_namespaced_job(