apache-airflow-providers-google 10.14.0rc1__py3-none-any.whl → 10.15.0rc1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (126) hide show
  1. airflow/providers/google/__init__.py +1 -1
  2. airflow/providers/google/ads/hooks/ads.py +1 -2
  3. airflow/providers/google/cloud/hooks/automl.py +13 -13
  4. airflow/providers/google/cloud/hooks/bigquery.py +208 -256
  5. airflow/providers/google/cloud/hooks/bigquery_dts.py +6 -6
  6. airflow/providers/google/cloud/hooks/bigtable.py +8 -8
  7. airflow/providers/google/cloud/hooks/cloud_batch.py +1 -1
  8. airflow/providers/google/cloud/hooks/cloud_build.py +19 -20
  9. airflow/providers/google/cloud/hooks/cloud_composer.py +4 -4
  10. airflow/providers/google/cloud/hooks/cloud_memorystore.py +10 -10
  11. airflow/providers/google/cloud/hooks/cloud_run.py +1 -1
  12. airflow/providers/google/cloud/hooks/cloud_sql.py +18 -19
  13. airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +3 -3
  14. airflow/providers/google/cloud/hooks/compute.py +16 -16
  15. airflow/providers/google/cloud/hooks/compute_ssh.py +1 -1
  16. airflow/providers/google/cloud/hooks/datacatalog.py +22 -22
  17. airflow/providers/google/cloud/hooks/dataflow.py +48 -49
  18. airflow/providers/google/cloud/hooks/dataform.py +16 -16
  19. airflow/providers/google/cloud/hooks/datafusion.py +15 -15
  20. airflow/providers/google/cloud/hooks/datapipeline.py +3 -3
  21. airflow/providers/google/cloud/hooks/dataplex.py +19 -19
  22. airflow/providers/google/cloud/hooks/dataprep.py +10 -10
  23. airflow/providers/google/cloud/hooks/dataproc.py +132 -14
  24. airflow/providers/google/cloud/hooks/dataproc_metastore.py +13 -13
  25. airflow/providers/google/cloud/hooks/datastore.py +3 -3
  26. airflow/providers/google/cloud/hooks/dlp.py +25 -25
  27. airflow/providers/google/cloud/hooks/gcs.py +39 -27
  28. airflow/providers/google/cloud/hooks/gdm.py +3 -3
  29. airflow/providers/google/cloud/hooks/kms.py +3 -3
  30. airflow/providers/google/cloud/hooks/kubernetes_engine.py +63 -48
  31. airflow/providers/google/cloud/hooks/life_sciences.py +13 -12
  32. airflow/providers/google/cloud/hooks/looker.py +8 -9
  33. airflow/providers/google/cloud/hooks/mlengine.py +12 -12
  34. airflow/providers/google/cloud/hooks/natural_language.py +2 -2
  35. airflow/providers/google/cloud/hooks/os_login.py +1 -1
  36. airflow/providers/google/cloud/hooks/pubsub.py +9 -9
  37. airflow/providers/google/cloud/hooks/secret_manager.py +1 -1
  38. airflow/providers/google/cloud/hooks/spanner.py +11 -11
  39. airflow/providers/google/cloud/hooks/speech_to_text.py +1 -1
  40. airflow/providers/google/cloud/hooks/stackdriver.py +7 -7
  41. airflow/providers/google/cloud/hooks/tasks.py +11 -11
  42. airflow/providers/google/cloud/hooks/text_to_speech.py +1 -1
  43. airflow/providers/google/cloud/hooks/translate.py +1 -1
  44. airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +13 -13
  45. airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +6 -6
  46. airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +45 -50
  47. airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +13 -13
  48. airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +9 -9
  49. airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +128 -11
  50. airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +10 -10
  51. airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +8 -8
  52. airflow/providers/google/cloud/hooks/video_intelligence.py +2 -2
  53. airflow/providers/google/cloud/hooks/vision.py +1 -1
  54. airflow/providers/google/cloud/hooks/workflows.py +10 -10
  55. airflow/providers/google/cloud/links/datafusion.py +12 -5
  56. airflow/providers/google/cloud/operators/bigquery.py +11 -11
  57. airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +3 -1
  58. airflow/providers/google/cloud/operators/dataflow.py +16 -16
  59. airflow/providers/google/cloud/operators/datafusion.py +9 -1
  60. airflow/providers/google/cloud/operators/dataproc.py +444 -69
  61. airflow/providers/google/cloud/operators/kubernetes_engine.py +6 -6
  62. airflow/providers/google/cloud/operators/life_sciences.py +10 -9
  63. airflow/providers/google/cloud/operators/mlengine.py +96 -96
  64. airflow/providers/google/cloud/operators/pubsub.py +2 -0
  65. airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +33 -3
  66. airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +59 -2
  67. airflow/providers/google/cloud/secrets/secret_manager.py +8 -7
  68. airflow/providers/google/cloud/sensors/bigquery.py +20 -16
  69. airflow/providers/google/cloud/sensors/cloud_composer.py +11 -8
  70. airflow/providers/google/cloud/sensors/dataproc_metastore.py +12 -2
  71. airflow/providers/google/cloud/sensors/gcs.py +8 -7
  72. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +1 -0
  73. airflow/providers/google/cloud/transfers/cassandra_to_gcs.py +4 -4
  74. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +1 -0
  75. airflow/providers/google/cloud/transfers/gcs_to_sftp.py +1 -1
  76. airflow/providers/google/cloud/transfers/mssql_to_gcs.py +1 -1
  77. airflow/providers/google/cloud/transfers/mysql_to_gcs.py +1 -1
  78. airflow/providers/google/cloud/transfers/oracle_to_gcs.py +1 -1
  79. airflow/providers/google/cloud/transfers/postgres_to_gcs.py +1 -1
  80. airflow/providers/google/cloud/transfers/presto_to_gcs.py +1 -1
  81. airflow/providers/google/cloud/transfers/s3_to_gcs.py +3 -3
  82. airflow/providers/google/cloud/transfers/sftp_to_gcs.py +1 -1
  83. airflow/providers/google/cloud/transfers/sql_to_gcs.py +3 -3
  84. airflow/providers/google/cloud/transfers/trino_to_gcs.py +1 -1
  85. airflow/providers/google/cloud/triggers/bigquery.py +12 -12
  86. airflow/providers/google/cloud/triggers/bigquery_dts.py +1 -1
  87. airflow/providers/google/cloud/triggers/cloud_batch.py +3 -1
  88. airflow/providers/google/cloud/triggers/cloud_build.py +2 -2
  89. airflow/providers/google/cloud/triggers/cloud_run.py +1 -1
  90. airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +6 -6
  91. airflow/providers/google/cloud/triggers/dataflow.py +3 -1
  92. airflow/providers/google/cloud/triggers/datafusion.py +2 -2
  93. airflow/providers/google/cloud/triggers/dataplex.py +2 -2
  94. airflow/providers/google/cloud/triggers/dataproc.py +34 -14
  95. airflow/providers/google/cloud/triggers/gcs.py +12 -8
  96. airflow/providers/google/cloud/triggers/kubernetes_engine.py +2 -2
  97. airflow/providers/google/cloud/triggers/mlengine.py +2 -2
  98. airflow/providers/google/cloud/triggers/pubsub.py +1 -1
  99. airflow/providers/google/cloud/triggers/vertex_ai.py +99 -0
  100. airflow/providers/google/cloud/utils/bigquery.py +2 -2
  101. airflow/providers/google/cloud/utils/credentials_provider.py +2 -2
  102. airflow/providers/google/cloud/utils/dataform.py +1 -1
  103. airflow/providers/google/cloud/utils/dataproc.py +25 -0
  104. airflow/providers/google/cloud/utils/field_validator.py +2 -2
  105. airflow/providers/google/cloud/utils/helpers.py +2 -2
  106. airflow/providers/google/cloud/utils/mlengine_operator_utils.py +1 -1
  107. airflow/providers/google/cloud/utils/mlengine_prediction_summary.py +1 -1
  108. airflow/providers/google/common/auth_backend/google_openid.py +2 -2
  109. airflow/providers/google/common/hooks/base_google.py +87 -23
  110. airflow/providers/google/common/hooks/discovery_api.py +2 -2
  111. airflow/providers/google/common/utils/id_token_credentials.py +5 -5
  112. airflow/providers/google/firebase/hooks/firestore.py +3 -3
  113. airflow/providers/google/get_provider_info.py +7 -2
  114. airflow/providers/google/leveldb/hooks/leveldb.py +4 -4
  115. airflow/providers/google/marketing_platform/hooks/analytics.py +11 -14
  116. airflow/providers/google/marketing_platform/hooks/campaign_manager.py +11 -11
  117. airflow/providers/google/marketing_platform/hooks/display_video.py +13 -13
  118. airflow/providers/google/marketing_platform/hooks/search_ads.py +4 -4
  119. airflow/providers/google/marketing_platform/operators/analytics.py +37 -32
  120. airflow/providers/google/suite/hooks/calendar.py +2 -2
  121. airflow/providers/google/suite/hooks/drive.py +7 -7
  122. airflow/providers/google/suite/hooks/sheets.py +8 -8
  123. {apache_airflow_providers_google-10.14.0rc1.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/METADATA +11 -11
  124. {apache_airflow_providers_google-10.14.0rc1.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/RECORD +126 -124
  125. {apache_airflow_providers_google-10.14.0rc1.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/WHEEL +0 -0
  126. {apache_airflow_providers_google-10.14.0rc1.dist-info → apache_airflow_providers_google-10.15.0rc1.dist-info}/entry_points.txt +0 -0
@@ -22,6 +22,7 @@ import warnings
22
22
  from functools import cached_property
23
23
  from typing import TYPE_CHECKING, Any, Sequence
24
24
 
25
+ from deprecated import deprecated
25
26
  from google.api_core.exceptions import AlreadyExists
26
27
  from google.cloud.container_v1.types import Cluster
27
28
 
@@ -510,13 +511,12 @@ class GKEStartPodOperator(KubernetesPodOperator):
510
511
  raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.")
511
512
 
512
513
  @staticmethod
514
+ @deprecated(
515
+ reason="Please use `fetch_cluster_info` instead to get the cluster info for connecting to it.",
516
+ category=AirflowProviderDeprecationWarning,
517
+ )
513
518
  def get_gke_config_file():
514
- warnings.warn(
515
- "The `get_gke_config_file` method is deprecated, "
516
- "please use `fetch_cluster_info` instead to get the cluster info for connecting to it.",
517
- AirflowProviderDeprecationWarning,
518
- stacklevel=2,
519
- )
519
+ pass
520
520
 
521
521
  @cached_property
522
522
  def cluster_hook(self) -> GKEHook:
@@ -18,9 +18,10 @@
18
18
  """Operators that interact with Google Cloud Life Sciences service."""
19
19
  from __future__ import annotations
20
20
 
21
- import warnings
22
21
  from typing import TYPE_CHECKING, Sequence
23
22
 
23
+ from deprecated import deprecated
24
+
24
25
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
25
26
  from airflow.providers.google.cloud.hooks.life_sciences import LifeSciencesHook
26
27
  from airflow.providers.google.cloud.links.life_sciences import LifeSciencesLink
@@ -30,6 +31,14 @@ if TYPE_CHECKING:
30
31
  from airflow.utils.context import Context
31
32
 
32
33
 
34
+ @deprecated(
35
+ reason=(
36
+ "Consider using Google Cloud Batch Operators instead."
37
+ "The Life Sciences API (beta) will be discontinued "
38
+ "on July 8, 2025 in favor of Google Cloud Batch."
39
+ ),
40
+ category=AirflowProviderDeprecationWarning,
41
+ )
33
42
  class LifeSciencesRunPipelineOperator(GoogleCloudBaseOperator):
34
43
  """
35
44
  Runs a Life Sciences Pipeline.
@@ -87,14 +96,6 @@ class LifeSciencesRunPipelineOperator(GoogleCloudBaseOperator):
87
96
  self._validate_inputs()
88
97
  self.impersonation_chain = impersonation_chain
89
98
 
90
- warnings.warn(
91
- """This operator is deprecated. Consider using Google Cloud Batch Operators instead.
92
- The Life Sciences API (beta) will be discontinued on July 8, 2025 in favor
93
- of Google Cloud Batch.""",
94
- AirflowProviderDeprecationWarning,
95
- stacklevel=3,
96
- )
97
-
98
99
  def _validate_inputs(self) -> None:
99
100
  if not self.body:
100
101
  raise AirflowException("The required parameter 'body' is missing")
@@ -22,9 +22,9 @@ from __future__ import annotations
22
22
  import logging
23
23
  import re
24
24
  import time
25
- import warnings
26
25
  from typing import TYPE_CHECKING, Any, Sequence
27
26
 
27
+ from deprecated import deprecated
28
28
  from googleapiclient.errors import HttpError
29
29
 
30
30
  from airflow.configuration import conf
@@ -78,6 +78,14 @@ def _normalize_mlengine_job_id(job_id: str) -> str:
78
78
  return cleansed_job_id
79
79
 
80
80
 
81
+ @deprecated(
82
+ reason=(
83
+ "This operator is deprecated. All the functionality of legacy "
84
+ "MLEngine and new features are available on the Vertex AI platform. "
85
+ "Please use `CreateBatchPredictionJobOperator`"
86
+ ),
87
+ category=AirflowProviderDeprecationWarning,
88
+ )
81
89
  class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator):
82
90
  """
83
91
  Start a Google Cloud ML Engine prediction job.
@@ -214,14 +222,6 @@ class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator):
214
222
  self._labels = labels
215
223
  self._impersonation_chain = impersonation_chain
216
224
 
217
- warnings.warn(
218
- "This operator is deprecated. All the functionality of legacy "
219
- "MLEngine and new features are available on the Vertex AI platform. "
220
- "Please use `CreateBatchPredictionJobOperator`",
221
- AirflowProviderDeprecationWarning,
222
- stacklevel=3,
223
- )
224
-
225
225
  if not self._project_id:
226
226
  raise AirflowException("Google Cloud project id is required.")
227
227
  if not self._job_id:
@@ -296,6 +296,13 @@ class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator):
296
296
  return finished_prediction_job["predictionOutput"]
297
297
 
298
298
 
299
+ @deprecated(
300
+ reason=(
301
+ "This operator is deprecated. Consider using operators for specific operations: "
302
+ "MLEngineCreateModelOperator, MLEngineGetModelOperator."
303
+ ),
304
+ category=AirflowProviderDeprecationWarning,
305
+ )
299
306
  class MLEngineManageModelOperator(GoogleCloudBaseOperator):
300
307
  """
301
308
  Operator for managing a Google Cloud ML Engine model.
@@ -345,14 +352,6 @@ class MLEngineManageModelOperator(GoogleCloudBaseOperator):
345
352
  **kwargs,
346
353
  ) -> None:
347
354
  super().__init__(**kwargs)
348
-
349
- warnings.warn(
350
- "This operator is deprecated. Consider using operators for specific operations: "
351
- "MLEngineCreateModelOperator, MLEngineGetModelOperator.",
352
- AirflowProviderDeprecationWarning,
353
- stacklevel=3,
354
- )
355
-
356
355
  self._project_id = project_id
357
356
  self._model = model
358
357
  self._operation = operation
@@ -372,6 +371,14 @@ class MLEngineManageModelOperator(GoogleCloudBaseOperator):
372
371
  raise ValueError(f"Unknown operation: {self._operation}")
373
372
 
374
373
 
374
+ @deprecated(
375
+ reason=(
376
+ "This operator is deprecated. All the functionality of legacy "
377
+ "MLEngine and new features are available on the Vertex AI platform. "
378
+ "Please use appropriate VertexAI operator."
379
+ ),
380
+ category=AirflowProviderDeprecationWarning,
381
+ )
375
382
  class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
376
383
  """
377
384
  Creates a new model.
@@ -422,14 +429,6 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
422
429
  self._gcp_conn_id = gcp_conn_id
423
430
  self._impersonation_chain = impersonation_chain
424
431
 
425
- warnings.warn(
426
- "This operator is deprecated. All the functionality of legacy "
427
- "MLEngine and new features are available on the Vertex AI platform. "
428
- "Please use appropriate VertexAI operator.",
429
- AirflowProviderDeprecationWarning,
430
- stacklevel=3,
431
- )
432
-
433
432
  def execute(self, context: Context):
434
433
  hook = MLEngineHook(
435
434
  gcp_conn_id=self._gcp_conn_id,
@@ -448,6 +447,14 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
448
447
  return hook.create_model(project_id=self._project_id, model=self._model)
449
448
 
450
449
 
450
+ @deprecated(
451
+ reason=(
452
+ "This operator is deprecated. All the functionality of legacy "
453
+ "MLEngine and new features are available on the Vertex AI platform. "
454
+ "Please use `GetModelOperator`"
455
+ ),
456
+ category=AirflowProviderDeprecationWarning,
457
+ )
451
458
  class MLEngineGetModelOperator(GoogleCloudBaseOperator):
452
459
  """
453
460
  Gets a particular model.
@@ -498,14 +505,6 @@ class MLEngineGetModelOperator(GoogleCloudBaseOperator):
498
505
  self._gcp_conn_id = gcp_conn_id
499
506
  self._impersonation_chain = impersonation_chain
500
507
 
501
- warnings.warn(
502
- "This operator is deprecated. All the functionality of legacy "
503
- "MLEngine and new features are available on the Vertex AI platform. "
504
- "Please use `GetModelOperator`",
505
- AirflowProviderDeprecationWarning,
506
- stacklevel=3,
507
- )
508
-
509
508
  def execute(self, context: Context):
510
509
  hook = MLEngineHook(
511
510
  gcp_conn_id=self._gcp_conn_id,
@@ -523,6 +522,14 @@ class MLEngineGetModelOperator(GoogleCloudBaseOperator):
523
522
  return hook.get_model(project_id=self._project_id, model_name=self._model_name)
524
523
 
525
524
 
525
+ @deprecated(
526
+ reason=(
527
+ "This operator is deprecated. All the functionality of legacy "
528
+ "MLEngine and new features are available on the Vertex AI platform. "
529
+ "Please use `DeleteModelOperator`"
530
+ ),
531
+ category=AirflowProviderDeprecationWarning,
532
+ )
526
533
  class MLEngineDeleteModelOperator(GoogleCloudBaseOperator):
527
534
  """
528
535
  Deletes a model.
@@ -579,14 +586,6 @@ class MLEngineDeleteModelOperator(GoogleCloudBaseOperator):
579
586
  self._gcp_conn_id = gcp_conn_id
580
587
  self._impersonation_chain = impersonation_chain
581
588
 
582
- warnings.warn(
583
- "This operator is deprecated. All the functionality of legacy "
584
- "MLEngine and new features are available on the Vertex AI platform. "
585
- "Please use `DeleteModelOperator`",
586
- AirflowProviderDeprecationWarning,
587
- stacklevel=3,
588
- )
589
-
590
589
  def execute(self, context: Context):
591
590
  hook = MLEngineHook(
592
591
  gcp_conn_id=self._gcp_conn_id,
@@ -606,6 +605,14 @@ class MLEngineDeleteModelOperator(GoogleCloudBaseOperator):
606
605
  )
607
606
 
608
607
 
608
+ @deprecated(
609
+ reason=(
610
+ "This operator is deprecated. Consider using operators for specific operations: "
611
+ "MLEngineCreateVersion, MLEngineSetDefaultVersion, "
612
+ "MLEngineListVersions, MLEngineDeleteVersion."
613
+ ),
614
+ category=AirflowProviderDeprecationWarning,
615
+ )
609
616
  class MLEngineManageVersionOperator(GoogleCloudBaseOperator):
610
617
  """
611
618
  Operator for managing a Google Cloud ML Engine version.
@@ -688,13 +695,6 @@ class MLEngineManageVersionOperator(GoogleCloudBaseOperator):
688
695
  self._gcp_conn_id = gcp_conn_id
689
696
  self._impersonation_chain = impersonation_chain
690
697
 
691
- warnings.warn(
692
- "This operator is deprecated. Consider using operators for specific operations: "
693
- "MLEngineCreateVersion, MLEngineSetDefaultVersion, MLEngineListVersions, MLEngineDeleteVersion.",
694
- AirflowProviderDeprecationWarning,
695
- stacklevel=3,
696
- )
697
-
698
698
  def execute(self, context: Context):
699
699
  if "name" not in self._version:
700
700
  self._version["name"] = self._version_name
@@ -724,6 +724,14 @@ class MLEngineManageVersionOperator(GoogleCloudBaseOperator):
724
724
  raise ValueError(f"Unknown operation: {self._operation}")
725
725
 
726
726
 
727
+ @deprecated(
728
+ reason=(
729
+ "This operator is deprecated. All the functionality of legacy "
730
+ "MLEngine and new features are available on the Vertex AI platform. "
731
+ "Please use parent_model parameter for VertexAI operators instead."
732
+ ),
733
+ category=AirflowProviderDeprecationWarning,
734
+ )
727
735
  class MLEngineCreateVersionOperator(GoogleCloudBaseOperator):
728
736
  """
729
737
  Creates a new version in the model.
@@ -779,14 +787,6 @@ class MLEngineCreateVersionOperator(GoogleCloudBaseOperator):
779
787
  self._impersonation_chain = impersonation_chain
780
788
  self._validate_inputs()
781
789
 
782
- warnings.warn(
783
- "This operator is deprecated. All the functionality of legacy "
784
- "MLEngine and new features are available on the Vertex AI platform. "
785
- "Please use parent_model parameter for VertexAI operators instead.",
786
- AirflowProviderDeprecationWarning,
787
- stacklevel=3,
788
- )
789
-
790
790
  def _validate_inputs(self):
791
791
  if not self._model_name:
792
792
  raise AirflowException("The model_name parameter could not be empty.")
@@ -815,6 +815,14 @@ class MLEngineCreateVersionOperator(GoogleCloudBaseOperator):
815
815
  )
816
816
 
817
817
 
818
+ @deprecated(
819
+ reason=(
820
+ "This operator is deprecated. All the functionality of legacy "
821
+ "MLEngine and new features are available on the Vertex AI platform. "
822
+ "Please use `SetDefaultVersionOnModelOperator`"
823
+ ),
824
+ category=AirflowProviderDeprecationWarning,
825
+ )
818
826
  class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator):
819
827
  """
820
828
  Sets a version in the model.
@@ -872,14 +880,6 @@ class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator):
872
880
  self._impersonation_chain = impersonation_chain
873
881
  self._validate_inputs()
874
882
 
875
- warnings.warn(
876
- "This operator is deprecated. All the functionality of legacy "
877
- "MLEngine and new features are available on the Vertex AI platform. "
878
- "Please use `SetDefaultVersionOnModelOperator` instead.",
879
- AirflowProviderDeprecationWarning,
880
- stacklevel=3,
881
- )
882
-
883
883
  def _validate_inputs(self):
884
884
  if not self._model_name:
885
885
  raise AirflowException("The model_name parameter could not be empty.")
@@ -908,6 +908,14 @@ class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator):
908
908
  )
909
909
 
910
910
 
911
+ @deprecated(
912
+ reason=(
913
+ "This operator is deprecated. All the functionality of legacy "
914
+ "MLEngine and new features are available on the Vertex AI platform. "
915
+ "Please use `ListModelVersionsOperator`"
916
+ ),
917
+ category=AirflowProviderDeprecationWarning,
918
+ )
911
919
  class MLEngineListVersionsOperator(GoogleCloudBaseOperator):
912
920
  """
913
921
  Lists all available versions of the model.
@@ -961,14 +969,6 @@ class MLEngineListVersionsOperator(GoogleCloudBaseOperator):
961
969
  self._impersonation_chain = impersonation_chain
962
970
  self._validate_inputs()
963
971
 
964
- warnings.warn(
965
- "This operator is deprecated. All the functionality of legacy "
966
- "MLEngine and new features are available on the Vertex AI platform. "
967
- "Please use `ListModelVersionsOperator` instead.",
968
- AirflowProviderDeprecationWarning,
969
- stacklevel=3,
970
- )
971
-
972
972
  def _validate_inputs(self):
973
973
  if not self._model_name:
974
974
  raise AirflowException("The model_name parameter could not be empty.")
@@ -994,6 +994,14 @@ class MLEngineListVersionsOperator(GoogleCloudBaseOperator):
994
994
  )
995
995
 
996
996
 
997
+ @deprecated(
998
+ reason=(
999
+ "This operator is deprecated. All the functionality of legacy "
1000
+ "MLEngine and new features are available on the Vertex AI platform. "
1001
+ "Please use `DeleteModelVersionOperator`"
1002
+ ),
1003
+ category=AirflowProviderDeprecationWarning,
1004
+ )
997
1005
  class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator):
998
1006
  """
999
1007
  Deletes the version from the model.
@@ -1051,14 +1059,6 @@ class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator):
1051
1059
  self._impersonation_chain = impersonation_chain
1052
1060
  self._validate_inputs()
1053
1061
 
1054
- warnings.warn(
1055
- "This operator is deprecated. All the functionality of legacy "
1056
- "MLEngine and new features are available on the Vertex AI platform. "
1057
- "Please use `DeleteModelVersionOperator` instead.",
1058
- AirflowProviderDeprecationWarning,
1059
- stacklevel=3,
1060
- )
1061
-
1062
1062
  def _validate_inputs(self):
1063
1063
  if not self._model_name:
1064
1064
  raise AirflowException("The model_name parameter could not be empty.")
@@ -1086,6 +1086,14 @@ class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator):
1086
1086
  )
1087
1087
 
1088
1088
 
1089
+ @deprecated(
1090
+ reason=(
1091
+ "This operator is deprecated. All the functionality of legacy "
1092
+ "MLEngine and new features are available on the Vertex AI platform. "
1093
+ "Please use `CreateCustomPythonPackageTrainingJobOperator`"
1094
+ ),
1095
+ category=AirflowProviderDeprecationWarning,
1096
+ )
1089
1097
  class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
1090
1098
  """
1091
1099
  Operator for launching a MLEngine training job.
@@ -1220,14 +1228,6 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
1220
1228
  self.deferrable = deferrable
1221
1229
  self.cancel_on_kill = cancel_on_kill
1222
1230
 
1223
- warnings.warn(
1224
- "This operator is deprecated. All the functionality of legacy "
1225
- "MLEngine and new features are available on the Vertex AI platform. "
1226
- "Please use `CreateCustomPythonPackageTrainingJobOperator` instead.",
1227
- AirflowProviderDeprecationWarning,
1228
- stacklevel=3,
1229
- )
1230
-
1231
1231
  custom = self._scale_tier is not None and self._scale_tier.upper() == "CUSTOM"
1232
1232
  custom_image = (
1233
1233
  custom
@@ -1428,6 +1428,14 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
1428
1428
  self.log.info("Skipping to cancel job: %s:%s.%s", self._project_id, self.job_id)
1429
1429
 
1430
1430
 
1431
+ @deprecated(
1432
+ reason=(
1433
+ "This operator is deprecated. All the functionality of legacy "
1434
+ "MLEngine and new features are available on the Vertex AI platform. "
1435
+ "Please use `CancelCustomTrainingJobOperator`"
1436
+ ),
1437
+ category=AirflowProviderDeprecationWarning,
1438
+ )
1431
1439
  class MLEngineTrainingCancelJobOperator(GoogleCloudBaseOperator):
1432
1440
  """
1433
1441
  Operator for cleaning up failed MLEngine training job.
@@ -1474,14 +1482,6 @@ class MLEngineTrainingCancelJobOperator(GoogleCloudBaseOperator):
1474
1482
  self._gcp_conn_id = gcp_conn_id
1475
1483
  self._impersonation_chain = impersonation_chain
1476
1484
 
1477
- warnings.warn(
1478
- "This operator is deprecated. All the functionality of legacy "
1479
- "MLEngine and new features are available on the Vertex AI platform. "
1480
- "Please use `CancelCustomTrainingJobOperator` instead.",
1481
- AirflowProviderDeprecationWarning,
1482
- stacklevel=3,
1483
- )
1484
-
1485
1485
  if not self._project_id:
1486
1486
  raise AirflowException("Google Cloud project id is required.")
1487
1487
 
@@ -781,6 +781,8 @@ class PubSubPullOperator(GoogleCloudBaseOperator):
781
781
  context: Context,
782
782
  ) -> list:
783
783
  """
784
+ Act as a default message callback.
785
+
784
786
  This method can be overridden by subclasses or by `messages_callback` constructor argument.
785
787
 
786
788
  This default implementation converts `ReceivedMessage` objects into JSON-serializable dicts.
@@ -445,9 +445,19 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
445
445
  self,
446
446
  *,
447
447
  command: Sequence[str] = [],
448
+ region: str,
449
+ parent_model: str | None = None,
450
+ impersonation_chain: str | Sequence[str] | None = None,
451
+ dataset_id: str | None = None,
448
452
  **kwargs,
449
453
  ) -> None:
450
- super().__init__(**kwargs)
454
+ super().__init__(
455
+ region=region,
456
+ parent_model=parent_model,
457
+ impersonation_chain=impersonation_chain,
458
+ dataset_id=dataset_id,
459
+ **kwargs,
460
+ )
451
461
  self.command = command
452
462
 
453
463
  def execute(self, context: Context):
@@ -816,9 +826,19 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
816
826
  *,
817
827
  python_package_gcs_uri: str,
818
828
  python_module_name: str,
829
+ region: str,
830
+ parent_model: str | None = None,
831
+ impersonation_chain: str | Sequence[str] | None = None,
832
+ dataset_id: str | None = None,
819
833
  **kwargs,
820
834
  ) -> None:
821
- super().__init__(**kwargs)
835
+ super().__init__(
836
+ region=region,
837
+ parent_model=parent_model,
838
+ impersonation_chain=impersonation_chain,
839
+ dataset_id=dataset_id,
840
+ **kwargs,
841
+ )
822
842
  self.python_package_gcs_uri = python_package_gcs_uri
823
843
  self.python_module_name = python_module_name
824
844
 
@@ -1190,9 +1210,19 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
1190
1210
  *,
1191
1211
  script_path: str,
1192
1212
  requirements: Sequence[str] | None = None,
1213
+ region: str,
1214
+ parent_model: str | None = None,
1215
+ impersonation_chain: str | Sequence[str] | None = None,
1216
+ dataset_id: str | None = None,
1193
1217
  **kwargs,
1194
1218
  ) -> None:
1195
- super().__init__(**kwargs)
1219
+ super().__init__(
1220
+ region=region,
1221
+ parent_model=parent_model,
1222
+ impersonation_chain=impersonation_chain,
1223
+ dataset_id=dataset_id,
1224
+ **kwargs,
1225
+ )
1196
1226
  self.requirements = requirements
1197
1227
  self.script_path = script_path
1198
1228
 
@@ -20,12 +20,14 @@
20
20
 
21
21
  from __future__ import annotations
22
22
 
23
- from typing import TYPE_CHECKING, Sequence
23
+ from typing import TYPE_CHECKING, Any, Sequence
24
24
 
25
25
  from google.api_core.exceptions import NotFound
26
26
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
27
27
  from google.cloud.aiplatform_v1.types import HyperparameterTuningJob
28
28
 
29
+ from airflow.configuration import conf
30
+ from airflow.exceptions import AirflowException
29
31
  from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
30
32
  HyperparameterTuningJobHook,
31
33
  )
@@ -34,6 +36,7 @@ from airflow.providers.google.cloud.links.vertex_ai import (
34
36
  VertexAITrainingLink,
35
37
  )
36
38
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
39
+ from airflow.providers.google.cloud.triggers.vertex_ai import CreateHyperparameterTuningJobTrigger
37
40
 
38
41
  if TYPE_CHECKING:
39
42
  from google.api_core.retry import Retry
@@ -124,7 +127,7 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
124
127
  `service_account` is required with provided `tensorboard`. For more information on configuring
125
128
  your service account please visit:
126
129
  https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
127
- :param sync: Whether to execute this method synchronously. If False, this method will unblock and it
130
+ :param sync: Whether to execute this method synchronously. If False, this method will unblock, and it
128
131
  will be executed in a concurrent Future.
129
132
  :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
130
133
  :param impersonation_chain: Optional service account to impersonate using short-term
@@ -135,6 +138,9 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
135
138
  If set as a sequence, the identities from the list must grant
136
139
  Service Account Token Creator IAM role to the directly preceding identity, with first
137
140
  account from the list granting this role to the originating account (templated).
141
+ :param deferrable: Run operator in the deferrable mode. Note that it requires calling the operator
142
+ with `sync=False` parameter.
143
+ :param poll_interval: Interval size which defines how often job status is checked in deferrable mode.
138
144
  """
139
145
 
140
146
  template_fields = [
@@ -177,6 +183,8 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
177
183
  # END: run param
178
184
  gcp_conn_id: str = "google_cloud_default",
179
185
  impersonation_chain: str | Sequence[str] | None = None,
186
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
187
+ poll_interval: int = 10,
180
188
  **kwargs,
181
189
  ) -> None:
182
190
  super().__init__(**kwargs)
@@ -209,8 +217,17 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
209
217
  self.gcp_conn_id = gcp_conn_id
210
218
  self.impersonation_chain = impersonation_chain
211
219
  self.hook: HyperparameterTuningJobHook | None = None
220
+ self.deferrable = deferrable
221
+ self.poll_interval = poll_interval
212
222
 
213
223
  def execute(self, context: Context):
224
+ if self.deferrable and self.sync:
225
+ raise AirflowException(
226
+ "Deferrable mode can be used only with sync=False option. "
227
+ "If you are willing to run the operator in deferrable mode, please, set sync=False. "
228
+ "Otherwise, disable deferrable mode `deferrable=False`."
229
+ )
230
+
214
231
  self.log.info("Creating Hyperparameter Tuning job")
215
232
  self.hook = HyperparameterTuningJobHook(
216
233
  gcp_conn_id=self.gcp_conn_id,
@@ -243,12 +260,26 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
243
260
  enable_web_access=self.enable_web_access,
244
261
  tensorboard=self.tensorboard,
245
262
  sync=self.sync,
263
+ wait_job_completed=not self.deferrable,
246
264
  )
247
265
 
248
266
  hyperparameter_tuning_job = result.to_dict()
249
267
  hyperparameter_tuning_job_id = self.hook.extract_hyperparameter_tuning_job_id(
250
268
  hyperparameter_tuning_job
251
269
  )
270
+ if self.deferrable:
271
+ self.defer(
272
+ trigger=CreateHyperparameterTuningJobTrigger(
273
+ conn_id=self.gcp_conn_id,
274
+ project_id=self.project_id,
275
+ location=self.region,
276
+ job_id=hyperparameter_tuning_job_id,
277
+ poll_interval=self.poll_interval,
278
+ impersonation_chain=self.impersonation_chain,
279
+ ),
280
+ method_name="execute_complete",
281
+ )
282
+
252
283
  self.log.info("Hyperparameter Tuning job was created. Job id: %s", hyperparameter_tuning_job_id)
253
284
 
254
285
  self.xcom_push(context, key="hyperparameter_tuning_job_id", value=hyperparameter_tuning_job_id)
@@ -262,6 +293,32 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
262
293
  if self.hook:
263
294
  self.hook.cancel_hyperparameter_tuning_job()
264
295
 
296
+ def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]:
297
+ if event and event["status"] == "error":
298
+ raise AirflowException(event["message"])
299
+ job: dict[str, Any] = event["job"]
300
+ self.log.info("Hyperparameter tuning job %s created and completed successfully.", job["name"])
301
+ hook = HyperparameterTuningJobHook(
302
+ gcp_conn_id=self.gcp_conn_id,
303
+ impersonation_chain=self.impersonation_chain,
304
+ )
305
+ job_id = hook.extract_hyperparameter_tuning_job_id(job)
306
+ self.xcom_push(
307
+ context,
308
+ key="hyperparameter_tuning_job_id",
309
+ value=job_id,
310
+ )
311
+ self.xcom_push(
312
+ context,
313
+ key="training_conf",
314
+ value={
315
+ "training_conf_id": job_id,
316
+ "region": self.region,
317
+ "project_id": self.project_id,
318
+ },
319
+ )
320
+ return event["job"]
321
+
265
322
 
266
323
  class GetHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
267
324
  """
@@ -18,9 +18,9 @@
18
18
  from __future__ import annotations
19
19
 
20
20
  import logging
21
- import warnings
22
21
  from typing import Sequence
23
22
 
23
+ from deprecated import deprecated
24
24
  from google.auth.exceptions import DefaultCredentialsError
25
25
 
26
26
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
@@ -159,6 +159,13 @@ class CloudSecretManagerBackend(BaseSecretsBackend, LoggingMixin):
159
159
 
160
160
  return self._get_secret(self.connections_prefix, conn_id)
161
161
 
162
+ @deprecated(
163
+ reason=(
164
+ "Method `CloudSecretManagerBackend.get_conn_uri` is deprecated and will be removed "
165
+ "in a future release. Please use method `get_conn_value` instead."
166
+ ),
167
+ category=AirflowProviderDeprecationWarning,
168
+ )
162
169
  def get_conn_uri(self, conn_id: str) -> str | None:
163
170
  """
164
171
  Return URI representation of Connection conn_id.
@@ -168,12 +175,6 @@ class CloudSecretManagerBackend(BaseSecretsBackend, LoggingMixin):
168
175
  :param conn_id: the connection id
169
176
  :return: deserialized Connection
170
177
  """
171
- warnings.warn(
172
- f"Method `{self.__class__.__name__}.get_conn_uri` is deprecated and will be removed "
173
- "in a future release. Please use method `get_conn_value` instead.",
174
- AirflowProviderDeprecationWarning,
175
- stacklevel=2,
176
- )
177
178
  return self.get_conn_value(conn_id)
178
179
 
179
180
  def get_variable(self, key: str) -> str | None: