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.
Files changed (257) hide show
  1. airflow/providers/google/3rd-party-licenses/NOTICE +2 -12
  2. airflow/providers/google/__init__.py +3 -3
  3. airflow/providers/google/ads/hooks/ads.py +39 -6
  4. airflow/providers/google/ads/operators/ads.py +2 -2
  5. airflow/providers/google/ads/transfers/ads_to_gcs.py +2 -2
  6. airflow/providers/google/assets/gcs.py +1 -11
  7. airflow/providers/google/cloud/bundles/__init__.py +16 -0
  8. airflow/providers/google/cloud/bundles/gcs.py +161 -0
  9. airflow/providers/google/cloud/hooks/alloy_db.py +1 -1
  10. airflow/providers/google/cloud/hooks/bigquery.py +176 -293
  11. airflow/providers/google/cloud/hooks/cloud_batch.py +1 -1
  12. airflow/providers/google/cloud/hooks/cloud_build.py +1 -1
  13. airflow/providers/google/cloud/hooks/cloud_composer.py +288 -15
  14. airflow/providers/google/cloud/hooks/cloud_logging.py +109 -0
  15. airflow/providers/google/cloud/hooks/cloud_memorystore.py +1 -1
  16. airflow/providers/google/cloud/hooks/cloud_run.py +18 -10
  17. airflow/providers/google/cloud/hooks/cloud_sql.py +102 -23
  18. airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +29 -7
  19. airflow/providers/google/cloud/hooks/compute.py +1 -1
  20. airflow/providers/google/cloud/hooks/compute_ssh.py +6 -2
  21. airflow/providers/google/cloud/hooks/datacatalog.py +10 -1
  22. airflow/providers/google/cloud/hooks/dataflow.py +72 -95
  23. airflow/providers/google/cloud/hooks/dataform.py +1 -1
  24. airflow/providers/google/cloud/hooks/datafusion.py +21 -19
  25. airflow/providers/google/cloud/hooks/dataplex.py +2 -2
  26. airflow/providers/google/cloud/hooks/dataprep.py +1 -1
  27. airflow/providers/google/cloud/hooks/dataproc.py +73 -72
  28. airflow/providers/google/cloud/hooks/dataproc_metastore.py +1 -1
  29. airflow/providers/google/cloud/hooks/dlp.py +1 -1
  30. airflow/providers/google/cloud/hooks/functions.py +1 -1
  31. airflow/providers/google/cloud/hooks/gcs.py +112 -15
  32. airflow/providers/google/cloud/hooks/gdm.py +1 -1
  33. airflow/providers/google/cloud/hooks/gen_ai.py +196 -0
  34. airflow/providers/google/cloud/hooks/kubernetes_engine.py +3 -3
  35. airflow/providers/google/cloud/hooks/looker.py +6 -2
  36. airflow/providers/google/cloud/hooks/managed_kafka.py +1 -1
  37. airflow/providers/google/cloud/hooks/mlengine.py +4 -3
  38. airflow/providers/google/cloud/hooks/pubsub.py +3 -0
  39. airflow/providers/google/cloud/hooks/secret_manager.py +102 -10
  40. airflow/providers/google/cloud/hooks/spanner.py +74 -9
  41. airflow/providers/google/cloud/hooks/stackdriver.py +11 -9
  42. airflow/providers/google/cloud/hooks/tasks.py +1 -1
  43. airflow/providers/google/cloud/hooks/translate.py +2 -2
  44. airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +2 -210
  45. airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +3 -3
  46. airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +28 -2
  47. airflow/providers/google/cloud/hooks/vertex_ai/experiment_service.py +202 -0
  48. airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +308 -8
  49. airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +79 -75
  50. airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +1 -1
  51. airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +1 -1
  52. airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +1 -1
  53. airflow/providers/google/cloud/hooks/vertex_ai/ray.py +223 -0
  54. airflow/providers/google/cloud/hooks/vision.py +3 -3
  55. airflow/providers/google/cloud/hooks/workflows.py +1 -1
  56. airflow/providers/google/cloud/links/alloy_db.py +0 -46
  57. airflow/providers/google/cloud/links/base.py +77 -13
  58. airflow/providers/google/cloud/links/bigquery.py +0 -47
  59. airflow/providers/google/cloud/links/bigquery_dts.py +0 -20
  60. airflow/providers/google/cloud/links/bigtable.py +0 -48
  61. airflow/providers/google/cloud/links/cloud_build.py +0 -73
  62. airflow/providers/google/cloud/links/cloud_functions.py +0 -33
  63. airflow/providers/google/cloud/links/cloud_memorystore.py +0 -58
  64. airflow/providers/google/cloud/links/{life_sciences.py → cloud_run.py} +5 -27
  65. airflow/providers/google/cloud/links/cloud_sql.py +0 -33
  66. airflow/providers/google/cloud/links/cloud_storage_transfer.py +17 -44
  67. airflow/providers/google/cloud/links/cloud_tasks.py +7 -26
  68. airflow/providers/google/cloud/links/compute.py +0 -58
  69. airflow/providers/google/cloud/links/data_loss_prevention.py +0 -169
  70. airflow/providers/google/cloud/links/datacatalog.py +23 -54
  71. airflow/providers/google/cloud/links/dataflow.py +0 -34
  72. airflow/providers/google/cloud/links/dataform.py +0 -64
  73. airflow/providers/google/cloud/links/datafusion.py +1 -96
  74. airflow/providers/google/cloud/links/dataplex.py +0 -154
  75. airflow/providers/google/cloud/links/dataprep.py +0 -24
  76. airflow/providers/google/cloud/links/dataproc.py +11 -95
  77. airflow/providers/google/cloud/links/datastore.py +0 -31
  78. airflow/providers/google/cloud/links/kubernetes_engine.py +9 -60
  79. airflow/providers/google/cloud/links/managed_kafka.py +0 -70
  80. airflow/providers/google/cloud/links/mlengine.py +0 -70
  81. airflow/providers/google/cloud/links/pubsub.py +0 -32
  82. airflow/providers/google/cloud/links/spanner.py +0 -33
  83. airflow/providers/google/cloud/links/stackdriver.py +0 -30
  84. airflow/providers/google/cloud/links/translate.py +17 -187
  85. airflow/providers/google/cloud/links/vertex_ai.py +28 -195
  86. airflow/providers/google/cloud/links/workflows.py +0 -52
  87. airflow/providers/google/cloud/log/gcs_task_handler.py +58 -22
  88. airflow/providers/google/cloud/log/stackdriver_task_handler.py +9 -6
  89. airflow/providers/google/cloud/openlineage/CloudStorageTransferJobFacet.json +68 -0
  90. airflow/providers/google/cloud/openlineage/CloudStorageTransferRunFacet.json +60 -0
  91. airflow/providers/google/cloud/openlineage/DataFusionRunFacet.json +32 -0
  92. airflow/providers/google/cloud/openlineage/facets.py +102 -1
  93. airflow/providers/google/cloud/openlineage/mixins.py +10 -8
  94. airflow/providers/google/cloud/openlineage/utils.py +15 -1
  95. airflow/providers/google/cloud/operators/alloy_db.py +71 -56
  96. airflow/providers/google/cloud/operators/bigquery.py +73 -636
  97. airflow/providers/google/cloud/operators/bigquery_dts.py +4 -6
  98. airflow/providers/google/cloud/operators/bigtable.py +37 -8
  99. airflow/providers/google/cloud/operators/cloud_base.py +21 -1
  100. airflow/providers/google/cloud/operators/cloud_batch.py +3 -3
  101. airflow/providers/google/cloud/operators/cloud_build.py +76 -33
  102. airflow/providers/google/cloud/operators/cloud_composer.py +129 -41
  103. airflow/providers/google/cloud/operators/cloud_logging_sink.py +341 -0
  104. airflow/providers/google/cloud/operators/cloud_memorystore.py +69 -43
  105. airflow/providers/google/cloud/operators/cloud_run.py +24 -6
  106. airflow/providers/google/cloud/operators/cloud_sql.py +8 -17
  107. airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +93 -12
  108. airflow/providers/google/cloud/operators/compute.py +9 -41
  109. airflow/providers/google/cloud/operators/datacatalog.py +157 -21
  110. airflow/providers/google/cloud/operators/dataflow.py +40 -16
  111. airflow/providers/google/cloud/operators/dataform.py +15 -5
  112. airflow/providers/google/cloud/operators/datafusion.py +42 -21
  113. airflow/providers/google/cloud/operators/dataplex.py +194 -110
  114. airflow/providers/google/cloud/operators/dataprep.py +1 -5
  115. airflow/providers/google/cloud/operators/dataproc.py +80 -36
  116. airflow/providers/google/cloud/operators/dataproc_metastore.py +97 -89
  117. airflow/providers/google/cloud/operators/datastore.py +23 -7
  118. airflow/providers/google/cloud/operators/dlp.py +6 -29
  119. airflow/providers/google/cloud/operators/functions.py +17 -8
  120. airflow/providers/google/cloud/operators/gcs.py +12 -9
  121. airflow/providers/google/cloud/operators/gen_ai.py +389 -0
  122. airflow/providers/google/cloud/operators/kubernetes_engine.py +62 -100
  123. airflow/providers/google/cloud/operators/looker.py +2 -2
  124. airflow/providers/google/cloud/operators/managed_kafka.py +108 -53
  125. airflow/providers/google/cloud/operators/natural_language.py +1 -1
  126. airflow/providers/google/cloud/operators/pubsub.py +68 -15
  127. airflow/providers/google/cloud/operators/spanner.py +26 -13
  128. airflow/providers/google/cloud/operators/speech_to_text.py +2 -3
  129. airflow/providers/google/cloud/operators/stackdriver.py +1 -9
  130. airflow/providers/google/cloud/operators/tasks.py +1 -12
  131. airflow/providers/google/cloud/operators/text_to_speech.py +2 -3
  132. airflow/providers/google/cloud/operators/translate.py +41 -17
  133. airflow/providers/google/cloud/operators/translate_speech.py +2 -3
  134. airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +39 -19
  135. airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +30 -10
  136. airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +55 -27
  137. airflow/providers/google/cloud/operators/vertex_ai/dataset.py +70 -8
  138. airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +43 -9
  139. airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py +435 -0
  140. airflow/providers/google/cloud/operators/vertex_ai/feature_store.py +532 -1
  141. airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +135 -115
  142. airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +12 -10
  143. airflow/providers/google/cloud/operators/vertex_ai/model_service.py +57 -11
  144. airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +31 -8
  145. airflow/providers/google/cloud/operators/vertex_ai/ray.py +393 -0
  146. airflow/providers/google/cloud/operators/video_intelligence.py +1 -1
  147. airflow/providers/google/cloud/operators/vision.py +2 -2
  148. airflow/providers/google/cloud/operators/workflows.py +18 -15
  149. airflow/providers/google/cloud/secrets/secret_manager.py +3 -2
  150. airflow/providers/google/cloud/sensors/bigquery.py +3 -3
  151. airflow/providers/google/cloud/sensors/bigquery_dts.py +2 -3
  152. airflow/providers/google/cloud/sensors/bigtable.py +11 -4
  153. airflow/providers/google/cloud/sensors/cloud_composer.py +533 -30
  154. airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +2 -3
  155. airflow/providers/google/cloud/sensors/dataflow.py +26 -10
  156. airflow/providers/google/cloud/sensors/dataform.py +2 -3
  157. airflow/providers/google/cloud/sensors/datafusion.py +4 -5
  158. airflow/providers/google/cloud/sensors/dataplex.py +2 -3
  159. airflow/providers/google/cloud/sensors/dataprep.py +2 -2
  160. airflow/providers/google/cloud/sensors/dataproc.py +2 -3
  161. airflow/providers/google/cloud/sensors/dataproc_metastore.py +2 -3
  162. airflow/providers/google/cloud/sensors/gcs.py +4 -5
  163. airflow/providers/google/cloud/sensors/looker.py +2 -3
  164. airflow/providers/google/cloud/sensors/pubsub.py +4 -5
  165. airflow/providers/google/cloud/sensors/tasks.py +2 -2
  166. airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py +2 -3
  167. airflow/providers/google/cloud/sensors/workflows.py +2 -3
  168. airflow/providers/google/cloud/transfers/adls_to_gcs.py +1 -1
  169. airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py +2 -2
  170. airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py +4 -3
  171. airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py +11 -8
  172. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +10 -5
  173. airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +7 -3
  174. airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +12 -1
  175. airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +24 -10
  176. airflow/providers/google/cloud/transfers/bigquery_to_sql.py +104 -5
  177. airflow/providers/google/cloud/transfers/calendar_to_gcs.py +1 -1
  178. airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +3 -3
  179. airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +4 -4
  180. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +21 -13
  181. airflow/providers/google/cloud/transfers/gcs_to_gcs.py +4 -3
  182. airflow/providers/google/cloud/transfers/gcs_to_local.py +6 -4
  183. airflow/providers/google/cloud/transfers/gcs_to_sftp.py +11 -5
  184. airflow/providers/google/cloud/transfers/gdrive_to_gcs.py +6 -2
  185. airflow/providers/google/cloud/transfers/gdrive_to_local.py +2 -2
  186. airflow/providers/google/cloud/transfers/http_to_gcs.py +193 -0
  187. airflow/providers/google/cloud/transfers/local_to_gcs.py +2 -2
  188. airflow/providers/google/cloud/transfers/mssql_to_gcs.py +1 -1
  189. airflow/providers/google/cloud/transfers/oracle_to_gcs.py +36 -11
  190. airflow/providers/google/cloud/transfers/postgres_to_gcs.py +42 -9
  191. airflow/providers/google/cloud/transfers/s3_to_gcs.py +13 -7
  192. airflow/providers/google/cloud/transfers/salesforce_to_gcs.py +2 -2
  193. airflow/providers/google/cloud/transfers/sftp_to_gcs.py +14 -5
  194. airflow/providers/google/cloud/transfers/sheets_to_gcs.py +3 -3
  195. airflow/providers/google/cloud/transfers/sql_to_gcs.py +10 -10
  196. airflow/providers/google/cloud/triggers/bigquery.py +76 -35
  197. airflow/providers/google/cloud/triggers/cloud_build.py +1 -1
  198. airflow/providers/google/cloud/triggers/cloud_composer.py +303 -47
  199. airflow/providers/google/cloud/triggers/cloud_run.py +3 -3
  200. airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +92 -2
  201. airflow/providers/google/cloud/triggers/dataflow.py +122 -0
  202. airflow/providers/google/cloud/triggers/datafusion.py +1 -1
  203. airflow/providers/google/cloud/triggers/dataplex.py +14 -2
  204. airflow/providers/google/cloud/triggers/dataproc.py +123 -53
  205. airflow/providers/google/cloud/triggers/kubernetes_engine.py +47 -28
  206. airflow/providers/google/cloud/triggers/mlengine.py +1 -1
  207. airflow/providers/google/cloud/triggers/pubsub.py +15 -19
  208. airflow/providers/google/cloud/triggers/vertex_ai.py +1 -1
  209. airflow/providers/google/cloud/utils/bigquery_get_data.py +1 -1
  210. airflow/providers/google/cloud/utils/credentials_provider.py +2 -2
  211. airflow/providers/google/cloud/utils/field_sanitizer.py +1 -1
  212. airflow/providers/google/cloud/utils/field_validator.py +2 -3
  213. airflow/providers/google/common/auth_backend/google_openid.py +4 -4
  214. airflow/providers/google/common/deprecated.py +2 -1
  215. airflow/providers/google/common/hooks/base_google.py +27 -9
  216. airflow/providers/google/common/hooks/operation_helpers.py +1 -1
  217. airflow/providers/google/common/links/storage.py +0 -22
  218. airflow/providers/google/common/utils/get_secret.py +31 -0
  219. airflow/providers/google/common/utils/id_token_credentials.py +3 -4
  220. airflow/providers/google/firebase/hooks/firestore.py +1 -1
  221. airflow/providers/google/firebase/operators/firestore.py +3 -3
  222. airflow/providers/google/get_provider_info.py +56 -52
  223. airflow/providers/google/go_module_utils.py +35 -3
  224. airflow/providers/google/leveldb/hooks/leveldb.py +27 -2
  225. airflow/providers/google/leveldb/operators/leveldb.py +2 -2
  226. airflow/providers/google/marketing_platform/hooks/campaign_manager.py +1 -1
  227. airflow/providers/google/marketing_platform/hooks/display_video.py +3 -109
  228. airflow/providers/google/marketing_platform/hooks/search_ads.py +1 -1
  229. airflow/providers/google/marketing_platform/links/analytics_admin.py +5 -14
  230. airflow/providers/google/marketing_platform/operators/analytics_admin.py +2 -3
  231. airflow/providers/google/marketing_platform/operators/campaign_manager.py +6 -6
  232. airflow/providers/google/marketing_platform/operators/display_video.py +28 -489
  233. airflow/providers/google/marketing_platform/operators/search_ads.py +2 -2
  234. airflow/providers/google/marketing_platform/sensors/campaign_manager.py +2 -2
  235. airflow/providers/google/marketing_platform/sensors/display_video.py +3 -64
  236. airflow/providers/google/suite/hooks/calendar.py +2 -2
  237. airflow/providers/google/suite/hooks/sheets.py +16 -2
  238. airflow/providers/google/suite/operators/sheets.py +8 -3
  239. airflow/providers/google/suite/sensors/drive.py +2 -2
  240. airflow/providers/google/suite/transfers/gcs_to_gdrive.py +3 -3
  241. airflow/providers/google/suite/transfers/gcs_to_sheets.py +1 -1
  242. airflow/providers/google/suite/transfers/local_to_drive.py +3 -3
  243. airflow/providers/google/suite/transfers/sql_to_sheets.py +5 -4
  244. airflow/providers/google/version_compat.py +15 -1
  245. {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/METADATA +90 -46
  246. apache_airflow_providers_google-19.3.0.dist-info/RECORD +331 -0
  247. apache_airflow_providers_google-19.3.0.dist-info/licenses/NOTICE +5 -0
  248. airflow/providers/google/cloud/hooks/automl.py +0 -673
  249. airflow/providers/google/cloud/hooks/life_sciences.py +0 -159
  250. airflow/providers/google/cloud/links/automl.py +0 -193
  251. airflow/providers/google/cloud/operators/automl.py +0 -1362
  252. airflow/providers/google/cloud/operators/life_sciences.py +0 -119
  253. airflow/providers/google/cloud/operators/mlengine.py +0 -112
  254. apache_airflow_providers_google-15.1.0rc1.dist-info/RECORD +0 -321
  255. {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/WHEEL +0 -0
  256. {apache_airflow_providers_google-15.1.0rc1.dist-info → apache_airflow_providers_google-19.3.0.dist-info}/entry_points.txt +0 -0
  257. {airflow/providers/google → apache_airflow_providers_google-19.3.0.dist-info/licenses}/LICENSE +0 -0
@@ -1,1362 +0,0 @@
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 Google AutoML operators."""
19
-
20
- from __future__ import annotations
21
-
22
- import ast
23
- from collections.abc import Sequence
24
- from functools import cached_property
25
- from typing import TYPE_CHECKING, cast
26
-
27
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.cloud.automl_v1beta1 import (
29
- ColumnSpec,
30
- Dataset,
31
- Model,
32
- PredictResponse,
33
- TableSpec,
34
- )
35
-
36
- from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
37
- from airflow.providers.google.cloud.hooks.automl import CloudAutoMLHook
38
- from airflow.providers.google.cloud.hooks.vertex_ai.prediction_service import PredictionServiceHook
39
- from airflow.providers.google.cloud.links.translate import (
40
- TranslationDatasetListLink,
41
- TranslationLegacyDatasetLink,
42
- TranslationLegacyModelLink,
43
- TranslationLegacyModelPredictLink,
44
- TranslationLegacyModelTrainLink,
45
- )
46
- from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
47
- from airflow.providers.google.common.deprecated import deprecated
48
- from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
49
-
50
- if TYPE_CHECKING:
51
- from google.api_core.retry import Retry
52
-
53
- from airflow.utils.context import Context
54
-
55
- MetaData = Sequence[tuple[str, str]]
56
-
57
-
58
- @deprecated(
59
- planned_removal_date="September 30, 2025",
60
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTabularTrainingJobOperator, "
61
- "airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLVideoTrainingJobOperator, "
62
- "airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLImageTrainingJobOperator, "
63
- "airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator, "
64
- "airflow.providers.google.cloud.operators.translate.TranslateCreateModelOperator",
65
- category=AirflowProviderDeprecationWarning,
66
- )
67
- class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
68
- """
69
- Creates Google Cloud AutoML model.
70
-
71
- .. warning::
72
- AutoMLTrainModelOperator for tables, video intelligence, vision and natural language has been deprecated
73
- and no longer available. Please use
74
- :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLTabularTrainingJobOperator`,
75
- :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLVideoTrainingJobOperator`,
76
- :class:`airflow.providers.google.cloud.operators.vertex_ai.auto_ml.CreateAutoMLImageTrainingJobOperator`,
77
- :class:`airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator`,
78
- :class:`airflow.providers.google.cloud.operators.translate.TranslateCreateModelOperator`.
79
- instead.
80
-
81
- .. seealso::
82
- For more information on how to use this operator, take a look at the guide:
83
- :ref:`howto/operator:AutoMLTrainModelOperator`
84
-
85
- :param model: Model definition.
86
- :param project_id: ID of the Google Cloud project where model will be created if None then
87
- default project_id is used.
88
- :param location: The location of the project.
89
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
90
- retried.
91
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
92
- `retry` is specified, the timeout applies to each individual attempt.
93
- :param metadata: Additional metadata that is provided to the method.
94
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
95
- :param impersonation_chain: Optional service account to impersonate using short-term
96
- credentials, or chained list of accounts required to get the access_token
97
- of the last account in the list, which will be impersonated in the request.
98
- If set as a string, the account must grant the originating account
99
- the Service Account Token Creator IAM role.
100
- If set as a sequence, the identities from the list must grant
101
- Service Account Token Creator IAM role to the directly preceding identity, with first
102
- account from the list granting this role to the originating account (templated).
103
- """
104
-
105
- template_fields: Sequence[str] = (
106
- "model",
107
- "location",
108
- "project_id",
109
- "impersonation_chain",
110
- )
111
- operator_extra_links = (
112
- TranslationLegacyModelTrainLink(),
113
- TranslationLegacyModelLink(),
114
- )
115
-
116
- def __init__(
117
- self,
118
- *,
119
- model: dict,
120
- location: str,
121
- project_id: str = PROVIDE_PROJECT_ID,
122
- metadata: MetaData = (),
123
- timeout: float | None = None,
124
- retry: Retry | _MethodDefault = DEFAULT,
125
- gcp_conn_id: str = "google_cloud_default",
126
- impersonation_chain: str | Sequence[str] | None = None,
127
- **kwargs,
128
- ) -> None:
129
- super().__init__(**kwargs)
130
- self.model = model
131
- self.location = location
132
- self.project_id = project_id
133
- self.metadata = metadata
134
- self.timeout = timeout
135
- self.retry = retry
136
- self.gcp_conn_id = gcp_conn_id
137
- self.impersonation_chain = impersonation_chain
138
-
139
- def execute(self, context: Context):
140
- hook = CloudAutoMLHook(
141
- gcp_conn_id=self.gcp_conn_id,
142
- impersonation_chain=self.impersonation_chain,
143
- )
144
- self.log.info("Creating model %s...", self.model["display_name"])
145
- operation = hook.create_model(
146
- model=self.model,
147
- location=self.location,
148
- project_id=self.project_id,
149
- retry=self.retry,
150
- timeout=self.timeout,
151
- metadata=self.metadata,
152
- )
153
- project_id = self.project_id or hook.project_id
154
- if project_id:
155
- TranslationLegacyModelTrainLink.persist(
156
- context=context, task_instance=self, project_id=project_id
157
- )
158
- operation_result = hook.wait_for_operation(timeout=self.timeout, operation=operation)
159
- result = Model.to_dict(operation_result)
160
- model_id = hook.extract_object_id(result)
161
- self.log.info("Model is created, model_id: %s", model_id)
162
-
163
- self.xcom_push(context, key="model_id", value=model_id)
164
- if project_id:
165
- TranslationLegacyModelLink.persist(
166
- context=context,
167
- task_instance=self,
168
- dataset_id=self.model["dataset_id"] or "-",
169
- model_id=model_id,
170
- project_id=project_id,
171
- )
172
- return result
173
-
174
-
175
- @deprecated(
176
- planned_removal_date="September 30, 2025",
177
- use_instead="airflow.providers.google.cloud.operators.translate.TranslateTextOperator",
178
- category=AirflowProviderDeprecationWarning,
179
- )
180
- class AutoMLPredictOperator(GoogleCloudBaseOperator):
181
- """
182
- Runs prediction operation on Google Cloud AutoML.
183
-
184
- .. warning::
185
- AutoMLPredictOperator for text, image, and video prediction has been deprecated.
186
- Please use endpoint_id param instead of model_id param.
187
-
188
- .. seealso::
189
- For more information on how to use this operator, take a look at the guide:
190
- :ref:`howto/operator:AutoMLPredictOperator`
191
-
192
- :param model_id: Name of the model requested to serve the batch prediction.
193
- :param endpoint_id: Name of the endpoint used for the prediction.
194
- :param payload: Name of the model used for the prediction.
195
- :param project_id: ID of the Google Cloud project where model is located if None then
196
- default project_id is used.
197
- :param location: The location of the project.
198
- :param operation_params: Additional domain-specific parameters for the predictions.
199
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
200
- retried.
201
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
202
- `retry` is specified, the timeout applies to each individual attempt.
203
- :param metadata: Additional metadata that is provided to the method.
204
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
205
- :param impersonation_chain: Optional service account to impersonate using short-term
206
- credentials, or chained list of accounts required to get the access_token
207
- of the last account in the list, which will be impersonated in the request.
208
- If set as a string, the account must grant the originating account
209
- the Service Account Token Creator IAM role.
210
- If set as a sequence, the identities from the list must grant
211
- Service Account Token Creator IAM role to the directly preceding identity, with first
212
- account from the list granting this role to the originating account (templated).
213
- """
214
-
215
- template_fields: Sequence[str] = (
216
- "model_id",
217
- "location",
218
- "project_id",
219
- "impersonation_chain",
220
- )
221
- operator_extra_links = (TranslationLegacyModelPredictLink(),)
222
-
223
- def __init__(
224
- self,
225
- *,
226
- model_id: str | None = None,
227
- endpoint_id: str | None = None,
228
- location: str,
229
- payload: dict,
230
- operation_params: dict[str, str] | None = None,
231
- instances: list[str] | None = None,
232
- project_id: str = PROVIDE_PROJECT_ID,
233
- metadata: MetaData = (),
234
- timeout: float | None = None,
235
- retry: Retry | _MethodDefault = DEFAULT,
236
- gcp_conn_id: str = "google_cloud_default",
237
- impersonation_chain: str | Sequence[str] | None = None,
238
- **kwargs,
239
- ) -> None:
240
- super().__init__(**kwargs)
241
-
242
- self.model_id = model_id
243
- self.endpoint_id = endpoint_id
244
- self.operation_params = operation_params # type: ignore
245
- self.instances = instances
246
- self.location = location
247
- self.project_id = project_id
248
- self.metadata = metadata
249
- self.timeout = timeout
250
- self.retry = retry
251
- self.payload = payload
252
- self.gcp_conn_id = gcp_conn_id
253
- self.impersonation_chain = impersonation_chain
254
-
255
- @cached_property
256
- def hook(self) -> CloudAutoMLHook | PredictionServiceHook:
257
- if self.model_id:
258
- return CloudAutoMLHook(
259
- gcp_conn_id=self.gcp_conn_id,
260
- impersonation_chain=self.impersonation_chain,
261
- )
262
- # endpoint_id defined
263
- return PredictionServiceHook(
264
- gcp_conn_id=self.gcp_conn_id,
265
- impersonation_chain=self.impersonation_chain,
266
- )
267
-
268
- @cached_property
269
- def model(self) -> Model | None:
270
- if self.model_id:
271
- hook = cast("CloudAutoMLHook", self.hook)
272
- return hook.get_model(
273
- model_id=self.model_id,
274
- location=self.location,
275
- project_id=self.project_id,
276
- retry=self.retry,
277
- timeout=self.timeout,
278
- metadata=self.metadata,
279
- )
280
- return None
281
-
282
- def execute(self, context: Context):
283
- if self.model_id is None and self.endpoint_id is None:
284
- raise AirflowException("You must specify model_id or endpoint_id!")
285
-
286
- hook = self.hook
287
- if self.model_id:
288
- result = hook.predict(
289
- model_id=self.model_id,
290
- payload=self.payload,
291
- location=self.location,
292
- project_id=self.project_id,
293
- params=self.operation_params,
294
- retry=self.retry,
295
- timeout=self.timeout,
296
- metadata=self.metadata,
297
- )
298
- else: # self.endpoint_id is defined
299
- result = hook.predict(
300
- endpoint_id=self.endpoint_id,
301
- instances=self.instances,
302
- payload=self.payload,
303
- location=self.location,
304
- project_id=self.project_id,
305
- parameters=self.operation_params,
306
- retry=self.retry,
307
- timeout=self.timeout,
308
- metadata=self.metadata,
309
- )
310
-
311
- project_id = self.project_id or hook.project_id
312
- dataset_id: str | None = self.model.dataset_id if self.model else None
313
- if project_id and self.model_id and dataset_id:
314
- TranslationLegacyModelPredictLink.persist(
315
- context=context,
316
- task_instance=self,
317
- model_id=self.model_id,
318
- dataset_id=dataset_id,
319
- project_id=project_id,
320
- )
321
- return PredictResponse.to_dict(result)
322
-
323
-
324
- @deprecated(
325
- planned_removal_date="September 30, 2025",
326
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.CreateDatasetOperator, "
327
- "airflow.providers.google.cloud.operators.translate.TranslateCreateDatasetOperator",
328
- category=AirflowProviderDeprecationWarning,
329
- )
330
- class AutoMLCreateDatasetOperator(GoogleCloudBaseOperator):
331
- """
332
- Creates a Google Cloud AutoML dataset.
333
-
334
- AutoMLCreateDatasetOperator for tables, video intelligence, vision and natural language has been
335
- deprecated and no longer available. Please use
336
- :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.CreateDatasetOperator`,
337
- :class:`airflow.providers.google.cloud.operators.translate.TranslateCreateDatasetOperator` instead.
338
-
339
- .. seealso::
340
- For more information on how to use this operator, take a look at the guide:
341
- :ref:`howto/operator:AutoMLCreateDatasetOperator`
342
-
343
- :param dataset: The dataset to create. If a dict is provided, it must be of the
344
- same form as the protobuf message Dataset.
345
- :param project_id: ID of the Google Cloud project where dataset is located if None then
346
- default project_id is used.
347
- :param location: The location of the project.
348
- :param params: Additional domain-specific parameters for the predictions.
349
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
350
- retried.
351
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
352
- `retry` is specified, the timeout applies to each individual attempt.
353
- :param metadata: Additional metadata that is provided to the method.
354
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
355
- :param impersonation_chain: Optional service account to impersonate using short-term
356
- credentials, or chained list of accounts required to get the access_token
357
- of the last account in the list, which will be impersonated in the request.
358
- If set as a string, the account must grant the originating account
359
- the Service Account Token Creator IAM role.
360
- If set as a sequence, the identities from the list must grant
361
- Service Account Token Creator IAM role to the directly preceding identity, with first
362
- account from the list granting this role to the originating account (templated).
363
- """
364
-
365
- template_fields: Sequence[str] = (
366
- "dataset",
367
- "location",
368
- "project_id",
369
- "impersonation_chain",
370
- )
371
- operator_extra_links = (TranslationLegacyDatasetLink(),)
372
-
373
- def __init__(
374
- self,
375
- *,
376
- dataset: dict,
377
- location: str,
378
- project_id: str = PROVIDE_PROJECT_ID,
379
- metadata: MetaData = (),
380
- timeout: float | None = None,
381
- retry: Retry | _MethodDefault = DEFAULT,
382
- gcp_conn_id: str = "google_cloud_default",
383
- impersonation_chain: str | Sequence[str] | None = None,
384
- **kwargs,
385
- ) -> None:
386
- super().__init__(**kwargs)
387
-
388
- self.dataset = dataset
389
- self.location = location
390
- self.project_id = project_id
391
- self.metadata = metadata
392
- self.timeout = timeout
393
- self.retry = retry
394
- self.gcp_conn_id = gcp_conn_id
395
- self.impersonation_chain = impersonation_chain
396
-
397
- def execute(self, context: Context):
398
- hook = CloudAutoMLHook(
399
- gcp_conn_id=self.gcp_conn_id,
400
- impersonation_chain=self.impersonation_chain,
401
- )
402
- self.log.info("Creating dataset %s...", self.dataset)
403
- result = hook.create_dataset(
404
- dataset=self.dataset,
405
- location=self.location,
406
- project_id=self.project_id,
407
- retry=self.retry,
408
- timeout=self.timeout,
409
- metadata=self.metadata,
410
- )
411
- result = Dataset.to_dict(result)
412
- dataset_id = hook.extract_object_id(result)
413
- self.log.info("Creating completed. Dataset id: %s", dataset_id)
414
-
415
- self.xcom_push(context, key="dataset_id", value=dataset_id)
416
- project_id = self.project_id or hook.project_id
417
- if project_id:
418
- TranslationLegacyDatasetLink.persist(
419
- context=context,
420
- task_instance=self,
421
- dataset_id=dataset_id,
422
- project_id=project_id,
423
- )
424
- return result
425
-
426
-
427
- @deprecated(
428
- planned_removal_date="September 30, 2025",
429
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.ImportDataOperator, "
430
- "airflow.providers.google.cloud.operators.translate.TranslateImportDataOperator",
431
- category=AirflowProviderDeprecationWarning,
432
- )
433
- class AutoMLImportDataOperator(GoogleCloudBaseOperator):
434
- """
435
- Imports data to a Google Cloud AutoML dataset.
436
-
437
- .. warning::
438
- AutoMLImportDataOperator for tables, video intelligence, vision and natural language has been deprecated
439
- and no longer available. Please use
440
- :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.ImportDataOperator` instead.
441
-
442
- .. seealso::
443
- For more information on how to use this operator, take a look at the guide:
444
- :ref:`howto/operator:AutoMLImportDataOperator`
445
-
446
- :param dataset_id: ID of dataset to be updated.
447
- :param input_config: The desired input location and its domain specific semantics, if any.
448
- If a dict is provided, it must be of the same form as the protobuf message InputConfig.
449
- :param project_id: ID of the Google Cloud project where dataset is located if None then
450
- default project_id is used.
451
- :param location: The location of the project.
452
- :param params: Additional domain-specific parameters for the predictions.
453
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
454
- retried.
455
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
456
- `retry` is specified, the timeout applies to each individual attempt.
457
- :param metadata: Additional metadata that is provided to the method.
458
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
459
- :param impersonation_chain: Optional service account to impersonate using short-term
460
- credentials, or chained list of accounts required to get the access_token
461
- of the last account in the list, which will be impersonated in the request.
462
- If set as a string, the account must grant the originating account
463
- the Service Account Token Creator IAM role.
464
- If set as a sequence, the identities from the list must grant
465
- Service Account Token Creator IAM role to the directly preceding identity, with first
466
- account from the list granting this role to the originating account (templated).
467
- """
468
-
469
- template_fields: Sequence[str] = (
470
- "dataset_id",
471
- "input_config",
472
- "location",
473
- "project_id",
474
- "impersonation_chain",
475
- )
476
- operator_extra_links = (TranslationLegacyDatasetLink(),)
477
-
478
- def __init__(
479
- self,
480
- *,
481
- dataset_id: str,
482
- location: str,
483
- input_config: dict,
484
- project_id: str = PROVIDE_PROJECT_ID,
485
- metadata: MetaData = (),
486
- timeout: float | None = None,
487
- retry: Retry | _MethodDefault = DEFAULT,
488
- gcp_conn_id: str = "google_cloud_default",
489
- impersonation_chain: str | Sequence[str] | None = None,
490
- **kwargs,
491
- ) -> None:
492
- super().__init__(**kwargs)
493
-
494
- self.dataset_id = dataset_id
495
- self.input_config = input_config
496
- self.location = location
497
- self.project_id = project_id
498
- self.metadata = metadata
499
- self.timeout = timeout
500
- self.retry = retry
501
- self.gcp_conn_id = gcp_conn_id
502
- self.impersonation_chain = impersonation_chain
503
-
504
- def execute(self, context: Context):
505
- hook = CloudAutoMLHook(
506
- gcp_conn_id=self.gcp_conn_id,
507
- impersonation_chain=self.impersonation_chain,
508
- )
509
- hook.get_dataset(
510
- dataset_id=self.dataset_id,
511
- location=self.location,
512
- project_id=self.project_id,
513
- retry=self.retry,
514
- timeout=self.timeout,
515
- metadata=self.metadata,
516
- )
517
- self.log.info("Importing data to dataset...")
518
- operation = hook.import_data(
519
- dataset_id=self.dataset_id,
520
- input_config=self.input_config,
521
- location=self.location,
522
- project_id=self.project_id,
523
- retry=self.retry,
524
- timeout=self.timeout,
525
- metadata=self.metadata,
526
- )
527
- hook.wait_for_operation(timeout=self.timeout, operation=operation)
528
- self.log.info("Import is completed")
529
- project_id = self.project_id or hook.project_id
530
- if project_id:
531
- TranslationLegacyDatasetLink.persist(
532
- context=context,
533
- task_instance=self,
534
- dataset_id=self.dataset_id,
535
- project_id=project_id,
536
- )
537
-
538
-
539
- @deprecated(
540
- planned_removal_date="September 30, 2025",
541
- category=AirflowProviderDeprecationWarning,
542
- reason="Shutdown of legacy version of AutoML Tables on March 31, 2024.",
543
- )
544
- class AutoMLTablesListColumnSpecsOperator(GoogleCloudBaseOperator):
545
- """
546
- Lists column specs in a table.
547
-
548
- .. warning::
549
- Operator AutoMLTablesListColumnSpecsOperator has been deprecated due to shutdown of
550
- a legacy version of AutoML Tables on March 31, 2024. For additional information
551
- see: https://cloud.google.com/automl-tables/docs/deprecations.
552
-
553
- .. seealso::
554
- For more information on how to use this operator, take a look at the guide:
555
- :ref:`howto/operator:AutoMLTablesListColumnSpecsOperator`
556
-
557
- :param dataset_id: Name of the dataset.
558
- :param table_spec_id: table_spec_id for path builder.
559
- :param field_mask: Mask specifying which fields to read. If a dict is provided, it must be of the same
560
- form as the protobuf message `google.cloud.automl_v1beta1.types.FieldMask`
561
- :param filter_: Filter expression, see go/filtering.
562
- :param page_size: The maximum number of resources contained in the
563
- underlying API response. If page streaming is performed per
564
- resource, this parameter does not affect the return value. If page
565
- streaming is performed per page, this determines the maximum number
566
- of resources in a page.
567
- :param project_id: ID of the Google Cloud project where dataset is located if None then
568
- default project_id is used.
569
- :param location: The location of the project.
570
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
571
- retried.
572
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
573
- `retry` is specified, the timeout applies to each individual attempt.
574
- :param metadata: Additional metadata that is provided to the method.
575
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
576
- :param impersonation_chain: Optional service account to impersonate using short-term
577
- credentials, or chained list of accounts required to get the access_token
578
- of the last account in the list, which will be impersonated in the request.
579
- If set as a string, the account must grant the originating account
580
- the Service Account Token Creator IAM role.
581
- If set as a sequence, the identities from the list must grant
582
- Service Account Token Creator IAM role to the directly preceding identity, with first
583
- account from the list granting this role to the originating account (templated).
584
- """
585
-
586
- template_fields: Sequence[str] = (
587
- "dataset_id",
588
- "table_spec_id",
589
- "field_mask",
590
- "filter_",
591
- "location",
592
- "project_id",
593
- "impersonation_chain",
594
- )
595
- operator_extra_links = (TranslationLegacyDatasetLink(),)
596
-
597
- def __init__(
598
- self,
599
- *,
600
- dataset_id: str,
601
- table_spec_id: str,
602
- location: str,
603
- field_mask: dict | None = None,
604
- filter_: str | None = None,
605
- page_size: int | None = None,
606
- project_id: str = PROVIDE_PROJECT_ID,
607
- metadata: MetaData = (),
608
- timeout: float | None = None,
609
- retry: Retry | _MethodDefault = DEFAULT,
610
- gcp_conn_id: str = "google_cloud_default",
611
- impersonation_chain: str | Sequence[str] | None = None,
612
- **kwargs,
613
- ) -> None:
614
- super().__init__(**kwargs)
615
- self.dataset_id = dataset_id
616
- self.table_spec_id = table_spec_id
617
- self.field_mask = field_mask
618
- self.filter_ = filter_
619
- self.page_size = page_size
620
- self.location = location
621
- self.project_id = project_id
622
- self.metadata = metadata
623
- self.timeout = timeout
624
- self.retry = retry
625
- self.gcp_conn_id = gcp_conn_id
626
- self.impersonation_chain = impersonation_chain
627
-
628
- def execute(self, context: Context):
629
- hook = CloudAutoMLHook(
630
- gcp_conn_id=self.gcp_conn_id,
631
- impersonation_chain=self.impersonation_chain,
632
- )
633
- self.log.info("Requesting column specs.")
634
- page_iterator = hook.list_column_specs(
635
- dataset_id=self.dataset_id,
636
- table_spec_id=self.table_spec_id,
637
- field_mask=self.field_mask,
638
- filter_=self.filter_,
639
- page_size=self.page_size,
640
- location=self.location,
641
- project_id=self.project_id,
642
- retry=self.retry,
643
- timeout=self.timeout,
644
- metadata=self.metadata,
645
- )
646
- result = [ColumnSpec.to_dict(spec) for spec in page_iterator]
647
- self.log.info("Columns specs obtained.")
648
- project_id = self.project_id or hook.project_id
649
- if project_id:
650
- TranslationLegacyDatasetLink.persist(
651
- context=context,
652
- task_instance=self,
653
- dataset_id=self.dataset_id,
654
- project_id=project_id,
655
- )
656
- return result
657
-
658
-
659
- @deprecated(
660
- planned_removal_date="September 30, 2025",
661
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.UpdateDatasetOperator",
662
- category=AirflowProviderDeprecationWarning,
663
- reason="Shutdown of legacy version of AutoML Tables on March 31, 2024.",
664
- )
665
- class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
666
- """
667
- Updates a dataset.
668
-
669
- .. warning::
670
- Operator AutoMLTablesUpdateDatasetOperator has been deprecated due to shutdown of
671
- a legacy version of AutoML Tables on March 31, 2024. For additional information
672
- see: https://cloud.google.com/automl-tables/docs/deprecations.
673
- Please use :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.UpdateDatasetOperator`
674
- instead.
675
-
676
- .. seealso::
677
- For more information on how to use this operator, take a look at the guide:
678
- :ref:`howto/operator:AutoMLTablesUpdateDatasetOperator`
679
-
680
- :param dataset: The dataset which replaces the resource on the server.
681
- If a dict is provided, it must be of the same form as the protobuf message Dataset.
682
- :param update_mask: The update mask applies to the resource. If a dict is provided, it must
683
- be of the same form as the protobuf message FieldMask.
684
- :param location: The location of the project.
685
- :param params: Additional domain-specific parameters for the predictions.
686
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
687
- retried.
688
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
689
- `retry` is specified, the timeout applies to each individual attempt.
690
- :param metadata: Additional metadata that is provided to the method.
691
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
692
- :param impersonation_chain: Optional service account to impersonate using short-term
693
- credentials, or chained list of accounts required to get the access_token
694
- of the last account in the list, which will be impersonated in the request.
695
- If set as a string, the account must grant the originating account
696
- the Service Account Token Creator IAM role.
697
- If set as a sequence, the identities from the list must grant
698
- Service Account Token Creator IAM role to the directly preceding identity, with first
699
- account from the list granting this role to the originating account (templated).
700
- """
701
-
702
- template_fields: Sequence[str] = (
703
- "dataset",
704
- "update_mask",
705
- "location",
706
- "impersonation_chain",
707
- )
708
- operator_extra_links = (TranslationLegacyDatasetLink(),)
709
-
710
- def __init__(
711
- self,
712
- *,
713
- dataset: dict,
714
- location: str,
715
- update_mask: dict | None = None,
716
- metadata: MetaData = (),
717
- timeout: float | None = None,
718
- retry: Retry | _MethodDefault = DEFAULT,
719
- gcp_conn_id: str = "google_cloud_default",
720
- impersonation_chain: str | Sequence[str] | None = None,
721
- **kwargs,
722
- ) -> None:
723
- super().__init__(**kwargs)
724
-
725
- self.dataset = dataset
726
- self.update_mask = update_mask
727
- self.location = location
728
- self.metadata = metadata
729
- self.timeout = timeout
730
- self.retry = retry
731
- self.gcp_conn_id = gcp_conn_id
732
- self.impersonation_chain = impersonation_chain
733
-
734
- def execute(self, context: Context):
735
- hook = CloudAutoMLHook(
736
- gcp_conn_id=self.gcp_conn_id,
737
- impersonation_chain=self.impersonation_chain,
738
- )
739
- self.log.info("Updating AutoML dataset %s.", self.dataset["name"])
740
- result = hook.update_dataset(
741
- dataset=self.dataset,
742
- update_mask=self.update_mask,
743
- retry=self.retry,
744
- timeout=self.timeout,
745
- metadata=self.metadata,
746
- )
747
- self.log.info("Dataset updated.")
748
- project_id = hook.project_id
749
- if project_id:
750
- TranslationLegacyDatasetLink.persist(
751
- context=context,
752
- task_instance=self,
753
- dataset_id=hook.extract_object_id(self.dataset),
754
- project_id=project_id,
755
- )
756
- return Dataset.to_dict(result)
757
-
758
-
759
- @deprecated(
760
- planned_removal_date="September 30, 2025",
761
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator",
762
- category=AirflowProviderDeprecationWarning,
763
- )
764
- class AutoMLGetModelOperator(GoogleCloudBaseOperator):
765
- """
766
- Get Google Cloud AutoML model.
767
-
768
- .. warning::
769
- AutoMLGetModelOperator for tables, video intelligence, vision and natural language has been deprecated
770
- and no longer available. Please use
771
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator` instead.
772
-
773
- .. seealso::
774
- For more information on how to use this operator, take a look at the guide:
775
- :ref:`howto/operator:AutoMLGetModelOperator`
776
-
777
- :param model_id: Name of the model requested to serve the prediction.
778
- :param project_id: ID of the Google Cloud project where model is located if None then
779
- default project_id is used.
780
- :param location: The location of the project.
781
- :param params: Additional domain-specific parameters for the predictions.
782
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
783
- retried.
784
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
785
- `retry` is specified, the timeout applies to each individual attempt.
786
- :param metadata: Additional metadata that is provided to the method.
787
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
788
- :param impersonation_chain: Optional service account to impersonate using short-term
789
- credentials, or chained list of accounts required to get the access_token
790
- of the last account in the list, which will be impersonated in the request.
791
- If set as a string, the account must grant the originating account
792
- the Service Account Token Creator IAM role.
793
- If set as a sequence, the identities from the list must grant
794
- Service Account Token Creator IAM role to the directly preceding identity, with first
795
- account from the list granting this role to the originating account (templated).
796
- """
797
-
798
- template_fields: Sequence[str] = (
799
- "model_id",
800
- "location",
801
- "project_id",
802
- "impersonation_chain",
803
- )
804
- operator_extra_links = (TranslationLegacyModelLink(),)
805
-
806
- def __init__(
807
- self,
808
- *,
809
- model_id: str,
810
- location: str,
811
- project_id: str = PROVIDE_PROJECT_ID,
812
- metadata: MetaData = (),
813
- timeout: float | None = None,
814
- retry: Retry | _MethodDefault = DEFAULT,
815
- gcp_conn_id: str = "google_cloud_default",
816
- impersonation_chain: str | Sequence[str] | None = None,
817
- **kwargs,
818
- ) -> None:
819
- super().__init__(**kwargs)
820
-
821
- self.model_id = model_id
822
- self.location = location
823
- self.project_id = project_id
824
- self.metadata = metadata
825
- self.timeout = timeout
826
- self.retry = retry
827
- self.gcp_conn_id = gcp_conn_id
828
- self.impersonation_chain = impersonation_chain
829
-
830
- def execute(self, context: Context):
831
- hook = CloudAutoMLHook(
832
- gcp_conn_id=self.gcp_conn_id,
833
- impersonation_chain=self.impersonation_chain,
834
- )
835
- result = hook.get_model(
836
- model_id=self.model_id,
837
- location=self.location,
838
- project_id=self.project_id,
839
- retry=self.retry,
840
- timeout=self.timeout,
841
- metadata=self.metadata,
842
- )
843
- model = Model.to_dict(result)
844
- project_id = self.project_id or hook.project_id
845
- if project_id:
846
- TranslationLegacyModelLink.persist(
847
- context=context,
848
- task_instance=self,
849
- dataset_id=model["dataset_id"],
850
- model_id=self.model_id,
851
- project_id=project_id,
852
- )
853
- return model
854
-
855
-
856
- @deprecated(
857
- planned_removal_date="September 30, 2025",
858
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator, "
859
- "airflow.providers.google.cloud.operators.translate.TranslateDeleteModelOperator",
860
- category=AirflowProviderDeprecationWarning,
861
- )
862
- class AutoMLDeleteModelOperator(GoogleCloudBaseOperator):
863
- """
864
- Delete Google Cloud AutoML model.
865
-
866
- .. warning::
867
- AutoMLDeleteModelOperator for tables, video intelligence, vision and natural language has been deprecated
868
- and no longer available. Please use
869
- :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator` instead.
870
-
871
- .. seealso::
872
- For more information on how to use this operator, take a look at the guide:
873
- :ref:`howto/operator:AutoMLDeleteModelOperator`
874
-
875
- :param model_id: Name of the model requested to serve the prediction.
876
- :param project_id: ID of the Google Cloud project where model is located if None then
877
- default project_id is used.
878
- :param location: The location of the project.
879
- :param params: Additional domain-specific parameters for the predictions.
880
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
881
- retried.
882
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
883
- `retry` is specified, the timeout applies to each individual attempt.
884
- :param metadata: Additional metadata that is provided to the method.
885
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
886
- :param impersonation_chain: Optional service account to impersonate using short-term
887
- credentials, or chained list of accounts required to get the access_token
888
- of the last account in the list, which will be impersonated in the request.
889
- If set as a string, the account must grant the originating account
890
- the Service Account Token Creator IAM role.
891
- If set as a sequence, the identities from the list must grant
892
- Service Account Token Creator IAM role to the directly preceding identity, with first
893
- account from the list granting this role to the originating account (templated).
894
- """
895
-
896
- template_fields: Sequence[str] = (
897
- "model_id",
898
- "location",
899
- "project_id",
900
- "impersonation_chain",
901
- )
902
-
903
- def __init__(
904
- self,
905
- *,
906
- model_id: str,
907
- location: str,
908
- project_id: str = PROVIDE_PROJECT_ID,
909
- metadata: MetaData = (),
910
- timeout: float | None = None,
911
- retry: Retry | _MethodDefault = DEFAULT,
912
- gcp_conn_id: str = "google_cloud_default",
913
- impersonation_chain: str | Sequence[str] | None = None,
914
- **kwargs,
915
- ) -> None:
916
- super().__init__(**kwargs)
917
-
918
- self.model_id = model_id
919
- self.location = location
920
- self.project_id = project_id
921
- self.metadata = metadata
922
- self.timeout = timeout
923
- self.retry = retry
924
- self.gcp_conn_id = gcp_conn_id
925
- self.impersonation_chain = impersonation_chain
926
-
927
- def execute(self, context: Context):
928
- hook = CloudAutoMLHook(
929
- gcp_conn_id=self.gcp_conn_id,
930
- impersonation_chain=self.impersonation_chain,
931
- )
932
- hook.get_model(
933
- model_id=self.model_id,
934
- location=self.location,
935
- project_id=self.project_id,
936
- retry=self.retry,
937
- timeout=self.timeout,
938
- metadata=self.metadata,
939
- )
940
- operation = hook.delete_model(
941
- model_id=self.model_id,
942
- location=self.location,
943
- project_id=self.project_id,
944
- retry=self.retry,
945
- timeout=self.timeout,
946
- metadata=self.metadata,
947
- )
948
- hook.wait_for_operation(timeout=self.timeout, operation=operation)
949
- self.log.info("Deletion is completed")
950
-
951
-
952
- @deprecated(
953
- planned_removal_date="September 30, 2025",
954
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.endpoint_service.DeployModelOperator",
955
- category=AirflowProviderDeprecationWarning,
956
- )
957
- class AutoMLDeployModelOperator(GoogleCloudBaseOperator):
958
- """
959
- Deploys a model; if a model is already deployed, deploying it with the same parameters has no effect.
960
-
961
- Deploying with different parameters (as e.g. changing node_number) will
962
- reset the deployment state without pausing the model_id's availability.
963
-
964
- Only applicable for Text Classification, Image Object Detection and Tables; all other
965
- domains manage deployment automatically.
966
-
967
- .. warning::
968
- Operator AutoMLDeployModelOperator has been deprecated due to shutdown of a legacy version
969
- of AutoML Natural Language, Vision, Video Intelligence on March 31, 2024.
970
- For additional information see: https://cloud.google.com/vision/automl/docs/deprecations .
971
- Please use :class:`airflow.providers.google.cloud.operators.vertex_ai.endpoint_service.DeployModelOperator`
972
- instead.
973
-
974
- .. seealso::
975
- For more information on how to use this operator, take a look at the guide:
976
- :ref:`howto/operator:AutoMLDeployModelOperator`
977
-
978
- :param model_id: Name of the model to be deployed.
979
- :param image_detection_metadata: Model deployment metadata specific to Image Object Detection.
980
- If a dict is provided, it must be of the same form as the protobuf message
981
- ImageObjectDetectionModelDeploymentMetadata
982
- :param project_id: ID of the Google Cloud project where model is located if None then
983
- default project_id is used.
984
- :param location: The location of the project.
985
- :param params: Additional domain-specific parameters for the predictions.
986
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
987
- retried.
988
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
989
- `retry` is specified, the timeout applies to each individual attempt.
990
- :param metadata: Additional metadata that is provided to the method.
991
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
992
- :param impersonation_chain: Optional service account to impersonate using short-term
993
- credentials, or chained list of accounts required to get the access_token
994
- of the last account in the list, which will be impersonated in the request.
995
- If set as a string, the account must grant the originating account
996
- the Service Account Token Creator IAM role.
997
- If set as a sequence, the identities from the list must grant
998
- Service Account Token Creator IAM role to the directly preceding identity, with first
999
- account from the list granting this role to the originating account (templated).
1000
- """
1001
-
1002
- template_fields: Sequence[str] = (
1003
- "model_id",
1004
- "location",
1005
- "project_id",
1006
- "impersonation_chain",
1007
- )
1008
-
1009
- def __init__(
1010
- self,
1011
- *,
1012
- model_id: str,
1013
- location: str,
1014
- project_id: str = PROVIDE_PROJECT_ID,
1015
- image_detection_metadata: dict | None = None,
1016
- metadata: Sequence[tuple[str, str]] = (),
1017
- timeout: float | None = None,
1018
- retry: Retry | _MethodDefault = DEFAULT,
1019
- gcp_conn_id: str = "google_cloud_default",
1020
- impersonation_chain: str | Sequence[str] | None = None,
1021
- **kwargs,
1022
- ) -> None:
1023
- super().__init__(**kwargs)
1024
-
1025
- self.model_id = model_id
1026
- self.image_detection_metadata = image_detection_metadata
1027
- self.location = location
1028
- self.project_id = project_id
1029
- self.metadata = metadata
1030
- self.timeout = timeout
1031
- self.retry = retry
1032
- self.gcp_conn_id = gcp_conn_id
1033
- self.impersonation_chain = impersonation_chain
1034
-
1035
- def execute(self, context: Context):
1036
- hook = CloudAutoMLHook(
1037
- gcp_conn_id=self.gcp_conn_id,
1038
- impersonation_chain=self.impersonation_chain,
1039
- )
1040
- self.log.info("Deploying model_id %s", self.model_id)
1041
- operation = hook.deploy_model(
1042
- model_id=self.model_id,
1043
- location=self.location,
1044
- project_id=self.project_id,
1045
- image_detection_metadata=self.image_detection_metadata,
1046
- retry=self.retry,
1047
- timeout=self.timeout,
1048
- metadata=self.metadata,
1049
- )
1050
- hook.wait_for_operation(timeout=self.timeout, operation=operation)
1051
- self.log.info("Model was deployed successfully.")
1052
-
1053
-
1054
- @deprecated(
1055
- planned_removal_date="September 30, 2025",
1056
- category=AirflowProviderDeprecationWarning,
1057
- reason="Shutdown of legacy version of AutoML Tables on March 31, 2024.",
1058
- )
1059
- class AutoMLTablesListTableSpecsOperator(GoogleCloudBaseOperator):
1060
- """
1061
- Lists table specs in a dataset.
1062
-
1063
- .. warning::
1064
- Operator AutoMLTablesListTableSpecsOperator has been deprecated due to shutdown of
1065
- a legacy version of AutoML Tables on March 31, 2024. For additional information
1066
- see: https://cloud.google.com/automl-tables/docs/deprecations.
1067
-
1068
- .. seealso::
1069
- For more information on how to use this operator, take a look at the guide:
1070
- :ref:`howto/operator:AutoMLTablesListTableSpecsOperator`
1071
-
1072
- :param dataset_id: Name of the dataset.
1073
- :param filter_: Filter expression, see go/filtering.
1074
- :param page_size: The maximum number of resources contained in the
1075
- underlying API response. If page streaming is performed per
1076
- resource, this parameter does not affect the return value. If page
1077
- streaming is performed per-page, this determines the maximum number
1078
- of resources in a page.
1079
- :param project_id: ID of the Google Cloud project if None then
1080
- default project_id is used.
1081
- :param location: The location of the project.
1082
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
1083
- retried.
1084
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
1085
- `retry` is specified, the timeout applies to each individual attempt.
1086
- :param metadata: Additional metadata that is provided to the method.
1087
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
1088
- :param impersonation_chain: Optional service account to impersonate using short-term
1089
- credentials, or chained list of accounts required to get the access_token
1090
- of the last account in the list, which will be impersonated in the request.
1091
- If set as a string, the account must grant the originating account
1092
- the Service Account Token Creator IAM role.
1093
- If set as a sequence, the identities from the list must grant
1094
- Service Account Token Creator IAM role to the directly preceding identity, with first
1095
- account from the list granting this role to the originating account (templated).
1096
- """
1097
-
1098
- template_fields: Sequence[str] = (
1099
- "dataset_id",
1100
- "filter_",
1101
- "location",
1102
- "project_id",
1103
- "impersonation_chain",
1104
- )
1105
- operator_extra_links = (TranslationLegacyDatasetLink(),)
1106
-
1107
- def __init__(
1108
- self,
1109
- *,
1110
- dataset_id: str,
1111
- location: str,
1112
- page_size: int | None = None,
1113
- filter_: str | None = None,
1114
- project_id: str = PROVIDE_PROJECT_ID,
1115
- metadata: MetaData = (),
1116
- timeout: float | None = None,
1117
- retry: Retry | _MethodDefault = DEFAULT,
1118
- gcp_conn_id: str = "google_cloud_default",
1119
- impersonation_chain: str | Sequence[str] | None = None,
1120
- **kwargs,
1121
- ) -> None:
1122
- super().__init__(**kwargs)
1123
- self.dataset_id = dataset_id
1124
- self.filter_ = filter_
1125
- self.page_size = page_size
1126
- self.location = location
1127
- self.project_id = project_id
1128
- self.metadata = metadata
1129
- self.timeout = timeout
1130
- self.retry = retry
1131
- self.gcp_conn_id = gcp_conn_id
1132
- self.impersonation_chain = impersonation_chain
1133
-
1134
- def execute(self, context: Context):
1135
- hook = CloudAutoMLHook(
1136
- gcp_conn_id=self.gcp_conn_id,
1137
- impersonation_chain=self.impersonation_chain,
1138
- )
1139
- self.log.info("Requesting table specs for %s.", self.dataset_id)
1140
- page_iterator = hook.list_table_specs(
1141
- dataset_id=self.dataset_id,
1142
- filter_=self.filter_,
1143
- page_size=self.page_size,
1144
- location=self.location,
1145
- project_id=self.project_id,
1146
- retry=self.retry,
1147
- timeout=self.timeout,
1148
- metadata=self.metadata,
1149
- )
1150
- result = [TableSpec.to_dict(spec) for spec in page_iterator]
1151
- self.log.info(result)
1152
- self.log.info("Table specs obtained.")
1153
- project_id = self.project_id or hook.project_id
1154
- if project_id:
1155
- TranslationLegacyDatasetLink.persist(
1156
- context=context,
1157
- task_instance=self,
1158
- dataset_id=self.dataset_id,
1159
- project_id=project_id,
1160
- )
1161
- return result
1162
-
1163
-
1164
- @deprecated(
1165
- planned_removal_date="September 30, 2025",
1166
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.ListDatasetsOperator, "
1167
- "airflow.providers.google.cloud.operators.translate.TranslateDatasetsListOperator",
1168
- category=AirflowProviderDeprecationWarning,
1169
- )
1170
- class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
1171
- """
1172
- Lists AutoML Datasets in project.
1173
-
1174
- .. warning::
1175
- AutoMLListDatasetOperator for tables, video intelligence, vision and natural language has been deprecated
1176
- and no longer available. Please use
1177
- :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.ListDatasetsOperator` instead.
1178
-
1179
- .. seealso::
1180
- For more information on how to use this operator, take a look at the guide:
1181
- :ref:`howto/operator:AutoMLListDatasetOperator`
1182
-
1183
- :param project_id: ID of the Google Cloud project where datasets are located if None then
1184
- default project_id is used.
1185
- :param location: The location of the project.
1186
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
1187
- retried.
1188
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
1189
- `retry` is specified, the timeout applies to each individual attempt.
1190
- :param metadata: Additional metadata that is provided to the method.
1191
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
1192
- :param impersonation_chain: Optional service account to impersonate using short-term
1193
- credentials, or chained list of accounts required to get the access_token
1194
- of the last account in the list, which will be impersonated in the request.
1195
- If set as a string, the account must grant the originating account
1196
- the Service Account Token Creator IAM role.
1197
- If set as a sequence, the identities from the list must grant
1198
- Service Account Token Creator IAM role to the directly preceding identity, with first
1199
- account from the list granting this role to the originating account (templated).
1200
- """
1201
-
1202
- template_fields: Sequence[str] = (
1203
- "location",
1204
- "project_id",
1205
- "impersonation_chain",
1206
- )
1207
- operator_extra_links = (TranslationDatasetListLink(),)
1208
-
1209
- def __init__(
1210
- self,
1211
- *,
1212
- location: str,
1213
- project_id: str = PROVIDE_PROJECT_ID,
1214
- metadata: MetaData = (),
1215
- timeout: float | None = None,
1216
- retry: Retry | _MethodDefault = DEFAULT,
1217
- gcp_conn_id: str = "google_cloud_default",
1218
- impersonation_chain: str | Sequence[str] | None = None,
1219
- **kwargs,
1220
- ) -> None:
1221
- super().__init__(**kwargs)
1222
- self.location = location
1223
- self.project_id = project_id
1224
- self.metadata = metadata
1225
- self.timeout = timeout
1226
- self.retry = retry
1227
- self.gcp_conn_id = gcp_conn_id
1228
- self.impersonation_chain = impersonation_chain
1229
-
1230
- def execute(self, context: Context):
1231
- hook = CloudAutoMLHook(
1232
- gcp_conn_id=self.gcp_conn_id,
1233
- impersonation_chain=self.impersonation_chain,
1234
- )
1235
- self.log.info("Requesting datasets")
1236
- page_iterator = hook.list_datasets(
1237
- location=self.location,
1238
- project_id=self.project_id,
1239
- retry=self.retry,
1240
- timeout=self.timeout,
1241
- metadata=self.metadata,
1242
- )
1243
- result = []
1244
- for dataset in page_iterator:
1245
- result.append(Dataset.to_dict(dataset))
1246
- self.log.info("Datasets obtained.")
1247
-
1248
- self.xcom_push(
1249
- context,
1250
- key="dataset_id_list",
1251
- value=[hook.extract_object_id(d) for d in result],
1252
- )
1253
- project_id = self.project_id or hook.project_id
1254
- if project_id:
1255
- TranslationDatasetListLink.persist(context=context, task_instance=self, project_id=project_id)
1256
- return result
1257
-
1258
-
1259
- @deprecated(
1260
- planned_removal_date="September 30, 2025",
1261
- use_instead="airflow.providers.google.cloud.operators.vertex_ai.dataset.ListDatasetsOperator, "
1262
- "airflow.providers.google.cloud.operators.translate.TranslateDatasetsListOperator",
1263
- category=AirflowProviderDeprecationWarning,
1264
- )
1265
- class AutoMLDeleteDatasetOperator(GoogleCloudBaseOperator):
1266
- """
1267
- Deletes a dataset and all of its contents.
1268
-
1269
- AutoMLDeleteDatasetOperator for tables, video intelligence, vision and natural language has been
1270
- deprecated and no longer available. Please use
1271
- :class:`airflow.providers.google.cloud.operators.vertex_ai.dataset.DeleteDatasetOperator` instead.
1272
-
1273
- .. seealso::
1274
- For more information on how to use this operator, take a look at the guide:
1275
- :ref:`howto/operator:AutoMLDeleteDatasetOperator`
1276
-
1277
- :param dataset_id: Name of the dataset_id, list of dataset_id or string of dataset_id
1278
- coma separated to be deleted.
1279
- :param project_id: ID of the Google Cloud project where dataset is located if None then
1280
- default project_id is used.
1281
- :param location: The location of the project.
1282
- :param retry: A retry object used to retry requests. If `None` is specified, requests will not be
1283
- retried.
1284
- :param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if
1285
- `retry` is specified, the timeout applies to each individual attempt.
1286
- :param metadata: Additional metadata that is provided to the method.
1287
- :param gcp_conn_id: The connection ID to use to connect to Google Cloud.
1288
- :param impersonation_chain: Optional service account to impersonate using short-term
1289
- credentials, or chained list of accounts required to get the access_token
1290
- of the last account in the list, which will be impersonated in the request.
1291
- If set as a string, the account must grant the originating account
1292
- the Service Account Token Creator IAM role.
1293
- If set as a sequence, the identities from the list must grant
1294
- Service Account Token Creator IAM role to the directly preceding identity, with first
1295
- account from the list granting this role to the originating account (templated).
1296
- """
1297
-
1298
- template_fields: Sequence[str] = (
1299
- "dataset_id",
1300
- "location",
1301
- "project_id",
1302
- "impersonation_chain",
1303
- )
1304
-
1305
- def __init__(
1306
- self,
1307
- *,
1308
- dataset_id: str | list[str],
1309
- location: str,
1310
- project_id: str = PROVIDE_PROJECT_ID,
1311
- metadata: MetaData = (),
1312
- timeout: float | None = None,
1313
- retry: Retry | _MethodDefault = DEFAULT,
1314
- gcp_conn_id: str = "google_cloud_default",
1315
- impersonation_chain: str | Sequence[str] | None = None,
1316
- **kwargs,
1317
- ) -> None:
1318
- super().__init__(**kwargs)
1319
-
1320
- self.dataset_id = dataset_id
1321
- self.location = location
1322
- self.project_id = project_id
1323
- self.metadata = metadata
1324
- self.timeout = timeout
1325
- self.retry = retry
1326
- self.gcp_conn_id = gcp_conn_id
1327
- self.impersonation_chain = impersonation_chain
1328
-
1329
- @staticmethod
1330
- def _parse_dataset_id(dataset_id: str | list[str]) -> list[str]:
1331
- if not isinstance(dataset_id, str):
1332
- return dataset_id
1333
- try:
1334
- return ast.literal_eval(dataset_id)
1335
- except (SyntaxError, ValueError):
1336
- return dataset_id.split(",")
1337
-
1338
- def execute(self, context: Context):
1339
- hook = CloudAutoMLHook(
1340
- gcp_conn_id=self.gcp_conn_id,
1341
- impersonation_chain=self.impersonation_chain,
1342
- )
1343
- hook.get_dataset(
1344
- dataset_id=self.dataset_id,
1345
- location=self.location,
1346
- project_id=self.project_id,
1347
- retry=self.retry,
1348
- timeout=self.timeout,
1349
- metadata=self.metadata,
1350
- )
1351
- dataset_id_list = self._parse_dataset_id(self.dataset_id)
1352
- for dataset_id in dataset_id_list:
1353
- self.log.info("Deleting dataset %s", dataset_id)
1354
- hook.delete_dataset(
1355
- dataset_id=dataset_id,
1356
- location=self.location,
1357
- project_id=self.project_id,
1358
- retry=self.retry,
1359
- timeout=self.timeout,
1360
- metadata=self.metadata,
1361
- )
1362
- self.log.info("Dataset deleted.")