Skip to content

Commit 804548d

Browse files
michalslowikowski00turbaszek
michalslowikowski00
and
turbaszek
authored
Add Dataprep operators (#10304)
Add DataprepGetJobGroupOperator and DataprepRunJobGroupOperator for Dataprep service. Co-authored-by: Tomek Urbaszek <tomasz.urbaszek@polidea.com>
1 parent f40ac9b commit 804548d

8 files changed

Lines changed: 469 additions & 67 deletions

File tree

‎airflow/models/connection.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
),
5252
"cassandra": ("airflow.providers.apache.cassandra.hooks.cassandra.CassandraHook", "cassandra_conn_id"),
5353
"cloudant": ("airflow.providers.cloudant.hooks.cloudant.CloudantHook", "cloudant_conn_id"),
54-
"dataprep": ("airflow.providers.google.cloud.hooks.dataprep.GoogleDataprepHook", "dataprep_conn_id"),
54+
"dataprep": ("airflow.providers.google.cloud.hooks.dataprep.GoogleDataprepHook", "dataprep_default"),
5555
"docker": ("airflow.providers.docker.hooks.docker.DockerHook", "docker_conn_id"),
5656
"elasticsearch": (
5757
"airflow.providers.elasticsearch.hooks.elasticsearch.ElasticsearchHook",

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

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,56 @@
1717
"""
1818
Example Airflow DAG that shows how to use Google Dataprep.
1919
"""
20+
import os
2021

2122
from airflow import models
22-
from airflow.providers.google.cloud.operators.dataprep import DataprepGetJobsForJobGroupOperator
23+
from airflow.providers.google.cloud.operators.dataprep import (
24+
DataprepGetJobGroupOperator,
25+
DataprepGetJobsForJobGroupOperator,
26+
DataprepRunJobGroupOperator,
27+
)
2328
from airflow.utils import dates
2429

25-
JOB_ID = 6269792
30+
DATAPREP_JOB_ID = int(os.environ.get('DATAPREP_JOB_ID', 12345677))
31+
DATAPREP_JOB_RECIPE_ID = int(os.environ.get('DATAPREP_JOB_RECIPE_ID', 12345677))
32+
DATAPREP_BUCKET = os.environ.get("DATAPREP_BUCKET", "gs://afl-sql/name@email.com")
33+
34+
DATA = {
35+
"wrangledDataset": {"id": DATAPREP_JOB_RECIPE_ID},
36+
"overrides": {
37+
"execution": "dataflow",
38+
"profiler": False,
39+
"writesettings": [
40+
{
41+
"path": DATAPREP_BUCKET,
42+
"action": "create",
43+
"format": "csv",
44+
"compression": "none",
45+
"header": False,
46+
"asSingleFile": False,
47+
}
48+
],
49+
},
50+
}
51+
2652

2753
with models.DAG(
28-
"example_dataprep", schedule_interval=None, start_date=dates.days_ago(1) # Override to match your needs
54+
"example_dataprep", schedule_interval=None, start_date=dates.days_ago(1), # Override to match your needs
2955
) as dag:
56+
# [START how_to_dataprep_run_job_group_operator]
57+
run_job_group = DataprepRunJobGroupOperator(task_id="run_job_group", body_request=DATA)
58+
# [END how_to_dataprep_run_job_group_operator]
3059

3160
# [START how_to_dataprep_get_jobs_for_job_group_operator]
3261
get_jobs_for_job_group = DataprepGetJobsForJobGroupOperator(
33-
task_id="get_jobs_for_job_group", job_id=JOB_ID
62+
task_id="get_jobs_for_job_group", job_id=DATAPREP_JOB_ID
3463
)
3564
# [END how_to_dataprep_get_jobs_for_job_group_operator]
65+
66+
# [START how_to_dataprep_get_job_group_operator]
67+
get_job_group = DataprepGetJobGroupOperator(
68+
task_id="get_job_group", job_group_id=DATAPREP_JOB_ID, embed="", include_deleted=False,
69+
)
70+
# [END how_to_dataprep_get_job_group_operator]
71+
72+
run_job_group >> [get_jobs_for_job_group, get_job_group]

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

Lines changed: 59 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,14 @@
1818
"""
1919
This module contains Google Dataprep hook.
2020
"""
21+
import json
22+
import os
2123
from typing import Any, Dict
2224

2325
import requests
26+
from requests import HTTPError
2427
from tenacity import retry, stop_after_attempt, wait_exponential
2528

26-
from airflow import AirflowException
2729
from airflow.hooks.base_hook import BaseHook
2830

2931

@@ -37,10 +39,13 @@ class GoogleDataprepHook(BaseHook):
3739
3840
"""
3941

40-
def __init__(self, dataprep_conn_id: str = "dataprep_conn_id") -> None:
42+
def __init__(self, dataprep_conn_id: str = "dataprep_default") -> None:
4143
super().__init__()
4244
self.dataprep_conn_id = dataprep_conn_id
43-
self._url = "https://api.clouddataprep.com/v4/jobGroups"
45+
conn = self.get_connection(self.dataprep_conn_id)
46+
extra_dejson = conn.extra_dejson
47+
self._token = extra_dejson.get("extra__dataprep__token")
48+
self._base_url = extra_dejson.get("extra__dataprep__base_url", "https://api.clouddataprep.com")
4449

4550
@property
4651
def _headers(self) -> Dict[str, str]:
@@ -50,26 +55,63 @@ def _headers(self) -> Dict[str, str]:
5055
}
5156
return headers
5257

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-
6458
@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, max=10))
6559
def get_jobs_for_job_group(self, job_id: int) -> Dict[str, Any]:
6660
"""
6761
Get information about the batch jobs within a Cloud Dataprep job.
6862
69-
:param job_id The ID of the job that will be fetched.
63+
:param job_id: The ID of the job that will be fetched
7064
:type job_id: int
7165
"""
72-
url: str = f"{self._url}/{job_id}/jobs"
66+
67+
endpoint_path = f"v4/jobGroups/{job_id}/jobs"
68+
url: str = os.path.join(self._base_url, endpoint_path)
7369
response = requests.get(url, headers=self._headers)
74-
response.raise_for_status()
70+
self._raise_for_status(response)
71+
return response.json()
72+
73+
@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, max=10))
74+
def get_job_group(self, job_group_id: int, embed: str, include_deleted: bool) -> Dict[str, Any]:
75+
"""
76+
Get the specified job group.
77+
A job group is a job that is executed from a specific node in a flow.
78+
79+
:param job_group_id: The ID of the job that will be fetched
80+
:type job_group_id: int
81+
:param embed: Comma-separated list of objects to pull in as part of the response
82+
:type embed: str
83+
:param include_deleted: if set to "true", will include deleted objects
84+
:type include_deleted: bool
85+
"""
86+
87+
params: Dict[str, Any] = {"embed": embed, "includeDeleted": include_deleted}
88+
endpoint_path = f"v4/jobGroups/{job_group_id}"
89+
url: str = os.path.join(self._base_url, endpoint_path)
90+
response = requests.get(url, headers=self._headers, params=params)
91+
self._raise_for_status(response)
92+
return response.json()
93+
94+
@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, max=10))
95+
def run_job_group(self, body_request: dict) -> Dict[str, Any]:
96+
"""
97+
Creates a ``jobGroup``, which launches the specified job as the authenticated user.
98+
This performs the same action as clicking on the Run Job button in the application.
99+
To get recipe_id please follow the Dataprep API documentation
100+
https://clouddataprep.com/documentation/api#operation/runJobGroup
101+
102+
:param body_request: The identifier for the recipe you would like to run.
103+
:type body_request: dict
104+
"""
105+
106+
endpoint_path = "v4/jobGroups"
107+
url: str = os.path.join(self._base_url, endpoint_path)
108+
response = requests.post(url, headers=self._headers, data=json.dumps(body_request))
109+
self._raise_for_status(response)
75110
return response.json()
111+
112+
def _raise_for_status(self, response: requests.models.Response) -> None:
113+
try:
114+
response.raise_for_status()
115+
except HTTPError:
116+
self.log.error(response.json().get('exception'))
117+
raise

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

Lines changed: 73 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,20 +35,90 @@ class DataprepGetJobsForJobGroupOperator(BaseOperator):
3535
For more information on how to use this operator, take a look at the guide:
3636
:ref:`howto/operator:DataprepGetJobsForJobGroupOperator`
3737
38-
3938
:param job_id The ID of the job that will be requests
4039
:type job_id: int
4140
"""
4241

4342
template_fields = ("job_id",)
4443

4544
@apply_defaults
46-
def __init__(self, *, job_id: int, **kwargs) -> None:
45+
def __init__(self, *, dataprep_conn_id: str = "dataprep_default", job_id: int, **kwargs) -> None:
4746
super().__init__(**kwargs)
47+
self.dataprep_conn_id = (dataprep_conn_id,)
4848
self.job_id = job_id
4949

5050
def execute(self, context: Dict):
5151
self.log.info("Fetching data for job with id: %d ...", self.job_id)
52-
hook = GoogleDataprepHook(dataprep_conn_id="dataprep_conn_id")
52+
hook = GoogleDataprepHook(dataprep_conn_id="dataprep_default",)
5353
response = hook.get_jobs_for_job_group(job_id=self.job_id)
5454
return response
55+
56+
57+
class DataprepGetJobGroupOperator(BaseOperator):
58+
"""
59+
Get the specified job group.
60+
A job group is a job that is executed from a specific node in a flow.
61+
API documentation https://clouddataprep.com/documentation/api#section/Overview
62+
63+
.. seealso::
64+
For more information on how to use this operator, take a look at the guide:
65+
:ref:`howto/operator:DataprepGetJobGroupOperator`
66+
67+
:param job_group_id: The ID of the job that will be requests
68+
:type job_group_id: int
69+
:param embed: Comma-separated list of objects to pull in as part of the response
70+
:type embed: string
71+
:param include_deleted: if set to "true", will include deleted objects
72+
:type include_deleted: bool
73+
"""
74+
75+
template_fields = ("job_group_id", "embed")
76+
77+
@apply_defaults
78+
def __init__(
79+
self,
80+
*,
81+
dataprep_conn_id: str = "dataprep_default",
82+
job_group_id: int,
83+
embed: str,
84+
include_deleted: bool,
85+
**kwargs,
86+
) -> None:
87+
super().__init__(**kwargs)
88+
self.dataprep_conn_id: str = dataprep_conn_id
89+
self.job_group_id = job_group_id
90+
self.embed = embed
91+
self.include_deleted = include_deleted
92+
93+
def execute(self, context: Dict):
94+
self.log.info("Fetching data for job with id: %d ...", self.job_group_id)
95+
hook = GoogleDataprepHook(dataprep_conn_id=self.dataprep_conn_id)
96+
response = hook.get_job_group(
97+
job_group_id=self.job_group_id, embed=self.embed, include_deleted=self.include_deleted,
98+
)
99+
return response
100+
101+
102+
class DataprepRunJobGroupOperator(BaseOperator):
103+
"""
104+
Create a ``jobGroup``, which launches the specified job as the authenticated user.
105+
This performs the same action as clicking on the Run Job button in the application.
106+
To get recipe_id please follow the Dataprep API documentation
107+
https://clouddataprep.com/documentation/api#operation/runJobGroup
108+
109+
:param recipe_id: The identifier for the recipe you would like to run.
110+
:type recipe_id: int
111+
"""
112+
113+
template_fields = ("body_request",)
114+
115+
def __init__(self, *, dataprep_conn_id: str = "dataprep_default", body_request: dict, **kwargs) -> None:
116+
super().__init__(**kwargs)
117+
self.body_request = body_request
118+
self.dataprep_conn_id = dataprep_conn_id
119+
120+
def execute(self, context: None):
121+
self.log.info("Creating a job...")
122+
hook = GoogleDataprepHook(dataprep_conn_id=self.dataprep_conn_id)
123+
response = hook.run_job_group(body_request=self.body_request)
124+
return response

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

Lines changed: 62 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,30 @@
1717
1818
Google Dataprep Operators
1919
=========================
20-
`Google Dataprep API documentation is available here <https://cloud.google.com/dataprep/docs/html/API-Reference_145281441>`__
20+
Dataprep is the intelligent cloud data service to visually explore, clean, and prepare data for analysis and machine learning.
21+
Service can be use to explore and transform raw data from disparate and/or large datasets into clean and structured data for further analysis and processing.
22+
Dataprep Job is an internal object encoding the information necessary to run a part of a Cloud Dataprep job group.
23+
For more information about the service visit `Google Dataprep API documentation <https://cloud.google.com/dataprep/docs/html/API-Reference_145281441>`_
24+
25+
Before you begin
26+
^^^^^^^^^^^^^^^^
27+
Before using Dataprep within Airflow you need to authenticate your account with TOKEN.
28+
To get connection Dataprep with Airflow you need Dataprep token. Please follow Dataprep `instructions <https://clouddataprep.com/documentation/api#section/Authentication>`_ to do it.
29+
30+
TOKEN should be added to the Connection in Airflow in JSON format.
31+
You can check `how to do such connection <https://airflow.readthedocs.io/en/stable/howto/connection/index.html#editing-a-connection-with-the-ui>`_.
32+
33+
The DataprepRunJobGroupOperator will run specified job. Operator required a recipe id. To identify the recipe id please use `API documentation for runJobGroup <https://clouddataprep.com/documentation/api#operation/runJobGroup>`_
34+
E.g. if the URL is /flows/10?recipe=7, the recipe id is 7. The recipe cannot be created via this operator. It can be created only via UI which is available `here <https://clouddataprep.com/>`_.
35+
Some of parameters can be override by DAG's body request. How to do it is shown in example dag.
36+
37+
See following example:
38+
Set values for these fields:
39+
.. code-block::
40+
41+
Conn Id: "your_conn_id"
42+
Extra: {"extra__dataprep__token": "TOKEN",
43+
"extra__dataprep__base_url": "https://api.clouddataprep.com"}
2144
2245
.. contents::
2346
:depth: 1
@@ -28,33 +51,58 @@ Prerequisite Tasks
2851

2952
.. include:: /howto/operator/google/_partials/prerequisite_tasks.rst
3053

54+
.. _howto/operator:DataprepRunJobGroupOperator:
55+
56+
Run Job Group
57+
^^^^^^^^^^^^^
58+
59+
Operator task is to create a job group, which launches the specified job as the authenticated user.
60+
This performs the same action as clicking on the Run Job button in the application.
61+
62+
To get information about jobs within a Cloud Dataprep job use:
63+
:class:`~airflow.providers.google.cloud.operators.dataprep.DataprepRunJobGroupOperator`
64+
65+
Example usage:
66+
67+
.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_dataprep.py
68+
:language: python
69+
:dedent: 4
70+
:start-after: [START how_to_dataprep_run_job_group_operator]
71+
:end-before: [END how_to_dataprep_run_job_group_operator]
72+
3173
.. _howto/operator:DataprepGetJobsForJobGroupOperator:
3274

3375
Get Jobs For Job Group
3476
^^^^^^^^^^^^^^^^^^^^^^
3577

78+
Operator task is to get information about the batch jobs within a Cloud Dataprep job.
79+
3680
To get information about jobs within a Cloud Dataprep job use:
3781
:class:`~airflow.providers.google.cloud.operators.dataprep.DataprepGetJobsForJobGroupOperator`
3882

39-
To get connection Dataprep with Airflow you need Dataprep token.
40-
Please follow Dataprep instructions.
41-
https://clouddataprep.com/documentation/api#section/Authentication
83+
Example usage:
84+
85+
.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_dataprep.py
86+
:language: python
87+
:dedent: 4
88+
:start-after: [START how_to_dataprep_get_jobs_for_job_group_operator]
89+
:end-before: [END how_to_dataprep_get_jobs_for_job_group_operator]
90+
91+
.. _howto/operator:DataprepGetJobGroupOperator:
4292

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
93+
Get Job Group
94+
^^^^^^^^^^^^^
4695

47-
See following example:
48-
Set values for these fields:
49-
.. code-block::
96+
Operator task is to get the specified job group.
97+
A job group is a job that is executed from a specific node in a flow.
5098

51-
Conn Id: "your_conn_id"
52-
Extra: "{\"token\": \"TOKEN\"}
99+
To get information about jobs within a Cloud Dataprep job use:
100+
:class:`~airflow.providers.google.cloud.operators.dataprep.DataprepGetJobGroupOperator`
53101

54102
Example usage:
55103

56104
.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_dataprep.py
57105
:language: python
58106
: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]
107+
:start-after: [START how_to_dataprep_get_job_group_operator]
108+
:end-before: [END how_to_dataprep_get_job_group_operator]

0 commit comments

Comments
 (0)