Add CloudBatchHook and operators (#32606)

apache · Aug 18, 2023 · c6a014a · c6a014a
1 parent 2c0fa0c
commit c6a014a
Show file tree

Hide file tree

Showing 11 changed files with 1,694 additions and 0 deletions.
diff --git a/airflow/providers/google/cloud/hooks/cloud_batch.py b/airflow/providers/google/cloud/hooks/cloud_batch.py
@@ -0,0 +1,204 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import itertools
+import json
+from time import sleep
+from typing import Iterable, Sequence
+
+from google.api_core import operation # type: ignore
+from google.cloud.batch import ListJobsRequest, ListTasksRequest
+from google.cloud.batch_v1 import (
+ BatchServiceAsyncClient,
+ BatchServiceClient,
+ CreateJobRequest,
+ Job,
+ JobStatus,
+ Task,
+)
+from google.cloud.batch_v1.services.batch_service import pagers
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+
+
+class CloudBatchHook(GoogleBaseHook):
+ """
+ Hook for the Google Cloud Batch service.
+
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :param impersonation_chain: Optional service account to impersonate using short-term
+ credentials, or chained list of accounts required to get the access_token
+ of the last account in the list, which will be impersonated in the request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding identity, with first
+ account from the list granting this role to the originating account.
+ """
+
+ def __init__(
+ self,
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ ) -> None:
+ super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain)
+ self._client: BatchServiceClient | None = None
+
+ def get_conn(self):
+ """
+ Retrieves connection to GCE Batch.
+
+ :return: Google Batch Service client object.
+ """
+ if self._client is None:
+ self._client = BatchServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
+ return self._client
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def submit_batch_job(
+ self, job_name: str, job: Job, region: str, project_id: str = PROVIDE_PROJECT_ID
+ ) -> Job:
+ if isinstance(job, dict):
+ job = Job.from_json(json.dumps(job))
+
+ create_request = CreateJobRequest()
+ create_request.job = job
+ create_request.job_id = job_name
+ create_request.parent = f"projects/{project_id}/locations/{region}"
+
+ return self.get_conn().create_job(create_request)
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def delete_job(
+ self, job_name: str, region: str, project_id: str = PROVIDE_PROJECT_ID
+ ) -> operation.Operation:
+ return self.get_conn().delete_job(name=f"projects/{project_id}/locations/{region}/jobs/{job_name}")
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_jobs(
+ self,
+ region: str,
+ project_id: str = PROVIDE_PROJECT_ID,
+ filter: str | None = None,
+ limit: int | None = None,
+ ) -> Iterable[Job]:
+
+ if limit is not None and limit < 0:
+ raise AirflowException("The limit for the list jobs request should be greater or equal to zero")
+
+ list_jobs_request: ListJobsRequest = ListJobsRequest(
+ parent=f"projects/{project_id}/locations/{region}", filter=filter
+ )
+
+ jobs: pagers.ListJobsPager = self.get_conn().list_jobs(request=list_jobs_request)
+
+ return list(itertools.islice(jobs, limit))
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def list_tasks(
+ self,
+ region: str,
+ job_name: str,
+ project_id: str = PROVIDE_PROJECT_ID,
+ group_name: str = "group0",
+ filter: str | None = None,
+ limit: int | None = None,
+ ) -> Iterable[Task]:
+
+ if limit is not None and limit < 0:
+ raise AirflowException("The limit for the list tasks request should be greater or equal to zero")
+
+ list_tasks_request: ListTasksRequest = ListTasksRequest(
+ parent=f"projects/{project_id}/locations/{region}/jobs/{job_name}/taskGroups/{group_name}",
+ filter=filter,
+ )
+
+ tasks: pagers.ListTasksPager = self.get_conn().list_tasks(request=list_tasks_request)
+
+ return list(itertools.islice(tasks, limit))
+
+ def wait_for_job(
+ self, job_name: str, polling_period_seconds: float = 10, timeout: float | None = None
+ ) -> Job:
+ client = self.get_conn()
+ while timeout is None or timeout > 0:
+ try:
+ job = client.get_job(name=f"{job_name}")
+ status: JobStatus.State = job.status.state
+ if (
+ status == JobStatus.State.SUCCEEDED
+ or status == JobStatus.State.FAILED
+ or status == JobStatus.State.DELETION_IN_PROGRESS
+ ):
+ return job
+ else:
+ sleep(polling_period_seconds)
+ except Exception as e:
+ self.log.exception("Exception occurred while checking for job completion.")
+ raise e
+
+ if timeout is not None:
+ timeout -= polling_period_seconds
+
+ raise AirflowException(f"Job with name [{job_name}] timed out")
+
+ def get_job(self, job_name) -> Job:
+ return self.get_conn().get_job(name=job_name)
+
+
+class CloudBatchAsyncHook(GoogleBaseHook):
+ """
+ Async hook for the Google Cloud Batch service.
+
+ :param gcp_conn_id: The connection ID to use when fetching connection info.
+ :param impersonation_chain: Optional service account to impersonate using short-term
+ credentials, or chained list of accounts required to get the access_token
+ of the last account in the list, which will be impersonated in the request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding identity, with first
+ account from the list granting this role to the originating account.
+ """
+
+ def __init__(
+ self,
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+
+ self._client: BatchServiceAsyncClient | None = None
+ super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain)
+
+ def get_conn(self):
+ if self._client is None:
+ self._client = BatchServiceAsyncClient(
+ credentials=self.get_credentials(), client_info=CLIENT_INFO
+ )
+
+ return self._client
+
+ async def get_batch_job(
+ self,
+ job_name: str,
+ ) -> Job:
+ client = self.get_conn()
+ return await client.get_job(name=f"{job_name}")