Added DataprepGetJobsForJobGroupOperator (#10246)

apache · Aug 9, 2020 · ef08831 · ef08831
1 parent 06a1836
commit ef08831
Show file tree

Hide file tree

Showing 8 changed files with 366 additions and 0 deletions.
diff --git a/airflow/models/connection.py b/airflow/models/connection.py
@@ -50,6 +50,7 @@
  ),
  "cassandra": ("airflow.providers.apache.cassandra.hooks.cassandra.CassandraHook", "cassandra_conn_id"),
  "cloudant": ("airflow.providers.cloudant.hooks.cloudant.CloudantHook", "cloudant_conn_id"),
+ "dataprep": ("airflow.providers.google.cloud.hooks.dataprep.GoogleDataprepHook", "dataprep_conn_id"),
  "docker": ("airflow.providers.docker.hooks.docker.DockerHook", "docker_conn_id"),
  "elasticsearch": (
  "airflow.providers.elasticsearch.hooks.elasticsearch.ElasticsearchHook",

diff --git a/airflow/providers/google/cloud/example_dags/example_dataprep.py b/airflow/providers/google/cloud/example_dags/example_dataprep.py
@@ -0,0 +1,37 @@
+# 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.
+"""
+Example Airflow DAG that shows how to use Google Dataprep.
+"""
+
+from airflow import models
+from airflow.providers.google.cloud.operators.dataprep import DataprepGetJobsForJobGroupOperator
+from airflow.utils import dates
+
+JOB_ID = 6269792
+
+with models.DAG(
+ "example_dataprep",
+ schedule_interval=None, # Override to match your needs
+ start_date=dates.days_ago(1)
+) as dag:
+
+ # [START how_to_dataprep_get_jobs_for_job_group_operator]
+ get_jobs_for_job_group = DataprepGetJobsForJobGroupOperator(
+ task_id="get_jobs_for_job_group", job_id=JOB_ID
+ )
+ # [END how_to_dataprep_get_jobs_for_job_group_operator]
diff --git a/airflow/providers/google/cloud/hooks/dataprep.py b/airflow/providers/google/cloud/hooks/dataprep.py
@@ -0,0 +1,75 @@
+#
+# 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 Dataprep hook.
+"""
+from typing import Any, Dict
+
+import requests
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+from airflow import AirflowException
+from airflow.hooks.base_hook import BaseHook
+
+
+class GoogleDataprepHook(BaseHook):
+ """
+ Hook for connection with Dataprep API.
+ To get connection Dataprep with Airflow you need Dataprep token.
+ https://clouddataprep.com/documentation/api#section/Authentication
+
+ It should be added to the Connection in Airflow in JSON format.
+
+ """
+
+ def __init__(self, dataprep_conn_id: str = "dataprep_conn_id") -> None:
+ super().__init__()
+ self.dataprep_conn_id = dataprep_conn_id
+ self._url = "https://api.clouddataprep.com/v4/jobGroups"
+
+ @property
+ def _headers(self) -> Dict[str, str]:
+ headers = {
+ "Content-Type": "application/json",
+ "Authorization": f"Bearer {self._token}",
+ }
+ return headers
+
+ @property
+ def _token(self) -> str:
+ conn = self.get_connection(self.dataprep_conn_id)
+ token = conn.extra_dejson.get("token")
+ if token is None:
+ raise AirflowException(
+ "Dataprep token is missing or has invalid format. "
+ "Please make sure that Dataprep token is added to the Airflow Connections."
+ )
+ return token
+
+ @retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, max=10))
+ def get_jobs_for_job_group(self, job_id: int) -> Dict[str, Any]:
+ """
+ Get information about the batch jobs within a Cloud Dataprep job.
+
+ :param job_id The ID of the job that will be fetched.
+ :type job_id: int
+ """
+ url: str = f"{self._url}/{job_id}/jobs"
+ response = requests.get(url, headers=self._headers)
+ response.raise_for_status()
+ return response.json()
diff --git a/airflow/providers/google/cloud/operators/dataprep.py b/airflow/providers/google/cloud/operators/dataprep.py
@@ -0,0 +1,56 @@
+#
+# 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 a Google Dataprep operator.
+"""
+
+from typing import Dict
+
+from airflow.models import BaseOperator
+from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook
+from airflow.utils.decorators import apply_defaults
+
+
+class DataprepGetJobsForJobGroupOperator(BaseOperator):
+ """
+ Get information about the batch jobs within a Cloud Dataprep job.
+ API documentation https://clouddataprep.com/documentation/api#section/Overview
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:DataprepGetJobsForJobGroupOperator`
+
+
+ :param job_id The ID of the job that will be requests
+ :type job_id: int
+ """
+
+ template_fields = ("job_id",)
+
+ @apply_defaults
+ def __init__(
+ self, *, job_id: int, **kwargs
+ ) -> None:
+ super().__init__(**kwargs)
+ self.job_id = job_id
+
+ def execute(self, context: Dict):
+ self.log.info("Fetching data for job with id: %d ...", self.job_id)
+ hook = GoogleDataprepHook(dataprep_conn_id="dataprep_conn_id")
+ response = hook.get_jobs_for_job_group(job_id=self.job_id)
+ return response
diff --git a/docs/howto/operator/google/cloud/dataprep.rst b/docs/howto/operator/google/cloud/dataprep.rst
@@ -0,0 +1,60 @@
+ .. 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.
+
+Google Dataprep Operators
+=========================
+`Google Dataprep API documentation is available here <https://cloud.google.com/dataprep/docs/html/API-Reference_145281441>`__
+
+.. contents::
+ :depth: 1
+ :local:
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include:: /howto/operator/google/_partials/prerequisite_tasks.rst
+
+.. _howto/operator:DataprepGetJobsForJobGroupOperator:
+
+Get Jobs For Job Group
+^^^^^^^^^^^^^^^^^^^^^^
+
+To get information about jobs within a Cloud Dataprep job use:
+:class:`~airflow.providers.google.cloud.operators.dataprep.DataprepGetJobsForJobGroupOperator`
+
+To get connection Dataprep with Airflow you need Dataprep token.
+Please follow Dataprep instructions.
+https://clouddataprep.com/documentation/api#section/Authentication
+
+It should be added to the Connection in Airflow in JSON format.
+Her you can check how to do such connection:
+https://airflow.readthedocs.io/en/stable/howto/connection/index.html#editing-a-connection-with-the-ui
+
+See following example:
+Set values for these fields:
+.. code-block::
+
+ Conn Id: "your_conn_id"
+ Extra: "{\"token\": \"TOKEN\"}
+
+Example usage:
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_dataprep.py
+ :language: python
+ :dedent: 4
+ :start-after: [START how_to_dataprep_get_jobs_for_job_group_operator]
+ :end-before: [END how_to_dataprep_get_jobs_for_job_group_operator]
diff --git a/docs/operators-and-hooks-ref.rst b/docs/operators-and-hooks-ref.rst
@@ -720,6 +720,12 @@ These integrations allow you to perform various operations within the Google Clo
  - :mod:`airflow.providers.google.cloud.operators.dataflow`
  -
 
+ * - `Dataprep <https://cloud.google.com/dataprep/>`__
+ - :doc:`How to use <howto/operator/google/cloud/dataprep>`
+ - :mod:`airflow.providers.google.cloud.hooks.dataprep`
+ - :mod:`airflow.providers.google.cloud.operators.dataprep`
+ -
+
  * - `Dataproc <https://cloud.google.com/dataproc/>`__
  - :doc:`How to use <howto/operator/google/cloud/dataproc>`
  - :mod:`airflow.providers.google.cloud.hooks.dataproc`

diff --git a/tests/providers/google/cloud/hooks/test_dataprep.py b/tests/providers/google/cloud/hooks/test_dataprep.py
@@ -0,0 +1,97 @@
+#
+# 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 unittest import mock
+
+import pytest
+from mock import patch
+from requests import HTTPError
+from tenacity import RetryError
+
+from airflow.providers.google.cloud.hooks import dataprep
+
+JOB_ID = 1234567
+URL = "https://api.clouddataprep.com/v4/jobGroups"
+TOKEN = "1111"
+EXTRA = {"token": TOKEN}
+
+
+@pytest.fixture(scope="class")
+def mock_hook():
+ with mock.patch("airflow.hooks.base_hook.BaseHook.get_connection") as conn:
+ hook = dataprep.GoogleDataprepHook(dataprep_conn_id="dataprep_conn_id")
+ conn.return_value.extra_dejson = EXTRA
+ yield hook
+
+
+class TestGoogleDataprepHook:
+ def test_get_token(self, mock_hook):
+ assert mock_hook._token == TOKEN
+
+ @patch("airflow.providers.google.cloud.hooks.dataprep.requests.get")
+ def test_mock_should_be_called_once_with_params(self, mock_get_request, mock_hook):
+ mock_hook.get_jobs_for_job_group(job_id=JOB_ID)
+ mock_get_request.assert_called_once_with(
+ f"{URL}/{JOB_ID}/jobs",
+ headers={
+ "Content-Type": "application/json",
+ "Authorization": f"Bearer {TOKEN}",
+ },
+ )
+
+ @patch(
+ "airflow.providers.google.cloud.hooks.dataprep.requests.get",
+ side_effect=[HTTPError(), mock.MagicMock()],
+ )
+ def test_should_pass_after_retry(self, mock_get_request, mock_hook):
+ mock_hook.get_jobs_for_job_group(JOB_ID)
+ assert mock_get_request.call_count == 2
+
+ @patch(
+ "airflow.providers.google.cloud.hooks.dataprep.requests.get",
+ side_effect=[mock.MagicMock(), HTTPError()],
+ )
+ def test_should_not_retry_after_success(self, mock_get_request, mock_hook):
+ mock_hook.get_jobs_for_job_group.retry.sleep = mock.Mock() # pylint: disable=no-member
+ mock_hook.get_jobs_for_job_group(JOB_ID)
+ assert mock_get_request.call_count == 1
+
+ @patch(
+ "airflow.providers.google.cloud.hooks.dataprep.requests.get",
+ side_effect=[
+ HTTPError(),
+ HTTPError(),
+ HTTPError(),
+ HTTPError(),
+ mock.MagicMock(),
+ ],
+ )
+ def test_should_retry_after_four_errors(self, mock_get_request, mock_hook):
+ mock_hook.get_jobs_for_job_group.retry.sleep = mock.Mock() # pylint: disable=no-member
+ mock_hook.get_jobs_for_job_group(JOB_ID)
+ assert mock_get_request.call_count == 5
+
+ @patch(
+ "airflow.providers.google.cloud.hooks.dataprep.requests.get",
+ side_effect=[HTTPError(), HTTPError(), HTTPError(), HTTPError(), HTTPError()],
+ )
+ def test_raise_error_after_five_calls(self, mock_get_request, mock_hook):
+ with pytest.raises(RetryError) as err:
+ mock_hook.get_jobs_for_job_group.retry.sleep = mock.Mock() # pylint: disable=no-member
+ mock_hook.get_jobs_for_job_group(JOB_ID)
+ assert "HTTPError" in str(err)
+ assert mock_get_request.call_count == 5
diff --git a/tests/providers/google/cloud/operators/test_dataprep.py b/tests/providers/google/cloud/operators/test_dataprep.py
@@ -0,0 +1,34 @@
+#
+# 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 unittest import TestCase, mock
+
+from airflow.providers.google.cloud.operators.dataprep import DataprepGetJobsForJobGroupOperator
+
+JOB_ID = 143
+TASK_ID = "dataprep_job"
+
+
+class TestDataprepGetJobsForJobGroupOperator(TestCase):
+ @mock.patch(
+ "airflow.providers.google.cloud.operators.dataprep.GoogleDataprepHook"
+ )
+ def test_execute(self, hook_mock):
+ op = DataprepGetJobsForJobGroupOperator(job_id=JOB_ID, task_id=TASK_ID)
+ op.execute(context={})
+ hook_mock.assert_called_once_with(dataprep_conn_id='dataprep_conn_id')
+ hook_mock.return_value.get_jobs_for_job_group.assert_called_once_with(job_id=JOB_ID)