apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.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/3rd-party-licenses/LICENSES.txt +14 -0
- airflow/providers/google/3rd-party-licenses/NOTICE +5 -0
- airflow/providers/google/__init__.py +3 -3
- airflow/providers/google/_vendor/__init__.py +0 -0
- airflow/providers/google/_vendor/json_merge_patch.py +91 -0
- airflow/providers/google/ads/hooks/ads.py +52 -43
- airflow/providers/google/ads/operators/ads.py +2 -2
- airflow/providers/google/ads/transfers/ads_to_gcs.py +3 -19
- airflow/providers/google/assets/gcs.py +1 -11
- airflow/providers/google/cloud/_internal_client/secret_manager_client.py +3 -2
- airflow/providers/google/cloud/bundles/gcs.py +161 -0
- airflow/providers/google/cloud/hooks/alloy_db.py +2 -3
- airflow/providers/google/cloud/hooks/bigquery.py +195 -318
- airflow/providers/google/cloud/hooks/bigquery_dts.py +8 -8
- airflow/providers/google/cloud/hooks/bigtable.py +3 -2
- airflow/providers/google/cloud/hooks/cloud_batch.py +8 -9
- airflow/providers/google/cloud/hooks/cloud_build.py +6 -65
- airflow/providers/google/cloud/hooks/cloud_composer.py +292 -24
- airflow/providers/google/cloud/hooks/cloud_logging.py +109 -0
- airflow/providers/google/cloud/hooks/cloud_memorystore.py +4 -3
- airflow/providers/google/cloud/hooks/cloud_run.py +20 -11
- airflow/providers/google/cloud/hooks/cloud_sql.py +136 -64
- airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +35 -15
- airflow/providers/google/cloud/hooks/compute.py +7 -6
- airflow/providers/google/cloud/hooks/compute_ssh.py +7 -4
- airflow/providers/google/cloud/hooks/datacatalog.py +12 -3
- airflow/providers/google/cloud/hooks/dataflow.py +87 -242
- airflow/providers/google/cloud/hooks/dataform.py +9 -14
- airflow/providers/google/cloud/hooks/datafusion.py +7 -9
- airflow/providers/google/cloud/hooks/dataplex.py +13 -12
- airflow/providers/google/cloud/hooks/dataprep.py +2 -2
- airflow/providers/google/cloud/hooks/dataproc.py +76 -74
- airflow/providers/google/cloud/hooks/dataproc_metastore.py +4 -3
- airflow/providers/google/cloud/hooks/dlp.py +5 -4
- airflow/providers/google/cloud/hooks/gcs.py +144 -33
- airflow/providers/google/cloud/hooks/gen_ai.py +196 -0
- airflow/providers/google/cloud/hooks/kms.py +3 -2
- airflow/providers/google/cloud/hooks/kubernetes_engine.py +22 -17
- airflow/providers/google/cloud/hooks/looker.py +6 -1
- airflow/providers/google/cloud/hooks/managed_kafka.py +227 -3
- airflow/providers/google/cloud/hooks/mlengine.py +7 -8
- airflow/providers/google/cloud/hooks/natural_language.py +3 -2
- airflow/providers/google/cloud/hooks/os_login.py +3 -2
- airflow/providers/google/cloud/hooks/pubsub.py +6 -6
- airflow/providers/google/cloud/hooks/secret_manager.py +105 -12
- airflow/providers/google/cloud/hooks/spanner.py +75 -10
- airflow/providers/google/cloud/hooks/speech_to_text.py +3 -2
- airflow/providers/google/cloud/hooks/stackdriver.py +18 -18
- airflow/providers/google/cloud/hooks/tasks.py +4 -3
- airflow/providers/google/cloud/hooks/text_to_speech.py +3 -2
- airflow/providers/google/cloud/hooks/translate.py +8 -17
- airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +8 -222
- airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +9 -15
- airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +33 -283
- airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +5 -12
- airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +6 -12
- airflow/providers/google/cloud/hooks/vertex_ai/experiment_service.py +202 -0
- airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +311 -10
- airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +79 -75
- airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +7 -13
- airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +8 -12
- airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +6 -12
- airflow/providers/google/cloud/hooks/vertex_ai/prediction_service.py +3 -2
- airflow/providers/google/cloud/hooks/vertex_ai/ray.py +223 -0
- airflow/providers/google/cloud/hooks/video_intelligence.py +3 -2
- airflow/providers/google/cloud/hooks/vision.py +7 -7
- airflow/providers/google/cloud/hooks/workflows.py +4 -3
- airflow/providers/google/cloud/links/alloy_db.py +0 -46
- airflow/providers/google/cloud/links/base.py +77 -7
- 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 -46
- 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 -90
- 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 -89
- airflow/providers/google/cloud/links/datastore.py +0 -31
- airflow/providers/google/cloud/links/kubernetes_engine.py +11 -61
- airflow/providers/google/cloud/links/managed_kafka.py +11 -51
- 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 +166 -118
- airflow/providers/google/cloud/log/stackdriver_task_handler.py +14 -9
- 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 +141 -40
- airflow/providers/google/cloud/openlineage/mixins.py +14 -13
- airflow/providers/google/cloud/openlineage/utils.py +19 -3
- airflow/providers/google/cloud/operators/alloy_db.py +76 -61
- airflow/providers/google/cloud/operators/bigquery.py +104 -667
- airflow/providers/google/cloud/operators/bigquery_dts.py +12 -12
- airflow/providers/google/cloud/operators/bigtable.py +38 -7
- airflow/providers/google/cloud/operators/cloud_base.py +22 -1
- airflow/providers/google/cloud/operators/cloud_batch.py +18 -18
- airflow/providers/google/cloud/operators/cloud_build.py +80 -36
- airflow/providers/google/cloud/operators/cloud_composer.py +157 -71
- airflow/providers/google/cloud/operators/cloud_logging_sink.py +341 -0
- airflow/providers/google/cloud/operators/cloud_memorystore.py +74 -46
- airflow/providers/google/cloud/operators/cloud_run.py +39 -20
- airflow/providers/google/cloud/operators/cloud_sql.py +46 -61
- airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +92 -14
- airflow/providers/google/cloud/operators/compute.py +18 -50
- airflow/providers/google/cloud/operators/datacatalog.py +167 -29
- airflow/providers/google/cloud/operators/dataflow.py +38 -15
- airflow/providers/google/cloud/operators/dataform.py +19 -7
- airflow/providers/google/cloud/operators/datafusion.py +43 -43
- airflow/providers/google/cloud/operators/dataplex.py +212 -126
- airflow/providers/google/cloud/operators/dataprep.py +1 -5
- airflow/providers/google/cloud/operators/dataproc.py +134 -207
- airflow/providers/google/cloud/operators/dataproc_metastore.py +102 -84
- airflow/providers/google/cloud/operators/datastore.py +22 -6
- airflow/providers/google/cloud/operators/dlp.py +24 -45
- airflow/providers/google/cloud/operators/functions.py +21 -14
- airflow/providers/google/cloud/operators/gcs.py +15 -12
- airflow/providers/google/cloud/operators/gen_ai.py +389 -0
- airflow/providers/google/cloud/operators/kubernetes_engine.py +115 -106
- airflow/providers/google/cloud/operators/looker.py +1 -1
- airflow/providers/google/cloud/operators/managed_kafka.py +362 -40
- airflow/providers/google/cloud/operators/natural_language.py +5 -3
- airflow/providers/google/cloud/operators/pubsub.py +69 -21
- airflow/providers/google/cloud/operators/spanner.py +53 -45
- airflow/providers/google/cloud/operators/speech_to_text.py +5 -4
- airflow/providers/google/cloud/operators/stackdriver.py +5 -11
- airflow/providers/google/cloud/operators/tasks.py +6 -15
- airflow/providers/google/cloud/operators/text_to_speech.py +4 -3
- airflow/providers/google/cloud/operators/translate.py +46 -20
- airflow/providers/google/cloud/operators/translate_speech.py +4 -3
- airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +44 -34
- airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +34 -12
- airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +62 -53
- airflow/providers/google/cloud/operators/vertex_ai/dataset.py +75 -11
- airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +48 -12
- 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 -116
- airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +16 -12
- airflow/providers/google/cloud/operators/vertex_ai/model_service.py +62 -14
- airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +35 -10
- airflow/providers/google/cloud/operators/vertex_ai/ray.py +393 -0
- airflow/providers/google/cloud/operators/video_intelligence.py +5 -3
- airflow/providers/google/cloud/operators/vision.py +7 -5
- airflow/providers/google/cloud/operators/workflows.py +24 -19
- airflow/providers/google/cloud/secrets/secret_manager.py +2 -1
- airflow/providers/google/cloud/sensors/bigquery.py +2 -2
- airflow/providers/google/cloud/sensors/bigquery_dts.py +6 -4
- airflow/providers/google/cloud/sensors/bigtable.py +14 -6
- airflow/providers/google/cloud/sensors/cloud_composer.py +535 -33
- airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +6 -5
- airflow/providers/google/cloud/sensors/dataflow.py +27 -10
- airflow/providers/google/cloud/sensors/dataform.py +2 -2
- airflow/providers/google/cloud/sensors/datafusion.py +4 -4
- airflow/providers/google/cloud/sensors/dataplex.py +7 -5
- airflow/providers/google/cloud/sensors/dataprep.py +2 -2
- airflow/providers/google/cloud/sensors/dataproc.py +10 -9
- airflow/providers/google/cloud/sensors/dataproc_metastore.py +4 -3
- airflow/providers/google/cloud/sensors/gcs.py +22 -21
- airflow/providers/google/cloud/sensors/looker.py +5 -5
- airflow/providers/google/cloud/sensors/pubsub.py +20 -20
- airflow/providers/google/cloud/sensors/tasks.py +2 -2
- airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +2 -2
- airflow/providers/google/cloud/sensors/workflows.py +6 -4
- 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 +2 -2
- airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py +11 -8
- airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +14 -13
- 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 +18 -22
- airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +4 -5
- airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +45 -38
- airflow/providers/google/cloud/transfers/gcs_to_gcs.py +2 -2
- airflow/providers/google/cloud/transfers/gcs_to_local.py +5 -3
- airflow/providers/google/cloud/transfers/gcs_to_sftp.py +10 -4
- 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 +44 -12
- airflow/providers/google/cloud/transfers/s3_to_gcs.py +12 -6
- airflow/providers/google/cloud/transfers/salesforce_to_gcs.py +2 -2
- airflow/providers/google/cloud/transfers/sftp_to_gcs.py +36 -14
- 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 +75 -34
- airflow/providers/google/cloud/triggers/bigquery_dts.py +2 -1
- airflow/providers/google/cloud/triggers/cloud_batch.py +2 -1
- airflow/providers/google/cloud/triggers/cloud_build.py +3 -2
- airflow/providers/google/cloud/triggers/cloud_composer.py +303 -47
- airflow/providers/google/cloud/triggers/cloud_run.py +2 -2
- airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +96 -5
- airflow/providers/google/cloud/triggers/dataflow.py +125 -2
- airflow/providers/google/cloud/triggers/datafusion.py +1 -1
- airflow/providers/google/cloud/triggers/dataplex.py +16 -3
- airflow/providers/google/cloud/triggers/dataproc.py +124 -53
- airflow/providers/google/cloud/triggers/kubernetes_engine.py +46 -28
- airflow/providers/google/cloud/triggers/mlengine.py +1 -1
- airflow/providers/google/cloud/triggers/pubsub.py +17 -20
- airflow/providers/google/cloud/triggers/vertex_ai.py +8 -7
- airflow/providers/google/cloud/utils/bigquery.py +5 -7
- airflow/providers/google/cloud/utils/bigquery_get_data.py +1 -1
- airflow/providers/google/cloud/utils/credentials_provider.py +4 -3
- airflow/providers/google/cloud/utils/dataform.py +1 -1
- airflow/providers/google/cloud/utils/external_token_supplier.py +0 -1
- airflow/providers/google/cloud/utils/field_validator.py +1 -2
- airflow/providers/google/cloud/utils/validators.py +43 -0
- airflow/providers/google/common/auth_backend/google_openid.py +26 -9
- airflow/providers/google/common/consts.py +2 -1
- airflow/providers/google/common/deprecated.py +2 -1
- airflow/providers/google/common/hooks/base_google.py +40 -43
- airflow/providers/google/common/hooks/operation_helpers.py +78 -0
- 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 +4 -5
- airflow/providers/google/firebase/operators/firestore.py +2 -2
- airflow/providers/google/get_provider_info.py +61 -216
- airflow/providers/google/go_module_utils.py +35 -3
- airflow/providers/google/leveldb/hooks/leveldb.py +30 -6
- airflow/providers/google/leveldb/operators/leveldb.py +2 -2
- airflow/providers/google/marketing_platform/hooks/analytics_admin.py +3 -2
- 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 +4 -5
- airflow/providers/google/marketing_platform/operators/analytics_admin.py +7 -6
- airflow/providers/google/marketing_platform/operators/campaign_manager.py +5 -5
- 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 +4 -64
- airflow/providers/google/suite/hooks/calendar.py +1 -1
- airflow/providers/google/suite/hooks/drive.py +2 -2
- airflow/providers/google/suite/hooks/sheets.py +15 -1
- 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 +2 -2
- 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-14.0.0.dist-info → apache_airflow_providers_google-19.1.0rc1.dist-info}/METADATA +117 -72
- apache_airflow_providers_google-19.1.0rc1.dist-info/RECORD +331 -0
- {apache_airflow_providers_google-14.0.0.dist-info → apache_airflow_providers_google-19.1.0rc1.dist-info}/WHEEL +1 -1
- apache_airflow_providers_google-19.1.0rc1.dist-info/licenses/NOTICE +5 -0
- airflow/providers/google/cloud/example_dags/example_cloud_task.py +0 -54
- airflow/providers/google/cloud/hooks/automl.py +0 -679
- 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 -1360
- airflow/providers/google/cloud/operators/life_sciences.py +0 -119
- airflow/providers/google/cloud/operators/mlengine.py +0 -1515
- airflow/providers/google/cloud/utils/mlengine_operator_utils.py +0 -273
- apache_airflow_providers_google-14.0.0.dist-info/RECORD +0 -318
- /airflow/providers/google/cloud/{example_dags → bundles}/__init__.py +0 -0
- {apache_airflow_providers_google-14.0.0.dist-info → apache_airflow_providers_google-19.1.0rc1.dist-info}/entry_points.txt +0 -0
- {airflow/providers/google → apache_airflow_providers_google-19.1.0rc1.dist-info/licenses}/LICENSE +0 -0
|
@@ -23,9 +23,10 @@ from collections.abc import AsyncIterator, Sequence
|
|
|
23
23
|
from functools import cached_property
|
|
24
24
|
from typing import Any
|
|
25
25
|
|
|
26
|
+
from google.cloud.pubsub_v1.types import ReceivedMessage
|
|
27
|
+
|
|
26
28
|
from airflow.providers.google.cloud.hooks.pubsub import PubSubAsyncHook
|
|
27
29
|
from airflow.triggers.base import BaseTrigger, TriggerEvent
|
|
28
|
-
from google.cloud.pubsub_v1.types import ReceivedMessage
|
|
29
30
|
|
|
30
31
|
|
|
31
32
|
class PubsubPullTrigger(BaseTrigger):
|
|
@@ -84,27 +85,23 @@ class PubsubPullTrigger(BaseTrigger):
|
|
|
84
85
|
},
|
|
85
86
|
)
|
|
86
87
|
|
|
87
|
-
async def run(self) -> AsyncIterator[TriggerEvent]:
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
await self.message_acknowledgement(pulled_messages)
|
|
88
|
+
async def run(self) -> AsyncIterator[TriggerEvent]:
|
|
89
|
+
while True:
|
|
90
|
+
if pulled_messages := await self.hook.pull(
|
|
91
|
+
project_id=self.project_id,
|
|
92
|
+
subscription=self.subscription,
|
|
93
|
+
max_messages=self.max_messages,
|
|
94
|
+
return_immediately=True,
|
|
95
|
+
):
|
|
96
|
+
if self.ack_messages:
|
|
97
|
+
await self.message_acknowledgement(pulled_messages)
|
|
98
98
|
|
|
99
|
-
|
|
99
|
+
messages_json = [ReceivedMessage.to_dict(m) for m in pulled_messages]
|
|
100
100
|
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
except Exception as e:
|
|
106
|
-
yield TriggerEvent({"status": "error", "message": str(e)})
|
|
107
|
-
return
|
|
101
|
+
yield TriggerEvent({"status": "success", "message": messages_json})
|
|
102
|
+
return
|
|
103
|
+
self.log.info("Sleeping for %s seconds.", self.poke_interval)
|
|
104
|
+
await asyncio.sleep(self.poke_interval)
|
|
108
105
|
|
|
109
106
|
async def message_acknowledgement(self, pulled_messages):
|
|
110
107
|
await self.hook.acknowledge(
|
|
@@ -20,6 +20,14 @@ from collections.abc import AsyncIterator, Sequence
|
|
|
20
20
|
from functools import cached_property
|
|
21
21
|
from typing import TYPE_CHECKING, Any
|
|
22
22
|
|
|
23
|
+
from google.cloud.aiplatform_v1 import (
|
|
24
|
+
BatchPredictionJob,
|
|
25
|
+
HyperparameterTuningJob,
|
|
26
|
+
JobState,
|
|
27
|
+
PipelineState,
|
|
28
|
+
types,
|
|
29
|
+
)
|
|
30
|
+
|
|
23
31
|
from airflow.exceptions import AirflowException
|
|
24
32
|
from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobAsyncHook
|
|
25
33
|
from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobAsyncHook
|
|
@@ -28,13 +36,6 @@ from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job im
|
|
|
28
36
|
)
|
|
29
37
|
from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobAsyncHook
|
|
30
38
|
from airflow.triggers.base import BaseTrigger, TriggerEvent
|
|
31
|
-
from google.cloud.aiplatform_v1 import (
|
|
32
|
-
BatchPredictionJob,
|
|
33
|
-
HyperparameterTuningJob,
|
|
34
|
-
JobState,
|
|
35
|
-
PipelineState,
|
|
36
|
-
types,
|
|
37
|
-
)
|
|
38
39
|
|
|
39
40
|
if TYPE_CHECKING:
|
|
40
41
|
from proto import Message
|
|
@@ -27,16 +27,15 @@ def bq_cast(string_field: str, bq_type: str) -> None | int | float | bool | str:
|
|
|
27
27
|
"""
|
|
28
28
|
if string_field is None:
|
|
29
29
|
return None
|
|
30
|
-
|
|
30
|
+
if bq_type == "INTEGER":
|
|
31
31
|
return int(string_field)
|
|
32
|
-
|
|
32
|
+
if bq_type in ("FLOAT", "TIMESTAMP"):
|
|
33
33
|
return float(string_field)
|
|
34
|
-
|
|
34
|
+
if bq_type == "BOOLEAN":
|
|
35
35
|
if string_field not in ["true", "false"]:
|
|
36
36
|
raise ValueError(f"{string_field} must have value 'true' or 'false'")
|
|
37
37
|
return string_field == "true"
|
|
38
|
-
|
|
39
|
-
return string_field
|
|
38
|
+
return string_field
|
|
40
39
|
|
|
41
40
|
|
|
42
41
|
def convert_job_id(job_id: str | list[str], project_id: str, location: str | None) -> Any:
|
|
@@ -51,5 +50,4 @@ def convert_job_id(job_id: str | list[str], project_id: str, location: str | Non
|
|
|
51
50
|
location = location or "US"
|
|
52
51
|
if isinstance(job_id, list):
|
|
53
52
|
return [f"{project_id}:{location}:{i}" for i in job_id]
|
|
54
|
-
|
|
55
|
-
return f"{project_id}:{location}:{job_id}"
|
|
53
|
+
return f"{project_id}:{location}:{job_id}"
|
|
@@ -23,9 +23,9 @@ from google.cloud.bigquery.table import Row, RowIterator
|
|
|
23
23
|
|
|
24
24
|
if TYPE_CHECKING:
|
|
25
25
|
from collections.abc import Iterator
|
|
26
|
-
from logging import Logger
|
|
27
26
|
|
|
28
27
|
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
|
|
28
|
+
from airflow.sdk.types import Logger
|
|
29
29
|
|
|
30
30
|
|
|
31
31
|
def bigquery_get_data(
|
|
@@ -29,6 +29,10 @@ from urllib.parse import urlencode
|
|
|
29
29
|
|
|
30
30
|
import google.auth
|
|
31
31
|
import google.oauth2.service_account
|
|
32
|
+
from google.auth import impersonated_credentials
|
|
33
|
+
from google.auth.credentials import AnonymousCredentials, Credentials
|
|
34
|
+
from google.auth.environment_vars import CREDENTIALS, LEGACY_PROJECT, PROJECT
|
|
35
|
+
|
|
32
36
|
from airflow.exceptions import AirflowException
|
|
33
37
|
from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
|
|
34
38
|
from airflow.providers.google.cloud.utils.external_token_supplier import (
|
|
@@ -36,9 +40,6 @@ from airflow.providers.google.cloud.utils.external_token_supplier import (
|
|
|
36
40
|
)
|
|
37
41
|
from airflow.utils.log.logging_mixin import LoggingMixin
|
|
38
42
|
from airflow.utils.process_utils import patch_environ
|
|
39
|
-
from google.auth import impersonated_credentials # type: ignore[attr-defined]
|
|
40
|
-
from google.auth.credentials import AnonymousCredentials, Credentials
|
|
41
|
-
from google.auth.environment_vars import CREDENTIALS, LEGACY_PROJECT, PROJECT
|
|
42
43
|
|
|
43
44
|
log = logging.getLogger(__name__)
|
|
44
45
|
|
|
@@ -202,7 +202,7 @@ def make_initialization_workspace_flow(
|
|
|
202
202
|
def define_default_location(region: str) -> DataformLocations:
|
|
203
203
|
if "us" in region:
|
|
204
204
|
return DataformLocations.US
|
|
205
|
-
|
|
205
|
+
if "europe" in region:
|
|
206
206
|
return DataformLocations.EUROPE
|
|
207
207
|
|
|
208
208
|
regions_mapping: Mapping[str, DataformLocations] = {}
|
|
@@ -134,8 +134,7 @@ Here are the guidelines that you should follow to make validation forward-compat
|
|
|
134
134
|
from __future__ import annotations
|
|
135
135
|
|
|
136
136
|
import re
|
|
137
|
-
from collections.abc import Sequence
|
|
138
|
-
from typing import Callable
|
|
137
|
+
from collections.abc import Callable, Sequence
|
|
139
138
|
|
|
140
139
|
from airflow.exceptions import AirflowException
|
|
141
140
|
from airflow.utils.log.logging_mixin import LoggingMixin
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Licensed to the Apache Software Foundation (ASF) under one
|
|
3
|
+
# or more contributor license agreements. See the NOTICE file
|
|
4
|
+
# distributed with this work for additional information
|
|
5
|
+
# regarding copyright ownership. The ASF licenses this file
|
|
6
|
+
# to you under the Apache License, Version 2.0 (the
|
|
7
|
+
# "License"); you may not use this file except in compliance
|
|
8
|
+
# with the License. You may obtain a copy of the License at
|
|
9
|
+
#
|
|
10
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
11
|
+
#
|
|
12
|
+
# Unless required by applicable law or agreed to in writing,
|
|
13
|
+
# software distributed under the License is distributed on an
|
|
14
|
+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
15
|
+
# KIND, either express or implied. See the License for the
|
|
16
|
+
# specific language governing permissions and limitations
|
|
17
|
+
# under the License.
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import json
|
|
21
|
+
from json import JSONDecodeError
|
|
22
|
+
|
|
23
|
+
from wtforms.validators import ValidationError
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class ValidJson:
|
|
27
|
+
"""
|
|
28
|
+
Validates data is valid JSON.
|
|
29
|
+
|
|
30
|
+
:param message:
|
|
31
|
+
Error message to raise in case of a validation error.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(self, message=None):
|
|
35
|
+
self.message = message
|
|
36
|
+
|
|
37
|
+
def __call__(self, form, field):
|
|
38
|
+
if field.data:
|
|
39
|
+
try:
|
|
40
|
+
json.loads(field.data)
|
|
41
|
+
except JSONDecodeError as ex:
|
|
42
|
+
message = self.message or f"JSON Validation Error: {ex}"
|
|
43
|
+
raise ValidationError(message=field.gettext(message.format(field.data)))
|
|
@@ -20,24 +20,36 @@
|
|
|
20
20
|
from __future__ import annotations
|
|
21
21
|
|
|
22
22
|
import logging
|
|
23
|
+
from collections.abc import Callable
|
|
23
24
|
from functools import wraps
|
|
24
|
-
from typing import
|
|
25
|
+
from typing import TypeVar, cast
|
|
25
26
|
|
|
26
|
-
from flask import Response, current_app, request as flask_request # type: ignore
|
|
27
|
-
|
|
28
|
-
import google
|
|
29
27
|
import google.auth.transport.requests
|
|
30
28
|
import google.oauth2.id_token
|
|
31
|
-
|
|
32
|
-
|
|
29
|
+
|
|
30
|
+
try:
|
|
31
|
+
from flask import Response, current_app, request as flask_request
|
|
32
|
+
except ImportError:
|
|
33
|
+
raise ImportError(
|
|
34
|
+
"Google requires FAB provider to be installed in order to use this auth backend. "
|
|
35
|
+
"Please install the FAB provider by running: "
|
|
36
|
+
"pip install apache-airflow-providers-google[fab]"
|
|
37
|
+
)
|
|
33
38
|
from google.auth import exceptions
|
|
34
39
|
from google.auth.transport.requests import AuthorizedSession
|
|
35
40
|
from google.oauth2 import service_account
|
|
36
41
|
|
|
42
|
+
from airflow.configuration import conf
|
|
43
|
+
from airflow.exceptions import AirflowProviderDeprecationWarning
|
|
44
|
+
from airflow.providers.google.common.deprecated import deprecated
|
|
45
|
+
from airflow.providers.google.common.utils.id_token_credentials import get_default_id_token_credentials
|
|
46
|
+
|
|
37
47
|
log = logging.getLogger(__name__)
|
|
38
48
|
|
|
39
49
|
_GOOGLE_ISSUERS = ("accounts.google.com", "https://accounts.google.com")
|
|
40
|
-
AUDIENCE = conf.get(
|
|
50
|
+
AUDIENCE = conf.get(
|
|
51
|
+
"api", "google_oauth2_audience", fallback="project-id-random-value.apps.googleusercontent.com"
|
|
52
|
+
)
|
|
41
53
|
|
|
42
54
|
|
|
43
55
|
def create_client_session():
|
|
@@ -52,6 +64,11 @@ def create_client_session():
|
|
|
52
64
|
return AuthorizedSession(credentials=id_token_credentials)
|
|
53
65
|
|
|
54
66
|
|
|
67
|
+
@deprecated(
|
|
68
|
+
planned_removal_release="apache-airflow-providers-google==15.0.0",
|
|
69
|
+
reason="Auth backends are not supported on Airflow 3, and this entire module will be removed",
|
|
70
|
+
category=AirflowProviderDeprecationWarning,
|
|
71
|
+
)
|
|
55
72
|
def init_app(_):
|
|
56
73
|
"""Initialize authentication."""
|
|
57
74
|
|
|
@@ -104,7 +121,7 @@ def _lookup_user(user_email: str):
|
|
|
104
121
|
|
|
105
122
|
|
|
106
123
|
def _set_current_user(user):
|
|
107
|
-
current_app.appbuilder.sm.lm._update_request_context_with_user(user=user)
|
|
124
|
+
current_app.appbuilder.sm.lm._update_request_context_with_user(user=user)
|
|
108
125
|
|
|
109
126
|
|
|
110
127
|
T = TypeVar("T", bound=Callable)
|
|
@@ -137,4 +154,4 @@ def requires_authentication(function: T):
|
|
|
137
154
|
|
|
138
155
|
return function(*args, **kwargs)
|
|
139
156
|
|
|
140
|
-
return cast(T, decorated)
|
|
157
|
+
return cast("T", decorated)
|
|
@@ -16,9 +16,10 @@
|
|
|
16
16
|
# under the License.
|
|
17
17
|
from __future__ import annotations
|
|
18
18
|
|
|
19
|
-
from airflow import version
|
|
20
19
|
from google.api_core.gapic_v1.client_info import ClientInfo
|
|
21
20
|
|
|
21
|
+
from airflow import version
|
|
22
|
+
|
|
22
23
|
GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME = "execute_complete"
|
|
23
24
|
|
|
24
25
|
CLIENT_INFO = ClientInfo(client_library_version="airflow_v" + version.version)
|
|
@@ -18,8 +18,9 @@ from __future__ import annotations
|
|
|
18
18
|
|
|
19
19
|
import inspect
|
|
20
20
|
import re
|
|
21
|
+
from collections.abc import Callable
|
|
21
22
|
from datetime import date, datetime
|
|
22
|
-
from typing import Any
|
|
23
|
+
from typing import Any
|
|
23
24
|
|
|
24
25
|
from deprecated import deprecated as standard_deprecated
|
|
25
26
|
from deprecated.classic import ClassicAdapter
|
|
@@ -26,44 +26,39 @@ import json
|
|
|
26
26
|
import logging
|
|
27
27
|
import os
|
|
28
28
|
import tempfile
|
|
29
|
-
from collections.abc import Generator, Sequence
|
|
29
|
+
from collections.abc import Callable, Generator, Sequence
|
|
30
30
|
from contextlib import ExitStack, contextmanager
|
|
31
31
|
from subprocess import check_output
|
|
32
|
-
from typing import TYPE_CHECKING, Any,
|
|
32
|
+
from typing import TYPE_CHECKING, Any, TypeVar, cast
|
|
33
33
|
|
|
34
|
+
import google.auth
|
|
34
35
|
import google_auth_httplib2
|
|
35
36
|
import requests
|
|
36
37
|
import tenacity
|
|
37
38
|
from asgiref.sync import sync_to_async
|
|
38
39
|
from gcloud.aio.auth.token import Token, TokenResponse
|
|
40
|
+
from google.api_core.exceptions import Forbidden, ResourceExhausted, TooManyRequests
|
|
41
|
+
from google.auth import _cloud_sdk, compute_engine
|
|
42
|
+
from google.auth.environment_vars import CLOUD_SDK_CONFIG_DIR, CREDENTIALS
|
|
43
|
+
from google.auth.exceptions import RefreshError
|
|
44
|
+
from google.auth.transport import _http_client
|
|
39
45
|
from googleapiclient import discovery
|
|
40
46
|
from googleapiclient.errors import HttpError
|
|
41
47
|
from googleapiclient.http import MediaIoBaseDownload, build_http, set_user_agent
|
|
42
48
|
from requests import Session
|
|
43
49
|
|
|
44
|
-
import google.auth
|
|
45
|
-
import google.oauth2.service_account
|
|
46
50
|
from airflow import version
|
|
47
|
-
from airflow.exceptions import AirflowException
|
|
48
|
-
from airflow.
|
|
51
|
+
from airflow.exceptions import AirflowException
|
|
52
|
+
from airflow.providers.common.compat.sdk import BaseHook
|
|
49
53
|
from airflow.providers.google.cloud.utils.credentials_provider import (
|
|
50
54
|
_get_scopes,
|
|
51
55
|
_get_target_principal_and_delegates,
|
|
52
56
|
get_credentials_and_project_id,
|
|
53
57
|
)
|
|
54
|
-
from airflow.providers.google.common.consts import CLIENT_INFO
|
|
55
|
-
from airflow.providers.google.common.deprecated import deprecated
|
|
56
58
|
from airflow.utils.process_utils import patch_environ
|
|
57
|
-
from google.api_core.exceptions import Forbidden, ResourceExhausted, TooManyRequests
|
|
58
|
-
from google.auth import _cloud_sdk, compute_engine # type: ignore[attr-defined]
|
|
59
|
-
from google.auth.environment_vars import CLOUD_SDK_CONFIG_DIR, CREDENTIALS
|
|
60
|
-
from google.auth.exceptions import RefreshError
|
|
61
|
-
from google.auth.transport import _http_client
|
|
62
59
|
|
|
63
60
|
if TYPE_CHECKING:
|
|
64
61
|
from aiohttp import ClientSession
|
|
65
|
-
|
|
66
|
-
from google.api_core.gapic_v1.client_info import ClientInfo
|
|
67
62
|
from google.auth.credentials import Credentials
|
|
68
63
|
|
|
69
64
|
log = logging.getLogger(__name__)
|
|
@@ -154,13 +149,16 @@ class retry_if_temporary_refresh_credentials(tenacity.retry_if_exception):
|
|
|
154
149
|
# This allows the 'project_id' argument to be of type str instead of str | None,
|
|
155
150
|
# making it easier to type hint the function body without dealing with the None
|
|
156
151
|
# case that can never happen at runtime.
|
|
157
|
-
PROVIDE_PROJECT_ID: str = cast(str, None)
|
|
152
|
+
PROVIDE_PROJECT_ID: str = cast("str", None)
|
|
158
153
|
|
|
159
154
|
T = TypeVar("T", bound=Callable)
|
|
160
155
|
RT = TypeVar("RT")
|
|
161
156
|
|
|
157
|
+
# Sentinel value to distinguish "parameter not provided" from "parameter explicitly set to a value"
|
|
158
|
+
_UNSET = object()
|
|
159
|
+
|
|
162
160
|
|
|
163
|
-
def get_field(extras: dict, field_name: str):
|
|
161
|
+
def get_field(extras: dict, field_name: str) -> str | None:
|
|
164
162
|
"""Get field from extra, first checking short name, then for backcompat we check for prefixed name."""
|
|
165
163
|
if field_name.startswith("extra__"):
|
|
166
164
|
raise ValueError(
|
|
@@ -168,9 +166,11 @@ def get_field(extras: dict, field_name: str):
|
|
|
168
166
|
"when using this method."
|
|
169
167
|
)
|
|
170
168
|
if field_name in extras:
|
|
171
|
-
|
|
169
|
+
value = extras[field_name]
|
|
170
|
+
return None if value == "" else value
|
|
172
171
|
prefixed_name = f"extra__google_cloud_platform__{field_name}"
|
|
173
|
-
|
|
172
|
+
value = extras.get(prefixed_name)
|
|
173
|
+
return None if value == "" else value
|
|
174
174
|
|
|
175
175
|
|
|
176
176
|
class GoogleBaseHook(BaseHook):
|
|
@@ -410,7 +410,22 @@ class GoogleBaseHook(BaseHook):
|
|
|
410
410
|
custom UI elements to the hook page, which allow admins to specify
|
|
411
411
|
service_account, key_path, etc. They get formatted as shown below.
|
|
412
412
|
"""
|
|
413
|
-
|
|
413
|
+
# New behavior: If default is _UNSET, parameter was not provided
|
|
414
|
+
# Check connection extras first, return None if not found (caller handles default)
|
|
415
|
+
if default is _UNSET:
|
|
416
|
+
if hasattr(self, "extras"):
|
|
417
|
+
value = get_field(self.extras, f)
|
|
418
|
+
if value is not None:
|
|
419
|
+
return value
|
|
420
|
+
return None
|
|
421
|
+
|
|
422
|
+
# Old behavior (for backward compatibility):
|
|
423
|
+
# Check connection extras, but properly handle False values
|
|
424
|
+
if hasattr(self, "extras"):
|
|
425
|
+
value = get_field(self.extras, f)
|
|
426
|
+
if value is not None:
|
|
427
|
+
return value
|
|
428
|
+
return default
|
|
414
429
|
|
|
415
430
|
@property
|
|
416
431
|
def project_id(self) -> str:
|
|
@@ -443,24 +458,6 @@ class GoogleBaseHook(BaseHook):
|
|
|
443
458
|
f"Please check the connection configuration."
|
|
444
459
|
)
|
|
445
460
|
|
|
446
|
-
@property
|
|
447
|
-
@deprecated(
|
|
448
|
-
planned_removal_date="March 01, 2025",
|
|
449
|
-
use_instead="airflow.providers.google.common.consts.CLIENT_INFO",
|
|
450
|
-
category=AirflowProviderDeprecationWarning,
|
|
451
|
-
)
|
|
452
|
-
def client_info(self) -> ClientInfo:
|
|
453
|
-
"""
|
|
454
|
-
Return client information used to generate a user-agent for API calls.
|
|
455
|
-
|
|
456
|
-
It allows for better errors tracking.
|
|
457
|
-
|
|
458
|
-
This object is only used by the google-cloud-* libraries that are built specifically for
|
|
459
|
-
the Google Cloud. It is not supported by The Google APIs Python Client that use Discovery
|
|
460
|
-
based APIs.
|
|
461
|
-
"""
|
|
462
|
-
return CLIENT_INFO
|
|
463
|
-
|
|
464
461
|
@property
|
|
465
462
|
def scopes(self) -> Sequence[str]:
|
|
466
463
|
"""
|
|
@@ -500,7 +497,7 @@ class GoogleBaseHook(BaseHook):
|
|
|
500
497
|
"after": tenacity.after_log(log, logging.DEBUG),
|
|
501
498
|
}
|
|
502
499
|
default_kwargs.update(**kwargs)
|
|
503
|
-
return cast(T, tenacity.retry(*args, **default_kwargs)(func))
|
|
500
|
+
return cast("T", tenacity.retry(*args, **default_kwargs)(func))
|
|
504
501
|
|
|
505
502
|
return decorator
|
|
506
503
|
|
|
@@ -518,7 +515,7 @@ class GoogleBaseHook(BaseHook):
|
|
|
518
515
|
"after": tenacity.after_log(log, logging.DEBUG),
|
|
519
516
|
}
|
|
520
517
|
default_kwargs.update(**kwargs)
|
|
521
|
-
return cast(T, tenacity.retry(*args, **default_kwargs)(func))
|
|
518
|
+
return cast("T", tenacity.retry(*args, **default_kwargs)(func))
|
|
522
519
|
|
|
523
520
|
return decorator
|
|
524
521
|
|
|
@@ -570,7 +567,7 @@ class GoogleBaseHook(BaseHook):
|
|
|
570
567
|
with self.provide_gcp_credential_file_as_context():
|
|
571
568
|
return func(self, *args, **kwargs)
|
|
572
569
|
|
|
573
|
-
return cast(T, wrapper)
|
|
570
|
+
return cast("T", wrapper)
|
|
574
571
|
|
|
575
572
|
@contextmanager
|
|
576
573
|
def provide_gcp_credential_file_as_context(self) -> Generator[str | None, None, None]:
|
|
@@ -719,7 +716,7 @@ class _CredentialsToken(Token):
|
|
|
719
716
|
scopes: Sequence[str] | None = None,
|
|
720
717
|
) -> None:
|
|
721
718
|
_scopes: list[str] | None = list(scopes) if scopes else None
|
|
722
|
-
super().__init__(session=cast(Session, session), scopes=_scopes)
|
|
719
|
+
super().__init__(session=cast("Session", session), scopes=_scopes)
|
|
723
720
|
self.credentials = credentials
|
|
724
721
|
self.project = project
|
|
725
722
|
|
|
@@ -744,7 +741,7 @@ class _CredentialsToken(Token):
|
|
|
744
741
|
async def refresh(self, *, timeout: int) -> TokenResponse:
|
|
745
742
|
await sync_to_async(self.credentials.refresh)(google.auth.transport.requests.Request())
|
|
746
743
|
|
|
747
|
-
self.access_token = cast(str, self.credentials.token)
|
|
744
|
+
self.access_token = cast("str", self.credentials.token)
|
|
748
745
|
self.access_token_duration = 3600
|
|
749
746
|
self.access_token_acquired_at = self._now()
|
|
750
747
|
return TokenResponse(value=self.access_token, expires_in=self.access_token_duration)
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
#
|
|
2
|
+
# Licensed to the Apache Software Foundation (ASF) under one
|
|
3
|
+
# or more contributor license agreements. See the NOTICE file
|
|
4
|
+
# distributed with this work for additional information
|
|
5
|
+
# regarding copyright ownership. The ASF licenses this file
|
|
6
|
+
# to you under the Apache License, Version 2.0 (the
|
|
7
|
+
# "License"); you may not use this file except in compliance
|
|
8
|
+
# with the License. You may obtain a copy of the License at
|
|
9
|
+
#
|
|
10
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
11
|
+
#
|
|
12
|
+
# Unless required by applicable law or agreed to in writing,
|
|
13
|
+
# software distributed under the License is distributed on an
|
|
14
|
+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
15
|
+
# KIND, either express or implied. See the License for the
|
|
16
|
+
# specific language governing permissions and limitations
|
|
17
|
+
# under the License.
|
|
18
|
+
"""This module contains a helper class to work with `google.api_core.operation.Operation` object."""
|
|
19
|
+
|
|
20
|
+
from __future__ import annotations
|
|
21
|
+
|
|
22
|
+
from typing import TYPE_CHECKING
|
|
23
|
+
|
|
24
|
+
from google.api_core.exceptions import GoogleAPICallError
|
|
25
|
+
|
|
26
|
+
from airflow.exceptions import AirflowException
|
|
27
|
+
|
|
28
|
+
if TYPE_CHECKING:
|
|
29
|
+
from google.api_core.operation import Operation
|
|
30
|
+
from google.api_core.retry import Retry
|
|
31
|
+
from proto import Message
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class OperationHelper:
|
|
35
|
+
"""Helper class to work with `operation.Operation` objects."""
|
|
36
|
+
|
|
37
|
+
@staticmethod
|
|
38
|
+
def wait_for_operation_result(
|
|
39
|
+
operation: Operation,
|
|
40
|
+
timeout: int | None = None,
|
|
41
|
+
polling: Retry | None = None,
|
|
42
|
+
retry: Retry | None = None,
|
|
43
|
+
) -> Message:
|
|
44
|
+
"""
|
|
45
|
+
Wait for long-lasting operation result to be retrieved.
|
|
46
|
+
|
|
47
|
+
For advance usage please check the docs on:
|
|
48
|
+
:class:`google.api_core.future.polling.PollingFuture`
|
|
49
|
+
:class:`google.api_core.retry.Retry`
|
|
50
|
+
|
|
51
|
+
:param operation: The initial operation to get result from.
|
|
52
|
+
:param timeout: How long (in seconds) to wait for the operation to complete.
|
|
53
|
+
If None, wait indefinitely. Overrides polling.timeout if both specified.
|
|
54
|
+
:param polling: How often and for how long to call polling RPC periodically.
|
|
55
|
+
:param retry: How to retry the operation polling if error occurs.
|
|
56
|
+
"""
|
|
57
|
+
try:
|
|
58
|
+
return operation.result(timeout=timeout, polling=polling, retry=retry)
|
|
59
|
+
except GoogleAPICallError as ex:
|
|
60
|
+
raise AirflowException("Google API error on operation result call") from ex
|
|
61
|
+
except Exception:
|
|
62
|
+
error = operation.exception(timeout=timeout)
|
|
63
|
+
raise AirflowException(error)
|
|
64
|
+
|
|
65
|
+
def wait_for_operation(
|
|
66
|
+
self,
|
|
67
|
+
operation: Operation,
|
|
68
|
+
timeout: float | int | None = None,
|
|
69
|
+
):
|
|
70
|
+
"""
|
|
71
|
+
Legacy method name wrapper.
|
|
72
|
+
|
|
73
|
+
Intended to use with existing hooks/operators, until the proper deprecation and replacement provided.
|
|
74
|
+
"""
|
|
75
|
+
if isinstance(timeout, float):
|
|
76
|
+
timeout = int(timeout)
|
|
77
|
+
|
|
78
|
+
return self.wait_for_operation_result(operation=operation, timeout=timeout)
|
|
@@ -18,18 +18,12 @@
|
|
|
18
18
|
|
|
19
19
|
from __future__ import annotations
|
|
20
20
|
|
|
21
|
-
from typing import TYPE_CHECKING
|
|
22
|
-
|
|
23
21
|
from airflow.providers.google.cloud.links.base import BaseGoogleLink
|
|
24
22
|
|
|
25
23
|
BASE_LINK = "https://console.cloud.google.com"
|
|
26
24
|
GCS_STORAGE_LINK = BASE_LINK + "/storage/browser/{uri};tab=objects?project={project_id}"
|
|
27
25
|
GCS_FILE_DETAILS_LINK = BASE_LINK + "/storage/browser/_details/{uri};tab=live_object?project={project_id}"
|
|
28
26
|
|
|
29
|
-
if TYPE_CHECKING:
|
|
30
|
-
from airflow.models import BaseOperator
|
|
31
|
-
from airflow.utils.context import Context
|
|
32
|
-
|
|
33
27
|
|
|
34
28
|
class StorageLink(BaseGoogleLink):
|
|
35
29
|
"""Helper class for constructing GCS Storage link."""
|
|
@@ -38,14 +32,6 @@ class StorageLink(BaseGoogleLink):
|
|
|
38
32
|
key = "storage_conf"
|
|
39
33
|
format_str = GCS_STORAGE_LINK
|
|
40
34
|
|
|
41
|
-
@staticmethod
|
|
42
|
-
def persist(context: Context, task_instance, uri: str, project_id: str | None):
|
|
43
|
-
task_instance.xcom_push(
|
|
44
|
-
context=context,
|
|
45
|
-
key=StorageLink.key,
|
|
46
|
-
value={"uri": uri, "project_id": project_id},
|
|
47
|
-
)
|
|
48
|
-
|
|
49
35
|
|
|
50
36
|
class FileDetailsLink(BaseGoogleLink):
|
|
51
37
|
"""Helper class for constructing GCS file details link."""
|
|
@@ -53,11 +39,3 @@ class FileDetailsLink(BaseGoogleLink):
|
|
|
53
39
|
name = "GCS File Details"
|
|
54
40
|
key = "file_details"
|
|
55
41
|
format_str = GCS_FILE_DETAILS_LINK
|
|
56
|
-
|
|
57
|
-
@staticmethod
|
|
58
|
-
def persist(context: Context, task_instance: BaseOperator, uri: str, project_id: str | None):
|
|
59
|
-
task_instance.xcom_push(
|
|
60
|
-
context=context,
|
|
61
|
-
key=FileDetailsLink.key,
|
|
62
|
-
value={"uri": uri, "project_id": project_id},
|
|
63
|
-
)
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# Licensed to the Apache Software Foundation (ASF) under one
|
|
2
|
+
# or more contributor license agreements. See the NOTICE file
|
|
3
|
+
# distributed with this work for additional information
|
|
4
|
+
# regarding copyright ownership. The ASF licenses this file
|
|
5
|
+
# to you under the Apache License, Version 2.0 (the
|
|
6
|
+
# "License"); you may not use this file except in compliance
|
|
7
|
+
# with the License. You may obtain a copy of the License at
|
|
8
|
+
#
|
|
9
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
10
|
+
#
|
|
11
|
+
# Unless required by applicable law or agreed to in writing,
|
|
12
|
+
# software distributed under the License is distributed on an
|
|
13
|
+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
14
|
+
# KIND, either express or implied. See the License for the
|
|
15
|
+
# specific language governing permissions and limitations
|
|
16
|
+
# under the License.
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
from google.cloud.exceptions import NotFound
|
|
21
|
+
|
|
22
|
+
from airflow.providers.google.cloud.hooks.secret_manager import (
|
|
23
|
+
GoogleCloudSecretManagerHook,
|
|
24
|
+
)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def get_secret(secret_id: str) -> str:
|
|
28
|
+
hook = GoogleCloudSecretManagerHook()
|
|
29
|
+
if hook.secret_exists(secret_id=secret_id):
|
|
30
|
+
return hook.access_secret(secret_id=secret_id).payload.data.decode()
|
|
31
|
+
raise NotFound("The secret '%s' not found", secret_id)
|