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.
- airflow/providers/google/__init__.py +3 -3
- airflow/providers/google/assets/gcs.py +1 -7
- airflow/providers/google/cloud/hooks/alloy_db.py +289 -0
- airflow/providers/google/cloud/hooks/cloud_batch.py +13 -5
- airflow/providers/google/cloud/hooks/dataproc.py +7 -3
- airflow/providers/google/cloud/hooks/dataproc_metastore.py +41 -22
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +7 -38
- airflow/providers/google/cloud/hooks/translate.py +355 -0
- airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +147 -0
- airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +10 -0
- airflow/providers/google/cloud/links/alloy_db.py +55 -0
- airflow/providers/google/cloud/links/translate.py +98 -0
- airflow/providers/google/cloud/log/stackdriver_task_handler.py +1 -5
- airflow/providers/google/cloud/openlineage/mixins.py +4 -12
- airflow/providers/google/cloud/openlineage/utils.py +200 -22
- airflow/providers/google/cloud/operators/alloy_db.py +459 -0
- airflow/providers/google/cloud/operators/automl.py +55 -44
- airflow/providers/google/cloud/operators/bigquery.py +60 -15
- airflow/providers/google/cloud/operators/dataproc.py +12 -0
- airflow/providers/google/cloud/operators/gcs.py +5 -14
- airflow/providers/google/cloud/operators/kubernetes_engine.py +377 -705
- airflow/providers/google/cloud/operators/mlengine.py +41 -31
- airflow/providers/google/cloud/operators/translate.py +586 -1
- airflow/providers/google/cloud/operators/vertex_ai/feature_store.py +163 -0
- airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +5 -0
- airflow/providers/google/cloud/sensors/dataproc.py +2 -2
- airflow/providers/google/cloud/sensors/vertex_ai/__init__.py +16 -0
- airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +112 -0
- airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +6 -11
- airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +3 -0
- airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +3 -0
- airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +5 -10
- airflow/providers/google/cloud/transfers/gcs_to_gcs.py +3 -15
- airflow/providers/google/cloud/transfers/gcs_to_local.py +9 -0
- airflow/providers/google/cloud/transfers/local_to_gcs.py +41 -6
- airflow/providers/google/cloud/transfers/s3_to_gcs.py +15 -0
- airflow/providers/google/get_provider_info.py +30 -18
- airflow/providers/google/version_compat.py +36 -0
- {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/METADATA +16 -18
- {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/RECORD +42 -37
- airflow/providers/google/cloud/hooks/datapipeline.py +0 -71
- airflow/providers/google/cloud/openlineage/BigQueryErrorRunFacet.json +0 -30
- airflow/providers/google/cloud/operators/datapipeline.py +0 -63
- {apache_airflow_providers_google-11.0.0.dist-info → apache_airflow_providers_google-12.0.0.dist-info}/WHEEL +0 -0
- {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
|
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
|
105
|
-
self._ssl_ca_cert
|
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
|
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 ``
|
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
|
-
|
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
|
146
|
-
|
147
|
-
:param
|
148
|
-
|
149
|
-
:param
|
150
|
-
:param
|
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
|
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
|
-
"
|
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
|
-
|
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.
|
201
|
-
self.log.error("One of (project_id,
|
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 =
|
209
|
-
name=self.
|
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.
|
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
|
288
|
-
|
289
|
-
|
290
|
-
:param
|
291
|
-
|
292
|
-
|
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
|
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
|
-
"
|
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
|
-
|
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
|
-
|
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
|
-
|
464
|
+
KubernetesEngineClusterLink.persist(context=context, task_instance=self, cluster=self.body)
|
465
|
+
|
408
466
|
try:
|
409
|
-
|
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=
|
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
|
-
|
436
|
-
|
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.
|
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
|
476
|
-
|
477
|
-
:param cluster_name: The
|
478
|
-
:param
|
479
|
-
:param
|
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
|
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
|
-
|
491
|
-
|
492
|
-
|
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.
|
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.
|
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
|
-
|
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
|
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
|
-
{"
|
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
|
697
|
-
"
|
698
|
-
"be deleted in
|
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
|
-
@
|
727
|
-
|
728
|
-
|
729
|
-
|
730
|
-
|
731
|
-
|
732
|
-
|
733
|
-
|
734
|
-
|
735
|
-
|
736
|
-
|
737
|
-
|
738
|
-
|
739
|
-
|
740
|
-
|
741
|
-
|
742
|
-
|
743
|
-
|
744
|
-
|
745
|
-
|
746
|
-
|
747
|
-
|
748
|
-
|
749
|
-
|
750
|
-
def
|
751
|
-
|
752
|
-
|
753
|
-
|
754
|
-
|
755
|
-
|
756
|
-
"
|
757
|
-
|
758
|
-
|
759
|
-
|
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.
|
778
|
-
ssl_ca_cert=self.
|
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=
|
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="
|
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
|
-
{"
|
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.
|
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.
|
935
|
-
ssl_ca_cert=self.
|
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
|
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
|
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
|
1071
|
-
|
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
|
1077
|
-
:param
|
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
|
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
|
-
|
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.
|
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
|
-
|
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
|
1050
|
+
# All Kubernetes parameters (except config_file) are also valid for the GKECreateCustomResourceOperator.
|
1228
1051
|
if self.config_file:
|
1229
|
-
raise AirflowException(
|
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
|
-
|
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
|
1120
|
+
# All Kubernetes parameters (except config_file) are also valid for the GKEDeleteCustomResourceOperator.
|
1330
1121
|
if self.config_file:
|
1331
|
-
raise AirflowException(
|
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(
|
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
|
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
|
-
|
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.
|
1381
|
-
|
1382
|
-
|
1383
|
-
|
1384
|
-
|
1385
|
-
|
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
|
-
|
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
|
-
|
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
|
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
|
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
|
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
|
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(
|