apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.3.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/3rd-party-licenses/NOTICE +2 -12
- airflow/providers/google/__init__.py +3 -3
- airflow/providers/google/ads/hooks/ads.py +39 -6
- airflow/providers/google/ads/operators/ads.py +2 -2
- airflow/providers/google/ads/transfers/ads_to_gcs.py +2 -2
- airflow/providers/google/assets/gcs.py +1 -11
- airflow/providers/google/cloud/bundles/__init__.py +16 -0
- airflow/providers/google/cloud/bundles/gcs.py +161 -0
- airflow/providers/google/cloud/hooks/alloy_db.py +1 -1
- airflow/providers/google/cloud/hooks/bigquery.py +176 -293
- airflow/providers/google/cloud/hooks/cloud_batch.py +1 -1
- airflow/providers/google/cloud/hooks/cloud_build.py +1 -1
- airflow/providers/google/cloud/hooks/cloud_composer.py +288 -15
- airflow/providers/google/cloud/hooks/cloud_logging.py +109 -0
- airflow/providers/google/cloud/hooks/cloud_memorystore.py +1 -1
- airflow/providers/google/cloud/hooks/cloud_run.py +18 -10
- airflow/providers/google/cloud/hooks/cloud_sql.py +102 -23
- airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +29 -7
- airflow/providers/google/cloud/hooks/compute.py +1 -1
- airflow/providers/google/cloud/hooks/compute_ssh.py +6 -2
- airflow/providers/google/cloud/hooks/datacatalog.py +10 -1
- airflow/providers/google/cloud/hooks/dataflow.py +72 -95
- airflow/providers/google/cloud/hooks/dataform.py +1 -1
- airflow/providers/google/cloud/hooks/datafusion.py +21 -19
- airflow/providers/google/cloud/hooks/dataplex.py +2 -2
- airflow/providers/google/cloud/hooks/dataprep.py +1 -1
- airflow/providers/google/cloud/hooks/dataproc.py +73 -72
- airflow/providers/google/cloud/hooks/dataproc_metastore.py +1 -1
- airflow/providers/google/cloud/hooks/dlp.py +1 -1
- airflow/providers/google/cloud/hooks/functions.py +1 -1
- airflow/providers/google/cloud/hooks/gcs.py +112 -15
- airflow/providers/google/cloud/hooks/gdm.py +1 -1
- airflow/providers/google/cloud/hooks/gen_ai.py +196 -0
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +3 -3
- airflow/providers/google/cloud/hooks/looker.py +6 -2
- airflow/providers/google/cloud/hooks/managed_kafka.py +1 -1
- airflow/providers/google/cloud/hooks/mlengine.py +4 -3
- airflow/providers/google/cloud/hooks/pubsub.py +3 -0
- airflow/providers/google/cloud/hooks/secret_manager.py +102 -10
- airflow/providers/google/cloud/hooks/spanner.py +74 -9
- airflow/providers/google/cloud/hooks/stackdriver.py +11 -9
- airflow/providers/google/cloud/hooks/tasks.py +1 -1
- airflow/providers/google/cloud/hooks/translate.py +2 -2
- airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +2 -210
- airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +3 -3
- airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +28 -2
- airflow/providers/google/cloud/hooks/vertex_ai/experiment_service.py +202 -0
- airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +308 -8
- airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +79 -75
- airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +1 -1
- airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +1 -1
- airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +1 -1
- airflow/providers/google/cloud/hooks/vertex_ai/ray.py +223 -0
- airflow/providers/google/cloud/hooks/vision.py +3 -3
- airflow/providers/google/cloud/hooks/workflows.py +1 -1
- airflow/providers/google/cloud/links/alloy_db.py +0 -46
- airflow/providers/google/cloud/links/base.py +77 -13
- airflow/providers/google/cloud/links/bigquery.py +0 -47
- airflow/providers/google/cloud/links/bigquery_dts.py +0 -20
- airflow/providers/google/cloud/links/bigtable.py +0 -48
- airflow/providers/google/cloud/links/cloud_build.py +0 -73
- airflow/providers/google/cloud/links/cloud_functions.py +0 -33
- airflow/providers/google/cloud/links/cloud_memorystore.py +0 -58
- airflow/providers/google/cloud/links/{life_sciences.py → cloud_run.py} +5 -27
- airflow/providers/google/cloud/links/cloud_sql.py +0 -33
- airflow/providers/google/cloud/links/cloud_storage_transfer.py +17 -44
- airflow/providers/google/cloud/links/cloud_tasks.py +7 -26
- airflow/providers/google/cloud/links/compute.py +0 -58
- airflow/providers/google/cloud/links/data_loss_prevention.py +0 -169
- airflow/providers/google/cloud/links/datacatalog.py +23 -54
- airflow/providers/google/cloud/links/dataflow.py +0 -34
- airflow/providers/google/cloud/links/dataform.py +0 -64
- airflow/providers/google/cloud/links/datafusion.py +1 -96
- airflow/providers/google/cloud/links/dataplex.py +0 -154
- airflow/providers/google/cloud/links/dataprep.py +0 -24
- airflow/providers/google/cloud/links/dataproc.py +11 -95
- airflow/providers/google/cloud/links/datastore.py +0 -31
- airflow/providers/google/cloud/links/kubernetes_engine.py +9 -60
- airflow/providers/google/cloud/links/managed_kafka.py +0 -70
- airflow/providers/google/cloud/links/mlengine.py +0 -70
- airflow/providers/google/cloud/links/pubsub.py +0 -32
- airflow/providers/google/cloud/links/spanner.py +0 -33
- airflow/providers/google/cloud/links/stackdriver.py +0 -30
- airflow/providers/google/cloud/links/translate.py +17 -187
- airflow/providers/google/cloud/links/vertex_ai.py +28 -195
- airflow/providers/google/cloud/links/workflows.py +0 -52
- airflow/providers/google/cloud/log/gcs_task_handler.py +58 -22
- airflow/providers/google/cloud/log/stackdriver_task_handler.py +9 -6
- airflow/providers/google/cloud/openlineage/CloudStorageTransferJobFacet.json +68 -0
- airflow/providers/google/cloud/openlineage/CloudStorageTransferRunFacet.json +60 -0
- airflow/providers/google/cloud/openlineage/DataFusionRunFacet.json +32 -0
- airflow/providers/google/cloud/openlineage/facets.py +102 -1
- airflow/providers/google/cloud/openlineage/mixins.py +10 -8
- airflow/providers/google/cloud/openlineage/utils.py +15 -1
- airflow/providers/google/cloud/operators/alloy_db.py +71 -56
- airflow/providers/google/cloud/operators/bigquery.py +73 -636
- airflow/providers/google/cloud/operators/bigquery_dts.py +4 -6
- airflow/providers/google/cloud/operators/bigtable.py +37 -8
- airflow/providers/google/cloud/operators/cloud_base.py +21 -1
- airflow/providers/google/cloud/operators/cloud_batch.py +3 -3
- airflow/providers/google/cloud/operators/cloud_build.py +76 -33
- airflow/providers/google/cloud/operators/cloud_composer.py +129 -41
- airflow/providers/google/cloud/operators/cloud_logging_sink.py +341 -0
- airflow/providers/google/cloud/operators/cloud_memorystore.py +69 -43
- airflow/providers/google/cloud/operators/cloud_run.py +24 -6
- airflow/providers/google/cloud/operators/cloud_sql.py +8 -17
- airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +93 -12
- airflow/providers/google/cloud/operators/compute.py +9 -41
- airflow/providers/google/cloud/operators/datacatalog.py +157 -21
- airflow/providers/google/cloud/operators/dataflow.py +40 -16
- airflow/providers/google/cloud/operators/dataform.py +15 -5
- airflow/providers/google/cloud/operators/datafusion.py +42 -21
- airflow/providers/google/cloud/operators/dataplex.py +194 -110
- airflow/providers/google/cloud/operators/dataprep.py +1 -5
- airflow/providers/google/cloud/operators/dataproc.py +80 -36
- airflow/providers/google/cloud/operators/dataproc_metastore.py +97 -89
- airflow/providers/google/cloud/operators/datastore.py +23 -7
- airflow/providers/google/cloud/operators/dlp.py +6 -29
- airflow/providers/google/cloud/operators/functions.py +17 -8
- airflow/providers/google/cloud/operators/gcs.py +12 -9
- airflow/providers/google/cloud/operators/gen_ai.py +389 -0
- airflow/providers/google/cloud/operators/kubernetes_engine.py +62 -100
- airflow/providers/google/cloud/operators/looker.py +2 -2
- airflow/providers/google/cloud/operators/managed_kafka.py +108 -53
- airflow/providers/google/cloud/operators/natural_language.py +1 -1
- airflow/providers/google/cloud/operators/pubsub.py +68 -15
- airflow/providers/google/cloud/operators/spanner.py +26 -13
- airflow/providers/google/cloud/operators/speech_to_text.py +2 -3
- airflow/providers/google/cloud/operators/stackdriver.py +1 -9
- airflow/providers/google/cloud/operators/tasks.py +1 -12
- airflow/providers/google/cloud/operators/text_to_speech.py +2 -3
- airflow/providers/google/cloud/operators/translate.py +41 -17
- airflow/providers/google/cloud/operators/translate_speech.py +2 -3
- airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +39 -19
- airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +30 -10
- airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +55 -27
- airflow/providers/google/cloud/operators/vertex_ai/dataset.py +70 -8
- airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +43 -9
- airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py +435 -0
- airflow/providers/google/cloud/operators/vertex_ai/feature_store.py +532 -1
- airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +135 -115
- airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +12 -10
- airflow/providers/google/cloud/operators/vertex_ai/model_service.py +57 -11
- airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +31 -8
- airflow/providers/google/cloud/operators/vertex_ai/ray.py +393 -0
- airflow/providers/google/cloud/operators/video_intelligence.py +1 -1
- airflow/providers/google/cloud/operators/vision.py +2 -2
- airflow/providers/google/cloud/operators/workflows.py +18 -15
- airflow/providers/google/cloud/secrets/secret_manager.py +3 -2
- airflow/providers/google/cloud/sensors/bigquery.py +3 -3
- airflow/providers/google/cloud/sensors/bigquery_dts.py +2 -3
- airflow/providers/google/cloud/sensors/bigtable.py +11 -4
- airflow/providers/google/cloud/sensors/cloud_composer.py +533 -30
- airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +2 -3
- airflow/providers/google/cloud/sensors/dataflow.py +26 -10
- airflow/providers/google/cloud/sensors/dataform.py +2 -3
- airflow/providers/google/cloud/sensors/datafusion.py +4 -5
- airflow/providers/google/cloud/sensors/dataplex.py +2 -3
- airflow/providers/google/cloud/sensors/dataprep.py +2 -2
- airflow/providers/google/cloud/sensors/dataproc.py +2 -3
- airflow/providers/google/cloud/sensors/dataproc_metastore.py +2 -3
- airflow/providers/google/cloud/sensors/gcs.py +4 -5
- airflow/providers/google/cloud/sensors/looker.py +2 -3
- airflow/providers/google/cloud/sensors/pubsub.py +4 -5
- airflow/providers/google/cloud/sensors/tasks.py +2 -2
- airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +2 -3
- airflow/providers/google/cloud/sensors/workflows.py +2 -3
- airflow/providers/google/cloud/transfers/adls_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +2 -2
- airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +4 -3
- airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py +11 -8
- airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +10 -5
- airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +7 -3
- airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +12 -1
- airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +24 -10
- airflow/providers/google/cloud/transfers/bigquery_to_sql.py +104 -5
- airflow/providers/google/cloud/transfers/calendar_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +3 -3
- airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +4 -4
- airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +21 -13
- airflow/providers/google/cloud/transfers/gcs_to_gcs.py +4 -3
- airflow/providers/google/cloud/transfers/gcs_to_local.py +6 -4
- airflow/providers/google/cloud/transfers/gcs_to_sftp.py +11 -5
- airflow/providers/google/cloud/transfers/gdrive_to_gcs.py +6 -2
- airflow/providers/google/cloud/transfers/gdrive_to_local.py +2 -2
- airflow/providers/google/cloud/transfers/http_to_gcs.py +193 -0
- airflow/providers/google/cloud/transfers/local_to_gcs.py +2 -2
- airflow/providers/google/cloud/transfers/mssql_to_gcs.py +1 -1
- airflow/providers/google/cloud/transfers/oracle_to_gcs.py +36 -11
- airflow/providers/google/cloud/transfers/postgres_to_gcs.py +42 -9
- airflow/providers/google/cloud/transfers/s3_to_gcs.py +13 -7
- airflow/providers/google/cloud/transfers/salesforce_to_gcs.py +2 -2
- airflow/providers/google/cloud/transfers/sftp_to_gcs.py +14 -5
- airflow/providers/google/cloud/transfers/sheets_to_gcs.py +3 -3
- airflow/providers/google/cloud/transfers/sql_to_gcs.py +10 -10
- airflow/providers/google/cloud/triggers/bigquery.py +76 -35
- airflow/providers/google/cloud/triggers/cloud_build.py +1 -1
- airflow/providers/google/cloud/triggers/cloud_composer.py +303 -47
- airflow/providers/google/cloud/triggers/cloud_run.py +3 -3
- airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +92 -2
- airflow/providers/google/cloud/triggers/dataflow.py +122 -0
- airflow/providers/google/cloud/triggers/datafusion.py +1 -1
- airflow/providers/google/cloud/triggers/dataplex.py +14 -2
- airflow/providers/google/cloud/triggers/dataproc.py +123 -53
- airflow/providers/google/cloud/triggers/kubernetes_engine.py +47 -28
- airflow/providers/google/cloud/triggers/mlengine.py +1 -1
- airflow/providers/google/cloud/triggers/pubsub.py +15 -19
- airflow/providers/google/cloud/triggers/vertex_ai.py +1 -1
- airflow/providers/google/cloud/utils/bigquery_get_data.py +1 -1
- airflow/providers/google/cloud/utils/credentials_provider.py +2 -2
- airflow/providers/google/cloud/utils/field_sanitizer.py +1 -1
- airflow/providers/google/cloud/utils/field_validator.py +2 -3
- airflow/providers/google/common/auth_backend/google_openid.py +4 -4
- airflow/providers/google/common/deprecated.py +2 -1
- airflow/providers/google/common/hooks/base_google.py +27 -9
- airflow/providers/google/common/hooks/operation_helpers.py +1 -1
- airflow/providers/google/common/links/storage.py +0 -22
- airflow/providers/google/common/utils/get_secret.py +31 -0
- airflow/providers/google/common/utils/id_token_credentials.py +3 -4
- airflow/providers/google/firebase/hooks/firestore.py +1 -1
- airflow/providers/google/firebase/operators/firestore.py +3 -3
- airflow/providers/google/get_provider_info.py +56 -52
- airflow/providers/google/go_module_utils.py +35 -3
- airflow/providers/google/leveldb/hooks/leveldb.py +27 -2
- airflow/providers/google/leveldb/operators/leveldb.py +2 -2
- airflow/providers/google/marketing_platform/hooks/campaign_manager.py +1 -1
- airflow/providers/google/marketing_platform/hooks/display_video.py +3 -109
- airflow/providers/google/marketing_platform/hooks/search_ads.py +1 -1
- airflow/providers/google/marketing_platform/links/analytics_admin.py +5 -14
- airflow/providers/google/marketing_platform/operators/analytics_admin.py +2 -3
- airflow/providers/google/marketing_platform/operators/campaign_manager.py +6 -6
- airflow/providers/google/marketing_platform/operators/display_video.py +28 -489
- airflow/providers/google/marketing_platform/operators/search_ads.py +2 -2
- airflow/providers/google/marketing_platform/sensors/campaign_manager.py +2 -2
- airflow/providers/google/marketing_platform/sensors/display_video.py +3 -64
- airflow/providers/google/suite/hooks/calendar.py +2 -2
- airflow/providers/google/suite/hooks/sheets.py +16 -2
- airflow/providers/google/suite/operators/sheets.py +8 -3
- airflow/providers/google/suite/sensors/drive.py +2 -2
- airflow/providers/google/suite/transfers/gcs_to_gdrive.py +3 -3
- airflow/providers/google/suite/transfers/gcs_to_sheets.py +1 -1
- airflow/providers/google/suite/transfers/local_to_drive.py +3 -3
- airflow/providers/google/suite/transfers/sql_to_sheets.py +5 -4
- airflow/providers/google/version_compat.py +15 -1
- {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/METADATA +90 -46
- apache_airflow_providers_google-19.3.0.dist-info/RECORD +331 -0
- apache_airflow_providers_google-19.3.0.dist-info/licenses/NOTICE +5 -0
- airflow/providers/google/cloud/hooks/automl.py +0 -673
- airflow/providers/google/cloud/hooks/life_sciences.py +0 -159
- airflow/providers/google/cloud/links/automl.py +0 -193
- airflow/providers/google/cloud/operators/automl.py +0 -1362
- airflow/providers/google/cloud/operators/life_sciences.py +0 -119
- airflow/providers/google/cloud/operators/mlengine.py +0 -112
- apache_airflow_providers_google-15.1.0rc1.dist-info/RECORD +0 -321
- {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/WHEEL +0 -0
- {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/entry_points.txt +0 -0
- {airflow/providers/google → apache_airflow_providers_google-19.3.0.dist-info/licenses}/LICENSE +0 -0
|
@@ -19,12 +19,12 @@
|
|
|
19
19
|
|
|
20
20
|
from __future__ import annotations
|
|
21
21
|
|
|
22
|
-
from collections.abc import Sequence
|
|
22
|
+
from collections.abc import Callable, Sequence
|
|
23
23
|
from functools import cached_property
|
|
24
|
-
from typing import TYPE_CHECKING, Any
|
|
24
|
+
from typing import TYPE_CHECKING, Any
|
|
25
25
|
|
|
26
26
|
from airflow.configuration import conf
|
|
27
|
-
from airflow.
|
|
27
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, PokeReturnValue
|
|
28
28
|
from airflow.providers.google.cloud.hooks.dataflow import (
|
|
29
29
|
DEFAULT_DATAFLOW_LOCATION,
|
|
30
30
|
DataflowHook,
|
|
@@ -37,10 +37,9 @@ from airflow.providers.google.cloud.triggers.dataflow import (
|
|
|
37
37
|
DataflowJobStatusTrigger,
|
|
38
38
|
)
|
|
39
39
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
40
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
41
40
|
|
|
42
41
|
if TYPE_CHECKING:
|
|
43
|
-
from airflow.
|
|
42
|
+
from airflow.providers.common.compat.sdk import Context
|
|
44
43
|
|
|
45
44
|
|
|
46
45
|
class DataflowJobStatusSensor(BaseSensorOperator):
|
|
@@ -342,7 +341,7 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
|
|
|
342
341
|
self.deferrable = deferrable
|
|
343
342
|
self.poll_interval = poll_interval
|
|
344
343
|
|
|
345
|
-
def poke(self, context: Context) -> bool:
|
|
344
|
+
def poke(self, context: Context) -> PokeReturnValue | bool:
|
|
346
345
|
if self.fail_on_terminal_state:
|
|
347
346
|
job = self.hook.get_job(
|
|
348
347
|
job_id=self.job_id,
|
|
@@ -359,8 +358,17 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
|
|
|
359
358
|
project_id=self.project_id,
|
|
360
359
|
location=self.location,
|
|
361
360
|
)
|
|
361
|
+
result = result if self.callback is None else self.callback(result)
|
|
362
|
+
|
|
363
|
+
if isinstance(result, PokeReturnValue):
|
|
364
|
+
return result
|
|
362
365
|
|
|
363
|
-
|
|
366
|
+
if bool(result):
|
|
367
|
+
return PokeReturnValue(
|
|
368
|
+
is_done=True,
|
|
369
|
+
xcom_value=result,
|
|
370
|
+
)
|
|
371
|
+
return False
|
|
364
372
|
|
|
365
373
|
def execute(self, context: Context) -> Any:
|
|
366
374
|
"""Airflow runs this method on the worker and defers using the trigger."""
|
|
@@ -464,7 +472,7 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
|
|
|
464
472
|
self.deferrable = deferrable
|
|
465
473
|
self.poll_interval = poll_interval
|
|
466
474
|
|
|
467
|
-
def poke(self, context: Context) -> bool:
|
|
475
|
+
def poke(self, context: Context) -> PokeReturnValue | bool:
|
|
468
476
|
if self.fail_on_terminal_state:
|
|
469
477
|
job = self.hook.get_job(
|
|
470
478
|
job_id=self.job_id,
|
|
@@ -481,8 +489,16 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
|
|
|
481
489
|
project_id=self.project_id,
|
|
482
490
|
location=self.location,
|
|
483
491
|
)
|
|
484
|
-
|
|
485
|
-
|
|
492
|
+
result = result if self.callback is None else self.callback(result)
|
|
493
|
+
if isinstance(result, PokeReturnValue):
|
|
494
|
+
return result
|
|
495
|
+
|
|
496
|
+
if bool(result):
|
|
497
|
+
return PokeReturnValue(
|
|
498
|
+
is_done=True,
|
|
499
|
+
xcom_value=result,
|
|
500
|
+
)
|
|
501
|
+
return False
|
|
486
502
|
|
|
487
503
|
def execute(self, context: Context) -> Any:
|
|
488
504
|
"""Airflow runs this method on the worker and defers using the trigger."""
|
|
@@ -22,12 +22,11 @@ from __future__ import annotations
|
|
|
22
22
|
from collections.abc import Iterable, Sequence
|
|
23
23
|
from typing import TYPE_CHECKING
|
|
24
24
|
|
|
25
|
-
from airflow.
|
|
25
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
26
26
|
from airflow.providers.google.cloud.hooks.dataform import DataformHook
|
|
27
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
28
27
|
|
|
29
28
|
if TYPE_CHECKING:
|
|
30
|
-
from airflow.
|
|
29
|
+
from airflow.providers.common.compat.sdk import Context
|
|
31
30
|
|
|
32
31
|
|
|
33
32
|
class DataformWorkflowInvocationStateSensor(BaseSensorOperator):
|
|
@@ -22,13 +22,12 @@ from __future__ import annotations
|
|
|
22
22
|
from collections.abc import Iterable, Sequence
|
|
23
23
|
from typing import TYPE_CHECKING
|
|
24
24
|
|
|
25
|
-
from airflow.
|
|
25
|
+
from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException, BaseSensorOperator
|
|
26
26
|
from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
|
|
27
27
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
28
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
29
28
|
|
|
30
29
|
if TYPE_CHECKING:
|
|
31
|
-
from airflow.
|
|
30
|
+
from airflow.providers.common.compat.sdk import Context
|
|
32
31
|
|
|
33
32
|
|
|
34
33
|
class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
|
|
@@ -110,7 +109,7 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
|
|
|
110
109
|
pipeline_id=self.pipeline_id,
|
|
111
110
|
namespace=self.namespace,
|
|
112
111
|
)
|
|
113
|
-
pipeline_status = pipeline_workflow
|
|
112
|
+
pipeline_status = pipeline_workflow.get("status")
|
|
114
113
|
except AirflowNotFoundException:
|
|
115
114
|
message = "Specified Pipeline ID was not found."
|
|
116
115
|
raise AirflowException(message)
|
|
@@ -127,4 +126,4 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
|
|
|
127
126
|
self.log.debug(
|
|
128
127
|
"Current status of the pipeline workflow for %s: %s.", self.pipeline_id, pipeline_status
|
|
129
128
|
)
|
|
130
|
-
return pipeline_status in self.expected_statuses
|
|
129
|
+
return pipeline_status is not None and pipeline_status in self.expected_statuses
|
|
@@ -25,19 +25,18 @@ from typing import TYPE_CHECKING
|
|
|
25
25
|
if TYPE_CHECKING:
|
|
26
26
|
from google.api_core.retry import Retry
|
|
27
27
|
|
|
28
|
-
from airflow.
|
|
28
|
+
from airflow.providers.common.compat.sdk import Context
|
|
29
29
|
|
|
30
30
|
from google.api_core.exceptions import GoogleAPICallError
|
|
31
31
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
|
32
32
|
from google.cloud.dataplex_v1.types import DataScanJob
|
|
33
33
|
|
|
34
|
-
from airflow.
|
|
34
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
35
35
|
from airflow.providers.google.cloud.hooks.dataplex import (
|
|
36
36
|
AirflowDataQualityScanException,
|
|
37
37
|
AirflowDataQualityScanResultTimeoutException,
|
|
38
38
|
DataplexHook,
|
|
39
39
|
)
|
|
40
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
41
40
|
|
|
42
41
|
|
|
43
42
|
class TaskState:
|
|
@@ -22,11 +22,11 @@ from __future__ import annotations
|
|
|
22
22
|
from collections.abc import Sequence
|
|
23
23
|
from typing import TYPE_CHECKING
|
|
24
24
|
|
|
25
|
+
from airflow.providers.common.compat.sdk import BaseSensorOperator
|
|
25
26
|
from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook, JobGroupStatuses
|
|
26
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
27
27
|
|
|
28
28
|
if TYPE_CHECKING:
|
|
29
|
-
from airflow.
|
|
29
|
+
from airflow.providers.common.compat.sdk import Context
|
|
30
30
|
|
|
31
31
|
|
|
32
32
|
class DataprepJobGroupIsFinishedSensor(BaseSensorOperator):
|
|
@@ -26,13 +26,12 @@ from typing import TYPE_CHECKING
|
|
|
26
26
|
from google.api_core.exceptions import ServerError
|
|
27
27
|
from google.cloud.dataproc_v1.types import Batch, JobStatus
|
|
28
28
|
|
|
29
|
-
from airflow.
|
|
29
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
30
30
|
from airflow.providers.google.cloud.hooks.dataproc import DataprocHook
|
|
31
31
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
32
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
33
32
|
|
|
34
33
|
if TYPE_CHECKING:
|
|
35
|
-
from airflow.
|
|
34
|
+
from airflow.providers.common.compat.sdk import Context
|
|
36
35
|
|
|
37
36
|
|
|
38
37
|
class DataprocJobSensor(BaseSensorOperator):
|
|
@@ -20,15 +20,14 @@ from __future__ import annotations
|
|
|
20
20
|
from collections.abc import Sequence
|
|
21
21
|
from typing import TYPE_CHECKING
|
|
22
22
|
|
|
23
|
-
from airflow.
|
|
23
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
24
24
|
from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
|
|
25
25
|
from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
|
|
26
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
27
26
|
|
|
28
27
|
if TYPE_CHECKING:
|
|
29
28
|
from google.api_core.operation import Operation
|
|
30
29
|
|
|
31
|
-
from airflow.
|
|
30
|
+
from airflow.providers.common.compat.sdk import Context
|
|
32
31
|
|
|
33
32
|
|
|
34
33
|
class MetastoreHivePartitionSensor(BaseSensorOperator):
|
|
@@ -21,14 +21,14 @@ from __future__ import annotations
|
|
|
21
21
|
|
|
22
22
|
import os
|
|
23
23
|
import textwrap
|
|
24
|
-
from collections.abc import Sequence
|
|
24
|
+
from collections.abc import Callable, Sequence
|
|
25
25
|
from datetime import datetime, timedelta
|
|
26
|
-
from typing import TYPE_CHECKING, Any
|
|
26
|
+
from typing import TYPE_CHECKING, Any
|
|
27
27
|
|
|
28
28
|
from google.cloud.storage.retry import DEFAULT_RETRY
|
|
29
29
|
|
|
30
30
|
from airflow.configuration import conf
|
|
31
|
-
from airflow.
|
|
31
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, poke_mode_only
|
|
32
32
|
from airflow.providers.google.cloud.hooks.gcs import GCSHook
|
|
33
33
|
from airflow.providers.google.cloud.triggers.gcs import (
|
|
34
34
|
GCSBlobTrigger,
|
|
@@ -36,12 +36,11 @@ from airflow.providers.google.cloud.triggers.gcs import (
|
|
|
36
36
|
GCSPrefixBlobTrigger,
|
|
37
37
|
GCSUploadSessionTrigger,
|
|
38
38
|
)
|
|
39
|
-
from airflow.sensors.base import BaseSensorOperator, poke_mode_only
|
|
40
39
|
|
|
41
40
|
if TYPE_CHECKING:
|
|
42
41
|
from google.api_core.retry import Retry
|
|
43
42
|
|
|
44
|
-
from airflow.
|
|
43
|
+
from airflow.providers.common.compat.sdk import Context
|
|
45
44
|
|
|
46
45
|
|
|
47
46
|
class GCSObjectExistenceSensor(BaseSensorOperator):
|
|
@@ -21,12 +21,11 @@ from __future__ import annotations
|
|
|
21
21
|
|
|
22
22
|
from typing import TYPE_CHECKING
|
|
23
23
|
|
|
24
|
-
from airflow.
|
|
24
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
25
25
|
from airflow.providers.google.cloud.hooks.looker import JobStatus, LookerHook
|
|
26
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
27
26
|
|
|
28
27
|
if TYPE_CHECKING:
|
|
29
|
-
from airflow.
|
|
28
|
+
from airflow.providers.common.compat.sdk import Context
|
|
30
29
|
|
|
31
30
|
|
|
32
31
|
class LookerCheckPdtBuildSensor(BaseSensorOperator):
|
|
@@ -19,21 +19,20 @@
|
|
|
19
19
|
|
|
20
20
|
from __future__ import annotations
|
|
21
21
|
|
|
22
|
-
from collections.abc import Sequence
|
|
22
|
+
from collections.abc import Callable, Sequence
|
|
23
23
|
from datetime import timedelta
|
|
24
|
-
from typing import TYPE_CHECKING, Any
|
|
24
|
+
from typing import TYPE_CHECKING, Any
|
|
25
25
|
|
|
26
26
|
from google.cloud import pubsub_v1
|
|
27
27
|
from google.cloud.pubsub_v1.types import ReceivedMessage
|
|
28
28
|
|
|
29
29
|
from airflow.configuration import conf
|
|
30
|
-
from airflow.
|
|
30
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
31
31
|
from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
|
|
32
32
|
from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
|
|
33
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
34
33
|
|
|
35
34
|
if TYPE_CHECKING:
|
|
36
|
-
from airflow.
|
|
35
|
+
from airflow.providers.common.compat.sdk import Context
|
|
37
36
|
|
|
38
37
|
|
|
39
38
|
class PubSubMessageTransformException(AirflowException):
|
|
@@ -22,12 +22,12 @@ from __future__ import annotations
|
|
|
22
22
|
from collections.abc import Sequence
|
|
23
23
|
from typing import TYPE_CHECKING
|
|
24
24
|
|
|
25
|
+
from airflow.providers.common.compat.sdk import BaseSensorOperator
|
|
25
26
|
from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
|
|
26
27
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
27
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
28
28
|
|
|
29
29
|
if TYPE_CHECKING:
|
|
30
|
-
from airflow.
|
|
30
|
+
from airflow.providers.common.compat.sdk import Context
|
|
31
31
|
|
|
32
32
|
|
|
33
33
|
class TaskQueueEmptySensor(BaseSensorOperator):
|
|
@@ -23,12 +23,11 @@ import time
|
|
|
23
23
|
from collections.abc import Sequence
|
|
24
24
|
from typing import TYPE_CHECKING
|
|
25
25
|
|
|
26
|
-
from airflow.
|
|
26
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
27
27
|
from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
|
|
28
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
29
28
|
|
|
30
29
|
if TYPE_CHECKING:
|
|
31
|
-
from airflow.
|
|
30
|
+
from airflow.providers.common.compat.sdk import Context
|
|
32
31
|
|
|
33
32
|
|
|
34
33
|
class FeatureViewSyncSensor(BaseSensorOperator):
|
|
@@ -22,15 +22,14 @@ from typing import TYPE_CHECKING
|
|
|
22
22
|
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
|
|
23
23
|
from google.cloud.workflows.executions_v1beta import Execution
|
|
24
24
|
|
|
25
|
-
from airflow.
|
|
25
|
+
from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
|
|
26
26
|
from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook
|
|
27
27
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
28
|
-
from airflow.sensors.base import BaseSensorOperator
|
|
29
28
|
|
|
30
29
|
if TYPE_CHECKING:
|
|
31
30
|
from google.api_core.retry import Retry
|
|
32
31
|
|
|
33
|
-
from airflow.
|
|
32
|
+
from airflow.providers.common.compat.sdk import Context
|
|
34
33
|
|
|
35
34
|
|
|
36
35
|
class WorkflowExecutionSensor(BaseSensorOperator):
|
|
@@ -35,7 +35,7 @@ except ModuleNotFoundError as e:
|
|
|
35
35
|
raise AirflowOptionalProviderFeatureException(e)
|
|
36
36
|
|
|
37
37
|
if TYPE_CHECKING:
|
|
38
|
-
from airflow.
|
|
38
|
+
from airflow.providers.common.compat.sdk import Context
|
|
39
39
|
|
|
40
40
|
|
|
41
41
|
class ADLSToGCSOperator(ADLSListOperator):
|
|
@@ -21,8 +21,8 @@ import tempfile
|
|
|
21
21
|
from collections.abc import Sequence
|
|
22
22
|
from typing import TYPE_CHECKING
|
|
23
23
|
|
|
24
|
-
from airflow.models import BaseOperator
|
|
25
24
|
from airflow.providers.google.cloud.hooks.gcs import GCSHook
|
|
25
|
+
from airflow.providers.google.version_compat import BaseOperator
|
|
26
26
|
|
|
27
27
|
try:
|
|
28
28
|
from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
|
|
@@ -32,7 +32,7 @@ except ModuleNotFoundError as e:
|
|
|
32
32
|
raise AirflowOptionalProviderFeatureException(e)
|
|
33
33
|
|
|
34
34
|
if TYPE_CHECKING:
|
|
35
|
-
from airflow.
|
|
35
|
+
from airflow.providers.common.compat.sdk import Context
|
|
36
36
|
|
|
37
37
|
|
|
38
38
|
class AzureBlobStorageToGCSOperator(BaseOperator):
|
|
@@ -22,9 +22,10 @@ from collections.abc import Sequence
|
|
|
22
22
|
from tempfile import NamedTemporaryFile
|
|
23
23
|
from typing import TYPE_CHECKING
|
|
24
24
|
|
|
25
|
-
from airflow.exceptions import
|
|
26
|
-
from airflow.
|
|
25
|
+
from airflow.exceptions import AirflowProviderDeprecationWarning
|
|
26
|
+
from airflow.providers.common.compat.sdk import AirflowException
|
|
27
27
|
from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url, gcs_object_is_directory
|
|
28
|
+
from airflow.providers.google.version_compat import BaseOperator
|
|
28
29
|
|
|
29
30
|
try:
|
|
30
31
|
from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook
|
|
@@ -34,7 +35,7 @@ except ModuleNotFoundError as e:
|
|
|
34
35
|
raise AirflowOptionalProviderFeatureException(e)
|
|
35
36
|
|
|
36
37
|
if TYPE_CHECKING:
|
|
37
|
-
from airflow.
|
|
38
|
+
from airflow.providers.common.compat.sdk import Context
|
|
38
39
|
|
|
39
40
|
|
|
40
41
|
class AzureFileShareToGCSOperator(BaseOperator):
|
|
@@ -22,12 +22,13 @@ from __future__ import annotations
|
|
|
22
22
|
from collections.abc import Sequence
|
|
23
23
|
from typing import TYPE_CHECKING
|
|
24
24
|
|
|
25
|
-
from airflow.models import BaseOperator
|
|
26
25
|
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
|
|
27
26
|
from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
|
|
27
|
+
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
28
|
+
from airflow.providers.google.version_compat import BaseOperator
|
|
28
29
|
|
|
29
30
|
if TYPE_CHECKING:
|
|
30
|
-
from airflow.
|
|
31
|
+
from airflow.providers.common.compat.sdk import Context
|
|
31
32
|
|
|
32
33
|
|
|
33
34
|
class BigQueryToBigQueryOperator(BaseOperator):
|
|
@@ -73,6 +74,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
|
|
|
73
74
|
If set as a sequence, the identities from the list must grant
|
|
74
75
|
Service Account Token Creator IAM role to the directly preceding identity, with first
|
|
75
76
|
account from the list granting this role to the originating account (templated).
|
|
77
|
+
:param project_id: Google Cloud Project where the job is running
|
|
76
78
|
"""
|
|
77
79
|
|
|
78
80
|
template_fields: Sequence[str] = (
|
|
@@ -93,6 +95,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
|
|
|
93
95
|
write_disposition: str = "WRITE_EMPTY",
|
|
94
96
|
create_disposition: str = "CREATE_IF_NEEDED",
|
|
95
97
|
gcp_conn_id: str = "google_cloud_default",
|
|
98
|
+
project_id: str = PROVIDE_PROJECT_ID,
|
|
96
99
|
labels: dict | None = None,
|
|
97
100
|
encryption_configuration: dict | None = None,
|
|
98
101
|
location: str | None = None,
|
|
@@ -112,6 +115,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
|
|
|
112
115
|
self.impersonation_chain = impersonation_chain
|
|
113
116
|
self.hook: BigQueryHook | None = None
|
|
114
117
|
self._job_conf: dict = {}
|
|
118
|
+
self.project_id = project_id
|
|
115
119
|
|
|
116
120
|
def _prepare_job_configuration(self):
|
|
117
121
|
self.source_project_dataset_tables = (
|
|
@@ -124,7 +128,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
|
|
|
124
128
|
for source_project_dataset_table in self.source_project_dataset_tables:
|
|
125
129
|
source_project, source_dataset, source_table = self.hook.split_tablename(
|
|
126
130
|
table_input=source_project_dataset_table,
|
|
127
|
-
default_project_id=self.
|
|
131
|
+
default_project_id=self.project_id,
|
|
128
132
|
var_name="source_project_dataset_table",
|
|
129
133
|
)
|
|
130
134
|
source_project_dataset_tables_fixup.append(
|
|
@@ -133,7 +137,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
|
|
|
133
137
|
|
|
134
138
|
destination_project, destination_dataset, destination_table = self.hook.split_tablename(
|
|
135
139
|
table_input=self.destination_project_dataset_table,
|
|
136
|
-
default_project_id=self.
|
|
140
|
+
default_project_id=self.project_id,
|
|
137
141
|
)
|
|
138
142
|
configuration = {
|
|
139
143
|
"copy": {
|
|
@@ -168,18 +172,17 @@ class BigQueryToBigQueryOperator(BaseOperator):
|
|
|
168
172
|
impersonation_chain=self.impersonation_chain,
|
|
169
173
|
)
|
|
170
174
|
|
|
171
|
-
if not self.
|
|
172
|
-
|
|
175
|
+
if not self.project_id:
|
|
176
|
+
self.project_id = self.hook.project_id
|
|
173
177
|
|
|
174
178
|
configuration = self._prepare_job_configuration()
|
|
175
179
|
self._job_conf = self.hook.insert_job(
|
|
176
|
-
configuration=configuration, project_id=self.
|
|
180
|
+
configuration=configuration, project_id=self.project_id
|
|
177
181
|
).to_api_repr()
|
|
178
182
|
|
|
179
183
|
dest_table_info = self._job_conf["configuration"]["copy"]["destinationTable"]
|
|
180
184
|
BigQueryTableLink.persist(
|
|
181
185
|
context=context,
|
|
182
|
-
task_instance=self,
|
|
183
186
|
dataset_id=dest_table_info["datasetId"],
|
|
184
187
|
project_id=dest_table_info["projectId"],
|
|
185
188
|
table_id=dest_table_info["tableId"],
|
|
@@ -26,18 +26,18 @@ from google.api_core.exceptions import Conflict
|
|
|
26
26
|
from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
|
|
27
27
|
|
|
28
28
|
from airflow.configuration import conf
|
|
29
|
-
from airflow.
|
|
30
|
-
from airflow.models import BaseOperator
|
|
29
|
+
from airflow.providers.common.compat.sdk import AirflowException
|
|
31
30
|
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
|
|
32
31
|
from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
|
|
33
32
|
from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
|
|
34
33
|
from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
|
|
34
|
+
from airflow.providers.google.version_compat import BaseOperator
|
|
35
35
|
from airflow.utils.helpers import merge_dicts
|
|
36
36
|
|
|
37
37
|
if TYPE_CHECKING:
|
|
38
38
|
from google.api_core.retry import Retry
|
|
39
39
|
|
|
40
|
-
from airflow.
|
|
40
|
+
from airflow.providers.common.compat.sdk import Context
|
|
41
41
|
|
|
42
42
|
|
|
43
43
|
class BigQueryToGCSOperator(BaseOperator):
|
|
@@ -87,6 +87,7 @@ class BigQueryToGCSOperator(BaseOperator):
|
|
|
87
87
|
:param reattach_states: Set of BigQuery job's states in case of which we should reattach
|
|
88
88
|
to the job. Should be other than final states.
|
|
89
89
|
:param deferrable: Run operator in the deferrable mode
|
|
90
|
+
:return: URIs for the objects created in Google Cloud Storage
|
|
90
91
|
"""
|
|
91
92
|
|
|
92
93
|
template_fields: Sequence[str] = (
|
|
@@ -215,8 +216,9 @@ class BigQueryToGCSOperator(BaseOperator):
|
|
|
215
216
|
job_id=self.job_id,
|
|
216
217
|
dag_id=self.dag_id,
|
|
217
218
|
task_id=self.task_id,
|
|
218
|
-
logical_date=
|
|
219
|
+
logical_date=None,
|
|
219
220
|
configuration=configuration,
|
|
221
|
+
run_after=hook.get_run_after_or_logical_date(context),
|
|
220
222
|
force_rerun=self.force_rerun,
|
|
221
223
|
)
|
|
222
224
|
|
|
@@ -254,7 +256,6 @@ class BigQueryToGCSOperator(BaseOperator):
|
|
|
254
256
|
dataset_id, project_id, table_id = conf["datasetId"], conf["projectId"], conf["tableId"]
|
|
255
257
|
BigQueryTableLink.persist(
|
|
256
258
|
context=context,
|
|
257
|
-
task_instance=self,
|
|
258
259
|
dataset_id=dataset_id,
|
|
259
260
|
project_id=project_id,
|
|
260
261
|
table_id=table_id,
|
|
@@ -275,6 +276,8 @@ class BigQueryToGCSOperator(BaseOperator):
|
|
|
275
276
|
else:
|
|
276
277
|
job.result(timeout=self.result_timeout, retry=self.result_retry)
|
|
277
278
|
|
|
279
|
+
return self.destination_cloud_storage_uris
|
|
280
|
+
|
|
278
281
|
def execute_complete(self, context: Context, event: dict[str, Any]):
|
|
279
282
|
"""
|
|
280
283
|
Return immediately and relies on trigger to throw a success event. Callback for the trigger.
|
|
@@ -291,6 +294,8 @@ class BigQueryToGCSOperator(BaseOperator):
|
|
|
291
294
|
# Save job_id as an attribute to be later used by listeners
|
|
292
295
|
self.job_id = event.get("job_id")
|
|
293
296
|
|
|
297
|
+
return self.destination_cloud_storage_uris
|
|
298
|
+
|
|
294
299
|
def get_openlineage_facets_on_complete(self, task_instance):
|
|
295
300
|
"""Implement on_complete as we will include final BQ job id."""
|
|
296
301
|
from airflow.providers.common.compat.openlineage.facet import (
|
|
@@ -21,6 +21,7 @@ from __future__ import annotations
|
|
|
21
21
|
|
|
22
22
|
import warnings
|
|
23
23
|
from collections.abc import Sequence
|
|
24
|
+
from functools import cached_property
|
|
24
25
|
from typing import TYPE_CHECKING
|
|
25
26
|
|
|
26
27
|
from airflow.exceptions import AirflowProviderDeprecationWarning
|
|
@@ -29,7 +30,7 @@ from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToS
|
|
|
29
30
|
from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
|
|
30
31
|
|
|
31
32
|
if TYPE_CHECKING:
|
|
32
|
-
from airflow.
|
|
33
|
+
from airflow.providers.common.compat.sdk import Context
|
|
33
34
|
|
|
34
35
|
|
|
35
36
|
class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
|
|
@@ -94,14 +95,17 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
|
|
|
94
95
|
self.mssql_conn_id = mssql_conn_id
|
|
95
96
|
self.source_project_dataset_table = source_project_dataset_table
|
|
96
97
|
|
|
97
|
-
|
|
98
|
+
@cached_property
|
|
99
|
+
def mssql_hook(self) -> MsSqlHook:
|
|
98
100
|
return MsSqlHook(schema=self.database, mssql_conn_id=self.mssql_conn_id)
|
|
99
101
|
|
|
102
|
+
def get_sql_hook(self) -> MsSqlHook:
|
|
103
|
+
return self.mssql_hook
|
|
104
|
+
|
|
100
105
|
def persist_links(self, context: Context) -> None:
|
|
101
106
|
project_id, dataset_id, table_id = self.source_project_dataset_table.split(".")
|
|
102
107
|
BigQueryTableLink.persist(
|
|
103
108
|
context=context,
|
|
104
|
-
task_instance=self,
|
|
105
109
|
dataset_id=dataset_id,
|
|
106
110
|
project_id=project_id,
|
|
107
111
|
table_id=table_id,
|
|
@@ -21,6 +21,7 @@ from __future__ import annotations
|
|
|
21
21
|
|
|
22
22
|
import warnings
|
|
23
23
|
from collections.abc import Sequence
|
|
24
|
+
from functools import cached_property
|
|
24
25
|
|
|
25
26
|
from airflow.exceptions import AirflowProviderDeprecationWarning
|
|
26
27
|
from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
|
|
@@ -76,5 +77,15 @@ class BigQueryToMySqlOperator(BigQueryToSqlBaseOperator):
|
|
|
76
77
|
)
|
|
77
78
|
self.mysql_conn_id = mysql_conn_id
|
|
78
79
|
|
|
79
|
-
|
|
80
|
+
@cached_property
|
|
81
|
+
def mysql_hook(self) -> MySqlHook:
|
|
80
82
|
return MySqlHook(schema=self.database, mysql_conn_id=self.mysql_conn_id)
|
|
83
|
+
|
|
84
|
+
def get_sql_hook(self) -> MySqlHook:
|
|
85
|
+
return self.mysql_hook
|
|
86
|
+
|
|
87
|
+
def execute(self, context):
|
|
88
|
+
# Set source_project_dataset_table here, after hooks are initialized and project_id is available
|
|
89
|
+
project_id = self.bigquery_hook.project_id
|
|
90
|
+
self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
|
|
91
|
+
return super().execute(context)
|
|
@@ -19,15 +19,19 @@
|
|
|
19
19
|
|
|
20
20
|
from __future__ import annotations
|
|
21
21
|
|
|
22
|
+
from functools import cached_property
|
|
22
23
|
from typing import TYPE_CHECKING
|
|
23
24
|
|
|
25
|
+
from psycopg2.extensions import register_adapter
|
|
26
|
+
from psycopg2.extras import Json
|
|
27
|
+
|
|
24
28
|
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
|
|
25
29
|
from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
|
|
26
30
|
from airflow.providers.google.cloud.utils.bigquery_get_data import bigquery_get_data
|
|
27
31
|
from airflow.providers.postgres.hooks.postgres import PostgresHook
|
|
28
32
|
|
|
29
33
|
if TYPE_CHECKING:
|
|
30
|
-
from airflow.
|
|
34
|
+
from airflow.providers.common.compat.sdk import Context
|
|
31
35
|
|
|
32
36
|
|
|
33
37
|
class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
|
|
@@ -75,26 +79,36 @@ class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
|
|
|
75
79
|
self.postgres_conn_id = postgres_conn_id
|
|
76
80
|
self.replace_index = replace_index
|
|
77
81
|
|
|
78
|
-
|
|
82
|
+
@cached_property
|
|
83
|
+
def postgres_hook(self) -> PostgresHook:
|
|
84
|
+
register_adapter(list, Json)
|
|
85
|
+
register_adapter(dict, Json)
|
|
79
86
|
return PostgresHook(database=self.database, postgres_conn_id=self.postgres_conn_id)
|
|
80
87
|
|
|
88
|
+
def get_sql_hook(self) -> PostgresHook:
|
|
89
|
+
return self.postgres_hook
|
|
90
|
+
|
|
81
91
|
def execute(self, context: Context) -> None:
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
92
|
+
if not self.bigquery_hook:
|
|
93
|
+
self.bigquery_hook = BigQueryHook(
|
|
94
|
+
gcp_conn_id=self.gcp_conn_id,
|
|
95
|
+
location=self.location,
|
|
96
|
+
impersonation_chain=self.impersonation_chain,
|
|
97
|
+
)
|
|
98
|
+
# Set source_project_dataset_table here, after hooks are initialized and project_id is available
|
|
99
|
+
project_id = self.bigquery_hook.project_id
|
|
100
|
+
self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
|
|
101
|
+
|
|
87
102
|
self.persist_links(context)
|
|
88
|
-
sql_hook: PostgresHook = self.get_sql_hook()
|
|
89
103
|
for rows in bigquery_get_data(
|
|
90
104
|
self.log,
|
|
91
105
|
self.dataset_id,
|
|
92
106
|
self.table_id,
|
|
93
|
-
|
|
107
|
+
self.bigquery_hook,
|
|
94
108
|
self.batch_size,
|
|
95
109
|
self.selected_fields,
|
|
96
110
|
):
|
|
97
|
-
|
|
111
|
+
self.postgres_hook.insert_rows(
|
|
98
112
|
table=self.target_table_name,
|
|
99
113
|
rows=rows,
|
|
100
114
|
target_fields=self.selected_fields,
|