apache-airflow-providers-google 10.14.0rc2__py3-none-any.whl → 10.15.0rc1__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/automl.py +13 -13
- airflow/providers/google/cloud/hooks/bigquery.py +193 -246
- airflow/providers/google/cloud/hooks/bigquery_dts.py +6 -6
- airflow/providers/google/cloud/hooks/bigtable.py +8 -8
- airflow/providers/google/cloud/hooks/cloud_batch.py +1 -1
- airflow/providers/google/cloud/hooks/cloud_build.py +19 -20
- airflow/providers/google/cloud/hooks/cloud_composer.py +4 -4
- airflow/providers/google/cloud/hooks/cloud_memorystore.py +10 -10
- airflow/providers/google/cloud/hooks/cloud_run.py +1 -1
- airflow/providers/google/cloud/hooks/cloud_sql.py +17 -17
- airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +3 -3
- airflow/providers/google/cloud/hooks/compute.py +16 -16
- airflow/providers/google/cloud/hooks/compute_ssh.py +1 -1
- airflow/providers/google/cloud/hooks/datacatalog.py +22 -22
- airflow/providers/google/cloud/hooks/dataflow.py +48 -49
- airflow/providers/google/cloud/hooks/dataform.py +16 -16
- airflow/providers/google/cloud/hooks/datafusion.py +15 -15
- airflow/providers/google/cloud/hooks/datapipeline.py +3 -3
- airflow/providers/google/cloud/hooks/dataplex.py +19 -19
- airflow/providers/google/cloud/hooks/dataprep.py +8 -8
- airflow/providers/google/cloud/hooks/dataproc.py +88 -0
- airflow/providers/google/cloud/hooks/dataproc_metastore.py +13 -13
- airflow/providers/google/cloud/hooks/datastore.py +3 -3
- airflow/providers/google/cloud/hooks/dlp.py +25 -25
- airflow/providers/google/cloud/hooks/gcs.py +25 -23
- airflow/providers/google/cloud/hooks/gdm.py +3 -3
- airflow/providers/google/cloud/hooks/kms.py +3 -3
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +63 -48
- airflow/providers/google/cloud/hooks/life_sciences.py +13 -12
- airflow/providers/google/cloud/hooks/looker.py +7 -7
- airflow/providers/google/cloud/hooks/mlengine.py +12 -12
- airflow/providers/google/cloud/hooks/natural_language.py +2 -2
- airflow/providers/google/cloud/hooks/os_login.py +1 -1
- airflow/providers/google/cloud/hooks/pubsub.py +9 -9
- airflow/providers/google/cloud/hooks/secret_manager.py +1 -1
- airflow/providers/google/cloud/hooks/spanner.py +11 -11
- airflow/providers/google/cloud/hooks/speech_to_text.py +1 -1
- airflow/providers/google/cloud/hooks/stackdriver.py +7 -7
- airflow/providers/google/cloud/hooks/tasks.py +11 -11
- airflow/providers/google/cloud/hooks/text_to_speech.py +1 -1
- airflow/providers/google/cloud/hooks/translate.py +1 -1
- airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +13 -13
- airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +6 -6
- airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +45 -50
- airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +13 -13
- airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +9 -9
- airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +128 -11
- airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +10 -10
- airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +8 -8
- airflow/providers/google/cloud/hooks/video_intelligence.py +2 -2
- airflow/providers/google/cloud/hooks/vision.py +1 -1
- airflow/providers/google/cloud/hooks/workflows.py +10 -10
- airflow/providers/google/cloud/links/datafusion.py +12 -5
- airflow/providers/google/cloud/operators/bigquery.py +9 -11
- airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +3 -1
- airflow/providers/google/cloud/operators/dataflow.py +16 -16
- airflow/providers/google/cloud/operators/datafusion.py +9 -1
- airflow/providers/google/cloud/operators/dataproc.py +298 -65
- airflow/providers/google/cloud/operators/kubernetes_engine.py +6 -6
- airflow/providers/google/cloud/operators/life_sciences.py +10 -9
- airflow/providers/google/cloud/operators/mlengine.py +96 -96
- airflow/providers/google/cloud/operators/pubsub.py +2 -0
- airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +33 -3
- airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +59 -2
- airflow/providers/google/cloud/secrets/secret_manager.py +8 -7
- airflow/providers/google/cloud/sensors/bigquery.py +20 -16
- airflow/providers/google/cloud/sensors/cloud_composer.py +11 -8
- airflow/providers/google/cloud/sensors/gcs.py +8 -7
- airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +4 -4
- airflow/providers/google/cloud/transfers/gcs_to_sftp.py +1 -1
- airflow/providers/google/cloud/transfers/mssql_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/mysql_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/oracle_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/postgres_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/presto_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/s3_to_gcs.py +3 -3
- airflow/providers/google/cloud/transfers/sftp_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/sql_to_gcs.py +3 -3
- airflow/providers/google/cloud/transfers/trino_to_gcs.py +1 -1
- airflow/providers/google/cloud/triggers/bigquery.py +12 -12
- airflow/providers/google/cloud/triggers/bigquery_dts.py +1 -1
- airflow/providers/google/cloud/triggers/cloud_batch.py +3 -1
- airflow/providers/google/cloud/triggers/cloud_build.py +2 -2
- airflow/providers/google/cloud/triggers/cloud_run.py +1 -1
- airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +6 -6
- airflow/providers/google/cloud/triggers/dataflow.py +3 -1
- airflow/providers/google/cloud/triggers/datafusion.py +2 -2
- airflow/providers/google/cloud/triggers/dataplex.py +2 -2
- airflow/providers/google/cloud/triggers/dataproc.py +2 -2
- airflow/providers/google/cloud/triggers/gcs.py +12 -8
- airflow/providers/google/cloud/triggers/kubernetes_engine.py +2 -2
- airflow/providers/google/cloud/triggers/mlengine.py +2 -2
- airflow/providers/google/cloud/triggers/pubsub.py +1 -1
- airflow/providers/google/cloud/triggers/vertex_ai.py +99 -0
- airflow/providers/google/cloud/utils/bigquery.py +2 -2
- airflow/providers/google/cloud/utils/credentials_provider.py +2 -2
- airflow/providers/google/cloud/utils/dataform.py +1 -1
- airflow/providers/google/cloud/utils/field_validator.py +2 -2
- airflow/providers/google/cloud/utils/helpers.py +2 -2
- airflow/providers/google/cloud/utils/mlengine_operator_utils.py +1 -1
- airflow/providers/google/cloud/utils/mlengine_prediction_summary.py +1 -1
- airflow/providers/google/common/auth_backend/google_openid.py +2 -2
- airflow/providers/google/common/hooks/base_google.py +29 -22
- airflow/providers/google/common/hooks/discovery_api.py +2 -2
- airflow/providers/google/common/utils/id_token_credentials.py +5 -5
- airflow/providers/google/firebase/hooks/firestore.py +3 -3
- airflow/providers/google/get_provider_info.py +7 -2
- airflow/providers/google/leveldb/hooks/leveldb.py +2 -2
- airflow/providers/google/marketing_platform/hooks/analytics.py +11 -14
- airflow/providers/google/marketing_platform/hooks/campaign_manager.py +11 -11
- airflow/providers/google/marketing_platform/hooks/display_video.py +13 -13
- airflow/providers/google/marketing_platform/hooks/search_ads.py +4 -4
- airflow/providers/google/marketing_platform/operators/analytics.py +37 -32
- airflow/providers/google/suite/hooks/calendar.py +2 -2
- airflow/providers/google/suite/hooks/drive.py +7 -7
- airflow/providers/google/suite/hooks/sheets.py +8 -8
- {apache_airflow_providers_google-10.14.0rc2.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/METADATA +11 -11
- {apache_airflow_providers_google-10.14.0rc2.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/RECORD +121 -120
- {apache_airflow_providers_google-10.14.0rc2.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/WHEEL +0 -0
- {apache_airflow_providers_google-10.14.0rc2.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/entry_points.txt +0 -0
@@ -19,9 +19,10 @@
|
|
19
19
|
|
20
20
|
from __future__ import annotations
|
21
21
|
|
22
|
-
import warnings
|
23
22
|
from typing import TYPE_CHECKING, Any, Sequence
|
24
23
|
|
24
|
+
from deprecated import deprecated
|
25
|
+
|
25
26
|
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
|
26
27
|
from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerExecutionTrigger
|
27
28
|
from airflow.sensors.base import BaseSensorOperator
|
@@ -30,6 +31,15 @@ if TYPE_CHECKING:
|
|
30
31
|
from airflow.utils.context import Context
|
31
32
|
|
32
33
|
|
34
|
+
@deprecated(
|
35
|
+
reason=(
|
36
|
+
"The `CloudComposerEnvironmentSensor` operator is deprecated. "
|
37
|
+
"You can achieve the same functionality "
|
38
|
+
"by using operators in deferrable or non-deferrable mode, since every operator for Cloud "
|
39
|
+
"Composer will wait for the operation to complete."
|
40
|
+
),
|
41
|
+
category=AirflowProviderDeprecationWarning,
|
42
|
+
)
|
33
43
|
class CloudComposerEnvironmentSensor(BaseSensorOperator):
|
34
44
|
"""
|
35
45
|
Check the status of the Cloud Composer Environment task.
|
@@ -65,13 +75,6 @@ class CloudComposerEnvironmentSensor(BaseSensorOperator):
|
|
65
75
|
pooling_period_seconds: int = 30,
|
66
76
|
**kwargs,
|
67
77
|
):
|
68
|
-
warnings.warn(
|
69
|
-
f"The `{self.__class__.__name__}` operator is deprecated. You can achieve the same functionality "
|
70
|
-
f"by using operators in deferrable or non-deferrable mode, since every operator for Cloud "
|
71
|
-
f"Composer will wait for the operation to complete.",
|
72
|
-
AirflowProviderDeprecationWarning,
|
73
|
-
stacklevel=2,
|
74
|
-
)
|
75
78
|
super().__init__(**kwargs)
|
76
79
|
self.project_id = project_id
|
77
80
|
self.region = region
|
@@ -20,10 +20,10 @@ from __future__ import annotations
|
|
20
20
|
|
21
21
|
import os
|
22
22
|
import textwrap
|
23
|
-
import warnings
|
24
23
|
from datetime import datetime, timedelta
|
25
24
|
from typing import TYPE_CHECKING, Any, Callable, Sequence
|
26
25
|
|
26
|
+
from deprecated import deprecated
|
27
27
|
from google.cloud.storage.retry import DEFAULT_RETRY
|
28
28
|
|
29
29
|
from airflow.configuration import conf
|
@@ -142,6 +142,13 @@ class GCSObjectExistenceSensor(BaseSensorOperator):
|
|
142
142
|
return event["message"]
|
143
143
|
|
144
144
|
|
145
|
+
@deprecated(
|
146
|
+
reason=(
|
147
|
+
"Class `GCSObjectExistenceAsyncSensor` is deprecated and will be removed in a future release. "
|
148
|
+
"Please use `GCSObjectExistenceSensor` and set `deferrable` attribute to `True` instead"
|
149
|
+
),
|
150
|
+
category=AirflowProviderDeprecationWarning,
|
151
|
+
)
|
145
152
|
class GCSObjectExistenceAsyncSensor(GCSObjectExistenceSensor):
|
146
153
|
"""
|
147
154
|
Checks for the existence of a file in Google Cloud Storage.
|
@@ -165,12 +172,6 @@ class GCSObjectExistenceAsyncSensor(GCSObjectExistenceSensor):
|
|
165
172
|
"""
|
166
173
|
|
167
174
|
def __init__(self, **kwargs: Any) -> None:
|
168
|
-
warnings.warn(
|
169
|
-
"Class `GCSObjectExistenceAsyncSensor` is deprecated and will be removed in a future release. "
|
170
|
-
"Please use `GCSObjectExistenceSensor` and set `deferrable` attribute to `True` instead",
|
171
|
-
AirflowProviderDeprecationWarning,
|
172
|
-
stacklevel=2,
|
173
|
-
)
|
174
175
|
super().__init__(deferrable=True, **kwargs)
|
175
176
|
|
176
177
|
|
@@ -218,7 +218,7 @@ class CassandraToGCSOperator(BaseOperator):
|
|
218
218
|
|
219
219
|
def _write_local_schema_file(self, cursor):
|
220
220
|
"""
|
221
|
-
|
221
|
+
Take a cursor, and writes the BigQuery schema for the results to a local file system.
|
222
222
|
|
223
223
|
:return: A dictionary where key is a filename to be used as an object
|
224
224
|
name in GCS, and values are file handles to local files that
|
@@ -253,7 +253,7 @@ class CassandraToGCSOperator(BaseOperator):
|
|
253
253
|
)
|
254
254
|
|
255
255
|
def generate_data_dict(self, names: Iterable[str], values: Any) -> dict[str, Any]:
|
256
|
-
"""
|
256
|
+
"""Generate data structure that will be stored as file in GCS."""
|
257
257
|
return {n: self.convert_value(v) for n, v in zip(names, values)}
|
258
258
|
|
259
259
|
def convert_value(self, value: Any | None) -> Any | None:
|
@@ -285,12 +285,12 @@ class CassandraToGCSOperator(BaseOperator):
|
|
285
285
|
raise AirflowException(f"Unexpected value: {value}")
|
286
286
|
|
287
287
|
def convert_array_types(self, value: list[Any] | SortedSet) -> list[Any]:
|
288
|
-
"""
|
288
|
+
"""Map convert_value over array."""
|
289
289
|
return [self.convert_value(nested_value) for nested_value in value]
|
290
290
|
|
291
291
|
def convert_user_type(self, value: Any) -> dict[str, Any]:
|
292
292
|
"""
|
293
|
-
|
293
|
+
Convert a user type to RECORD that contains n fields, where n is the number of attributes.
|
294
294
|
|
295
295
|
Each element in the user type class will be converted to its corresponding data type in BQ.
|
296
296
|
"""
|
@@ -176,7 +176,7 @@ class GCSToSFTPOperator(BaseOperator):
|
|
176
176
|
source_object: str,
|
177
177
|
destination_path: str,
|
178
178
|
) -> None:
|
179
|
-
"""
|
179
|
+
"""Copy single object."""
|
180
180
|
self.log.info(
|
181
181
|
"Executing copy of gs://%s/%s to %s",
|
182
182
|
self.source_bucket,
|
@@ -67,7 +67,7 @@ class MySQLToGCSOperator(BaseSQLToGCSOperator):
|
|
67
67
|
self.ensure_utc = ensure_utc
|
68
68
|
|
69
69
|
def query(self):
|
70
|
-
"""
|
70
|
+
"""Query mysql and returns a cursor to the results."""
|
71
71
|
mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id)
|
72
72
|
conn = mysql.get_conn()
|
73
73
|
cursor = conn.cursor()
|
@@ -62,7 +62,7 @@ class OracleToGCSOperator(BaseSQLToGCSOperator):
|
|
62
62
|
self.oracle_conn_id = oracle_conn_id
|
63
63
|
|
64
64
|
def query(self):
|
65
|
-
"""
|
65
|
+
"""Query Oracle and returns a cursor to the results."""
|
66
66
|
oracle = OracleHook(oracle_conn_id=self.oracle_conn_id)
|
67
67
|
conn = oracle.get_conn()
|
68
68
|
cursor = conn.cursor()
|
@@ -115,7 +115,7 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
|
|
115
115
|
return f"{self.dag_id}__{self.task_id}__{uuid.uuid4()}" if self.use_server_side_cursor else None
|
116
116
|
|
117
117
|
def query(self):
|
118
|
-
"""
|
118
|
+
"""Query Postgres and returns a cursor to the results."""
|
119
119
|
hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
|
120
120
|
conn = hook.get_conn()
|
121
121
|
cursor = conn.cursor(name=self._unique_name())
|
@@ -181,7 +181,7 @@ class PrestoToGCSOperator(BaseSQLToGCSOperator):
|
|
181
181
|
self.presto_conn_id = presto_conn_id
|
182
182
|
|
183
183
|
def query(self):
|
184
|
-
"""
|
184
|
+
"""Query presto and returns a cursor to the results."""
|
185
185
|
presto = PrestoHook(presto_conn_id=self.presto_conn_id)
|
186
186
|
conn = presto.get_conn()
|
187
187
|
cursor = conn.cursor()
|
@@ -219,7 +219,7 @@ class S3ToGCSOperator(S3ListOperator):
|
|
219
219
|
|
220
220
|
def s3_to_gcs_object(self, s3_object: str) -> str:
|
221
221
|
"""
|
222
|
-
|
222
|
+
Transform S3 path to GCS path according to the operator's logic.
|
223
223
|
|
224
224
|
If apply_gcs_prefix == True then <s3_prefix><content> => <gcs_prefix><content>
|
225
225
|
If apply_gcs_prefix == False then <s3_prefix><content> => <gcs_prefix><s3_prefix><content>
|
@@ -233,7 +233,7 @@ class S3ToGCSOperator(S3ListOperator):
|
|
233
233
|
|
234
234
|
def gcs_to_s3_object(self, gcs_object: str) -> str:
|
235
235
|
"""
|
236
|
-
|
236
|
+
Transform GCS path to S3 path according to the operator's logic.
|
237
237
|
|
238
238
|
If apply_gcs_prefix == True then <gcs_prefix><content> => <s3_prefix><content>
|
239
239
|
If apply_gcs_prefix == False then <gcs_prefix><s3_prefix><content> => <s3_prefix><content>
|
@@ -261,7 +261,7 @@ class S3ToGCSOperator(S3ListOperator):
|
|
261
261
|
self.log.info("All done, uploaded %d files to Google Cloud Storage", len(s3_objects))
|
262
262
|
|
263
263
|
def transfer_files_async(self, files: list[str], gcs_hook: GCSHook, s3_hook: S3Hook) -> None:
|
264
|
-
"""
|
264
|
+
"""Submit Google Cloud Storage Transfer Service job to copy files from AWS S3 to GCS."""
|
265
265
|
if not len(files):
|
266
266
|
raise ValueError("List of transferring files cannot be empty")
|
267
267
|
job_names = self.submit_transfer_jobs(files=files, gcs_hook=gcs_hook, s3_hook=s3_hook)
|
@@ -147,7 +147,7 @@ class SFTPToGCSOperator(BaseOperator):
|
|
147
147
|
source_path: str,
|
148
148
|
destination_object: str,
|
149
149
|
) -> None:
|
150
|
-
"""
|
150
|
+
"""Copy single object."""
|
151
151
|
self.log.info(
|
152
152
|
"Executing copy of %s to gs://%s/%s",
|
153
153
|
source_path,
|
@@ -228,7 +228,7 @@ class BaseSQLToGCSOperator(BaseOperator):
|
|
228
228
|
|
229
229
|
def _write_local_data_files(self, cursor):
|
230
230
|
"""
|
231
|
-
|
231
|
+
Take a cursor, and writes results to a local file.
|
232
232
|
|
233
233
|
:return: A dictionary where keys are filenames to be used as object
|
234
234
|
names in GCS, and values are file handles to local files that
|
@@ -348,7 +348,7 @@ class BaseSQLToGCSOperator(BaseOperator):
|
|
348
348
|
yield file_to_upload
|
349
349
|
|
350
350
|
def _get_file_to_upload(self, file_mime_type, file_no):
|
351
|
-
"""
|
351
|
+
"""Return a dictionary that represents the file to upload."""
|
352
352
|
tmp_file_handle = NamedTemporaryFile(mode="w", encoding="utf-8", delete=True)
|
353
353
|
return (
|
354
354
|
{
|
@@ -435,7 +435,7 @@ class BaseSQLToGCSOperator(BaseOperator):
|
|
435
435
|
|
436
436
|
def _write_local_schema_file(self, cursor):
|
437
437
|
"""
|
438
|
-
|
438
|
+
Take a cursor, and writes the BigQuery schema for the results to a local file system.
|
439
439
|
|
440
440
|
Schema for database will be read from cursor if not specified.
|
441
441
|
|
@@ -181,7 +181,7 @@ class TrinoToGCSOperator(BaseSQLToGCSOperator):
|
|
181
181
|
self.trino_conn_id = trino_conn_id
|
182
182
|
|
183
183
|
def query(self):
|
184
|
-
"""
|
184
|
+
"""Query trino and returns a cursor to the results."""
|
185
185
|
trino = TrinoHook(trino_conn_id=self.trino_conn_id)
|
186
186
|
conn = trino.get_conn()
|
187
187
|
cursor = conn.cursor()
|
@@ -68,7 +68,7 @@ class BigQueryInsertJobTrigger(BaseTrigger):
|
|
68
68
|
self.impersonation_chain = impersonation_chain
|
69
69
|
|
70
70
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
71
|
-
"""
|
71
|
+
"""Serialize BigQueryInsertJobTrigger arguments and classpath."""
|
72
72
|
return (
|
73
73
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryInsertJobTrigger",
|
74
74
|
{
|
@@ -83,7 +83,7 @@ class BigQueryInsertJobTrigger(BaseTrigger):
|
|
83
83
|
)
|
84
84
|
|
85
85
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
86
|
-
"""
|
86
|
+
"""Get current job execution status and yields a TriggerEvent."""
|
87
87
|
hook = self._get_async_hook()
|
88
88
|
try:
|
89
89
|
while True:
|
@@ -119,7 +119,7 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
|
|
119
119
|
"""BigQueryCheckTrigger run on the trigger worker."""
|
120
120
|
|
121
121
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
122
|
-
"""
|
122
|
+
"""Serialize BigQueryCheckTrigger arguments and classpath."""
|
123
123
|
return (
|
124
124
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryCheckTrigger",
|
125
125
|
{
|
@@ -134,7 +134,7 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
|
|
134
134
|
)
|
135
135
|
|
136
136
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
137
|
-
"""
|
137
|
+
"""Get current job execution status and yields a TriggerEvent."""
|
138
138
|
hook = self._get_async_hook()
|
139
139
|
try:
|
140
140
|
while True:
|
@@ -193,7 +193,7 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
|
|
193
193
|
self.as_dict = as_dict
|
194
194
|
|
195
195
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
196
|
-
"""
|
196
|
+
"""Serialize BigQueryInsertJobTrigger arguments and classpath."""
|
197
197
|
return (
|
198
198
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryGetDataTrigger",
|
199
199
|
{
|
@@ -209,7 +209,7 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
|
|
209
209
|
)
|
210
210
|
|
211
211
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
212
|
-
"""
|
212
|
+
"""Get current job execution status and yields a TriggerEvent with response data."""
|
213
213
|
hook = self._get_async_hook()
|
214
214
|
try:
|
215
215
|
while True:
|
@@ -307,7 +307,7 @@ class BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
|
|
307
307
|
self.ignore_zero = ignore_zero
|
308
308
|
|
309
309
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
310
|
-
"""
|
310
|
+
"""Serialize BigQueryCheckTrigger arguments and classpath."""
|
311
311
|
return (
|
312
312
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryIntervalCheckTrigger",
|
313
313
|
{
|
@@ -326,7 +326,7 @@ class BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
|
|
326
326
|
)
|
327
327
|
|
328
328
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
329
|
-
"""
|
329
|
+
"""Get current job execution status and yields a TriggerEvent."""
|
330
330
|
hook = self._get_async_hook()
|
331
331
|
try:
|
332
332
|
while True:
|
@@ -452,7 +452,7 @@ class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
|
|
452
452
|
self.tolerance = tolerance
|
453
453
|
|
454
454
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
455
|
-
"""
|
455
|
+
"""Serialize BigQueryValueCheckTrigger arguments and classpath."""
|
456
456
|
return (
|
457
457
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryValueCheckTrigger",
|
458
458
|
{
|
@@ -470,7 +470,7 @@ class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
|
|
470
470
|
)
|
471
471
|
|
472
472
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
473
|
-
"""
|
473
|
+
"""Get current job execution status and yields a TriggerEvent."""
|
474
474
|
hook = self._get_async_hook()
|
475
475
|
try:
|
476
476
|
while True:
|
@@ -536,7 +536,7 @@ class BigQueryTableExistenceTrigger(BaseTrigger):
|
|
536
536
|
self.impersonation_chain = impersonation_chain
|
537
537
|
|
538
538
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
539
|
-
"""
|
539
|
+
"""Serialize BigQueryTableExistenceTrigger arguments and classpath."""
|
540
540
|
return (
|
541
541
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryTableExistenceTrigger",
|
542
542
|
{
|
@@ -623,7 +623,7 @@ class BigQueryTablePartitionExistenceTrigger(BigQueryTableExistenceTrigger):
|
|
623
623
|
self.partition_id = partition_id
|
624
624
|
|
625
625
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
626
|
-
"""
|
626
|
+
"""Serialize BigQueryTablePartitionExistenceTrigger arguments and classpath."""
|
627
627
|
return (
|
628
628
|
"airflow.providers.google.cloud.triggers.bigquery.BigQueryTablePartitionExistenceTrigger",
|
629
629
|
{
|
@@ -66,7 +66,7 @@ class BigQueryDataTransferRunTrigger(BaseTrigger):
|
|
66
66
|
self.impersonation_chain = impersonation_chain
|
67
67
|
|
68
68
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
69
|
-
"""
|
69
|
+
"""Serialize class arguments and classpath."""
|
70
70
|
return (
|
71
71
|
"airflow.providers.google.cloud.triggers.bigquery_dts.BigQueryDataTransferRunTrigger",
|
72
72
|
{
|
@@ -67,7 +67,7 @@ class CloudBatchJobFinishedTrigger(BaseTrigger):
|
|
67
67
|
self.impersonation_chain = impersonation_chain
|
68
68
|
|
69
69
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
70
|
-
"""
|
70
|
+
"""Serialize class arguments and classpath."""
|
71
71
|
return (
|
72
72
|
"airflow.providers.google.cloud.triggers.cloud_batch.CloudBatchJobFinishedTrigger",
|
73
73
|
{
|
@@ -83,6 +83,8 @@ class CloudBatchJobFinishedTrigger(BaseTrigger):
|
|
83
83
|
|
84
84
|
async def run(self) -> AsyncIterator[TriggerEvent]:
|
85
85
|
"""
|
86
|
+
Fetch job status or yield certain Events.
|
87
|
+
|
86
88
|
Main loop of the class in where it is fetching the job status and yields certain Event.
|
87
89
|
|
88
90
|
If the job has status success then it yields TriggerEvent with success status, if job has
|
@@ -62,7 +62,7 @@ class CloudBuildCreateBuildTrigger(BaseTrigger):
|
|
62
62
|
self.location = location
|
63
63
|
|
64
64
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
65
|
-
"""
|
65
|
+
"""Serialize CloudBuildCreateBuildTrigger arguments and classpath."""
|
66
66
|
return (
|
67
67
|
"airflow.providers.google.cloud.triggers.cloud_build.CloudBuildCreateBuildTrigger",
|
68
68
|
{
|
@@ -76,7 +76,7 @@ class CloudBuildCreateBuildTrigger(BaseTrigger):
|
|
76
76
|
)
|
77
77
|
|
78
78
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
79
|
-
"""
|
79
|
+
"""Get current build execution status and yields a TriggerEvent."""
|
80
80
|
hook = self._get_async_hook()
|
81
81
|
try:
|
82
82
|
while True:
|
@@ -81,7 +81,7 @@ class CloudRunJobFinishedTrigger(BaseTrigger):
|
|
81
81
|
self.impersonation_chain = impersonation_chain
|
82
82
|
|
83
83
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
84
|
-
"""
|
84
|
+
"""Serialize class arguments and classpath."""
|
85
85
|
return (
|
86
86
|
"airflow.providers.google.cloud.triggers.cloud_run.CloudRunJobFinishedTrigger",
|
87
87
|
{
|
@@ -18,7 +18,7 @@
|
|
18
18
|
from __future__ import annotations
|
19
19
|
|
20
20
|
import asyncio
|
21
|
-
from typing import Any, AsyncIterator
|
21
|
+
from typing import Any, AsyncIterator, Iterable
|
22
22
|
|
23
23
|
from google.api_core.exceptions import GoogleAPIError
|
24
24
|
from google.cloud.storage_transfer_v1.types import TransferOperation
|
@@ -46,7 +46,7 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
|
|
46
46
|
self.poll_interval = poll_interval
|
47
47
|
|
48
48
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
49
|
-
"""
|
49
|
+
"""Serialize StorageTransferJobsTrigger arguments and classpath."""
|
50
50
|
return (
|
51
51
|
f"{self.__class__.__module__ }.{self.__class__.__qualname__}",
|
52
52
|
{
|
@@ -57,7 +57,7 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
|
|
57
57
|
)
|
58
58
|
|
59
59
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
60
|
-
"""
|
60
|
+
"""Get current data storage transfer jobs and yields a TriggerEvent."""
|
61
61
|
async_hook: CloudDataTransferServiceAsyncHook = self.get_async_hook()
|
62
62
|
|
63
63
|
while True:
|
@@ -67,11 +67,11 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
|
|
67
67
|
jobs_pager = await async_hook.get_jobs(job_names=self.job_names)
|
68
68
|
jobs, awaitable_operations = [], []
|
69
69
|
async for job in jobs_pager:
|
70
|
-
|
70
|
+
awaitable_operation = async_hook.get_latest_operation(job)
|
71
71
|
jobs.append(job)
|
72
|
-
awaitable_operations.append(
|
72
|
+
awaitable_operations.append(awaitable_operation)
|
73
73
|
|
74
|
-
operations:
|
74
|
+
operations: Iterable[TransferOperation | None] = await asyncio.gather(*awaitable_operations)
|
75
75
|
|
76
76
|
for job, operation in zip(jobs, operations):
|
77
77
|
if operation is None:
|
@@ -69,7 +69,7 @@ class TemplateJobStartTrigger(BaseTrigger):
|
|
69
69
|
self.cancel_timeout = cancel_timeout
|
70
70
|
|
71
71
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
72
|
-
"""
|
72
|
+
"""Serialize class arguments and classpath."""
|
73
73
|
return (
|
74
74
|
"airflow.providers.google.cloud.triggers.dataflow.TemplateJobStartTrigger",
|
75
75
|
{
|
@@ -85,6 +85,8 @@ class TemplateJobStartTrigger(BaseTrigger):
|
|
85
85
|
|
86
86
|
async def run(self):
|
87
87
|
"""
|
88
|
+
Fetch job status or yield certain Events.
|
89
|
+
|
88
90
|
Main loop of the class in where it is fetching the job status and yields certain Event.
|
89
91
|
|
90
92
|
If the job has status success then it yields TriggerEvent with success status, if job has
|
@@ -71,7 +71,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
|
|
71
71
|
self.success_states = success_states
|
72
72
|
|
73
73
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
74
|
-
"""
|
74
|
+
"""Serialize DataFusionStartPipelineTrigger arguments and classpath."""
|
75
75
|
return (
|
76
76
|
"airflow.providers.google.cloud.triggers.datafusion.DataFusionStartPipelineTrigger",
|
77
77
|
{
|
@@ -86,7 +86,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
|
|
86
86
|
)
|
87
87
|
|
88
88
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
89
|
-
"""
|
89
|
+
"""Get current pipeline status and yields a TriggerEvent."""
|
90
90
|
hook = self._get_async_hook()
|
91
91
|
try:
|
92
92
|
while True:
|
@@ -105,7 +105,7 @@ class DataplexDataQualityJobTrigger(BaseTrigger):
|
|
105
105
|
yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": self._convert_to_dict(job)})
|
106
106
|
|
107
107
|
def _convert_to_dict(self, job: DataScanJob) -> dict:
|
108
|
-
"""
|
108
|
+
"""Return a representation of a DataScanJob instance as a dict."""
|
109
109
|
return DataScanJob.to_dict(job)
|
110
110
|
|
111
111
|
|
@@ -187,5 +187,5 @@ class DataplexDataProfileJobTrigger(BaseTrigger):
|
|
187
187
|
yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": self._convert_to_dict(job)})
|
188
188
|
|
189
189
|
def _convert_to_dict(self, job: DataScanJob) -> dict:
|
190
|
-
"""
|
190
|
+
"""Return a representation of a DataScanJob instance as a dict."""
|
191
191
|
return DataScanJob.to_dict(job)
|
@@ -182,7 +182,7 @@ class DataprocBatchTrigger(DataprocBaseTrigger):
|
|
182
182
|
self.batch_id = batch_id
|
183
183
|
|
184
184
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
185
|
-
"""
|
185
|
+
"""Serialize DataprocBatchTrigger arguments and classpath."""
|
186
186
|
return (
|
187
187
|
"airflow.providers.google.cloud.triggers.dataproc.DataprocBatchTrigger",
|
188
188
|
{
|
@@ -244,7 +244,7 @@ class DataprocDeleteClusterTrigger(DataprocBaseTrigger):
|
|
244
244
|
self.metadata = metadata
|
245
245
|
|
246
246
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
247
|
-
"""
|
247
|
+
"""Serialize DataprocDeleteClusterTrigger arguments and classpath."""
|
248
248
|
return (
|
249
249
|
"airflow.providers.google.cloud.triggers.dataproc.DataprocDeleteClusterTrigger",
|
250
250
|
{
|
@@ -60,7 +60,7 @@ class GCSBlobTrigger(BaseTrigger):
|
|
60
60
|
self.hook_params = hook_params
|
61
61
|
|
62
62
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
63
|
-
"""
|
63
|
+
"""Serialize GCSBlobTrigger arguments and classpath."""
|
64
64
|
return (
|
65
65
|
"airflow.providers.google.cloud.triggers.gcs.GCSBlobTrigger",
|
66
66
|
{
|
@@ -93,7 +93,7 @@ class GCSBlobTrigger(BaseTrigger):
|
|
93
93
|
|
94
94
|
async def _object_exists(self, hook: GCSAsyncHook, bucket_name: str, object_name: str) -> str:
|
95
95
|
"""
|
96
|
-
|
96
|
+
Check for the existence of a file in Google Cloud Storage.
|
97
97
|
|
98
98
|
:param bucket_name: The Google Cloud Storage bucket where the object is.
|
99
99
|
:param object_name: The name of the blob_name to check in the Google cloud
|
@@ -143,7 +143,7 @@ class GCSCheckBlobUpdateTimeTrigger(BaseTrigger):
|
|
143
143
|
self.hook_params = hook_params
|
144
144
|
|
145
145
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
146
|
-
"""
|
146
|
+
"""Serialize GCSCheckBlobUpdateTimeTrigger arguments and classpath."""
|
147
147
|
return (
|
148
148
|
"airflow.providers.google.cloud.triggers.gcs.GCSCheckBlobUpdateTimeTrigger",
|
149
149
|
{
|
@@ -181,7 +181,7 @@ class GCSCheckBlobUpdateTimeTrigger(BaseTrigger):
|
|
181
181
|
self, hook: GCSAsyncHook, bucket_name: str, object_name: str, target_date: datetime
|
182
182
|
) -> tuple[bool, dict[str, Any]]:
|
183
183
|
"""
|
184
|
-
|
184
|
+
Check if the object in the bucket is updated.
|
185
185
|
|
186
186
|
:param hook: GCSAsyncHook Hook class
|
187
187
|
:param bucket_name: The Google Cloud Storage bucket where the object is.
|
@@ -248,7 +248,7 @@ class GCSPrefixBlobTrigger(GCSBlobTrigger):
|
|
248
248
|
self.prefix = prefix
|
249
249
|
|
250
250
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
251
|
-
"""
|
251
|
+
"""Serialize GCSPrefixBlobTrigger arguments and classpath."""
|
252
252
|
return (
|
253
253
|
"airflow.providers.google.cloud.triggers.gcs.GCSPrefixBlobTrigger",
|
254
254
|
{
|
@@ -282,7 +282,7 @@ class GCSPrefixBlobTrigger(GCSBlobTrigger):
|
|
282
282
|
|
283
283
|
async def _list_blobs_with_prefix(self, hook: GCSAsyncHook, bucket_name: str, prefix: str) -> list[str]:
|
284
284
|
"""
|
285
|
-
|
285
|
+
Return names of blobs which match the given prefix for a given bucket.
|
286
286
|
|
287
287
|
:param hook: The async hook to use for listing the blobs
|
288
288
|
:param bucket_name: The Google Cloud Storage bucket where the object is.
|
@@ -344,7 +344,7 @@ class GCSUploadSessionTrigger(GCSPrefixBlobTrigger):
|
|
344
344
|
self.last_activity_time: datetime | None = None
|
345
345
|
|
346
346
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
347
|
-
"""
|
347
|
+
"""Serialize GCSUploadSessionTrigger arguments and classpath."""
|
348
348
|
return (
|
349
349
|
"airflow.providers.google.cloud.triggers.gcs.GCSUploadSessionTrigger",
|
350
350
|
{
|
@@ -377,7 +377,11 @@ class GCSUploadSessionTrigger(GCSPrefixBlobTrigger):
|
|
377
377
|
yield TriggerEvent({"status": "error", "message": str(e)})
|
378
378
|
|
379
379
|
def _get_time(self) -> datetime:
|
380
|
-
"""
|
380
|
+
"""
|
381
|
+
Get current local date and time.
|
382
|
+
|
383
|
+
This is just a wrapper of datetime.datetime.now to simplify mocking in the unittests.
|
384
|
+
"""
|
381
385
|
return datetime.now()
|
382
386
|
|
383
387
|
def _is_bucket_updated(self, current_objects: set[str]) -> dict[str, str]:
|
@@ -174,7 +174,7 @@ class GKEOperationTrigger(BaseTrigger):
|
|
174
174
|
self._hook: GKEAsyncHook | None = None
|
175
175
|
|
176
176
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
177
|
-
"""
|
177
|
+
"""Serialize GKEOperationTrigger arguments and classpath."""
|
178
178
|
return (
|
179
179
|
"airflow.providers.google.cloud.triggers.kubernetes_engine.GKEOperationTrigger",
|
180
180
|
{
|
@@ -188,7 +188,7 @@ class GKEOperationTrigger(BaseTrigger):
|
|
188
188
|
)
|
189
189
|
|
190
190
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
191
|
-
"""
|
191
|
+
"""Get operation status and yields corresponding event."""
|
192
192
|
hook = self._get_hook()
|
193
193
|
try:
|
194
194
|
while True:
|
@@ -69,7 +69,7 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
|
|
69
69
|
self.impersonation_chain = impersonation_chain
|
70
70
|
|
71
71
|
def serialize(self) -> tuple[str, dict[str, Any]]:
|
72
|
-
"""
|
72
|
+
"""Serialize MLEngineStartTrainingJobTrigger arguments and classpath."""
|
73
73
|
return (
|
74
74
|
"airflow.providers.google.cloud.triggers.mlengine.MLEngineStartTrainingJobTrigger",
|
75
75
|
{
|
@@ -89,7 +89,7 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
|
|
89
89
|
)
|
90
90
|
|
91
91
|
async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
|
92
|
-
"""
|
92
|
+
"""Get current job execution status and yields a TriggerEvent."""
|
93
93
|
hook = self._get_async_hook()
|
94
94
|
try:
|
95
95
|
while True:
|