apache-airflow-providers-google 10.16.0rc1__py3-none-any.whl → 10.17.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 (196) hide show
  1. airflow/providers/google/__init__.py +1 -1
  2. airflow/providers/google/ads/hooks/ads.py +5 -4
  3. airflow/providers/google/ads/operators/ads.py +1 -0
  4. airflow/providers/google/cloud/example_dags/example_cloud_sql_query.py +1 -0
  5. airflow/providers/google/cloud/example_dags/example_cloud_task.py +1 -0
  6. airflow/providers/google/cloud/example_dags/example_facebook_ads_to_gcs.py +1 -0
  7. airflow/providers/google/cloud/example_dags/example_looker.py +1 -0
  8. airflow/providers/google/cloud/example_dags/example_presto_to_gcs.py +1 -0
  9. airflow/providers/google/cloud/example_dags/example_salesforce_to_gcs.py +1 -0
  10. airflow/providers/google/cloud/fs/gcs.py +1 -2
  11. airflow/providers/google/cloud/hooks/automl.py +1 -0
  12. airflow/providers/google/cloud/hooks/bigquery.py +87 -24
  13. airflow/providers/google/cloud/hooks/bigquery_dts.py +1 -0
  14. airflow/providers/google/cloud/hooks/bigtable.py +1 -0
  15. airflow/providers/google/cloud/hooks/cloud_build.py +1 -0
  16. airflow/providers/google/cloud/hooks/cloud_memorystore.py +1 -0
  17. airflow/providers/google/cloud/hooks/cloud_sql.py +1 -0
  18. airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +9 -4
  19. airflow/providers/google/cloud/hooks/compute.py +1 -0
  20. airflow/providers/google/cloud/hooks/compute_ssh.py +2 -2
  21. airflow/providers/google/cloud/hooks/dataflow.py +6 -5
  22. airflow/providers/google/cloud/hooks/datafusion.py +1 -0
  23. airflow/providers/google/cloud/hooks/datapipeline.py +1 -0
  24. airflow/providers/google/cloud/hooks/dataplex.py +1 -0
  25. airflow/providers/google/cloud/hooks/dataprep.py +1 -0
  26. airflow/providers/google/cloud/hooks/dataproc.py +3 -2
  27. airflow/providers/google/cloud/hooks/dataproc_metastore.py +1 -0
  28. airflow/providers/google/cloud/hooks/datastore.py +1 -0
  29. airflow/providers/google/cloud/hooks/dlp.py +1 -0
  30. airflow/providers/google/cloud/hooks/functions.py +1 -0
  31. airflow/providers/google/cloud/hooks/gcs.py +12 -5
  32. airflow/providers/google/cloud/hooks/kms.py +1 -0
  33. airflow/providers/google/cloud/hooks/kubernetes_engine.py +178 -300
  34. airflow/providers/google/cloud/hooks/life_sciences.py +1 -0
  35. airflow/providers/google/cloud/hooks/looker.py +1 -0
  36. airflow/providers/google/cloud/hooks/mlengine.py +1 -0
  37. airflow/providers/google/cloud/hooks/natural_language.py +1 -0
  38. airflow/providers/google/cloud/hooks/os_login.py +1 -0
  39. airflow/providers/google/cloud/hooks/pubsub.py +1 -0
  40. airflow/providers/google/cloud/hooks/secret_manager.py +1 -0
  41. airflow/providers/google/cloud/hooks/spanner.py +1 -0
  42. airflow/providers/google/cloud/hooks/speech_to_text.py +1 -0
  43. airflow/providers/google/cloud/hooks/stackdriver.py +1 -0
  44. airflow/providers/google/cloud/hooks/text_to_speech.py +1 -0
  45. airflow/providers/google/cloud/hooks/translate.py +1 -0
  46. airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +1 -0
  47. airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +255 -3
  48. airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +1 -0
  49. airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +1 -0
  50. airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +1 -0
  51. airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +197 -0
  52. airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +9 -9
  53. airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +231 -12
  54. airflow/providers/google/cloud/hooks/video_intelligence.py +1 -0
  55. airflow/providers/google/cloud/hooks/vision.py +1 -0
  56. airflow/providers/google/cloud/links/automl.py +1 -0
  57. airflow/providers/google/cloud/links/bigquery.py +1 -0
  58. airflow/providers/google/cloud/links/bigquery_dts.py +1 -0
  59. airflow/providers/google/cloud/links/cloud_memorystore.py +1 -0
  60. airflow/providers/google/cloud/links/cloud_sql.py +1 -0
  61. airflow/providers/google/cloud/links/cloud_tasks.py +1 -0
  62. airflow/providers/google/cloud/links/compute.py +1 -0
  63. airflow/providers/google/cloud/links/datacatalog.py +1 -0
  64. airflow/providers/google/cloud/links/dataflow.py +1 -0
  65. airflow/providers/google/cloud/links/dataform.py +1 -0
  66. airflow/providers/google/cloud/links/datafusion.py +1 -0
  67. airflow/providers/google/cloud/links/dataplex.py +1 -0
  68. airflow/providers/google/cloud/links/dataproc.py +1 -0
  69. airflow/providers/google/cloud/links/kubernetes_engine.py +28 -0
  70. airflow/providers/google/cloud/links/mlengine.py +1 -0
  71. airflow/providers/google/cloud/links/pubsub.py +1 -0
  72. airflow/providers/google/cloud/links/spanner.py +1 -0
  73. airflow/providers/google/cloud/links/stackdriver.py +1 -0
  74. airflow/providers/google/cloud/links/workflows.py +1 -0
  75. airflow/providers/google/cloud/log/stackdriver_task_handler.py +18 -4
  76. airflow/providers/google/cloud/operators/automl.py +1 -0
  77. airflow/providers/google/cloud/operators/bigquery.py +21 -0
  78. airflow/providers/google/cloud/operators/bigquery_dts.py +1 -0
  79. airflow/providers/google/cloud/operators/bigtable.py +1 -0
  80. airflow/providers/google/cloud/operators/cloud_base.py +1 -0
  81. airflow/providers/google/cloud/operators/cloud_build.py +1 -0
  82. airflow/providers/google/cloud/operators/cloud_memorystore.py +1 -0
  83. airflow/providers/google/cloud/operators/cloud_sql.py +1 -0
  84. airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +11 -5
  85. airflow/providers/google/cloud/operators/compute.py +1 -0
  86. airflow/providers/google/cloud/operators/dataflow.py +1 -0
  87. airflow/providers/google/cloud/operators/datafusion.py +1 -0
  88. airflow/providers/google/cloud/operators/datapipeline.py +1 -0
  89. airflow/providers/google/cloud/operators/dataprep.py +1 -0
  90. airflow/providers/google/cloud/operators/dataproc.py +3 -2
  91. airflow/providers/google/cloud/operators/dataproc_metastore.py +1 -0
  92. airflow/providers/google/cloud/operators/datastore.py +1 -0
  93. airflow/providers/google/cloud/operators/functions.py +1 -0
  94. airflow/providers/google/cloud/operators/gcs.py +1 -0
  95. airflow/providers/google/cloud/operators/kubernetes_engine.py +600 -4
  96. airflow/providers/google/cloud/operators/life_sciences.py +1 -0
  97. airflow/providers/google/cloud/operators/looker.py +1 -0
  98. airflow/providers/google/cloud/operators/mlengine.py +283 -259
  99. airflow/providers/google/cloud/operators/natural_language.py +1 -0
  100. airflow/providers/google/cloud/operators/pubsub.py +1 -0
  101. airflow/providers/google/cloud/operators/spanner.py +1 -0
  102. airflow/providers/google/cloud/operators/speech_to_text.py +1 -0
  103. airflow/providers/google/cloud/operators/text_to_speech.py +1 -0
  104. airflow/providers/google/cloud/operators/translate.py +1 -0
  105. airflow/providers/google/cloud/operators/translate_speech.py +1 -0
  106. airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +14 -7
  107. airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +67 -13
  108. airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +26 -8
  109. airflow/providers/google/cloud/operators/vertex_ai/dataset.py +1 -0
  110. airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +306 -0
  111. airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +29 -48
  112. airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +52 -17
  113. airflow/providers/google/cloud/operators/video_intelligence.py +1 -0
  114. airflow/providers/google/cloud/operators/vision.py +1 -0
  115. airflow/providers/google/cloud/secrets/secret_manager.py +1 -0
  116. airflow/providers/google/cloud/sensors/bigquery.py +1 -0
  117. airflow/providers/google/cloud/sensors/bigquery_dts.py +1 -0
  118. airflow/providers/google/cloud/sensors/bigtable.py +1 -0
  119. airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +1 -0
  120. airflow/providers/google/cloud/sensors/dataflow.py +1 -0
  121. airflow/providers/google/cloud/sensors/dataform.py +1 -0
  122. airflow/providers/google/cloud/sensors/datafusion.py +1 -0
  123. airflow/providers/google/cloud/sensors/dataplex.py +1 -0
  124. airflow/providers/google/cloud/sensors/dataprep.py +1 -0
  125. airflow/providers/google/cloud/sensors/dataproc.py +1 -0
  126. airflow/providers/google/cloud/sensors/gcs.py +1 -0
  127. airflow/providers/google/cloud/sensors/looker.py +1 -0
  128. airflow/providers/google/cloud/sensors/pubsub.py +1 -0
  129. airflow/providers/google/cloud/sensors/tasks.py +1 -0
  130. airflow/providers/google/cloud/transfers/bigquery_to_bigquery.py +1 -0
  131. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +1 -0
  132. airflow/providers/google/cloud/transfers/bigquery_to_mssql.py +1 -0
  133. airflow/providers/google/cloud/transfers/bigquery_to_mysql.py +1 -0
  134. airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +1 -0
  135. airflow/providers/google/cloud/transfers/bigquery_to_sql.py +1 -0
  136. airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +1 -0
  137. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +3 -2
  138. airflow/providers/google/cloud/transfers/gcs_to_gcs.py +1 -0
  139. airflow/providers/google/cloud/transfers/gcs_to_sftp.py +1 -0
  140. airflow/providers/google/cloud/transfers/local_to_gcs.py +1 -0
  141. airflow/providers/google/cloud/transfers/mssql_to_gcs.py +1 -0
  142. airflow/providers/google/cloud/transfers/mysql_to_gcs.py +1 -0
  143. airflow/providers/google/cloud/transfers/postgres_to_gcs.py +19 -1
  144. airflow/providers/google/cloud/transfers/s3_to_gcs.py +3 -5
  145. airflow/providers/google/cloud/transfers/sftp_to_gcs.py +1 -0
  146. airflow/providers/google/cloud/transfers/sql_to_gcs.py +4 -2
  147. airflow/providers/google/cloud/triggers/bigquery.py +4 -3
  148. airflow/providers/google/cloud/triggers/cloud_batch.py +1 -1
  149. airflow/providers/google/cloud/triggers/cloud_run.py +1 -0
  150. airflow/providers/google/cloud/triggers/cloud_sql.py +2 -0
  151. airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +14 -2
  152. airflow/providers/google/cloud/triggers/dataplex.py +1 -0
  153. airflow/providers/google/cloud/triggers/dataproc.py +1 -0
  154. airflow/providers/google/cloud/triggers/kubernetes_engine.py +72 -2
  155. airflow/providers/google/cloud/triggers/mlengine.py +2 -0
  156. airflow/providers/google/cloud/triggers/pubsub.py +3 -3
  157. airflow/providers/google/cloud/triggers/vertex_ai.py +107 -15
  158. airflow/providers/google/cloud/utils/field_sanitizer.py +2 -1
  159. airflow/providers/google/cloud/utils/field_validator.py +1 -0
  160. airflow/providers/google/cloud/utils/helpers.py +1 -0
  161. airflow/providers/google/cloud/utils/mlengine_operator_utils.py +1 -0
  162. airflow/providers/google/cloud/utils/mlengine_prediction_summary.py +1 -0
  163. airflow/providers/google/cloud/utils/openlineage.py +1 -0
  164. airflow/providers/google/common/auth_backend/google_openid.py +1 -0
  165. airflow/providers/google/common/hooks/base_google.py +2 -1
  166. airflow/providers/google/common/hooks/discovery_api.py +1 -0
  167. airflow/providers/google/common/links/storage.py +1 -0
  168. airflow/providers/google/common/utils/id_token_credentials.py +1 -0
  169. airflow/providers/google/firebase/hooks/firestore.py +1 -0
  170. airflow/providers/google/get_provider_info.py +9 -3
  171. airflow/providers/google/go_module_utils.py +1 -0
  172. airflow/providers/google/leveldb/hooks/leveldb.py +8 -7
  173. airflow/providers/google/marketing_platform/example_dags/example_display_video.py +1 -0
  174. airflow/providers/google/marketing_platform/hooks/analytics_admin.py +1 -0
  175. airflow/providers/google/marketing_platform/hooks/campaign_manager.py +1 -0
  176. airflow/providers/google/marketing_platform/hooks/display_video.py +1 -0
  177. airflow/providers/google/marketing_platform/hooks/search_ads.py +1 -0
  178. airflow/providers/google/marketing_platform/operators/analytics.py +1 -0
  179. airflow/providers/google/marketing_platform/operators/analytics_admin.py +4 -2
  180. airflow/providers/google/marketing_platform/operators/campaign_manager.py +1 -0
  181. airflow/providers/google/marketing_platform/operators/display_video.py +1 -0
  182. airflow/providers/google/marketing_platform/operators/search_ads.py +1 -0
  183. airflow/providers/google/marketing_platform/sensors/campaign_manager.py +1 -0
  184. airflow/providers/google/marketing_platform/sensors/display_video.py +2 -1
  185. airflow/providers/google/marketing_platform/sensors/search_ads.py +1 -0
  186. airflow/providers/google/suite/hooks/calendar.py +1 -0
  187. airflow/providers/google/suite/hooks/drive.py +1 -0
  188. airflow/providers/google/suite/hooks/sheets.py +1 -0
  189. airflow/providers/google/suite/sensors/drive.py +1 -0
  190. airflow/providers/google/suite/transfers/gcs_to_gdrive.py +7 -0
  191. airflow/providers/google/suite/transfers/gcs_to_sheets.py +4 -1
  192. airflow/providers/google/suite/transfers/local_to_drive.py +1 -0
  193. {apache_airflow_providers_google-10.16.0rc1.dist-info → apache_airflow_providers_google-10.17.0rc1.dist-info}/METADATA +22 -17
  194. {apache_airflow_providers_google-10.16.0rc1.dist-info → apache_airflow_providers_google-10.17.0rc1.dist-info}/RECORD +196 -194
  195. {apache_airflow_providers_google-10.16.0rc1.dist-info → apache_airflow_providers_google-10.17.0rc1.dist-info}/WHEEL +0 -0
  196. {apache_airflow_providers_google-10.16.0rc1.dist-info → apache_airflow_providers_google-10.17.0rc1.dist-info}/entry_points.txt +0 -0
@@ -16,6 +16,7 @@
16
16
  # specific language governing permissions and limitations
17
17
  # under the License.
18
18
  """This module contains Google Kubernetes Engine operators."""
19
+
19
20
  from __future__ import annotations
20
21
 
21
22
  import re
@@ -28,14 +29,21 @@ import yaml
28
29
  from deprecated import deprecated
29
30
  from google.api_core.exceptions import AlreadyExists
30
31
  from google.cloud.container_v1.types import Cluster
32
+ from kubernetes.client import V1JobList
31
33
  from kubernetes.utils.create_from_yaml import FailToCreateError
34
+ from packaging.version import parse as parse_version
32
35
 
33
36
  from airflow.configuration import conf
34
37
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
35
- from airflow.providers.cncf.kubernetes.operators.job import KubernetesJobOperator
38
+ from airflow.providers.cncf.kubernetes.operators.job import KubernetesDeleteJobOperator, KubernetesJobOperator
36
39
  from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
40
+ from airflow.providers.cncf.kubernetes.operators.resource import (
41
+ KubernetesCreateResourceOperator,
42
+ KubernetesDeleteResourceOperator,
43
+ )
37
44
  from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
38
45
  from airflow.providers.google.cloud.hooks.kubernetes_engine import (
46
+ GKECustomResourceHook,
39
47
  GKEDeploymentHook,
40
48
  GKEHook,
41
49
  GKEJobHook,
@@ -45,9 +53,15 @@ from airflow.providers.google.cloud.links.kubernetes_engine import (
45
53
  KubernetesEngineClusterLink,
46
54
  KubernetesEngineJobLink,
47
55
  KubernetesEnginePodLink,
56
+ KubernetesEngineWorkloadsLink,
48
57
  )
49
58
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
50
- from airflow.providers.google.cloud.triggers.kubernetes_engine import GKEOperationTrigger, GKEStartPodTrigger
59
+ from airflow.providers.google.cloud.triggers.kubernetes_engine import (
60
+ GKEJobTrigger,
61
+ GKEOperationTrigger,
62
+ GKEStartPodTrigger,
63
+ )
64
+ from airflow.providers_manager import ProvidersManager
51
65
  from airflow.utils.timezone import utcnow
52
66
 
53
67
  if TYPE_CHECKING:
@@ -62,7 +76,7 @@ class GKEClusterAuthDetails:
62
76
  """
63
77
  Helper for fetching information about cluster for connecting.
64
78
 
65
- :param cluster_name: The name of the Google Kubernetes Engine cluster the pod should be spawned in.
79
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
66
80
  :param project_id: The Google Developers Console project id.
67
81
  :param use_internal_ip: Use the internal IP address as the endpoint.
68
82
  :param cluster_hook: airflow hook for working with kubernetes cluster.
@@ -82,7 +96,7 @@ class GKEClusterAuthDetails:
82
96
  self._cluster_url = None
83
97
  self._ssl_ca_cert = None
84
98
 
85
- def fetch_cluster_info(self) -> tuple[str, str | None]:
99
+ def fetch_cluster_info(self) -> tuple[str, str]:
86
100
  """Fetch cluster info for connecting to it."""
87
101
  cluster = self.cluster_hook.get_cluster(
88
102
  name=self.cluster_name,
@@ -533,6 +547,7 @@ class GKEStartKueueInsideClusterOperator(GoogleCloudBaseOperator):
533
547
  impersonation_chain=self.impersonation_chain,
534
548
  cluster_url=self._cluster_url,
535
549
  ssl_ca_cert=self._ssl_ca_cert,
550
+ enable_tcp_keepalive=True,
536
551
  )
537
552
 
538
553
  @staticmethod
@@ -735,6 +750,7 @@ class GKEStartPodOperator(KubernetesPodOperator):
735
750
  cluster_url=self._cluster_url,
736
751
  ssl_ca_cert=self._ssl_ca_cert,
737
752
  impersonation_chain=self.impersonation_chain,
753
+ enable_tcp_keepalive=True,
738
754
  )
739
755
  return hook
740
756
 
@@ -824,6 +840,9 @@ class GKEStartJobOperator(KubernetesJobOperator):
824
840
  Service Account Token Creator IAM role to the directly preceding identity, with first
825
841
  account from the list granting this role to the originating account (templated).
826
842
  :param location: The location param is region name.
843
+ :param deferrable: Run operator in the deferrable mode.
844
+ :param poll_interval: (Deferrable mode only) polling period in seconds to
845
+ check for the status of job.
827
846
  """
828
847
 
829
848
  template_fields: Sequence[str] = tuple(
@@ -840,6 +859,8 @@ class GKEStartJobOperator(KubernetesJobOperator):
840
859
  project_id: str | None = None,
841
860
  gcp_conn_id: str = "google_cloud_default",
842
861
  impersonation_chain: str | Sequence[str] | None = None,
862
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
863
+ job_poll_interval: float = 10.0,
843
864
  **kwargs,
844
865
  ) -> None:
845
866
  super().__init__(**kwargs)
@@ -849,6 +870,8 @@ class GKEStartJobOperator(KubernetesJobOperator):
849
870
  self.gcp_conn_id = gcp_conn_id
850
871
  self.impersonation_chain = impersonation_chain
851
872
  self.use_internal_ip = use_internal_ip
873
+ self.deferrable = deferrable
874
+ self.job_poll_interval = job_poll_interval
852
875
 
853
876
  self.job: V1Job | None = None
854
877
  self._ssl_ca_cert: str | None = None
@@ -890,6 +913,579 @@ class GKEStartJobOperator(KubernetesJobOperator):
890
913
 
891
914
  def execute(self, context: Context):
892
915
  """Execute process of creating Job."""
916
+ if self.deferrable:
917
+ kubernetes_provider = ProvidersManager().providers["apache-airflow-providers-cncf-kubernetes"]
918
+ kubernetes_provider_name = kubernetes_provider.data["package-name"]
919
+ kubernetes_provider_version = kubernetes_provider.version
920
+ min_version = "8.0.1"
921
+ if parse_version(kubernetes_provider_version) <= parse_version(min_version):
922
+ raise AirflowException(
923
+ "You are trying to use `GKEStartJobOperator` in deferrable mode with the provider "
924
+ f"package {kubernetes_provider_name}=={kubernetes_provider_version} which doesn't "
925
+ f"support this feature. Please upgrade it to version higher than {min_version}."
926
+ )
927
+
928
+ self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
929
+ cluster_name=self.cluster_name,
930
+ project_id=self.project_id,
931
+ use_internal_ip=self.use_internal_ip,
932
+ cluster_hook=self.cluster_hook,
933
+ ).fetch_cluster_info()
934
+
935
+ return super().execute(context)
936
+
937
+ def execute_deferrable(self):
938
+ self.defer(
939
+ trigger=GKEJobTrigger(
940
+ cluster_url=self._cluster_url,
941
+ ssl_ca_cert=self._ssl_ca_cert,
942
+ job_name=self.job.metadata.name, # type: ignore[union-attr]
943
+ job_namespace=self.job.metadata.namespace, # type: ignore[union-attr]
944
+ gcp_conn_id=self.gcp_conn_id,
945
+ poll_interval=self.job_poll_interval,
946
+ impersonation_chain=self.impersonation_chain,
947
+ ),
948
+ method_name="execute_complete",
949
+ )
950
+
951
+
952
+ class GKEDescribeJobOperator(GoogleCloudBaseOperator):
953
+ """
954
+ Retrieve information about Job by given name.
955
+
956
+ .. seealso::
957
+ For more information on how to use this operator, take a look at the guide:
958
+ :ref:`howto/operator:GKEDescribeJobOperator`
959
+
960
+ :param job_name: The name of the Job to delete
961
+ :param project_id: The Google Developers Console project id.
962
+ :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
963
+ resides.
964
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
965
+ :param namespace: The name of the Google Kubernetes Engine namespace.
966
+ :param use_internal_ip: Use the internal IP address as the endpoint.
967
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
968
+ :param impersonation_chain: Optional service account to impersonate using short-term
969
+ credentials, or chained list of accounts required to get the access_token
970
+ of the last account in the list, which will be impersonated in the request.
971
+ If set as a string, the account must grant the originating account
972
+ the Service Account Token Creator IAM role.
973
+ If set as a sequence, the identities from the list must grant
974
+ Service Account Token Creator IAM role to the directly preceding identity, with first
975
+ account from the list granting this role to the originating account (templated).
976
+ """
977
+
978
+ template_fields: Sequence[str] = (
979
+ "project_id",
980
+ "gcp_conn_id",
981
+ "job_name",
982
+ "namespace",
983
+ "cluster_name",
984
+ "location",
985
+ "impersonation_chain",
986
+ )
987
+ operator_extra_links = (KubernetesEngineJobLink(),)
988
+
989
+ def __init__(
990
+ self,
991
+ *,
992
+ job_name: str,
993
+ location: str,
994
+ namespace: str,
995
+ cluster_name: str,
996
+ project_id: str | None = None,
997
+ use_internal_ip: bool = False,
998
+ gcp_conn_id: str = "google_cloud_default",
999
+ impersonation_chain: str | Sequence[str] | None = None,
1000
+ **kwargs,
1001
+ ) -> None:
1002
+ super().__init__(**kwargs)
1003
+
1004
+ self.project_id = project_id
1005
+ self.gcp_conn_id = gcp_conn_id
1006
+ self.location = location
1007
+ self.job_name = job_name
1008
+ self.namespace = namespace
1009
+ self.cluster_name = cluster_name
1010
+ self.use_internal_ip = use_internal_ip
1011
+ self.impersonation_chain = impersonation_chain
1012
+
1013
+ self.job: V1Job | None = None
1014
+ self._ssl_ca_cert: str
1015
+ self._cluster_url: str
1016
+
1017
+ @cached_property
1018
+ def cluster_hook(self) -> GKEHook:
1019
+ return GKEHook(
1020
+ gcp_conn_id=self.gcp_conn_id,
1021
+ location=self.location,
1022
+ impersonation_chain=self.impersonation_chain,
1023
+ )
1024
+
1025
+ @cached_property
1026
+ def hook(self) -> GKEJobHook:
1027
+ self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1028
+ cluster_name=self.cluster_name,
1029
+ project_id=self.project_id,
1030
+ use_internal_ip=self.use_internal_ip,
1031
+ cluster_hook=self.cluster_hook,
1032
+ ).fetch_cluster_info()
1033
+
1034
+ return GKEJobHook(
1035
+ gcp_conn_id=self.gcp_conn_id,
1036
+ cluster_url=self._cluster_url,
1037
+ ssl_ca_cert=self._ssl_ca_cert,
1038
+ )
1039
+
1040
+ def execute(self, context: Context) -> None:
1041
+ self.job = self.hook.get_job(job_name=self.job_name, namespace=self.namespace)
1042
+ self.log.info(
1043
+ "Retrieved description of Job %s from cluster %s:\n %s",
1044
+ self.job_name,
1045
+ self.cluster_name,
1046
+ self.job,
1047
+ )
1048
+ KubernetesEngineJobLink.persist(context=context, task_instance=self)
1049
+ return None
1050
+
1051
+
1052
+ class GKEListJobsOperator(GoogleCloudBaseOperator):
1053
+ """
1054
+ Retrieve list of Jobs.
1055
+
1056
+ If namespace parameter is specified, the list of Jobs from dedicated
1057
+ namespace will be retrieved. If no namespace specified, it will output Jobs from all namespaces.
1058
+
1059
+ .. seealso::
1060
+ For more information on how to use this operator, take a look at the guide:
1061
+ :ref:`howto/operator:GKEListJobsOperator`
1062
+
1063
+ :param project_id: The Google Developers Console project id.
1064
+ :param location: The name of the Google Kubernetes Engine zone or region in which the cluster
1065
+ resides.
1066
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1067
+ :param namespace: The name of the Google Kubernetes Engine namespace.
1068
+ :param use_internal_ip: Use the internal IP address as the endpoint.
1069
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
1070
+ :param do_xcom_push: If set to True the result list of Jobs will be pushed to the task result.
1071
+ :param impersonation_chain: Optional service account to impersonate using short-term
1072
+ credentials, or chained list of accounts required to get the access_token
1073
+ of the last account in the list, which will be impersonated in the request.
1074
+ If set as a string, the account must grant the originating account
1075
+ the Service Account Token Creator IAM role.
1076
+ If set as a sequence, the identities from the list must grant
1077
+ Service Account Token Creator IAM role to the directly preceding identity, with first
1078
+ account from the list granting this role to the originating account (templated).
1079
+ """
1080
+
1081
+ template_fields: Sequence[str] = (
1082
+ "project_id",
1083
+ "gcp_conn_id",
1084
+ "namespace",
1085
+ "cluster_name",
1086
+ "location",
1087
+ "impersonation_chain",
1088
+ )
1089
+ operator_extra_links = (KubernetesEngineWorkloadsLink(),)
1090
+
1091
+ def __init__(
1092
+ self,
1093
+ *,
1094
+ location: str,
1095
+ cluster_name: str,
1096
+ namespace: str | None = None,
1097
+ project_id: str | None = None,
1098
+ use_internal_ip: bool = False,
1099
+ do_xcom_push: bool = True,
1100
+ gcp_conn_id: str = "google_cloud_default",
1101
+ impersonation_chain: str | Sequence[str] | None = None,
1102
+ **kwargs,
1103
+ ) -> None:
1104
+ super().__init__(**kwargs)
1105
+
1106
+ self.project_id = project_id
1107
+ self.gcp_conn_id = gcp_conn_id
1108
+ self.location = location
1109
+ self.namespace = namespace
1110
+ self.cluster_name = cluster_name
1111
+ self.use_internal_ip = use_internal_ip
1112
+ self.do_xcom_push = do_xcom_push
1113
+ self.impersonation_chain = impersonation_chain
1114
+
1115
+ self._ssl_ca_cert: str
1116
+ self._cluster_url: str
1117
+
1118
+ @cached_property
1119
+ def cluster_hook(self) -> GKEHook:
1120
+ return GKEHook(
1121
+ gcp_conn_id=self.gcp_conn_id,
1122
+ location=self.location,
1123
+ impersonation_chain=self.impersonation_chain,
1124
+ )
1125
+
1126
+ @cached_property
1127
+ def hook(self) -> GKEJobHook:
1128
+ self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1129
+ cluster_name=self.cluster_name,
1130
+ project_id=self.project_id,
1131
+ use_internal_ip=self.use_internal_ip,
1132
+ cluster_hook=self.cluster_hook,
1133
+ ).fetch_cluster_info()
1134
+
1135
+ return GKEJobHook(
1136
+ gcp_conn_id=self.gcp_conn_id,
1137
+ cluster_url=self._cluster_url,
1138
+ ssl_ca_cert=self._ssl_ca_cert,
1139
+ )
1140
+
1141
+ def execute(self, context: Context) -> dict:
1142
+ if self.namespace:
1143
+ jobs = self.hook.list_jobs_from_namespace(namespace=self.namespace)
1144
+ else:
1145
+ jobs = self.hook.list_jobs_all_namespaces()
1146
+ for job in jobs.items:
1147
+ self.log.info("Retrieved description of Job:\n %s", job)
1148
+ if self.do_xcom_push:
1149
+ ti = context["ti"]
1150
+ ti.xcom_push(key="jobs_list", value=V1JobList.to_dict(jobs))
1151
+ KubernetesEngineWorkloadsLink.persist(context=context, task_instance=self)
1152
+ return V1JobList.to_dict(jobs)
1153
+
1154
+
1155
+ class GKECreateCustomResourceOperator(KubernetesCreateResourceOperator):
1156
+ """
1157
+ Create a resource in the specified Google Kubernetes Engine cluster.
1158
+
1159
+ This Operator assumes that the system has gcloud installed and has configured a
1160
+ connection id with a service account.
1161
+
1162
+ .. seealso::
1163
+ For more detail about Kubernetes Engine authentication have a look at the reference:
1164
+ https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#internal_ip
1165
+
1166
+ .. seealso::
1167
+ For more information on how to use this operator, take a look at the guide:
1168
+ :ref:`howto/operator:GKECreateCustomResourceOperator`
1169
+
1170
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
1171
+ cluster resides, e.g. 'us-central1-a'
1172
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1173
+ :param use_internal_ip: Use the internal IP address as the endpoint.
1174
+ :param project_id: The Google Developers Console project id
1175
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
1176
+ users to specify a service account.
1177
+ :param impersonation_chain: Optional service account to impersonate using short-term
1178
+ credentials, or list of accounts required to get the access_token
1179
+ of the last account in the list, which will be impersonated in the request.
1180
+ If set as a string, the account must grant the originating account
1181
+ the Service Account Token Creator IAM role.
1182
+ If set as a sequence, the identities from the list must grant
1183
+ Service Account Token Creator IAM role to the directly preceding identity, with first
1184
+ account from the list granting this role to the originating account (templated).
1185
+ """
1186
+
1187
+ template_fields: Sequence[str] = tuple(
1188
+ {"project_id", "location", "cluster_name"} | set(KubernetesCreateResourceOperator.template_fields)
1189
+ )
1190
+
1191
+ def __init__(
1192
+ self,
1193
+ *,
1194
+ location: str,
1195
+ cluster_name: str,
1196
+ use_internal_ip: bool = False,
1197
+ project_id: str | None = None,
1198
+ gcp_conn_id: str = "google_cloud_default",
1199
+ impersonation_chain: str | Sequence[str] | None = None,
1200
+ **kwargs,
1201
+ ) -> None:
1202
+ super().__init__(**kwargs)
1203
+ self.project_id = project_id
1204
+ self.location = location
1205
+ self.cluster_name = cluster_name
1206
+ self.gcp_conn_id = gcp_conn_id
1207
+ self.impersonation_chain = impersonation_chain
1208
+ self.use_internal_ip = use_internal_ip
1209
+
1210
+ self._ssl_ca_cert: str | None = None
1211
+ self._cluster_url: str | None = None
1212
+
1213
+ if self.gcp_conn_id is None:
1214
+ raise AirflowException(
1215
+ "The gcp_conn_id parameter has become required. If you want to use Application Default "
1216
+ "Credentials (ADC) strategy for authorization, create an empty connection "
1217
+ "called `google_cloud_default`.",
1218
+ )
1219
+ # There is no need to manage the kube_config file, as it will be generated automatically.
1220
+ # All Kubernetes parameters (except config_file) are also valid for the GKEStartPodOperator.
1221
+ if self.config_file:
1222
+ raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.")
1223
+
1224
+ @cached_property
1225
+ def cluster_hook(self) -> GKEHook:
1226
+ return GKEHook(
1227
+ gcp_conn_id=self.gcp_conn_id,
1228
+ location=self.location,
1229
+ impersonation_chain=self.impersonation_chain,
1230
+ )
1231
+
1232
+ @cached_property
1233
+ def hook(self) -> GKECustomResourceHook:
1234
+ if self._cluster_url is None or self._ssl_ca_cert is None:
1235
+ raise AttributeError(
1236
+ "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
1237
+ "Try to use self.get_kube_creds method",
1238
+ )
1239
+ return GKECustomResourceHook(
1240
+ gcp_conn_id=self.gcp_conn_id,
1241
+ cluster_url=self._cluster_url,
1242
+ ssl_ca_cert=self._ssl_ca_cert,
1243
+ impersonation_chain=self.impersonation_chain,
1244
+ )
1245
+
1246
+ def execute(self, context: Context):
1247
+ """Execute process of creating Custom Resource."""
1248
+ self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1249
+ cluster_name=self.cluster_name,
1250
+ project_id=self.project_id,
1251
+ use_internal_ip=self.use_internal_ip,
1252
+ cluster_hook=self.cluster_hook,
1253
+ ).fetch_cluster_info()
1254
+ return super().execute(context)
1255
+
1256
+
1257
+ class GKEDeleteCustomResourceOperator(KubernetesDeleteResourceOperator):
1258
+ """
1259
+ Delete a resource in the specified Google Kubernetes Engine cluster.
1260
+
1261
+ This Operator assumes that the system has gcloud installed and has configured a
1262
+ connection id with a service account.
1263
+
1264
+ .. seealso::
1265
+ For more detail about Kubernetes Engine authentication have a look at the reference:
1266
+ https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#internal_ip
1267
+
1268
+ .. seealso::
1269
+ For more information on how to use this operator, take a look at the guide:
1270
+ :ref:`howto/operator:GKEDeleteCustomResourceOperator`
1271
+
1272
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
1273
+ cluster resides, e.g. 'us-central1-a'
1274
+ :param cluster_name: The name of the Google Kubernetes Engine cluster.
1275
+ :param use_internal_ip: Use the internal IP address as the endpoint.
1276
+ :param project_id: The Google Developers Console project id
1277
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
1278
+ users to specify a service account.
1279
+ :param impersonation_chain: Optional service account to impersonate using short-term
1280
+ credentials, or list of accounts required to get the access_token
1281
+ of the last account in the list, which will be impersonated in the request.
1282
+ If set as a string, the account must grant the originating account
1283
+ the Service Account Token Creator IAM role.
1284
+ If set as a sequence, the identities from the list must grant
1285
+ Service Account Token Creator IAM role to the directly preceding identity, with first
1286
+ account from the list granting this role to the originating account (templated).
1287
+ """
1288
+
1289
+ template_fields: Sequence[str] = tuple(
1290
+ {"project_id", "location", "cluster_name"} | set(KubernetesDeleteResourceOperator.template_fields)
1291
+ )
1292
+
1293
+ def __init__(
1294
+ self,
1295
+ *,
1296
+ location: str,
1297
+ cluster_name: str,
1298
+ use_internal_ip: bool = False,
1299
+ project_id: str | None = None,
1300
+ gcp_conn_id: str = "google_cloud_default",
1301
+ impersonation_chain: str | Sequence[str] | None = None,
1302
+ **kwargs,
1303
+ ) -> None:
1304
+ super().__init__(**kwargs)
1305
+ self.project_id = project_id
1306
+ self.location = location
1307
+ self.cluster_name = cluster_name
1308
+ self.gcp_conn_id = gcp_conn_id
1309
+ self.impersonation_chain = impersonation_chain
1310
+ self.use_internal_ip = use_internal_ip
1311
+
1312
+ self._ssl_ca_cert: str | None = None
1313
+ self._cluster_url: str | None = None
1314
+
1315
+ if self.gcp_conn_id is None:
1316
+ raise AirflowException(
1317
+ "The gcp_conn_id parameter has become required. If you want to use Application Default "
1318
+ "Credentials (ADC) strategy for authorization, create an empty connection "
1319
+ "called `google_cloud_default`.",
1320
+ )
1321
+ # There is no need to manage the kube_config file, as it will be generated automatically.
1322
+ # All Kubernetes parameters (except config_file) are also valid for the GKEStartPodOperator.
1323
+ if self.config_file:
1324
+ raise AirflowException("config_file is not an allowed parameter for the GKEStartPodOperator.")
1325
+
1326
+ @cached_property
1327
+ def cluster_hook(self) -> GKEHook:
1328
+ return GKEHook(
1329
+ gcp_conn_id=self.gcp_conn_id,
1330
+ location=self.location,
1331
+ impersonation_chain=self.impersonation_chain,
1332
+ )
1333
+
1334
+ @cached_property
1335
+ def hook(self) -> GKECustomResourceHook:
1336
+ if self._cluster_url is None or self._ssl_ca_cert is None:
1337
+ raise AttributeError(
1338
+ "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
1339
+ "Try to use self.get_kube_creds method",
1340
+ )
1341
+ return GKECustomResourceHook(
1342
+ gcp_conn_id=self.gcp_conn_id,
1343
+ cluster_url=self._cluster_url,
1344
+ ssl_ca_cert=self._ssl_ca_cert,
1345
+ impersonation_chain=self.impersonation_chain,
1346
+ )
1347
+
1348
+ def execute(self, context: Context):
1349
+ """Execute process of deleting Custom Resource."""
1350
+ self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
1351
+ cluster_name=self.cluster_name,
1352
+ project_id=self.project_id,
1353
+ use_internal_ip=self.use_internal_ip,
1354
+ cluster_hook=self.cluster_hook,
1355
+ ).fetch_cluster_info()
1356
+ return super().execute(context)
1357
+
1358
+
1359
+ class GKEStartKueueJobOperator(GKEStartJobOperator):
1360
+ """
1361
+ Executes a Kubernetes Job in Kueue in the specified Google Kubernetes Engine cluster.
1362
+
1363
+ :param queue_name: The name of the Queue in the cluster
1364
+ """
1365
+
1366
+ def __init__(
1367
+ self,
1368
+ *,
1369
+ queue_name: str,
1370
+ **kwargs,
1371
+ ) -> None:
1372
+ super().__init__(**kwargs)
1373
+ self.queue_name = queue_name
1374
+
1375
+ if self.suspend is False:
1376
+ raise AirflowException(
1377
+ "The `suspend` parameter can't be False. If you want to use Kueue for running Job"
1378
+ " in a Kubernetes cluster, set the `suspend` parameter to True.",
1379
+ )
1380
+ elif self.suspend is None:
1381
+ warnings.warn(
1382
+ f"You have not set parameter `suspend` in class {self.__class__.__name__}. "
1383
+ "For running a Job in Kueue the `suspend` parameter should set to True.",
1384
+ UserWarning,
1385
+ stacklevel=2,
1386
+ )
1387
+ self.suspend = True
1388
+ self.labels.update({"kueue.x-k8s.io/queue-name": queue_name})
1389
+ self.annotations.update({"kueue.x-k8s.io/queue-name": queue_name})
1390
+
1391
+
1392
+ class GKEDeleteJobOperator(KubernetesDeleteJobOperator):
1393
+ """
1394
+ Delete a Kubernetes job in the specified Google Kubernetes Engine cluster.
1395
+
1396
+ This Operator assumes that the system has gcloud installed and has configured a
1397
+ connection id with a service account.
1398
+
1399
+ The **minimum** required to define a cluster to create are the variables
1400
+ ``task_id``, ``project_id``, ``location``, ``cluster_name``, ``name``,
1401
+ ``namespace``
1402
+
1403
+ .. seealso::
1404
+ For more detail about Kubernetes Engine authentication have a look at the reference:
1405
+ https://cloud.google.com/kubernetes-engine/docs/how-to/cluster-access-for-kubectl#internal_ip
1406
+
1407
+ .. seealso::
1408
+ For more information on how to use this operator, take a look at the guide:
1409
+ :ref:`howto/operator:GKEDeleteJobOperator`
1410
+
1411
+ :param location: The name of the Google Kubernetes Engine zone or region in which the
1412
+ cluster resides, e.g. 'us-central1-a'
1413
+ :param cluster_name: The name of the Google Kubernetes Engine cluster
1414
+ :param use_internal_ip: Use the internal IP address as the endpoint.
1415
+ :param project_id: The Google Developers Console project id
1416
+ :param gcp_conn_id: The Google cloud connection id to use. This allows for
1417
+ users to specify a service account.
1418
+ :param impersonation_chain: Optional service account to impersonate using short-term
1419
+ credentials, or list of accounts required to get the access_token
1420
+ of the last account in the list, which will be impersonated in the request.
1421
+ If set as a string, the account must grant the originating account
1422
+ the Service Account Token Creator IAM role.
1423
+ If set as a sequence, the identities from the list must grant
1424
+ Service Account Token Creator IAM role to the directly preceding identity, with first
1425
+ account from the list granting this role to the originating account (templated).
1426
+ """
1427
+
1428
+ template_fields: Sequence[str] = tuple(
1429
+ {"project_id", "location", "cluster_name"} | set(KubernetesDeleteJobOperator.template_fields)
1430
+ )
1431
+
1432
+ def __init__(
1433
+ self,
1434
+ *,
1435
+ location: str,
1436
+ cluster_name: str,
1437
+ use_internal_ip: bool = False,
1438
+ project_id: str | None = None,
1439
+ gcp_conn_id: str = "google_cloud_default",
1440
+ impersonation_chain: str | Sequence[str] | None = None,
1441
+ **kwargs,
1442
+ ) -> None:
1443
+ super().__init__(**kwargs)
1444
+ self.project_id = project_id
1445
+ self.location = location
1446
+ self.cluster_name = cluster_name
1447
+ self.gcp_conn_id = gcp_conn_id
1448
+ self.impersonation_chain = impersonation_chain
1449
+ self.use_internal_ip = use_internal_ip
1450
+
1451
+ self._ssl_ca_cert: str | None = None
1452
+ self._cluster_url: str | None = None
1453
+
1454
+ if self.gcp_conn_id is None:
1455
+ raise AirflowException(
1456
+ "The gcp_conn_id parameter has become required. If you want to use Application Default "
1457
+ "Credentials (ADC) strategy for authorization, create an empty connection "
1458
+ "called `google_cloud_default`.",
1459
+ )
1460
+ # There is no need to manage the kube_config file, as it will be generated automatically.
1461
+ # All Kubernetes parameters (except config_file) are also valid for the GKEDeleteJobOperator.
1462
+ if self.config_file:
1463
+ raise AirflowException("config_file is not an allowed parameter for the GKEDeleteJobOperator.")
1464
+
1465
+ @cached_property
1466
+ def cluster_hook(self) -> GKEHook:
1467
+ return GKEHook(
1468
+ gcp_conn_id=self.gcp_conn_id,
1469
+ location=self.location,
1470
+ impersonation_chain=self.impersonation_chain,
1471
+ )
1472
+
1473
+ @cached_property
1474
+ def hook(self) -> GKEJobHook:
1475
+ if self._cluster_url is None or self._ssl_ca_cert is None:
1476
+ raise AttributeError(
1477
+ "Cluster url and ssl_ca_cert should be defined before using self.hook method. "
1478
+ "Try to use self.get_kube_creds method",
1479
+ )
1480
+
1481
+ return GKEJobHook(
1482
+ gcp_conn_id=self.gcp_conn_id,
1483
+ cluster_url=self._cluster_url,
1484
+ ssl_ca_cert=self._ssl_ca_cert,
1485
+ )
1486
+
1487
+ def execute(self, context: Context):
1488
+ """Execute process of deleting Job."""
893
1489
  self._cluster_url, self._ssl_ca_cert = GKEClusterAuthDetails(
894
1490
  cluster_name=self.cluster_name,
895
1491
  project_id=self.project_id,
@@ -16,6 +16,7 @@
16
16
  # specific language governing permissions and limitations
17
17
  # under the License.
18
18
  """Operators that interact with Google Cloud Life Sciences service."""
19
+
19
20
  from __future__ import annotations
20
21
 
21
22
  from typing import TYPE_CHECKING, Sequence