apache-airflow-providers-google 10.22.0rc1__py3-none-any.whl → 10.23.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 +1 -1
- airflow/providers/google/cloud/hooks/bigquery.py +91 -54
- airflow/providers/google/cloud/hooks/cloud_build.py +3 -2
- airflow/providers/google/cloud/hooks/dataflow.py +112 -47
- airflow/providers/google/cloud/hooks/datapipeline.py +3 -3
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +15 -26
- airflow/providers/google/cloud/hooks/life_sciences.py +5 -7
- airflow/providers/google/cloud/hooks/secret_manager.py +3 -3
- airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +28 -8
- airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +11 -6
- airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +214 -34
- airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +11 -4
- airflow/providers/google/cloud/links/automl.py +13 -22
- airflow/providers/google/cloud/log/gcs_task_handler.py +1 -2
- airflow/providers/google/cloud/operators/bigquery.py +6 -4
- airflow/providers/google/cloud/operators/dataflow.py +186 -4
- airflow/providers/google/cloud/operators/datafusion.py +3 -2
- airflow/providers/google/cloud/operators/datapipeline.py +5 -6
- airflow/providers/google/cloud/operators/dataproc.py +30 -33
- airflow/providers/google/cloud/operators/gcs.py +4 -4
- airflow/providers/google/cloud/operators/kubernetes_engine.py +16 -2
- airflow/providers/google/cloud/operators/life_sciences.py +5 -7
- airflow/providers/google/cloud/operators/mlengine.py +42 -65
- airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +18 -4
- airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +5 -5
- airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +280 -9
- airflow/providers/google/cloud/operators/vertex_ai/model_service.py +4 -0
- airflow/providers/google/cloud/secrets/secret_manager.py +3 -5
- airflow/providers/google/cloud/sensors/bigquery.py +8 -27
- airflow/providers/google/cloud/sensors/bigquery_dts.py +1 -4
- airflow/providers/google/cloud/sensors/cloud_composer.py +9 -14
- airflow/providers/google/cloud/sensors/dataflow.py +1 -25
- airflow/providers/google/cloud/sensors/dataform.py +1 -4
- airflow/providers/google/cloud/sensors/datafusion.py +1 -7
- airflow/providers/google/cloud/sensors/dataplex.py +1 -31
- airflow/providers/google/cloud/sensors/dataproc.py +1 -16
- airflow/providers/google/cloud/sensors/dataproc_metastore.py +1 -7
- airflow/providers/google/cloud/sensors/gcs.py +5 -27
- airflow/providers/google/cloud/sensors/looker.py +1 -13
- airflow/providers/google/cloud/sensors/pubsub.py +11 -5
- airflow/providers/google/cloud/sensors/workflows.py +1 -4
- airflow/providers/google/cloud/transfers/sftp_to_gcs.py +6 -0
- airflow/providers/google/cloud/triggers/dataflow.py +145 -1
- airflow/providers/google/cloud/triggers/kubernetes_engine.py +66 -3
- airflow/providers/google/common/deprecated.py +176 -0
- airflow/providers/google/common/hooks/base_google.py +3 -2
- airflow/providers/google/get_provider_info.py +8 -10
- airflow/providers/google/marketing_platform/hooks/analytics.py +4 -2
- airflow/providers/google/marketing_platform/hooks/search_ads.py +169 -30
- airflow/providers/google/marketing_platform/operators/analytics.py +16 -33
- airflow/providers/google/marketing_platform/operators/search_ads.py +217 -156
- airflow/providers/google/marketing_platform/sensors/display_video.py +1 -4
- {apache_airflow_providers_google-10.22.0rc1.dist-info → apache_airflow_providers_google-10.23.0.dist-info}/METADATA +25 -23
- {apache_airflow_providers_google-10.22.0rc1.dist-info → apache_airflow_providers_google-10.23.0.dist-info}/RECORD +56 -56
- airflow/providers/google/marketing_platform/sensors/search_ads.py +0 -92
- {apache_airflow_providers_google-10.22.0rc1.dist-info → apache_airflow_providers_google-10.23.0.dist-info}/WHEEL +0 -0
- {apache_airflow_providers_google-10.22.0rc1.dist-info → apache_airflow_providers_google-10.23.0.dist-info}/entry_points.txt +0 -0
@@ -23,7 +23,7 @@ from functools import cached_property
|
|
23
23
|
from typing import TYPE_CHECKING, Any, Callable, Sequence
|
24
24
|
|
25
25
|
from airflow.configuration import conf
|
26
|
-
from airflow.exceptions import AirflowException
|
26
|
+
from airflow.exceptions import AirflowException
|
27
27
|
from airflow.providers.google.cloud.hooks.dataflow import (
|
28
28
|
DEFAULT_DATAFLOW_LOCATION,
|
29
29
|
DataflowHook,
|
@@ -117,10 +117,7 @@ class DataflowJobStatusSensor(BaseSensorOperator):
|
|
117
117
|
if job_status in self.expected_statuses:
|
118
118
|
return True
|
119
119
|
elif job_status in DataflowJobStatus.TERMINAL_STATES:
|
120
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
121
120
|
message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
|
122
|
-
if self.soft_fail:
|
123
|
-
raise AirflowSkipException(message)
|
124
121
|
raise AirflowException(message)
|
125
122
|
|
126
123
|
return False
|
@@ -154,9 +151,6 @@ class DataflowJobStatusSensor(BaseSensorOperator):
|
|
154
151
|
if event["status"] == "success":
|
155
152
|
self.log.info(event["message"])
|
156
153
|
return True
|
157
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
158
|
-
if self.soft_fail:
|
159
|
-
raise AirflowSkipException(f"Sensor failed with the following message: {event['message']}.")
|
160
154
|
raise AirflowException(f"Sensor failed with the following message: {event['message']}")
|
161
155
|
|
162
156
|
@cached_property
|
@@ -235,10 +229,7 @@ class DataflowJobMetricsSensor(BaseSensorOperator):
|
|
235
229
|
)
|
236
230
|
job_status = job["currentState"]
|
237
231
|
if job_status in DataflowJobStatus.TERMINAL_STATES:
|
238
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
239
232
|
message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
|
240
|
-
if self.soft_fail:
|
241
|
-
raise AirflowSkipException(message)
|
242
233
|
raise AirflowException(message)
|
243
234
|
|
244
235
|
result = self.hook.fetch_job_metrics_by_id(
|
@@ -279,9 +270,6 @@ class DataflowJobMetricsSensor(BaseSensorOperator):
|
|
279
270
|
if event["status"] == "success":
|
280
271
|
self.log.info(event["message"])
|
281
272
|
return event["result"] if self.callback is None else self.callback(event["result"])
|
282
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
283
|
-
if self.soft_fail:
|
284
|
-
raise AirflowSkipException(f"Sensor failed with the following message: {event['message']}.")
|
285
273
|
raise AirflowException(f"Sensor failed with the following message: {event['message']}")
|
286
274
|
|
287
275
|
@cached_property
|
@@ -362,10 +350,7 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
|
|
362
350
|
)
|
363
351
|
job_status = job["currentState"]
|
364
352
|
if job_status in DataflowJobStatus.TERMINAL_STATES:
|
365
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
366
353
|
message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
|
367
|
-
if self.soft_fail:
|
368
|
-
raise AirflowSkipException(message)
|
369
354
|
raise AirflowException(message)
|
370
355
|
|
371
356
|
result = self.hook.fetch_job_messages_by_id(
|
@@ -407,9 +392,6 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
|
|
407
392
|
if event["status"] == "success":
|
408
393
|
self.log.info(event["message"])
|
409
394
|
return event["result"] if self.callback is None else self.callback(event["result"])
|
410
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
411
|
-
if self.soft_fail:
|
412
|
-
raise AirflowSkipException(f"Sensor failed with the following message: {event['message']}.")
|
413
395
|
raise AirflowException(f"Sensor failed with the following message: {event['message']}")
|
414
396
|
|
415
397
|
@cached_property
|
@@ -490,10 +472,7 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
|
|
490
472
|
)
|
491
473
|
job_status = job["currentState"]
|
492
474
|
if job_status in DataflowJobStatus.TERMINAL_STATES:
|
493
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
494
475
|
message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
|
495
|
-
if self.soft_fail:
|
496
|
-
raise AirflowSkipException(message)
|
497
476
|
raise AirflowException(message)
|
498
477
|
|
499
478
|
result = self.hook.fetch_job_autoscaling_events_by_id(
|
@@ -534,9 +513,6 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
|
|
534
513
|
if event["status"] == "success":
|
535
514
|
self.log.info(event["message"])
|
536
515
|
return event["result"] if self.callback is None else self.callback(event["result"])
|
537
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
538
|
-
if self.soft_fail:
|
539
|
-
raise AirflowSkipException(f"Sensor failed with the following message: {event['message']}.")
|
540
516
|
raise AirflowException(f"Sensor failed with the following message: {event['message']}")
|
541
517
|
|
542
518
|
@cached_property
|
@@ -21,7 +21,7 @@ from __future__ import annotations
|
|
21
21
|
|
22
22
|
from typing import TYPE_CHECKING, Iterable, Sequence
|
23
23
|
|
24
|
-
from airflow.exceptions import AirflowException
|
24
|
+
from airflow.exceptions import AirflowException
|
25
25
|
from airflow.providers.google.cloud.hooks.dataform import DataformHook
|
26
26
|
from airflow.sensors.base import BaseSensorOperator
|
27
27
|
|
@@ -96,13 +96,10 @@ class DataformWorkflowInvocationStateSensor(BaseSensorOperator):
|
|
96
96
|
workflow_status = workflow_invocation.state
|
97
97
|
if workflow_status is not None:
|
98
98
|
if self.failure_statuses and workflow_status in self.failure_statuses:
|
99
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
100
99
|
message = (
|
101
100
|
f"Workflow Invocation with id '{self.workflow_invocation_id}' "
|
102
101
|
f"state is: {workflow_status}. Terminating sensor..."
|
103
102
|
)
|
104
|
-
if self.soft_fail:
|
105
|
-
raise AirflowSkipException(message)
|
106
103
|
raise AirflowException(message)
|
107
104
|
|
108
105
|
return workflow_status in self.expected_statuses
|
@@ -21,7 +21,7 @@ from __future__ import annotations
|
|
21
21
|
|
22
22
|
from typing import TYPE_CHECKING, Iterable, Sequence
|
23
23
|
|
24
|
-
from airflow.exceptions import AirflowException, AirflowNotFoundException
|
24
|
+
from airflow.exceptions import AirflowException, AirflowNotFoundException
|
25
25
|
from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
|
26
26
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
27
27
|
from airflow.sensors.base import BaseSensorOperator
|
@@ -111,22 +111,16 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
|
|
111
111
|
)
|
112
112
|
pipeline_status = pipeline_workflow["status"]
|
113
113
|
except AirflowNotFoundException:
|
114
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
115
114
|
message = "Specified Pipeline ID was not found."
|
116
|
-
if self.soft_fail:
|
117
|
-
raise AirflowSkipException(message)
|
118
115
|
raise AirflowException(message)
|
119
116
|
except AirflowException:
|
120
117
|
pass # Because the pipeline may not be visible in system yet
|
121
118
|
if pipeline_status is not None:
|
122
119
|
if self.failure_statuses and pipeline_status in self.failure_statuses:
|
123
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
124
120
|
message = (
|
125
121
|
f"Pipeline with id '{self.pipeline_id}' state is: {pipeline_status}. "
|
126
122
|
f"Terminating sensor..."
|
127
123
|
)
|
128
|
-
if self.soft_fail:
|
129
|
-
raise AirflowSkipException(message)
|
130
124
|
raise AirflowException(message)
|
131
125
|
|
132
126
|
self.log.debug(
|
@@ -30,7 +30,7 @@ from google.api_core.exceptions import GoogleAPICallError
|
|
30
30
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
31
31
|
from google.cloud.dataplex_v1.types import DataScanJob
|
32
32
|
|
33
|
-
from airflow.exceptions import AirflowException
|
33
|
+
from airflow.exceptions import AirflowException
|
34
34
|
from airflow.providers.google.cloud.hooks.dataplex import (
|
35
35
|
AirflowDataQualityScanException,
|
36
36
|
AirflowDataQualityScanResultTimeoutException,
|
@@ -118,10 +118,7 @@ class DataplexTaskStateSensor(BaseSensorOperator):
|
|
118
118
|
task_status = task.state
|
119
119
|
|
120
120
|
if task_status == TaskState.DELETING:
|
121
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
122
121
|
message = f"Task is going to be deleted {self.dataplex_task_id}"
|
123
|
-
if self.soft_fail:
|
124
|
-
raise AirflowSkipException(message)
|
125
122
|
raise AirflowException(message)
|
126
123
|
|
127
124
|
self.log.info("Current status of the Dataplex task %s => %s", self.dataplex_task_id, task_status)
|
@@ -202,12 +199,9 @@ class DataplexDataQualityJobStatusSensor(BaseSensorOperator):
|
|
202
199
|
if self.result_timeout:
|
203
200
|
duration = self._duration()
|
204
201
|
if duration > self.result_timeout:
|
205
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
206
202
|
message = (
|
207
203
|
f"Timeout: Data Quality scan {self.job_id} is not ready after {self.result_timeout}s"
|
208
204
|
)
|
209
|
-
if self.soft_fail:
|
210
|
-
raise AirflowSkipException(message)
|
211
205
|
raise AirflowDataQualityScanResultTimeoutException(message)
|
212
206
|
|
213
207
|
hook = DataplexHook(
|
@@ -227,10 +221,7 @@ class DataplexDataQualityJobStatusSensor(BaseSensorOperator):
|
|
227
221
|
metadata=self.metadata,
|
228
222
|
)
|
229
223
|
except GoogleAPICallError as e:
|
230
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
231
224
|
message = f"Error occurred when trying to retrieve Data Quality scan job: {self.data_scan_id}"
|
232
|
-
if self.soft_fail:
|
233
|
-
raise AirflowSkipException(message, e)
|
234
225
|
raise AirflowException(message, e)
|
235
226
|
|
236
227
|
job_status = job.state
|
@@ -238,26 +229,17 @@ class DataplexDataQualityJobStatusSensor(BaseSensorOperator):
|
|
238
229
|
"Current status of the Dataplex Data Quality scan job %s => %s", self.job_id, job_status
|
239
230
|
)
|
240
231
|
if job_status == DataScanJob.State.FAILED:
|
241
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
242
232
|
message = f"Data Quality scan job failed: {self.job_id}"
|
243
|
-
if self.soft_fail:
|
244
|
-
raise AirflowSkipException(message)
|
245
233
|
raise AirflowException(message)
|
246
234
|
if job_status == DataScanJob.State.CANCELLED:
|
247
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
248
235
|
message = f"Data Quality scan job cancelled: {self.job_id}"
|
249
|
-
if self.soft_fail:
|
250
|
-
raise AirflowSkipException(message)
|
251
236
|
raise AirflowException(message)
|
252
237
|
if self.fail_on_dq_failure:
|
253
238
|
if job_status == DataScanJob.State.SUCCEEDED and not job.data_quality_result.passed:
|
254
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
255
239
|
message = (
|
256
240
|
f"Data Quality job {self.job_id} execution failed due to failure of its scanning "
|
257
241
|
f"rules: {self.data_scan_id}"
|
258
242
|
)
|
259
|
-
if self.soft_fail:
|
260
|
-
raise AirflowSkipException(message)
|
261
243
|
raise AirflowDataQualityScanException(message)
|
262
244
|
return job_status == DataScanJob.State.SUCCEEDED
|
263
245
|
|
@@ -330,12 +312,9 @@ class DataplexDataProfileJobStatusSensor(BaseSensorOperator):
|
|
330
312
|
if self.result_timeout:
|
331
313
|
duration = self._duration()
|
332
314
|
if duration > self.result_timeout:
|
333
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
334
315
|
message = (
|
335
316
|
f"Timeout: Data Profile scan {self.job_id} is not ready after {self.result_timeout}s"
|
336
317
|
)
|
337
|
-
if self.soft_fail:
|
338
|
-
raise AirflowSkipException(message)
|
339
318
|
raise AirflowDataQualityScanResultTimeoutException(message)
|
340
319
|
|
341
320
|
hook = DataplexHook(
|
@@ -355,10 +334,7 @@ class DataplexDataProfileJobStatusSensor(BaseSensorOperator):
|
|
355
334
|
metadata=self.metadata,
|
356
335
|
)
|
357
336
|
except GoogleAPICallError as e:
|
358
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
359
337
|
message = f"Error occurred when trying to retrieve Data Profile scan job: {self.data_scan_id}"
|
360
|
-
if self.soft_fail:
|
361
|
-
raise AirflowSkipException(message, e)
|
362
338
|
raise AirflowException(message, e)
|
363
339
|
|
364
340
|
job_status = job.state
|
@@ -366,15 +342,9 @@ class DataplexDataProfileJobStatusSensor(BaseSensorOperator):
|
|
366
342
|
"Current status of the Dataplex Data Profile scan job %s => %s", self.job_id, job_status
|
367
343
|
)
|
368
344
|
if job_status == DataScanJob.State.FAILED:
|
369
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
370
345
|
message = f"Data Profile scan job failed: {self.job_id}"
|
371
|
-
if self.soft_fail:
|
372
|
-
raise AirflowSkipException(message)
|
373
346
|
raise AirflowException(message)
|
374
347
|
if job_status == DataScanJob.State.CANCELLED:
|
375
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
376
348
|
message = f"Data Profile scan job cancelled: {self.job_id}"
|
377
|
-
if self.soft_fail:
|
378
|
-
raise AirflowSkipException(message)
|
379
349
|
raise AirflowException(message)
|
380
350
|
return job_status == DataScanJob.State.SUCCEEDED
|
@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Sequence
|
|
25
25
|
from google.api_core.exceptions import ServerError
|
26
26
|
from google.cloud.dataproc_v1.types import Batch, JobStatus
|
27
27
|
|
28
|
-
from airflow.exceptions import AirflowException
|
28
|
+
from airflow.exceptions import AirflowException
|
29
29
|
from airflow.providers.google.cloud.hooks.dataproc import DataprocHook
|
30
30
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
31
31
|
from airflow.sensors.base import BaseSensorOperator
|
@@ -85,13 +85,10 @@ class DataprocJobSensor(BaseSensorOperator):
|
|
85
85
|
duration = self._duration()
|
86
86
|
self.log.info("DURATION RUN: %f", duration)
|
87
87
|
if duration > self.wait_timeout:
|
88
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
89
88
|
message = (
|
90
89
|
f"Timeout: dataproc job {self.dataproc_job_id} "
|
91
90
|
f"is not ready after {self.wait_timeout}s"
|
92
91
|
)
|
93
|
-
if self.soft_fail:
|
94
|
-
raise AirflowSkipException(message)
|
95
92
|
raise AirflowException(message)
|
96
93
|
self.log.info("Retrying. Dataproc API returned server error when waiting for job: %s", err)
|
97
94
|
return False
|
@@ -100,20 +97,14 @@ class DataprocJobSensor(BaseSensorOperator):
|
|
100
97
|
|
101
98
|
state = job.status.state
|
102
99
|
if state == JobStatus.State.ERROR:
|
103
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
104
100
|
message = f"Job failed:\n{job}"
|
105
|
-
if self.soft_fail:
|
106
|
-
raise AirflowSkipException(message)
|
107
101
|
raise AirflowException(message)
|
108
102
|
elif state in {
|
109
103
|
JobStatus.State.CANCELLED,
|
110
104
|
JobStatus.State.CANCEL_PENDING,
|
111
105
|
JobStatus.State.CANCEL_STARTED,
|
112
106
|
}:
|
113
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
114
107
|
message = f"Job was cancelled:\n{job}"
|
115
|
-
if self.soft_fail:
|
116
|
-
raise AirflowSkipException(message)
|
117
108
|
raise AirflowException(message)
|
118
109
|
elif JobStatus.State.DONE == state:
|
119
110
|
self.log.debug("Job %s completed successfully.", self.dataproc_job_id)
|
@@ -185,19 +176,13 @@ class DataprocBatchSensor(BaseSensorOperator):
|
|
185
176
|
|
186
177
|
state = batch.state
|
187
178
|
if state == Batch.State.FAILED:
|
188
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
189
179
|
message = "Batch failed"
|
190
|
-
if self.soft_fail:
|
191
|
-
raise AirflowSkipException(message)
|
192
180
|
raise AirflowException(message)
|
193
181
|
elif state in {
|
194
182
|
Batch.State.CANCELLED,
|
195
183
|
Batch.State.CANCELLING,
|
196
184
|
}:
|
197
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
198
185
|
message = "Batch was cancelled."
|
199
|
-
if self.soft_fail:
|
200
|
-
raise AirflowSkipException(message)
|
201
186
|
raise AirflowException(message)
|
202
187
|
elif state == Batch.State.SUCCEEDED:
|
203
188
|
self.log.debug("Batch %s completed successfully.", self.batch_id)
|
@@ -19,7 +19,7 @@ from __future__ import annotations
|
|
19
19
|
|
20
20
|
from typing import TYPE_CHECKING, Sequence
|
21
21
|
|
22
|
-
from airflow.exceptions import AirflowException
|
22
|
+
from airflow.exceptions import AirflowException
|
23
23
|
from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
|
24
24
|
from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
|
25
25
|
from airflow.sensors.base import BaseSensorOperator
|
@@ -99,20 +99,14 @@ class MetastoreHivePartitionSensor(BaseSensorOperator):
|
|
99
99
|
impersonation_chain=self.impersonation_chain,
|
100
100
|
)
|
101
101
|
if not (manifest and isinstance(manifest, dict)):
|
102
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
103
102
|
message = (
|
104
103
|
f"Failed to extract result manifest. "
|
105
104
|
f"Expected not empty dict, but this was received: {manifest}"
|
106
105
|
)
|
107
|
-
if self.soft_fail:
|
108
|
-
raise AirflowSkipException(message)
|
109
106
|
raise AirflowException(message)
|
110
107
|
|
111
108
|
if manifest.get("status", {}).get("code") != 0:
|
112
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
113
109
|
message = f"Request failed: {manifest.get('message')}"
|
114
|
-
if self.soft_fail:
|
115
|
-
raise AirflowSkipException(message)
|
116
110
|
raise AirflowException(message)
|
117
111
|
|
118
112
|
# Extract actual query results
|
@@ -24,11 +24,10 @@ import textwrap
|
|
24
24
|
from datetime import datetime, timedelta
|
25
25
|
from typing import TYPE_CHECKING, Any, Callable, Sequence
|
26
26
|
|
27
|
-
from deprecated import deprecated
|
28
27
|
from google.cloud.storage.retry import DEFAULT_RETRY
|
29
28
|
|
30
29
|
from airflow.configuration import conf
|
31
|
-
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
|
30
|
+
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
|
32
31
|
from airflow.providers.google.cloud.hooks.gcs import GCSHook
|
33
32
|
from airflow.providers.google.cloud.triggers.gcs import (
|
34
33
|
GCSBlobTrigger,
|
@@ -36,6 +35,7 @@ from airflow.providers.google.cloud.triggers.gcs import (
|
|
36
35
|
GCSPrefixBlobTrigger,
|
37
36
|
GCSUploadSessionTrigger,
|
38
37
|
)
|
38
|
+
from airflow.providers.google.common.deprecated import deprecated
|
39
39
|
from airflow.sensors.base import BaseSensorOperator, poke_mode_only
|
40
40
|
|
41
41
|
if TYPE_CHECKING:
|
@@ -137,19 +137,15 @@ class GCSObjectExistenceSensor(BaseSensorOperator):
|
|
137
137
|
Relies on trigger to throw an exception, otherwise it assumes execution was successful.
|
138
138
|
"""
|
139
139
|
if event["status"] == "error":
|
140
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
141
|
-
if self.soft_fail:
|
142
|
-
raise AirflowSkipException(event["message"])
|
143
140
|
raise AirflowException(event["message"])
|
144
141
|
self.log.info("File %s was found in bucket %s.", self.object, self.bucket)
|
145
142
|
return True
|
146
143
|
|
147
144
|
|
148
145
|
@deprecated(
|
149
|
-
|
150
|
-
|
151
|
-
|
152
|
-
),
|
146
|
+
planned_removal_date="November 01, 2024",
|
147
|
+
use_instead="GCSObjectExistenceSensor",
|
148
|
+
instructions="Please use GCSObjectExistenceSensor and set deferrable attribute to True.",
|
153
149
|
category=AirflowProviderDeprecationWarning,
|
154
150
|
)
|
155
151
|
class GCSObjectExistenceAsyncSensor(GCSObjectExistenceSensor):
|
@@ -284,15 +280,9 @@ class GCSObjectUpdateSensor(BaseSensorOperator):
|
|
284
280
|
"Checking last updated time for object %s in bucket : %s", self.object, self.bucket
|
285
281
|
)
|
286
282
|
return event["message"]
|
287
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
288
|
-
if self.soft_fail:
|
289
|
-
raise AirflowSkipException(event["message"])
|
290
283
|
raise AirflowException(event["message"])
|
291
284
|
|
292
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
293
285
|
message = "No event received in trigger callback"
|
294
|
-
if self.soft_fail:
|
295
|
-
raise AirflowSkipException(message)
|
296
286
|
raise AirflowException(message)
|
297
287
|
|
298
288
|
|
@@ -382,9 +372,6 @@ class GCSObjectsWithPrefixExistenceSensor(BaseSensorOperator):
|
|
382
372
|
self.log.info("Resuming from trigger and checking status")
|
383
373
|
if event["status"] == "success":
|
384
374
|
return event["matches"]
|
385
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
386
|
-
if self.soft_fail:
|
387
|
-
raise AirflowSkipException(event["message"])
|
388
375
|
raise AirflowException(event["message"])
|
389
376
|
|
390
377
|
|
@@ -514,13 +501,10 @@ class GCSUploadSessionCompleteSensor(BaseSensorOperator):
|
|
514
501
|
)
|
515
502
|
return False
|
516
503
|
|
517
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
518
504
|
message = (
|
519
505
|
"Illegal behavior: objects were deleted in "
|
520
506
|
f"{os.path.join(self.bucket, self.prefix)} between pokes."
|
521
507
|
)
|
522
|
-
if self.soft_fail:
|
523
|
-
raise AirflowSkipException(message)
|
524
508
|
raise AirflowException(message)
|
525
509
|
|
526
510
|
if self.last_activity_time:
|
@@ -592,13 +576,7 @@ class GCSUploadSessionCompleteSensor(BaseSensorOperator):
|
|
592
576
|
if event:
|
593
577
|
if event["status"] == "success":
|
594
578
|
return event["message"]
|
595
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
596
|
-
if self.soft_fail:
|
597
|
-
raise AirflowSkipException(event["message"])
|
598
579
|
raise AirflowException(event["message"])
|
599
580
|
|
600
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
601
581
|
message = "No event received in trigger callback"
|
602
|
-
if self.soft_fail:
|
603
|
-
raise AirflowSkipException(message)
|
604
582
|
raise AirflowException(message)
|
@@ -21,7 +21,7 @@ from __future__ import annotations
|
|
21
21
|
|
22
22
|
from typing import TYPE_CHECKING
|
23
23
|
|
24
|
-
from airflow.exceptions import AirflowException
|
24
|
+
from airflow.exceptions import AirflowException
|
25
25
|
from airflow.providers.google.cloud.hooks.looker import JobStatus, LookerHook
|
26
26
|
from airflow.sensors.base import BaseSensorOperator
|
27
27
|
|
@@ -54,10 +54,7 @@ class LookerCheckPdtBuildSensor(BaseSensorOperator):
|
|
54
54
|
self.hook = LookerHook(looker_conn_id=self.looker_conn_id)
|
55
55
|
|
56
56
|
if not self.materialization_id:
|
57
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
58
57
|
message = "Invalid `materialization_id`."
|
59
|
-
if self.soft_fail:
|
60
|
-
raise AirflowSkipException(message)
|
61
58
|
raise AirflowException(message)
|
62
59
|
|
63
60
|
# materialization_id is templated var pulling output from start task
|
@@ -66,22 +63,13 @@ class LookerCheckPdtBuildSensor(BaseSensorOperator):
|
|
66
63
|
|
67
64
|
if status == JobStatus.ERROR.value:
|
68
65
|
msg = status_dict["message"]
|
69
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
70
66
|
message = f'PDT materialization job failed. Job id: {self.materialization_id}. Message:\n"{msg}"'
|
71
|
-
if self.soft_fail:
|
72
|
-
raise AirflowSkipException(message)
|
73
67
|
raise AirflowException(message)
|
74
68
|
elif status == JobStatus.CANCELLED.value:
|
75
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
76
69
|
message = f"PDT materialization job was cancelled. Job id: {self.materialization_id}."
|
77
|
-
if self.soft_fail:
|
78
|
-
raise AirflowSkipException(message)
|
79
70
|
raise AirflowException(message)
|
80
71
|
elif status == JobStatus.UNKNOWN.value:
|
81
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
82
72
|
message = f"PDT materialization job has unknown status. Job id: {self.materialization_id}."
|
83
|
-
if self.soft_fail:
|
84
|
-
raise AirflowSkipException(message)
|
85
73
|
raise AirflowException(message)
|
86
74
|
elif status == JobStatus.DONE.value:
|
87
75
|
self.log.debug(
|
@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence
|
|
25
25
|
from google.cloud.pubsub_v1.types import ReceivedMessage
|
26
26
|
|
27
27
|
from airflow.configuration import conf
|
28
|
-
from airflow.exceptions import AirflowException
|
28
|
+
from airflow.exceptions import AirflowException
|
29
29
|
from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
|
30
30
|
from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
|
31
31
|
from airflow.sensors.base import BaseSensorOperator
|
@@ -69,6 +69,13 @@ class PubSubPullSensor(BaseSensorOperator):
|
|
69
69
|
full subscription path.
|
70
70
|
:param max_messages: The maximum number of messages to retrieve per
|
71
71
|
PubSub pull request
|
72
|
+
:param return_immediately: If this field set to true, the system will
|
73
|
+
respond immediately even if it there are no messages available to
|
74
|
+
return in the ``Pull`` response. Otherwise, the system may wait
|
75
|
+
(for a bounded amount of time) until at least one message is available,
|
76
|
+
rather than returning no messages. Warning: setting this field to
|
77
|
+
``true`` is discouraged because it adversely impacts the performance
|
78
|
+
of ``Pull`` operations. We recommend that users do not set this field.
|
72
79
|
:param ack_messages: If True, each message will be acknowledged
|
73
80
|
immediately rather than by any downstream tasks
|
74
81
|
:param gcp_conn_id: The connection ID to use connecting to
|
@@ -102,6 +109,7 @@ class PubSubPullSensor(BaseSensorOperator):
|
|
102
109
|
project_id: str,
|
103
110
|
subscription: str,
|
104
111
|
max_messages: int = 5,
|
112
|
+
return_immediately: bool = True,
|
105
113
|
ack_messages: bool = False,
|
106
114
|
gcp_conn_id: str = "google_cloud_default",
|
107
115
|
messages_callback: Callable[[list[ReceivedMessage], Context], Any] | None = None,
|
@@ -115,6 +123,7 @@ class PubSubPullSensor(BaseSensorOperator):
|
|
115
123
|
self.project_id = project_id
|
116
124
|
self.subscription = subscription
|
117
125
|
self.max_messages = max_messages
|
126
|
+
self.return_immediately = return_immediately
|
118
127
|
self.ack_messages = ack_messages
|
119
128
|
self.messages_callback = messages_callback
|
120
129
|
self.impersonation_chain = impersonation_chain
|
@@ -132,7 +141,7 @@ class PubSubPullSensor(BaseSensorOperator):
|
|
132
141
|
project_id=self.project_id,
|
133
142
|
subscription=self.subscription,
|
134
143
|
max_messages=self.max_messages,
|
135
|
-
return_immediately=
|
144
|
+
return_immediately=self.return_immediately,
|
136
145
|
)
|
137
146
|
|
138
147
|
handle_messages = self.messages_callback or self._default_message_callback
|
@@ -175,9 +184,6 @@ class PubSubPullSensor(BaseSensorOperator):
|
|
175
184
|
self.log.info("Sensor pulls messages: %s", event["message"])
|
176
185
|
return event["message"]
|
177
186
|
self.log.info("Sensor failed: %s", event["message"])
|
178
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
179
|
-
if self.soft_fail:
|
180
|
-
raise AirflowSkipException(event["message"])
|
181
187
|
raise AirflowException(event["message"])
|
182
188
|
|
183
189
|
def _default_message_callback(
|
@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Sequence
|
|
21
21
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
22
22
|
from google.cloud.workflows.executions_v1beta import Execution
|
23
23
|
|
24
|
-
from airflow.exceptions import AirflowException
|
24
|
+
from airflow.exceptions import AirflowException
|
25
25
|
from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook
|
26
26
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
27
27
|
from airflow.sensors.base import BaseSensorOperator
|
@@ -101,13 +101,10 @@ class WorkflowExecutionSensor(BaseSensorOperator):
|
|
101
101
|
|
102
102
|
state = execution.state
|
103
103
|
if state in self.failure_states:
|
104
|
-
# TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
|
105
104
|
message = (
|
106
105
|
f"Execution {self.execution_id} for workflow {self.execution_id} "
|
107
106
|
f"failed and is in `{state}` state"
|
108
107
|
)
|
109
|
-
if self.soft_fail:
|
110
|
-
raise AirflowSkipException(message)
|
111
108
|
raise AirflowException(message)
|
112
109
|
|
113
110
|
if state in self.success_states:
|
@@ -133,6 +133,12 @@ class SFTPToGCSOperator(BaseOperator):
|
|
133
133
|
|
134
134
|
for file in files:
|
135
135
|
destination_path = file.replace(base_path, self.destination_path, 1)
|
136
|
+
# See issue: https://github.com/apache/airflow/issues/41763
|
137
|
+
# If the destination_path is not specified, it defaults to an empty string. As a result,
|
138
|
+
# replacing base_path with an empty string is ineffective, causing the destination_path to
|
139
|
+
# retain the "/" prefix, if it has.
|
140
|
+
if not self.destination_path:
|
141
|
+
destination_path = destination_path.lstrip("/")
|
136
142
|
self._copy_single_object(gcs_hook, sftp_hook, file, destination_path)
|
137
143
|
|
138
144
|
else:
|