Add deferrable mode to DataFusionStartPipelineOperator (#28690)

apache · Jan 20, 2023 · e926bb9 · e926bb9
1 parent 5fcdd32
commit e926bb9
Show file tree

Hide file tree

Showing 16 changed files with 1,313 additions and 215 deletions.
diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py
@@ -21,14 +21,20 @@
 import os
 from time import monotonic, sleep
 from typing import Any, Dict, Sequence
-from urllib.parse import quote, urlencode
+from urllib.parse import quote, urlencode, urljoin
 
 import google.auth
+from aiohttp import ClientSession
+from gcloud.aio.auth import AioSession, Token
 from google.api_core.retry import exponential_sleep_generator
 from googleapiclient.discovery import Resource, build
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
+from airflow.exceptions import AirflowException, AirflowNotFoundException
+from airflow.providers.google.common.hooks.base_google import (
+ PROVIDE_PROJECT_ID,
+ GoogleBaseAsyncHook,
+ GoogleBaseHook,
+)
 
 Operation = Dict[str, Any]
 
@@ -154,12 +160,14 @@ def _cdap_request(
 
  @staticmethod
  def _check_response_status_and_data(response, message: str) -> None:
- if response.status != 200:
+ if response.status == 404:
+ raise AirflowNotFoundException(message)
+ elif response.status != 200:
  raise AirflowException(message)
  if response.data is None:
  raise AirflowException(
  "Empty response received. Please, check for possible root "
- "causes of this behavior either in DAG code or on Cloud Datafusion side"
+ "causes of this behavior either in DAG code or on Cloud DataFusion side"
  )
 
  def get_conn(self) -> Resource:
@@ -418,7 +426,7 @@ def start_pipeline(
  :param pipeline_name: Your pipeline name.
  :param instance_url: Endpoint on which the REST APIs is accessible for the instance.
  :param runtime_args: Optional runtime JSON args to be passed to the pipeline
- :param namespace: f your pipeline belongs to a Basic edition instance, the namespace ID
+ :param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID
  is always default. If your pipeline belongs to an Enterprise edition instance, you
  can create a namespace.
  """
@@ -469,3 +477,88 @@ def stop_pipeline(self, pipeline_name: str, instance_url: str, namespace: str =
  self._check_response_status_and_data(
  response, f"Stopping a pipeline failed with code {response.status}"
  )
+
+
+class DataFusionAsyncHook(GoogleBaseAsyncHook):
+ """Class to get asynchronous hook for DataFusion"""
+
+ sync_hook_class = DataFusionHook
+ scopes = ["https://www.googleapis.com/auth/cloud-platform"]
+
+ @staticmethod
+ def _base_url(instance_url: str, namespace: str) -> str:
+ return urljoin(f"{instance_url}/", f"v3/namespaces/{quote(namespace)}/apps/")
+
+ async def _get_link(self, url: str, session):
+ async with Token(scopes=self.scopes) as token:
+ session_aio = AioSession(session)
+ headers = {
+ "Authorization": f"Bearer {await token.get()}",
+ }
+ try:
+ pipeline = await session_aio.get(url=url, headers=headers)
+ except AirflowException:
+ pass # Because the pipeline may not be visible in system yet
+
+ return pipeline
+
+ async def get_pipeline(
+ self,
+ instance_url: str,
+ namespace: str,
+ pipeline_name: str,
+ pipeline_id: str,
+ session,
+ ):
+ base_url_link = self._base_url(instance_url, namespace)
+ url = urljoin(
+ base_url_link, f"{quote(pipeline_name)}/workflows/DataPipelineWorkflow/runs/{quote(pipeline_id)}"
+ )
+ return await self._get_link(url=url, session=session)
+
+ async def get_pipeline_status(
+ self,
+ pipeline_name: str,
+ instance_url: str,
+ pipeline_id: str,
+ namespace: str = "default",
+ success_states: list[str] | None = None,
+ ) -> str:
+ """
+ Gets a Cloud Data Fusion pipeline status asynchronously.
+
+ :param pipeline_name: Your pipeline name.
+ :param instance_url: Endpoint on which the REST APIs is accessible for the instance.
+ :param pipeline_id: Unique pipeline ID associated with specific pipeline
+ :param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID
+ is always default. If your pipeline belongs to an Enterprise edition instance, you
+ can create a namespace.
+ :param success_states: If provided the operator will wait for pipeline to be in one of
+ the provided states.
+ """
+ success_states = success_states or SUCCESS_STATES
+ async with ClientSession() as session:
+ try:
+ pipeline = await self.get_pipeline(
+ instance_url=instance_url,
+ namespace=namespace,
+ pipeline_name=pipeline_name,
+ pipeline_id=pipeline_id,
+ session=session,
+ )
+ self.log.info("Response pipeline: %s", pipeline)
+ pipeline = await pipeline.json(content_type=None)
+ current_pipeline_state = pipeline["status"]
+
+ if current_pipeline_state in success_states:
+ pipeline_status = "success"
+ elif current_pipeline_state in FAILURE_STATES:
+ pipeline_status = "failed"
+ else:
+ pipeline_status = "pending"
+ except OSError:
+ pipeline_status = "pending"
+ except Exception as e:
+ self.log.info("Retrieving pipeline status finished with errors...")
+ pipeline_status = str(e)
+ return pipeline_status
diff --git a/airflow/providers/google/cloud/links/datafusion.py b/airflow/providers/google/cloud/links/datafusion.py
@@ -0,0 +1,130 @@
+#
+# 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.
+"""This module contains Google Compute Engine links."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+from airflow.models import BaseOperatorLink, XCom
+
+if TYPE_CHECKING:
+ from airflow.models import BaseOperator
+ from airflow.models.taskinstance import TaskInstanceKey
+ from airflow.utils.context import Context
+
+
+BASE_LINK = "https://console.cloud.google.com/data-fusion"
+DATAFUSION_INSTANCE_LINK = BASE_LINK + "/locations/{region}/instances/{instance_name}?project={project_id}"
+DATAFUSION_PIPELINES_LINK = "{uri}/cdap/ns/default/pipelines"
+DATAFUSION_PIPELINE_LINK = "{uri}/pipelines/ns/default/view/{pipeline_name}"
+
+
+class BaseGoogleLink(BaseOperatorLink):
+ """
+ Override the base logic to prevent adding 'https://console.cloud.google.com'
+ in front of every link where uri is used
+ """
+
+ name: ClassVar[str]
+ key: ClassVar[str]
+ format_str: ClassVar[str]
+
+ def get_link(
+ self,
+ operator: BaseOperator,
+ *,
+ ti_key: TaskInstanceKey,
+ ) -> str:
+ conf = XCom.get_value(key=self.key, ti_key=ti_key)
+ if not conf:
+ return ""
+ if self.format_str.startswith("http"):
+ return self.format_str.format(**conf)
+ return self.format_str.format(**conf)
+
+
+class DataFusionInstanceLink(BaseGoogleLink):
+ """Helper class for constructing Data Fusion Instance link"""
+
+ name = "Data Fusion Instance"
+ key = "instance_conf"
+ format_str = DATAFUSION_INSTANCE_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance: BaseOperator,
+ location: str,
+ instance_name: str,
+ project_id: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=DataFusionInstanceLink.key,
+ value={
+ "region": location,
+ "instance_name": instance_name,
+ "project_id": project_id,
+ },
+ )
+
+
+class DataFusionPipelineLink(BaseGoogleLink):
+ """Helper class for constructing Data Fusion Pipeline link"""
+
+ name = "Data Fusion Pipeline"
+ key = "pipeline_conf"
+ format_str = DATAFUSION_PIPELINE_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance: BaseOperator,
+ uri: str,
+ pipeline_name: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=DataFusionPipelineLink.key,
+ value={
+ "uri": uri,
+ "pipeline_name": pipeline_name,
+ },
+ )
+
+
+class DataFusionPipelinesLink(BaseGoogleLink):
+ """Helper class for constructing list of Data Fusion Pipelines link"""
+
+ name = "Data Fusion Pipelines List"
+ key = "pipelines_conf"
+ format_str = DATAFUSION_PIPELINES_LINK
+
+ @staticmethod
+ def persist(
+ context: Context,
+ task_instance: BaseOperator,
+ uri: str,
+ ):
+ task_instance.xcom_push(
+ context=context,
+ key=DataFusionPipelinesLink.key,
+ value={
+ "uri": uri,
+ },
+ )