Skip to content

Commit ef08831

Browse files
michalslowikowski00michalslowikowski00
authored
Added DataprepGetJobsForJobGroupOperator (#10246)
1 parent 06a1836 commit ef08831

8 files changed

Lines changed: 366 additions & 0 deletions

File tree

‎airflow/models/connection.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
),
5151
"cassandra": ("airflow.providers.apache.cassandra.hooks.cassandra.CassandraHook", "cassandra_conn_id"),
5252
"cloudant": ("airflow.providers.cloudant.hooks.cloudant.CloudantHook", "cloudant_conn_id"),
53+
"dataprep": ("airflow.providers.google.cloud.hooks.dataprep.GoogleDataprepHook", "dataprep_conn_id"),
5354
"docker": ("airflow.providers.docker.hooks.docker.DockerHook", "docker_conn_id"),
5455
"elasticsearch": (
5556
"airflow.providers.elasticsearch.hooks.elasticsearch.ElasticsearchHook",

‎airflow/providers/google/cloud/example_dags/example_dataprep.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
Example Airflow DAG that shows how to use Google Dataprep.
19+
"""
20+
21+
from airflow import models
22+
from airflow.providers.google.cloud.operators.dataprep import DataprepGetJobsForJobGroupOperator
23+
from airflow.utils import dates
24+
25+
JOB_ID = 6269792
26+
27+
with models.DAG(
28+
"example_dataprep",
29+
schedule_interval=None, # Override to match your needs
30+
start_date=dates.days_ago(1)
31+
) as dag:
32+
33+
# [START how_to_dataprep_get_jobs_for_job_group_operator]
34+
get_jobs_for_job_group = DataprepGetJobsForJobGroupOperator(
35+
task_id="get_jobs_for_job_group", job_id=JOB_ID
36+
)
37+
# [END how_to_dataprep_get_jobs_for_job_group_operator]

‎airflow/providers/google/cloud/hooks/dataprep.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
"""
19+
This module contains Google Dataprep hook.
20+
"""
21+
from typing import Any, Dict
22+
23+
import requests
24+
from tenacity import retry, stop_after_attempt, wait_exponential
25+
26+
from airflow import AirflowException
27+
from airflow.hooks.base_hook import BaseHook
28+
29+
30+
class GoogleDataprepHook(BaseHook):
31+
"""
32+
Hook for connection with Dataprep API.
33+
To get connection Dataprep with Airflow you need Dataprep token.
34+
https://clouddataprep.com/documentation/api#section/Authentication
35+
36+
It should be added to the Connection in Airflow in JSON format.
37+
38+
"""
39+
40+
def __init__(self, dataprep_conn_id: str = "dataprep_conn_id") -> None:
41+
super().__init__()
42+
self.dataprep_conn_id = dataprep_conn_id
43+
self._url = "https://api.clouddataprep.com/v4/jobGroups"
44+
45+
@property
46+
def _headers(self) -> Dict[str, str]:
47+
headers = {
48+
"Content-Type": "application/json",
49+
"Authorization": f"Bearer {self._token}",
50+
}
51+
return headers
52+
53+
@property
54+
def _token(self) -> str:
55+
conn = self.get_connection(self.dataprep_conn_id)
56+
token = conn.extra_dejson.get("token")
57+
if token is None:
58+
raise AirflowException(
59+
"Dataprep token is missing or has invalid format. "
60+
"Please make sure that Dataprep token is added to the Airflow Connections."
61+
)
62+
return token
63+
64+
@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, max=10))
65+
def get_jobs_for_job_group(self, job_id: int) -> Dict[str, Any]:
66+
"""
67+
Get information about the batch jobs within a Cloud Dataprep job.
68+
69+
:param job_id The ID of the job that will be fetched.
70+
:type job_id: int
71+
"""
72+
url: str = f"{self._url}/{job_id}/jobs"
73+
response = requests.get(url, headers=self._headers)
74+
response.raise_for_status()
75+
return response.json()

‎airflow/providers/google/cloud/operators/dataprep.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
"""
19+
This module contains a Google Dataprep operator.
20+
"""
21+
22+
from typing import Dict
23+
24+
from airflow.models import BaseOperator
25+
from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook
26+
from airflow.utils.decorators import apply_defaults
27+
28+
29+
class DataprepGetJobsForJobGroupOperator(BaseOperator):
30+
"""
31+
Get information about the batch jobs within a Cloud Dataprep job.
32+
API documentation https://clouddataprep.com/documentation/api#section/Overview
33+
34+
.. seealso::
35+
For more information on how to use this operator, take a look at the guide:
36+
:ref:`howto/operator:DataprepGetJobsForJobGroupOperator`
37+
38+
39+
:param job_id The ID of the job that will be requests
40+
:type job_id: int
41+
"""
42+
43+
template_fields = ("job_id",)
44+
45+
@apply_defaults
46+
def __init__(
47+
self, *, job_id: int, **kwargs
48+
) -> None:
49+
super().__init__(**kwargs)
50+
self.job_id = job_id
51+
52+
def execute(self, context: Dict):
53+
self.log.info("Fetching data for job with id: %d ...", self.job_id)
54+
hook = GoogleDataprepHook(dataprep_conn_id="dataprep_conn_id")
55+
response = hook.get_jobs_for_job_group(job_id=self.job_id)
56+
return response

‎docs/howto/operator/google/cloud/dataprep.rst‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
Google Dataprep Operators
19+
=========================
20+
`Google Dataprep API documentation is available here <https://cloud.google.com/dataprep/docs/html/API-Reference_145281441>`__
21+
22+
.. contents::
23+
:depth: 1
24+
:local:
25+
26+
Prerequisite Tasks
27+
^^^^^^^^^^^^^^^^^^
28+
29+
.. include:: /howto/operator/google/_partials/prerequisite_tasks.rst
30+
31+
.. _howto/operator:DataprepGetJobsForJobGroupOperator:
32+
33+
Get Jobs For Job Group
34+
^^^^^^^^^^^^^^^^^^^^^^
35+
36+
To get information about jobs within a Cloud Dataprep job use:
37+
:class:`~airflow.providers.google.cloud.operators.dataprep.DataprepGetJobsForJobGroupOperator`
38+
39+
To get connection Dataprep with Airflow you need Dataprep token.
40+
Please follow Dataprep instructions.
41+
https://clouddataprep.com/documentation/api#section/Authentication
42+
43+
It should be added to the Connection in Airflow in JSON format.
44+
Her you can check how to do such connection:
45+
https://airflow.readthedocs.io/en/stable/howto/connection/index.html#editing-a-connection-with-the-ui
46+
47+
See following example:
48+
Set values for these fields:
49+
.. code-block::
50+
51+
Conn Id: "your_conn_id"
52+
Extra: "{\"token\": \"TOKEN\"}
53+
54+
Example usage:
55+
56+
.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_dataprep.py
57+
:language: python
58+
:dedent: 4
59+
:start-after: [START how_to_dataprep_get_jobs_for_job_group_operator]
60+
:end-before: [END how_to_dataprep_get_jobs_for_job_group_operator]

‎docs/operators-and-hooks-ref.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -720,6 +720,12 @@ These integrations allow you to perform various operations within the Google Clo
720720
- :mod:`airflow.providers.google.cloud.operators.dataflow`
721721
-
722722

723+
* - `Dataprep <https://cloud.google.com/dataprep/>`__
724+
- :doc:`How to use <howto/operator/google/cloud/dataprep>`
725+
- :mod:`airflow.providers.google.cloud.hooks.dataprep`
726+
- :mod:`airflow.providers.google.cloud.operators.dataprep`
727+
-
728+
723729
* - `Dataproc <https://cloud.google.com/dataproc/>`__
724730
- :doc:`How to use <howto/operator/google/cloud/dataproc>`
725731
- :mod:`airflow.providers.google.cloud.hooks.dataproc`

‎tests/providers/google/cloud/hooks/test_dataprep.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
from unittest import mock
19+
20+
import pytest
21+
from mock import patch
22+
from requests import HTTPError
23+
from tenacity import RetryError
24+
25+
from airflow.providers.google.cloud.hooks import dataprep
26+
27+
JOB_ID = 1234567
28+
URL = "https://api.clouddataprep.com/v4/jobGroups"
29+
TOKEN = "1111"
30+
EXTRA = {"token": TOKEN}
31+
32+
33+
@pytest.fixture(scope="class")
34+
def mock_hook():
35+
with mock.patch("airflow.hooks.base_hook.BaseHook.get_connection") as conn:
36+
hook = dataprep.GoogleDataprepHook(dataprep_conn_id="dataprep_conn_id")
37+
conn.return_value.extra_dejson = EXTRA
38+
yield hook
39+
40+
41+
class TestGoogleDataprepHook:
42+
def test_get_token(self, mock_hook):
43+
assert mock_hook._token == TOKEN
44+
45+
@patch("airflow.providers.google.cloud.hooks.dataprep.requests.get")
46+
def test_mock_should_be_called_once_with_params(self, mock_get_request, mock_hook):
47+
mock_hook.get_jobs_for_job_group(job_id=JOB_ID)
48+
mock_get_request.assert_called_once_with(
49+
f"{URL}/{JOB_ID}/jobs",
50+
headers={
51+
"Content-Type": "application/json",
52+
"Authorization": f"Bearer {TOKEN}",
53+
},
54+
)
55+
56+
@patch(
57+
"airflow.providers.google.cloud.hooks.dataprep.requests.get",
58+
side_effect=[HTTPError(), mock.MagicMock()],
59+
)
60+
def test_should_pass_after_retry(self, mock_get_request, mock_hook):
61+
mock_hook.get_jobs_for_job_group(JOB_ID)
62+
assert mock_get_request.call_count == 2
63+
64+
@patch(
65+
"airflow.providers.google.cloud.hooks.dataprep.requests.get",
66+
side_effect=[mock.MagicMock(), HTTPError()],
67+
)
68+
def test_should_not_retry_after_success(self, mock_get_request, mock_hook):
69+
mock_hook.get_jobs_for_job_group.retry.sleep = mock.Mock() # pylint: disable=no-member
70+
mock_hook.get_jobs_for_job_group(JOB_ID)
71+
assert mock_get_request.call_count == 1
72+
73+
@patch(
74+
"airflow.providers.google.cloud.hooks.dataprep.requests.get",
75+
side_effect=[
76+
HTTPError(),
77+
HTTPError(),
78+
HTTPError(),
79+
HTTPError(),
80+
mock.MagicMock(),
81+
],
82+
)
83+
def test_should_retry_after_four_errors(self, mock_get_request, mock_hook):
84+
mock_hook.get_jobs_for_job_group.retry.sleep = mock.Mock() # pylint: disable=no-member
85+
mock_hook.get_jobs_for_job_group(JOB_ID)
86+
assert mock_get_request.call_count == 5
87+
88+
@patch(
89+
"airflow.providers.google.cloud.hooks.dataprep.requests.get",
90+
side_effect=[HTTPError(), HTTPError(), HTTPError(), HTTPError(), HTTPError()],
91+
)
92+
def test_raise_error_after_five_calls(self, mock_get_request, mock_hook):
93+
with pytest.raises(RetryError) as err:
94+
mock_hook.get_jobs_for_job_group.retry.sleep = mock.Mock() # pylint: disable=no-member
95+
mock_hook.get_jobs_for_job_group(JOB_ID)
96+
assert "HTTPError" in str(err)
97+
assert mock_get_request.call_count == 5

‎tests/providers/google/cloud/operators/test_dataprep.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
from unittest import TestCase, mock
19+
20+
from airflow.providers.google.cloud.operators.dataprep import DataprepGetJobsForJobGroupOperator
21+
22+
JOB_ID = 143
23+
TASK_ID = "dataprep_job"
24+
25+
26+
class TestDataprepGetJobsForJobGroupOperator(TestCase):
27+
@mock.patch(
28+
"airflow.providers.google.cloud.operators.dataprep.GoogleDataprepHook"
29+
)
30+
def test_execute(self, hook_mock):
31+
op = DataprepGetJobsForJobGroupOperator(job_id=JOB_ID, task_id=TASK_ID)
32+
op.execute(context={})
33+
hook_mock.assert_called_once_with(dataprep_conn_id='dataprep_conn_id')
34+
hook_mock.return_value.get_jobs_for_job_group.assert_called_once_with(job_id=JOB_ID)

0 commit comments

Comments
 (0)