Glue DataBrew operator

airflow/providers/amazon/aws/hooks/glue_databrew.py

@@ -0,0 +1,68 @@
+#
+# 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
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class GlueDataBrewHook(AwsBaseHook):
+    """
+    Interact with AWS DataBrew.
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    .. seealso::
+        - :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    def __init__(self, *args, **kwargs):
+        kwargs["client_type"] = "databrew"
+        super().__init__(*args, **kwargs)
+
+    def job_completion(self, job_name: str, run_id: str, delay: int = 10, max_attempts: int = 60) -> str:
+        """
+        Wait until Glue DataBrew job reaches terminal status.
+
+        :param job_name: The name of the job being processed during this run.
+        :param run_id: The unique identifier of the job run.
+        :param delay: Time in seconds to delay between polls
+        :param maxAttempts: Maximum number of attempts to poll for completion
+        :return: job status
+        """
+        self.get_waiter("job_complete").wait(
+            Name=job_name,
+            RunId=run_id,
+            WaiterConfig={"Delay": delay, "maxAttempts": max_attempts},
+        )
+
+        status = self.get_job_state(job_name, run_id)
+        return status
+
+    def get_job_state(self, job_name: str, run_id: str) -> str:
+        """
+        Get the status of a job run.
+
+        :param job_name: The name of the job being processed during this run.
+        :param run_id: The unique identifier of the job run.
+        :return: State of the job run.
+            'STARTING'|'RUNNING'|'STOPPING'|'STOPPED'|'SUCCEEDED'|'FAILED'|'TIMEOUT'
+        """
+        response = self.conn.describe_job_run(Name=job_name, RunId=run_id)
+        return response["State"]

airflow/providers/amazon/aws/operators/glue_databrew.py

@@ -0,0 +1,110 @@
+#
+# 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
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Sequence
+
+from airflow.configuration import conf
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.triggers.glue_databrew import GlueDataBrewJobCompleteTrigger
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class GlueDataBrewStartJobOperator(BaseOperator):
+    """
+    Start an AWS Glue DataBrew job.
+
+    AWS Glue DataBrew is a visual data preparation tool that makes it easier
+    for data analysts and data scientists to clean and normalize data
+    to prepare it for analytics and machine learning (ML).
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:GlueDataBrewStartJobOperator`
+
+    :param job_name: unique job name per AWS Account
+    :param wait_for_completion: Whether to wait for job run completion. (default: True)
+    :param deferrable: If True, the operator will wait asynchronously for the job to complete.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param delay: Time in seconds to wait between status checks. Default is 30.
+    :return: dictionary with key run_id and value of the resulting job's run_id.
+    """
+
+    template_fields: Sequence[str] = (
+        "job_name",
+        "wait_for_completion",
+        "delay",
+        "deferrable",
+    )
+
+    def __init__(
+        self,
+        job_name: str,
+        wait_for_completion: bool = True,
+        delay: int = 30,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.job_name = job_name
+        self.wait_for_completion = wait_for_completion
+        self.deferrable = deferrable
+        self.delay = delay
+        self.aws_conn_id = aws_conn_id
+
+    @cached_property
+    def hook(self) -> GlueDataBrewHook:
+        return GlueDataBrewHook(aws_conn_id=self.aws_conn_id)
+
+    def execute(self, context: Context):
+        job = self.hook.conn.start_job_run(Name=self.job_name)
+        run_id = job["RunId"]
+
+        self.log.info("AWS Glue DataBrew Job: %s. Run Id: %s submitted.", self.job_name, run_id)
+
+        if self.deferrable:
+            self.log.info("Deferring job %s with run_id %s", self.job_name, run_id)
+            self.defer(
+                trigger=GlueDataBrewJobCompleteTrigger(
+                    aws_conn_id=self.aws_conn_id, job_name=self.job_name, run_id=run_id, delay=self.delay
+                ),
+                method_name="execute_complete",
+            )
+
+        elif self.wait_for_completion:
+            self.log.info(
+                "Waiting for AWS Glue DataBrew Job: %s. Run Id: %s to complete.", self.job_name, run_id
+            )
+            status = self.hook.job_completion(job_name=self.job_name, delay=self.delay, run_id=run_id)
+            self.log.info("Glue DataBrew Job: %s status: %s", self.job_name, status)
+
+        return {"run_id": run_id}
+
+    def execute_complete(self, context: Context, event=None) -> dict[str, str]:
+        run_id = event.get("run_id", "")
+        status = event.get("status", "")
+
+        self.log.info("AWS Glue DataBrew runID: %s completed with status: %s", run_id, status)
+
+        return {"run_id": run_id}

airflow/providers/amazon/aws/triggers/glue_databrew.py

@@ -0,0 +1,59 @@
+# 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
+
+from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+
+class GlueDataBrewJobCompleteTrigger(AwsBaseWaiterTrigger):
+    """
+    Watches for a Glue DataBrew job, triggers when it finishes.
+
+    :param job_name: Glue DataBrew job name
+    :param run_id: the ID of the specific run to watch for that job
+    :param delay: Number of seconds to wait between two checks. Default is 10 seconds.
+    :param max_attempts: Maximum number of attempts to wait for the job to complete. Default is 60 attempts.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        job_name: str,
+        run_id: str,
+        aws_conn_id: str,
+        delay: int = 10,
+        max_attempts: int = 60,
+        **kwargs,
+    ):
+        super().__init__(
+            serialized_fields={"job_name": job_name, "run_id": run_id},
+            waiter_name="job_complete",
+            waiter_args={"Name": job_name, "RunId": run_id},
+            failure_message=f"Error while waiting for job {job_name} with run id {run_id} to complete",
+            status_message=f"Run id: {run_id}",
+            status_queries=["State"],
+            return_value=run_id,
+            return_key="run_id",
+            waiter_delay=delay,
+            waiter_max_attempts=max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> GlueDataBrewHook:
+        return GlueDataBrewHook(aws_conn_id=self.aws_conn_id)

airflow/providers/amazon/aws/waiters/databrew.json

@@ -0,0 +1,36 @@
+{
+    "version": 2,
+    "waiters": {
+        "job_complete": {
+            "operation": "DescribeJobRun",
+            "delay": 30,
+            "maxAttempts": 60,
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "STOPPED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "SUCCEEDED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "FAILED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "TIMEOUT",
+                    "state": "success"
+                }
+            ]
+        }
+    }
+}

airflow/providers/amazon/provider.yaml

@@ -288,6 +288,12 @@ integrations:
     how-to-guide:
       - /docs/apache-airflow-providers-amazon/operators/appflow.rst
     tags: [aws]
+  - integration-name: AWS Glue DataBrew
+    external-doc-url: https://docs.aws.amazon.com/databrew/latest/dg/what-is.html
+    how-to-guide:
+      - /docs/apache-airflow-providers-amazon/operators/glue_databrew.rst
+    logo: /integration-logos/aws/AWS-Glue-DataBrew_64.png
+    tags: [aws]
 
 operators:
   - integration-name: Amazon Athena
@@ -364,6 +370,9 @@ operators:
   - integration-name: Amazon Appflow
     python-modules:
       - airflow.providers.amazon.aws.operators.appflow
+  - integration-name: AWS Glue DataBrew
+    python-modules:
+      - airflow.providers.amazon.aws.operators.glue_databrew
 
 sensors:
   - integration-name: Amazon Athena
@@ -540,6 +549,9 @@ hooks:
   - integration-name: Amazon Appflow
     python-modules:
       - airflow.providers.amazon.aws.hooks.appflow
+  - integration-name: AWS Glue DataBrew
+    python-modules:
+      - airflow.providers.amazon.aws.hooks.glue_databrew
 
 triggers:
   - integration-name: Amazon Web Services
@@ -588,6 +600,9 @@ triggers:
   - integration-name: Amazon Simple Queue Service (SQS)
     python-modules:
       - airflow.providers.amazon.aws.triggers.sqs
+  - integration-name: AWS Glue DataBrew
+    python-modules:
+      - airflow.providers.amazon.aws.triggers.glue_databrew
 
 transfers:
   - source-integration-name: Amazon DynamoDB

docs/apache-airflow-providers-amazon/operators/glue.rst

@@ -105,3 +105,4 @@ Reference
 
 * `AWS boto3 library documentation for Glue <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/glue.html>`__
 * `Glue IAM Role creation <https://docs.aws.amazon.com/glue/latest/dg/create-an-iam-role.html>`__
+* `AWS boto3 library documentation for Glue DataBrew <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/databrew.html>`__

docs/apache-airflow-providers-amazon/operators/glue_databrew.rst

@@ -0,0 +1,53 @@
+ .. 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.
+
+=================
+AWS Glue DataBrew
+=================
+
+`AWS Glue DataBrew <https://aws.amazon.com/glue/features/databrew/>`__ is a visual data preparation tool
+that makes it easier for data analysts and data scientists to clean and normalize data to prepare it
+for analytics and machine learning (ML). You can choose from over 250 prebuilt transformations to automate
+data preparation tasks, all without the need to write any code. You can automate filtering anomalies, converting
+data to standard formats and correcting invalid values, and other tasks.
+After your data is ready, you can immediately use it for analytics and ML projects.
+
+Prerequisite Tasks
+------------------
+
+.. include:: ../_partials/prerequisite_tasks.rst
+
+Operators
+---------
+
+.. _howto/operator:GlueDataBrewStartJobOperator:
+
+Start an AWS Glue DataBrew job
+==============================
+
+To submit a new AWS Glue DataBrew job you can use :class:`~airflow.providers.amazon.aws.operators.glue_databrew.GlueDataBrewStartJobOperator`.
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_glue_databrew.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_glue_databrew_start]
+    :end-before: [END howto_operator_glue_databrew_start]
+
+Reference
+---------
+
+* `AWS boto3 library documentation for Glue DataBrew <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/databrew.html>`__