From 369a8ff63b55c47fc65fe72d65d6dc5b6230cc8d Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Fri, 6 Oct 2023 12:31:12 -0400
Subject: [PATCH 01/21] Glue DataBrew operator

---
 airflow/providers/amazon/aws/hooks/glue.py    |  52 +++++++++
 .../providers/amazon/aws/operators/glue.py    | 100 +++++++++++++++++-
 airflow/providers/amazon/aws/triggers/glue.py |  59 ++++++++++-
 .../amazon/aws/waiters/databrew.json          |  54 ++++++++++
 .../operators/glue.rst                        |  14 ++-
 tests/providers/amazon/aws/hooks/test_glue.py |  14 ++-
 .../aws/operators/test_glue_databrew.py       |  44 ++++++++
 .../amazon/aws/triggers/test_glue_databrew.py |  46 ++++++++
 .../amazon/aws/waiters/test_custom_waiters.py |  40 +++++++
 .../providers/amazon/aws/example_databrew.py  |  60 +++++++++++
 10 files changed, 477 insertions(+), 6 deletions(-)
 create mode 100644 airflow/providers/amazon/aws/waiters/databrew.json
 create mode 100644 tests/providers/amazon/aws/operators/test_glue_databrew.py
 create mode 100644 tests/providers/amazon/aws/triggers/test_glue_databrew.py
 create mode 100644 tests/system/providers/amazon/aws/example_databrew.py

diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py
index baf6780e07802..66ee83fc0a7fe 100644
--- a/airflow/providers/amazon/aws/hooks/glue.py
+++ b/airflow/providers/amazon/aws/hooks/glue.py
@@ -422,3 +422,55 @@ def create_or_update_glue_job(self) -> str | None:
             self.conn.create_job(**config)
 
         return self.job_name
+
+
+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`
+    """
+
+    TERMINAL_STATES = ["STOPPED", "SUCCEEDED", "FAILED", "TIMEOUT"]
+    SUCCESS_STATES = ["SUCCEEDED"]
+    FAILURE_STATES = ["FAILED", "TIMEOUT"]
+    RUNNING_STATES = ["STARTING", "RUNNING", "STOPPING"]
+
+    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, maxAttempts: 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": maxAttempts},
+        )
+
+        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.get_conn().describe_job_run(Name=job_name, RunId=run_id)
+        return response["State"]
diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py
index 97bcaba66abc7..d304f53880e0c 100644
--- a/airflow/providers/amazon/aws/operators/glue.py
+++ b/airflow/providers/amazon/aws/operators/glue.py
@@ -25,10 +25,10 @@
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
-from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
+from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink
-from airflow.providers.amazon.aws.triggers.glue import GlueJobCompleteTrigger
+from airflow.providers.amazon.aws.triggers.glue import GlueDataBrewJobCompleteTrigger, GlueJobCompleteTrigger
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -230,3 +230,99 @@ def on_kill(self):
             )
             if not response["SuccessfulSubmissions"]:
                 self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", self.job_name, self._job_run_id)
+
+
+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.
+    """
+
+    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 data_brew_hook(self) -> GlueDataBrewHook:
+        return GlueDataBrewHook(aws_conn_id=self.aws_conn_id)
+
+    def execute(self, context: Context):
+        resp = {}
+        resp["job_name"] = self.job_name
+
+        job = self.data_brew_hook.conn.start_job_run(Name=self.job_name)
+        run_id = job["RunId"]
+        resp["run_id"] = run_id
+
+        status = self.data_brew_hook.get_job_state(self.job_name, run_id)
+        resp["status"] = status
+
+        self.log.info(
+            "AWS Glue DataBrew Job: %s. Run Id: %s submitted. Status: %s", self.job_name, run_id, status
+        )
+
+        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.data_brew_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)
+            resp["status"] = status
+
+        return resp
+
+    def execute_complete(self, context: Context, event=None) -> dict[str, str]:
+        result = {}
+        result["job_name"] = event.get("jobName", "")
+        result["run_id"] = event.get("runId", "")
+        result["status"] = event.get("status", "")
+        self.log.info(
+            "AWS Glue DataBrew Job: %s runId: %s status: %s",
+            result["job_name"],
+            result["run_id"],
+            result["status"],
+        )
+        return result
diff --git a/airflow/providers/amazon/aws/triggers/glue.py b/airflow/providers/amazon/aws/triggers/glue.py
index b5a0a673d3e18..72e9c177d4916 100644
--- a/airflow/providers/amazon/aws/triggers/glue.py
+++ b/airflow/providers/amazon/aws/triggers/glue.py
@@ -21,7 +21,7 @@
 from functools import cached_property
 from typing import Any, AsyncIterator
 
-from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
+from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
 from airflow.providers.amazon.aws.hooks.glue_catalog import GlueCatalogHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
@@ -148,3 +148,60 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                     break
                 else:
                     await asyncio.sleep(self.waiter_delay)
+
+
+class GlueDataBrewJobCompleteTrigger(BaseTrigger):
+    """
+    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 maxAttempts: 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,
+        maxAttempts: int = 60,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.job_name = job_name
+        self.run_id = run_id
+        self.aws_conn_id = aws_conn_id
+        self.delay = delay
+        self.maxAttempts = maxAttempts
+
+    @property
+    def hook(self) -> GlueDataBrewHook:
+        return GlueDataBrewHook(aws_conn_id=self.aws_conn_id)
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            # dynamically generate the fully qualified name of the class
+            self.__class__.__module__ + "." + self.__class__.__qualname__,
+            {
+                "job_name": self.job_name,
+                "run_id": self.run_id,
+                "aws_conn_id": self.aws_conn_id,
+                "delay": self.delay,
+                "maxAttempts": self.maxAttempts,
+            },
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        async with self.hook.async_conn as client:
+            waiter = self.hook.get_waiter("job_complete", deferrable=True, client=client)
+            await waiter.wait(
+                Name=self.job_name,
+                RunId=self.run_id,
+                WaiterConfig={"Delay": self.delay, "maxAttempts": self.maxAttempts},
+            )
+
+        result = self.hook.get_job_state(self.job_name, self.run_id)
+        yield TriggerEvent({"status": result, "RunId": self.run_id, "JobName": self.job_name})
diff --git a/airflow/providers/amazon/aws/waiters/databrew.json b/airflow/providers/amazon/aws/waiters/databrew.json
new file mode 100644
index 0000000000000..a18dc23d23a47
--- /dev/null
+++ b/airflow/providers/amazon/aws/waiters/databrew.json
@@ -0,0 +1,54 @@
+{
+    "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"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "Starting",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "STOPPING",
+                    "state": "retry"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "State",
+                    "expected": "RUNNING",
+                    "state": "retry"
+                }
+            ]
+        }
+    }
+}
diff --git a/docs/apache-airflow-providers-amazon/operators/glue.rst b/docs/apache-airflow-providers-amazon/operators/glue.rst
index e582e9d415b67..5f90eb536194f 100644
--- a/docs/apache-airflow-providers-amazon/operators/glue.rst
+++ b/docs/apache-airflow-providers-amazon/operators/glue.rst
@@ -100,8 +100,18 @@ use :class:`~airflow.providers.amazon.aws.sensors.glue.GlueJobSensor`
     :start-after: [START howto_sensor_glue]
     :end-before: [END howto_sensor_glue]
 
+Start an AWS Glue DataBrew job
+==============================
+
+To submit a new AWS Glue DataBrew job you can use :class:`~airflow.providers.amazon.aws.operators.glue.GlueDataBrewStartJobOperator`.
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_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 <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>`__
diff --git a/tests/providers/amazon/aws/hooks/test_glue.py b/tests/providers/amazon/aws/hooks/test_glue.py
index 799547dd576d2..1943eb0c0d19c 100644
--- a/tests/providers/amazon/aws/hooks/test_glue.py
+++ b/tests/providers/amazon/aws/hooks/test_glue.py
@@ -28,7 +28,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
+from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
 
 if TYPE_CHECKING:
     from unittest.mock import MagicMock
@@ -477,3 +477,15 @@ async def test_async_job_completion_failure(self, get_state_mock: MagicMock):
             await hook.async_job_completion("job_name", "run_id")
 
         assert get_state_mock.call_count == 3
+
+
+class TestGlueDataBrewHook:
+    job_name = "test-databrew-job"
+    runId = "test12345"
+
+    @mock.patch.object(GlueDataBrewHook, "get_job_state")
+    def test_get_job_state(self, get_job_state_mock: MagicMock):
+        get_job_state_mock.return_value = "SUCCEEDED"
+        hook = GlueDataBrewHook()
+        result = hook.get_job_state(self.job_name, self.runId)
+        assert result == "SUCCEEDED"
diff --git a/tests/providers/amazon/aws/operators/test_glue_databrew.py b/tests/providers/amazon/aws/operators/test_glue_databrew.py
new file mode 100644
index 0000000000000..cc7bd36b433b3
--- /dev/null
+++ b/tests/providers/amazon/aws/operators/test_glue_databrew.py
@@ -0,0 +1,44 @@
+#
+# 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 unittest import mock
+
+import pytest
+from moto import mock_databrew
+
+from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook
+from airflow.providers.amazon.aws.operators.glue import GlueDataBrewStartJobOperator
+
+JOB_NAME = "test_job"
+
+
+@pytest.fixture
+def hook() -> GlueDataBrewHook:
+    with mock_databrew():
+        yield GlueDataBrewHook(aws_conn_id="aws_default")
+
+
+class TestGlueDataBrewOperator:
+    @mock.patch.object(GlueDataBrewHook, "get_waiter")
+    def test_start_job_wait_for_completion(self, mock_hook_get_waiter):
+        operator = GlueDataBrewStartJobOperator(
+            task_id="task_test", job_name=JOB_NAME, wait_for_completion=True
+        )
+        operator.execute(None)
+        mock_hook_get_waiter.assert_called_once_with("databrew_job_complete")
diff --git a/tests/providers/amazon/aws/triggers/test_glue_databrew.py b/tests/providers/amazon/aws/triggers/test_glue_databrew.py
new file mode 100644
index 0000000000000..87577c9f809e1
--- /dev/null
+++ b/tests/providers/amazon/aws/triggers/test_glue_databrew.py
@@ -0,0 +1,46 @@
+# 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
+
+import pytest
+
+from airflow.providers.amazon.aws.triggers.glue import GlueDataBrewJobCompleteTrigger
+
+TEST_JOB_NAME = "test_job_name"
+TEST_JOB_RUN_ID = "a1234"
+TEST_JOB_RUN_STATUS = "SUCCEEDED"
+
+
+@pytest.fixture
+def trigger():
+    yield GlueDataBrewJobCompleteTrigger(
+        aws_conn_id="aws_default", job_name=TEST_JOB_NAME, run_id=TEST_JOB_RUN_ID
+    )
+
+
+class TestGlueDataBrewJobCompleteTrigger:
+    def test_serialize(self, trigger):
+        class_path, args = trigger.serialize()
+
+        class_name = class_path.split(".")[-1]
+        clazz = globals()[class_name]
+        instance = clazz(**args)
+
+        class_path2, args2 = instance.serialize()
+
+        assert class_path == class_path2
+        assert args == args2
diff --git a/tests/providers/amazon/aws/waiters/test_custom_waiters.py b/tests/providers/amazon/aws/waiters/test_custom_waiters.py
index 19f9296b6af71..d7bd5b6ead442 100644
--- a/tests/providers/amazon/aws/waiters/test_custom_waiters.py
+++ b/tests/providers/amazon/aws/waiters/test_custom_waiters.py
@@ -33,6 +33,7 @@
 from airflow.providers.amazon.aws.hooks.ecs import EcsClusterStates, EcsHook, EcsTaskDefinitionStates
 from airflow.providers.amazon.aws.hooks.eks import EksHook
 from airflow.providers.amazon.aws.hooks.emr import EmrHook
+from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook
 from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter
 
 
@@ -424,3 +425,42 @@ def test_steps_failed(self, mock_list_steps):
                 StepIds=[self.STEP_ID1, self.STEP_ID2],
                 WaiterConfig={"Delay": 0.01, "MaxAttempts": 3},
             )
+
+
+class TestCustomDataBrewWaiters:
+    """Test waiters from ``amazon/aws/waiters/glue.json``."""
+
+    JOB_NAME = "test_job"
+    RUN_ID = "123"
+
+    @pytest.fixture(autouse=True)
+    def setup_test_cases(self, monkeypatch):
+        self.client = boto3.client("databrew", region_name="eu-west-3")
+        monkeypatch.setattr(GlueDataBrewHook, "conn", self.client)
+
+    def test_service_waiters(self):
+        hook_waiters = GlueDataBrewHook(aws_conn_id=None).list_waiters()
+        assert "job_complete" in hook_waiters
+
+    @pytest.fixture
+    def mock_describe_job_runs(self):
+        """Mock ``GlueDataBrewHook.Client.describe_job_run`` method."""
+        with mock.patch.object(self.client, "describe_job_run") as m:
+            yield m
+
+    @staticmethod
+    def describe_jobs(status: str):
+        """
+        Helper function for generate minimal DescribeJobRun response for a single job.
+        https://docs.aws.amazon.com/databrew/latest/dg/API_DescribeJobRun.html
+        """
+        return {"State": status}
+
+    def test_job_succeeded(self, mock_describe_job_runs):
+        """Test job succeeded"""
+        mock_describe_job_runs.side_effect = [
+            self.describe_jobs(GlueDataBrewHook.RUNNING_STATES[1]),
+            self.describe_jobs(GlueDataBrewHook.TERMINAL_STATES[1]),
+        ]
+        waiter = GlueDataBrewHook(aws_conn_id=None).get_waiter("job_complete")
+        waiter.wait(name=self.JOB_NAME, runId=self.RUN_ID, WaiterConfig={"Delay": 0.2, "MaxAttempts": 2})
diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
new file mode 100644
index 0000000000000..1767a2586710c
--- /dev/null
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -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.
+from __future__ import annotations
+
+import pendulum
+
+from airflow.decorators import dag
+from airflow.models.baseoperator import chain
+from airflow.providers.amazon.aws.operators.glue import (
+    GlueDataBrewStartJobOperator,
+)
+from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_databrew_dag"
+
+# Externally fetched variables:
+JOB_NAME = "JOB_NAME"
+
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(JOB_NAME).build()
+
+
+@dag(DAG_ID, schedule=None, start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)
+def run_job():
+    # [START howto_operator_glue_databrew_start]
+
+    GlueDataBrewStartJobOperator(task_id="startjob", deferrable=True, job_name=JOB_NAME, delay=15)
+
+    # [END howto_operator_glue_databrew_start]
+
+
+test_context = sys_test_context_task()
+
+
+chain(test_context, run_job)
+
+from tests.system.utils.watcher import watcher  # noqa: E402
+
+# This test needs watcher in order to properly mark success/failure
+# when "tearDown" task with trigger rule is part of the DAG
+list(dag.tasks) >> watcher()
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)

From e37ad7b8333c70f41d19e6bd71000a2101ac1268 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Mon, 9 Oct 2023 08:59:26 -0400
Subject: [PATCH 02/21] Update airflow/providers/amazon/aws/operators/glue.py

Co-authored-by: Hussein Awala <hussein@awala.fr>
---
 airflow/providers/amazon/aws/operators/glue.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py
index d304f53880e0c..bdc58d2d984f3 100644
--- a/airflow/providers/amazon/aws/operators/glue.py
+++ b/airflow/providers/amazon/aws/operators/glue.py
@@ -275,7 +275,7 @@ def __init__(
         self.aws_conn_id = aws_conn_id
 
     @cached_property
-    def data_brew_hook(self) -> GlueDataBrewHook:
+    def hook(self) -> GlueDataBrewHook:
         return GlueDataBrewHook(aws_conn_id=self.aws_conn_id)
 
     def execute(self, context: Context):

From 245dad7ccb03fd804ad37687703401158418ef85 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:28:12 -0400
Subject: [PATCH 03/21] Update
 tests/system/providers/amazon/aws/example_databrew.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 tests/system/providers/amazon/aws/example_databrew.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
index 1767a2586710c..4edcafed0502a 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -27,11 +27,7 @@
 
 DAG_ID = "example_databrew_dag"
 
-# Externally fetched variables:
-JOB_NAME = "JOB_NAME"
-
-
-sys_test_context_task = SystemTestContextBuilder().add_variable(JOB_NAME).build()
+sys_test_context_task = SystemTestContextBuilder().build()
 
 
 @dag(DAG_ID, schedule=None, start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)

From c00dfeeea95e38de667d19adb6391ba7d6527fcc Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:28:22 -0400
Subject: [PATCH 04/21] Update airflow/providers/amazon/aws/operators/glue.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 airflow/providers/amazon/aws/operators/glue.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py
index bdc58d2d984f3..40433e826b12f 100644
--- a/airflow/providers/amazon/aws/operators/glue.py
+++ b/airflow/providers/amazon/aws/operators/glue.py
@@ -315,10 +315,11 @@ def execute(self, context: Context):
         return resp
 
     def execute_complete(self, context: Context, event=None) -> dict[str, str]:
-        result = {}
-        result["job_name"] = event.get("jobName", "")
-        result["run_id"] = event.get("runId", "")
-        result["status"] = event.get("status", "")
+        result = {
+                "job_name": event.get("jobName", ""),
+                "run_id": event.get("runId", ""),
+                "status": event.get("status", ""),
+        }
         self.log.info(
             "AWS Glue DataBrew Job: %s runId: %s status: %s",
             result["job_name"],

From 76b0a3788da718e8b1d579b287e8a7b14acca5f1 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:28:50 -0400
Subject: [PATCH 05/21] Update
 tests/system/providers/amazon/aws/example_databrew.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 tests/system/providers/amazon/aws/example_databrew.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
index 4edcafed0502a..b7af1afd4f8ca 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -25,7 +25,7 @@
 )
 from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
 
-DAG_ID = "example_databrew_dag"
+DAG_ID = "example_databrew"
 
 sys_test_context_task = SystemTestContextBuilder().build()
 

From e3f3f12e5b1e52e3a80ae26b8ef2d50ebaba28a1 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:29:12 -0400
Subject: [PATCH 06/21] Update
 tests/system/providers/amazon/aws/example_databrew.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 .../providers/amazon/aws/example_databrew.py   | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
index b7af1afd4f8ca..c7ba9be68892a 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -38,17 +38,13 @@ def run_job():
 
     # [END howto_operator_glue_databrew_start]
 
-
-test_context = sys_test_context_task()
-
-
-chain(test_context, run_job)
-
-from tests.system.utils.watcher import watcher  # noqa: E402
-
-# This test needs watcher in order to properly mark success/failure
-# when "tearDown" task with trigger rule is part of the DAG
-list(dag.tasks) >> watcher()
+    chain(test_context, start_job)
+    
+    from tests.system.utils.watcher import watcher  # noqa: E402
+    
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
 
 from tests.system.utils import get_test_run  # noqa: E402
 

From 0cb7808ef8dffbad9a0bb8e4937a387ee2297574 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:29:36 -0400
Subject: [PATCH 07/21] Update
 tests/system/providers/amazon/aws/example_databrew.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 tests/system/providers/amazon/aws/example_databrew.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
index c7ba9be68892a..fce03eed666bd 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -32,10 +32,13 @@
 
 @dag(DAG_ID, schedule=None, start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)
 def run_job():
-    # [START howto_operator_glue_databrew_start]
+    test_context = sys_test_context_task()
+    env_id = test_context["ENV_ID"]
 
-    GlueDataBrewStartJobOperator(task_id="startjob", deferrable=True, job_name=JOB_NAME, delay=15)
+    job_name = f"{env_id}-databrew-job"
 
+    # [START howto_operator_glue_databrew_start]
+    start_job = GlueDataBrewStartJobOperator(task_id="startjob", deferrable=True, job_name=job_name, delay=15)
     # [END howto_operator_glue_databrew_start]
 
     chain(test_context, start_job)

From 872cf244b8020a6b215d9720868b76c3145f8e00 Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Fri, 6 Oct 2023 15:15:41 -0400
Subject: [PATCH 08/21] Cleanup unused code in hook

---
 airflow/providers/amazon/aws/hooks/glue.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py
index 66ee83fc0a7fe..978334568e01c 100644
--- a/airflow/providers/amazon/aws/hooks/glue.py
+++ b/airflow/providers/amazon/aws/hooks/glue.py
@@ -435,11 +435,6 @@ class GlueDataBrewHook(AwsBaseHook):
         - :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
     """
 
-    TERMINAL_STATES = ["STOPPED", "SUCCEEDED", "FAILED", "TIMEOUT"]
-    SUCCESS_STATES = ["SUCCEEDED"]
-    FAILURE_STATES = ["FAILED", "TIMEOUT"]
-    RUNNING_STATES = ["STARTING", "RUNNING", "STOPPING"]
-
     def __init__(self, *args, **kwargs):
         kwargs["client_type"] = "databrew"
         super().__init__(*args, **kwargs)

From 631c7b4645bf9a573d59f0c4c6acb4b5cb07302e Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:39:35 -0400
Subject: [PATCH 09/21] Update
 tests/system/providers/amazon/aws/example_databrew.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 tests/system/providers/amazon/aws/example_databrew.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
index fce03eed666bd..01fb254d21b3a 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -30,7 +30,7 @@
 sys_test_context_task = SystemTestContextBuilder().build()
 
 
-@dag(DAG_ID, schedule=None, start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)
+@dag(DAG_ID, schedule="@once", start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)
 def run_job():
     test_context = sys_test_context_task()
     env_id = test_context["ENV_ID"]

From 7baa4d48ce391aeb869248abab73dfc2f413ea73 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:40:08 -0400
Subject: [PATCH 10/21] Update airflow/providers/amazon/aws/triggers/glue.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 airflow/providers/amazon/aws/triggers/glue.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/triggers/glue.py b/airflow/providers/amazon/aws/triggers/glue.py
index 72e9c177d4916..38d8f0b200313 100644
--- a/airflow/providers/amazon/aws/triggers/glue.py
+++ b/airflow/providers/amazon/aws/triggers/glue.py
@@ -167,7 +167,7 @@ def __init__(
         run_id: str,
         aws_conn_id: str,
         delay: int = 10,
-        maxAttempts: int = 60,
+        max_attempts: int = 60,
         **kwargs,
     ):
         super().__init__(**kwargs)

From 65ac3b4c235d60e88a8d47c2d0c67426abf5b113 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:40:25 -0400
Subject: [PATCH 11/21] Update airflow/providers/amazon/aws/triggers/glue.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 airflow/providers/amazon/aws/triggers/glue.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/triggers/glue.py b/airflow/providers/amazon/aws/triggers/glue.py
index 38d8f0b200313..ac9069d0ed45b 100644
--- a/airflow/providers/amazon/aws/triggers/glue.py
+++ b/airflow/providers/amazon/aws/triggers/glue.py
@@ -157,7 +157,7 @@ class GlueDataBrewJobCompleteTrigger(BaseTrigger):
     :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 maxAttempts: Maximum number of attempts to wait for the job to complete. Default is 60 attempts.
+    :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.
     """
 

From 87babf8f7598ee99570b0a7856de2bed3c482cc6 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:40:43 -0400
Subject: [PATCH 12/21] Update airflow/providers/amazon/aws/triggers/glue.py

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 airflow/providers/amazon/aws/triggers/glue.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/triggers/glue.py b/airflow/providers/amazon/aws/triggers/glue.py
index ac9069d0ed45b..79976cea3de43 100644
--- a/airflow/providers/amazon/aws/triggers/glue.py
+++ b/airflow/providers/amazon/aws/triggers/glue.py
@@ -175,7 +175,7 @@ def __init__(
         self.run_id = run_id
         self.aws_conn_id = aws_conn_id
         self.delay = delay
-        self.maxAttempts = maxAttempts
+        self.max_attempts = max_attempts
 
     @property
     def hook(self) -> GlueDataBrewHook:

From 2d9ae13e5c281011682fd3203ed16e5c6150b7ab Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Tue, 10 Oct 2023 11:41:37 -0400
Subject: [PATCH 13/21] Update
 docs/apache-airflow-providers-amazon/operators/glue.rst

Co-authored-by: Vincent <97131062+vincbeck@users.noreply.github.com>
---
 docs/apache-airflow-providers-amazon/operators/glue.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/apache-airflow-providers-amazon/operators/glue.rst b/docs/apache-airflow-providers-amazon/operators/glue.rst
index 5f90eb536194f..50956aaa15e4b 100644
--- a/docs/apache-airflow-providers-amazon/operators/glue.rst
+++ b/docs/apache-airflow-providers-amazon/operators/glue.rst
@@ -100,6 +100,8 @@ use :class:`~airflow.providers.amazon.aws.sensors.glue.GlueJobSensor`
     :start-after: [START howto_sensor_glue]
     :end-before: [END howto_sensor_glue]
 
+.. _howto/sensor:GlueDataBrewStartJobOperator:
+
 Start an AWS Glue DataBrew job
 ==============================
 

From 5478451e7a53359efa0afded16321de659924940 Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Tue, 10 Oct 2023 13:21:50 -0400
Subject: [PATCH 14/21] Use AwsBaseWaiterTrigger in DataBrew trigger

---
 airflow/providers/amazon/aws/hooks/glue.py    |  4 +-
 .../providers/amazon/aws/operators/glue.py    | 14 +++---
 airflow/providers/amazon/aws/triggers/glue.py | 48 ++++++-------------
 .../operators/glue.rst                        |  2 +
 4 files changed, 25 insertions(+), 43 deletions(-)

diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py
index 978334568e01c..ae6c6bcff2e1c 100644
--- a/airflow/providers/amazon/aws/hooks/glue.py
+++ b/airflow/providers/amazon/aws/hooks/glue.py
@@ -439,7 +439,7 @@ 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, maxAttempts: int = 60) -> str:
+    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.
 
@@ -452,7 +452,7 @@ def job_completion(self, job_name: str, run_id: str, delay: int = 10, maxAttempt
         self.get_waiter("job_complete").wait(
             Name=job_name,
             RunId=run_id,
-            WaiterConfig={"Delay": delay, "maxAttempts": maxAttempts},
+            WaiterConfig={"Delay": delay, "maxAttempts": max_attempts},
         )
 
         status = self.get_job_state(job_name, run_id)
diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py
index 40433e826b12f..5c4264e0fb22f 100644
--- a/airflow/providers/amazon/aws/operators/glue.py
+++ b/airflow/providers/amazon/aws/operators/glue.py
@@ -282,11 +282,11 @@ def execute(self, context: Context):
         resp = {}
         resp["job_name"] = self.job_name
 
-        job = self.data_brew_hook.conn.start_job_run(Name=self.job_name)
+        job = self.hook.conn.start_job_run(Name=self.job_name)
         run_id = job["RunId"]
         resp["run_id"] = run_id
 
-        status = self.data_brew_hook.get_job_state(self.job_name, run_id)
+        status = self.hook.get_job_state(self.job_name, run_id)
         resp["status"] = status
 
         self.log.info(
@@ -306,9 +306,7 @@ def execute(self, context: Context):
             self.log.info(
                 "Waiting for AWS Glue DataBrew Job: %s. Run Id: %s to complete.", self.job_name, run_id
             )
-            status = self.data_brew_hook.job_completion(
-                job_name=self.job_name, delay=self.delay, run_id=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)
             resp["status"] = status
 
@@ -316,9 +314,9 @@ def execute(self, context: Context):
 
     def execute_complete(self, context: Context, event=None) -> dict[str, str]:
         result = {
-                "job_name": event.get("jobName", ""),
-                "run_id": event.get("runId", ""),
-                "status": event.get("status", ""),
+            "job_name": event.get("jobName", ""),
+            "run_id": event.get("runId", ""),
+            "status": event.get("status", ""),
         }
         self.log.info(
             "AWS Glue DataBrew Job: %s runId: %s status: %s",
diff --git a/airflow/providers/amazon/aws/triggers/glue.py b/airflow/providers/amazon/aws/triggers/glue.py
index 79976cea3de43..ee03d14b074de 100644
--- a/airflow/providers/amazon/aws/triggers/glue.py
+++ b/airflow/providers/amazon/aws/triggers/glue.py
@@ -23,6 +23,7 @@
 
 from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
 from airflow.providers.amazon.aws.hooks.glue_catalog import GlueCatalogHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -150,7 +151,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                     await asyncio.sleep(self.waiter_delay)
 
 
-class GlueDataBrewJobCompleteTrigger(BaseTrigger):
+class GlueDataBrewJobCompleteTrigger(AwsBaseWaiterTrigger):
     """
     Watches for a Glue DataBrew job, triggers when it finishes.
 
@@ -170,38 +171,19 @@ def __init__(
         max_attempts: int = 60,
         **kwargs,
     ):
-        super().__init__(**kwargs)
-        self.job_name = job_name
-        self.run_id = run_id
-        self.aws_conn_id = aws_conn_id
-        self.delay = delay
-        self.max_attempts = max_attempts
+        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,
+        )
 
-    @property
     def hook(self) -> GlueDataBrewHook:
         return GlueDataBrewHook(aws_conn_id=self.aws_conn_id)
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        return (
-            # dynamically generate the fully qualified name of the class
-            self.__class__.__module__ + "." + self.__class__.__qualname__,
-            {
-                "job_name": self.job_name,
-                "run_id": self.run_id,
-                "aws_conn_id": self.aws_conn_id,
-                "delay": self.delay,
-                "maxAttempts": self.maxAttempts,
-            },
-        )
-
-    async def run(self) -> AsyncIterator[TriggerEvent]:
-        async with self.hook.async_conn as client:
-            waiter = self.hook.get_waiter("job_complete", deferrable=True, client=client)
-            await waiter.wait(
-                Name=self.job_name,
-                RunId=self.run_id,
-                WaiterConfig={"Delay": self.delay, "maxAttempts": self.maxAttempts},
-            )
-
-        result = self.hook.get_job_state(self.job_name, self.run_id)
-        yield TriggerEvent({"status": result, "RunId": self.run_id, "JobName": self.job_name})
diff --git a/docs/apache-airflow-providers-amazon/operators/glue.rst b/docs/apache-airflow-providers-amazon/operators/glue.rst
index 50956aaa15e4b..a68bfb78107fd 100644
--- a/docs/apache-airflow-providers-amazon/operators/glue.rst
+++ b/docs/apache-airflow-providers-amazon/operators/glue.rst
@@ -116,4 +116,6 @@ To submit a new AWS Glue DataBrew job you can use :class:`~airflow.providers.ama
 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>`__

From 610259eebf9f3d6b33e94407664c3cd7a43f7d88 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Wed, 11 Oct 2023 08:29:44 -0400
Subject: [PATCH 15/21] Update airflow/providers/amazon/aws/operators/glue.py

Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com>
---
 airflow/providers/amazon/aws/operators/glue.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py
index 5c4264e0fb22f..3f8c5b901259f 100644
--- a/airflow/providers/amazon/aws/operators/glue.py
+++ b/airflow/providers/amazon/aws/operators/glue.py
@@ -233,7 +233,8 @@ def on_kill(self):
 
 
 class GlueDataBrewStartJobOperator(BaseOperator):
-    """Start an AWS Glue DataBrew job.
+    """
+    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

From 8ea4d63c027987cf87b0d49860bff14aaed06ae4 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Wed, 11 Oct 2023 08:30:01 -0400
Subject: [PATCH 16/21] Update
 tests/providers/amazon/aws/waiters/test_custom_waiters.py

Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com>
---
 tests/providers/amazon/aws/waiters/test_custom_waiters.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/providers/amazon/aws/waiters/test_custom_waiters.py b/tests/providers/amazon/aws/waiters/test_custom_waiters.py
index d7bd5b6ead442..6a9a23cb26a4b 100644
--- a/tests/providers/amazon/aws/waiters/test_custom_waiters.py
+++ b/tests/providers/amazon/aws/waiters/test_custom_waiters.py
@@ -452,6 +452,7 @@ def mock_describe_job_runs(self):
     def describe_jobs(status: str):
         """
         Helper function for generate minimal DescribeJobRun response for a single job.
+        
         https://docs.aws.amazon.com/databrew/latest/dg/API_DescribeJobRun.html
         """
         return {"State": status}

From 01fd683ed574dcd263f00ab22b388b0e4bf7cc50 Mon Sep 17 00:00:00 2001
From: ellisms <114107920+ellisms@users.noreply.github.com>
Date: Thu, 12 Oct 2023 06:52:27 -0400
Subject: [PATCH 17/21] Update
 airflow/providers/amazon/aws/waiters/databrew.json

Co-authored-by: Andrey Anshin <Andrey.Anshin@taragol.is>
---
 .../providers/amazon/aws/waiters/databrew.json | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/airflow/providers/amazon/aws/waiters/databrew.json b/airflow/providers/amazon/aws/waiters/databrew.json
index a18dc23d23a47..41372def5b294 100644
--- a/airflow/providers/amazon/aws/waiters/databrew.json
+++ b/airflow/providers/amazon/aws/waiters/databrew.json
@@ -29,24 +29,6 @@
                     "argument": "State",
                     "expected": "TIMEOUT",
                     "state": "success"
-                },
-                {
-                    "matcher": "path",
-                    "argument": "State",
-                    "expected": "Starting",
-                    "state": "retry"
-                },
-                {
-                    "matcher": "path",
-                    "argument": "State",
-                    "expected": "STOPPING",
-                    "state": "retry"
-                },
-                {
-                    "matcher": "path",
-                    "argument": "State",
-                    "expected": "RUNNING",
-                    "state": "retry"
                 }
             ]
         }

From 79b1e9d92b767a203dd4adeb1ec4880b4e9a9dc0 Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Thu, 12 Oct 2023 11:50:33 -0400
Subject: [PATCH 18/21] place DataBrew components in their own files. Other
 minor PR changes

---
 .../providers/amazon/aws/hooks/databrew.py    |  68 +++++++++++
 airflow/providers/amazon/aws/hooks/glue.py    |  47 --------
 .../amazon/aws/operators/databrew.py          | 110 ++++++++++++++++++
 .../providers/amazon/aws/operators/glue.py    | 100 +---------------
 .../providers/amazon/aws/triggers/databrew.py |  59 ++++++++++
 airflow/providers/amazon/aws/triggers/glue.py |  41 +------
 airflow/providers/amazon/provider.yaml        |   9 ++
 .../operators/databrew.rst                    |  53 +++++++++
 .../operators/glue.rst                        |  13 ---
 .../amazon/aws/hooks/test_databrew.py         |  38 ++++++
 tests/providers/amazon/aws/hooks/test_glue.py |  14 +--
 ...test_glue_databrew.py => test_databrew.py} |  24 +++-
 .../amazon/aws/waiters/test_custom_waiters.py |  41 -------
 .../amazon/aws/waiters/test_databrew.py       |  68 +++++++++++
 .../providers/amazon/aws/example_databrew.py  |  14 +--
 15 files changed, 435 insertions(+), 264 deletions(-)
 create mode 100644 airflow/providers/amazon/aws/hooks/databrew.py
 create mode 100644 airflow/providers/amazon/aws/operators/databrew.py
 create mode 100644 airflow/providers/amazon/aws/triggers/databrew.py
 create mode 100644 docs/apache-airflow-providers-amazon/operators/databrew.rst
 create mode 100644 tests/providers/amazon/aws/hooks/test_databrew.py
 rename tests/providers/amazon/aws/operators/{test_glue_databrew.py => test_databrew.py} (58%)
 create mode 100644 tests/providers/amazon/aws/waiters/test_databrew.py

diff --git a/airflow/providers/amazon/aws/hooks/databrew.py b/airflow/providers/amazon/aws/hooks/databrew.py
new file mode 100644
index 0000000000000..f6c7b3ebd6f28
--- /dev/null
+++ b/airflow/providers/amazon/aws/hooks/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"]
diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py
index ae6c6bcff2e1c..baf6780e07802 100644
--- a/airflow/providers/amazon/aws/hooks/glue.py
+++ b/airflow/providers/amazon/aws/hooks/glue.py
@@ -422,50 +422,3 @@ def create_or_update_glue_job(self) -> str | None:
             self.conn.create_job(**config)
 
         return self.job_name
-
-
-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.get_conn().describe_job_run(Name=job_name, RunId=run_id)
-        return response["State"]
diff --git a/airflow/providers/amazon/aws/operators/databrew.py b/airflow/providers/amazon/aws/operators/databrew.py
new file mode 100644
index 0000000000000..a6569a1801df7
--- /dev/null
+++ b/airflow/providers/amazon/aws/operators/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.databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.triggers.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 runId.
+    """
+
+    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}
diff --git a/airflow/providers/amazon/aws/operators/glue.py b/airflow/providers/amazon/aws/operators/glue.py
index 3f8c5b901259f..97bcaba66abc7 100644
--- a/airflow/providers/amazon/aws/operators/glue.py
+++ b/airflow/providers/amazon/aws/operators/glue.py
@@ -25,10 +25,10 @@
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
-from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
+from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.links.glue import GlueJobRunDetailsLink
-from airflow.providers.amazon.aws.triggers.glue import GlueDataBrewJobCompleteTrigger, GlueJobCompleteTrigger
+from airflow.providers.amazon.aws.triggers.glue import GlueJobCompleteTrigger
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -230,99 +230,3 @@ def on_kill(self):
             )
             if not response["SuccessfulSubmissions"]:
                 self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", self.job_name, self._job_run_id)
-
-
-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.
-    """
-
-    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):
-        resp = {}
-        resp["job_name"] = self.job_name
-
-        job = self.hook.conn.start_job_run(Name=self.job_name)
-        run_id = job["RunId"]
-        resp["run_id"] = run_id
-
-        status = self.hook.get_job_state(self.job_name, run_id)
-        resp["status"] = status
-
-        self.log.info(
-            "AWS Glue DataBrew Job: %s. Run Id: %s submitted. Status: %s", self.job_name, run_id, status
-        )
-
-        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)
-            resp["status"] = status
-
-        return resp
-
-    def execute_complete(self, context: Context, event=None) -> dict[str, str]:
-        result = {
-            "job_name": event.get("jobName", ""),
-            "run_id": event.get("runId", ""),
-            "status": event.get("status", ""),
-        }
-        self.log.info(
-            "AWS Glue DataBrew Job: %s runId: %s status: %s",
-            result["job_name"],
-            result["run_id"],
-            result["status"],
-        )
-        return result
diff --git a/airflow/providers/amazon/aws/triggers/databrew.py b/airflow/providers/amazon/aws/triggers/databrew.py
new file mode 100644
index 0000000000000..697df713a3235
--- /dev/null
+++ b/airflow/providers/amazon/aws/triggers/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.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)
diff --git a/airflow/providers/amazon/aws/triggers/glue.py b/airflow/providers/amazon/aws/triggers/glue.py
index ee03d14b074de..b5a0a673d3e18 100644
--- a/airflow/providers/amazon/aws/triggers/glue.py
+++ b/airflow/providers/amazon/aws/triggers/glue.py
@@ -21,9 +21,8 @@
 from functools import cached_property
 from typing import Any, AsyncIterator
 
-from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
+from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
 from airflow.providers.amazon.aws.hooks.glue_catalog import GlueCatalogHook
-from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -149,41 +148,3 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                     break
                 else:
                     await asyncio.sleep(self.waiter_delay)
-
-
-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)
diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml
index 2f5796b64a2cf..bfb32c1e5b22a 100644
--- a/airflow/providers/amazon/provider.yaml
+++ b/airflow/providers/amazon/provider.yaml
@@ -364,6 +364,9 @@ operators:
   - integration-name: Amazon Appflow
     python-modules:
       - airflow.providers.amazon.aws.operators.appflow
+  - integration-name: Amazon Glue DataBrew
+    python-modules:
+      - airflow.providers.amazon.aws.operators.databrew
 
 sensors:
   - integration-name: Amazon Athena
@@ -540,6 +543,9 @@ hooks:
   - integration-name: Amazon Appflow
     python-modules:
       - airflow.providers.amazon.aws.hooks.appflow
+  - integration-name: Amazon Glue DataBrew
+    python-modules:
+      - airflow.providers.amazon.aws.hooks.databrew
 
 triggers:
   - integration-name: Amazon Web Services
@@ -588,6 +594,9 @@ triggers:
   - integration-name: Amazon Simple Queue Service (SQS)
     python-modules:
       - airflow.providers.amazon.aws.triggers.sqs
+  - integration-name: Amazon Glue DataBrew
+    python-modules:
+      - airflow.providers.amazon.aws.triggers.databrew
 
 transfers:
   - source-integration-name: Amazon DynamoDB
diff --git a/docs/apache-airflow-providers-amazon/operators/databrew.rst b/docs/apache-airflow-providers-amazon/operators/databrew.rst
new file mode 100644
index 0000000000000..aa55d1be3ca50
--- /dev/null
+++ b/docs/apache-airflow-providers-amazon/operators/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.GlueDataBrewStartJobOperator`.
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_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>`__
diff --git a/docs/apache-airflow-providers-amazon/operators/glue.rst b/docs/apache-airflow-providers-amazon/operators/glue.rst
index a68bfb78107fd..ddd21205c1673 100644
--- a/docs/apache-airflow-providers-amazon/operators/glue.rst
+++ b/docs/apache-airflow-providers-amazon/operators/glue.rst
@@ -100,19 +100,6 @@ use :class:`~airflow.providers.amazon.aws.sensors.glue.GlueJobSensor`
     :start-after: [START howto_sensor_glue]
     :end-before: [END howto_sensor_glue]
 
-.. _howto/sensor: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.GlueDataBrewStartJobOperator`.
-
-.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_databrew.py
-    :language: python
-    :dedent: 4
-    :start-after: [START howto_operator_glue_databrew_start]
-    :end-before: [END howto_operator_glue_databrew_start]
-
 Reference
 ---------
 
diff --git a/tests/providers/amazon/aws/hooks/test_databrew.py b/tests/providers/amazon/aws/hooks/test_databrew.py
new file mode 100644
index 0000000000000..9b5153edd5fad
--- /dev/null
+++ b/tests/providers/amazon/aws/hooks/test_databrew.py
@@ -0,0 +1,38 @@
+#
+# 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 typing import TYPE_CHECKING
+from unittest import mock
+
+from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
+
+if TYPE_CHECKING:
+    from unittest.mock import MagicMock
+
+
+class TestGlueDataBrewHook:
+    job_name = "test-databrew-job"
+    runId = "test12345"
+
+    @mock.patch.object(GlueDataBrewHook, "get_job_state")
+    def test_get_job_state(self, get_job_state_mock: MagicMock):
+        get_job_state_mock.return_value = "SUCCEEDED"
+        hook = GlueDataBrewHook()
+        result = hook.get_job_state(self.job_name, self.runId)
+        assert result == "SUCCEEDED"
diff --git a/tests/providers/amazon/aws/hooks/test_glue.py b/tests/providers/amazon/aws/hooks/test_glue.py
index 1943eb0c0d19c..799547dd576d2 100644
--- a/tests/providers/amazon/aws/hooks/test_glue.py
+++ b/tests/providers/amazon/aws/hooks/test_glue.py
@@ -28,7 +28,7 @@
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook, GlueJobHook
+from airflow.providers.amazon.aws.hooks.glue import GlueJobHook
 
 if TYPE_CHECKING:
     from unittest.mock import MagicMock
@@ -477,15 +477,3 @@ async def test_async_job_completion_failure(self, get_state_mock: MagicMock):
             await hook.async_job_completion("job_name", "run_id")
 
         assert get_state_mock.call_count == 3
-
-
-class TestGlueDataBrewHook:
-    job_name = "test-databrew-job"
-    runId = "test12345"
-
-    @mock.patch.object(GlueDataBrewHook, "get_job_state")
-    def test_get_job_state(self, get_job_state_mock: MagicMock):
-        get_job_state_mock.return_value = "SUCCEEDED"
-        hook = GlueDataBrewHook()
-        result = hook.get_job_state(self.job_name, self.runId)
-        assert result == "SUCCEEDED"
diff --git a/tests/providers/amazon/aws/operators/test_glue_databrew.py b/tests/providers/amazon/aws/operators/test_databrew.py
similarity index 58%
rename from tests/providers/amazon/aws/operators/test_glue_databrew.py
rename to tests/providers/amazon/aws/operators/test_databrew.py
index cc7bd36b433b3..07071ddc04070 100644
--- a/tests/providers/amazon/aws/operators/test_glue_databrew.py
+++ b/tests/providers/amazon/aws/operators/test_databrew.py
@@ -22,8 +22,8 @@
 import pytest
 from moto import mock_databrew
 
-from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook
-from airflow.providers.amazon.aws.operators.glue import GlueDataBrewStartJobOperator
+from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.operators.databrew import GlueDataBrewStartJobOperator
 
 JOB_NAME = "test_job"
 
@@ -35,10 +35,24 @@ def hook() -> GlueDataBrewHook:
 
 
 class TestGlueDataBrewOperator:
+    @mock.patch.object(GlueDataBrewHook, "conn")
     @mock.patch.object(GlueDataBrewHook, "get_waiter")
-    def test_start_job_wait_for_completion(self, mock_hook_get_waiter):
+    def test_start_job_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+        TEST_RUN_ID = "12345"
         operator = GlueDataBrewStartJobOperator(
-            task_id="task_test", job_name=JOB_NAME, wait_for_completion=True
+            task_id="task_test", job_name=JOB_NAME, wait_for_completion=True, aws_conn_id="aws_default"
         )
+        mock_conn.start_job_run(mock.MagicMock(), return_value=TEST_RUN_ID)
         operator.execute(None)
-        mock_hook_get_waiter.assert_called_once_with("databrew_job_complete")
+        mock_hook_get_waiter.assert_called_once_with("job_complete")
+
+    @mock.patch.object(GlueDataBrewHook, "conn")
+    @mock.patch.object(GlueDataBrewHook, "get_waiter")
+    def test_start_job_no_wait(self, mock_hook_get_waiter, mock_conn):
+        TEST_RUN_ID = "12345"
+        operator = GlueDataBrewStartJobOperator(
+            task_id="task_test", job_name=JOB_NAME, wait_for_completion=False, aws_conn_id="aws_default"
+        )
+        mock_conn.start_job_run(mock.MagicMock(), return_value=TEST_RUN_ID)
+        operator.execute(None)
+        mock_hook_get_waiter.assert_not_called()
diff --git a/tests/providers/amazon/aws/waiters/test_custom_waiters.py b/tests/providers/amazon/aws/waiters/test_custom_waiters.py
index 6a9a23cb26a4b..19f9296b6af71 100644
--- a/tests/providers/amazon/aws/waiters/test_custom_waiters.py
+++ b/tests/providers/amazon/aws/waiters/test_custom_waiters.py
@@ -33,7 +33,6 @@
 from airflow.providers.amazon.aws.hooks.ecs import EcsClusterStates, EcsHook, EcsTaskDefinitionStates
 from airflow.providers.amazon.aws.hooks.eks import EksHook
 from airflow.providers.amazon.aws.hooks.emr import EmrHook
-from airflow.providers.amazon.aws.hooks.glue import GlueDataBrewHook
 from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter
 
 
@@ -425,43 +424,3 @@ def test_steps_failed(self, mock_list_steps):
                 StepIds=[self.STEP_ID1, self.STEP_ID2],
                 WaiterConfig={"Delay": 0.01, "MaxAttempts": 3},
             )
-
-
-class TestCustomDataBrewWaiters:
-    """Test waiters from ``amazon/aws/waiters/glue.json``."""
-
-    JOB_NAME = "test_job"
-    RUN_ID = "123"
-
-    @pytest.fixture(autouse=True)
-    def setup_test_cases(self, monkeypatch):
-        self.client = boto3.client("databrew", region_name="eu-west-3")
-        monkeypatch.setattr(GlueDataBrewHook, "conn", self.client)
-
-    def test_service_waiters(self):
-        hook_waiters = GlueDataBrewHook(aws_conn_id=None).list_waiters()
-        assert "job_complete" in hook_waiters
-
-    @pytest.fixture
-    def mock_describe_job_runs(self):
-        """Mock ``GlueDataBrewHook.Client.describe_job_run`` method."""
-        with mock.patch.object(self.client, "describe_job_run") as m:
-            yield m
-
-    @staticmethod
-    def describe_jobs(status: str):
-        """
-        Helper function for generate minimal DescribeJobRun response for a single job.
-        
-        https://docs.aws.amazon.com/databrew/latest/dg/API_DescribeJobRun.html
-        """
-        return {"State": status}
-
-    def test_job_succeeded(self, mock_describe_job_runs):
-        """Test job succeeded"""
-        mock_describe_job_runs.side_effect = [
-            self.describe_jobs(GlueDataBrewHook.RUNNING_STATES[1]),
-            self.describe_jobs(GlueDataBrewHook.TERMINAL_STATES[1]),
-        ]
-        waiter = GlueDataBrewHook(aws_conn_id=None).get_waiter("job_complete")
-        waiter.wait(name=self.JOB_NAME, runId=self.RUN_ID, WaiterConfig={"Delay": 0.2, "MaxAttempts": 2})
diff --git a/tests/providers/amazon/aws/waiters/test_databrew.py b/tests/providers/amazon/aws/waiters/test_databrew.py
new file mode 100644
index 0000000000000..2a0fad503edad
--- /dev/null
+++ b/tests/providers/amazon/aws/waiters/test_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 unittest import mock
+
+import boto3
+import pytest
+
+from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
+
+RUNNING_STATES = ["STARTING", "RUNNING", "STOPPING"]
+TERMINAL_STATES = ["STOPPED", "SUCCEEDED", "FAILED"]
+
+
+class TestCustomDataBrewWaiters:
+    """Test waiters from ``amazon/aws/waiters/glue.json``."""
+
+    JOB_NAME = "test_job"
+    RUN_ID = "123"
+
+    @pytest.fixture(autouse=True)
+    def setup_test_cases(self, monkeypatch):
+        self.client = boto3.client("databrew", region_name="eu-west-3")
+        monkeypatch.setattr(GlueDataBrewHook, "conn", self.client)
+
+    def test_service_waiters(self):
+        hook_waiters = GlueDataBrewHook(aws_conn_id=None).list_waiters()
+        assert "job_complete" in hook_waiters
+
+    @pytest.fixture
+    def mock_describe_job_runs(self):
+        """Mock ``GlueDataBrewHook.Client.describe_job_run`` method."""
+        with mock.patch.object(self.client, "describe_job_run") as m:
+            yield m
+
+    @staticmethod
+    def describe_jobs(status: str):
+        """
+        Helper function for generate minimal DescribeJobRun response for a single job.
+
+        https://docs.aws.amazon.com/databrew/latest/dg/API_DescribeJobRun.html
+        """
+        return {"State": status}
+
+    def test_job_succeeded(self, mock_describe_job_runs):
+        """Test job succeeded"""
+        mock_describe_job_runs.side_effect = [
+            self.describe_jobs(RUNNING_STATES[1]),
+            self.describe_jobs(TERMINAL_STATES[1]),
+        ]
+        waiter = GlueDataBrewHook(aws_conn_id=None).get_waiter("job_complete")
+        waiter.wait(name=self.JOB_NAME, runId=self.RUN_ID, WaiterConfig={"Delay": 0.2, "MaxAttempts": 2})
diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_databrew.py
index 01fb254d21b3a..7d35df98498ad 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_databrew.py
@@ -18,9 +18,9 @@
 
 import pendulum
 
-from airflow.decorators import dag
 from airflow.models.baseoperator import chain
-from airflow.providers.amazon.aws.operators.glue import (
+from airflow.models.dag import DAG
+from airflow.providers.amazon.aws.operators.databrew import (
     GlueDataBrewStartJobOperator,
 )
 from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
@@ -30,8 +30,7 @@
 sys_test_context_task = SystemTestContextBuilder().build()
 
 
-@dag(DAG_ID, schedule="@once", start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False)
-def run_job():
+with DAG(DAG_ID, schedule="@once", start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False) as dag:
     test_context = sys_test_context_task()
     env_id = test_context["ENV_ID"]
 
@@ -42,13 +41,14 @@ def run_job():
     # [END howto_operator_glue_databrew_start]
 
     chain(test_context, start_job)
-    
-    from tests.system.utils.watcher import watcher  # noqa: E402
-    
+
+    from tests.system.utils.watcher import watcher
+
     # This test needs watcher in order to properly mark success/failure
     # when "tearDown" task with trigger rule is part of the DAG
     list(dag.tasks) >> watcher()
 
+
 from tests.system.utils import get_test_run  # noqa: E402
 
 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)

From 9863199db5f4e66bdc5f616266934b0f00c2d2c8 Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Thu, 12 Oct 2023 13:02:10 -0400
Subject: [PATCH 19/21] Pipeline failure fixes

---
 airflow/providers/amazon/aws/operators/databrew.py              | 2 +-
 docs/apache-airflow-providers-amazon/operators/databrew.rst     | 2 +-
 .../aws/triggers/{test_glue_databrew.py => test_databrew.py}    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)
 rename tests/providers/amazon/aws/triggers/{test_glue_databrew.py => test_databrew.py} (94%)

diff --git a/airflow/providers/amazon/aws/operators/databrew.py b/airflow/providers/amazon/aws/operators/databrew.py
index a6569a1801df7..5cd0ee4b5dd93 100644
--- a/airflow/providers/amazon/aws/operators/databrew.py
+++ b/airflow/providers/amazon/aws/operators/databrew.py
@@ -47,7 +47,7 @@ class GlueDataBrewStartJobOperator(BaseOperator):
         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 runId.
+    :return: dictionary with key run_id and value of the resulting job's run_id.
     """
 
     template_fields: Sequence[str] = (
diff --git a/docs/apache-airflow-providers-amazon/operators/databrew.rst b/docs/apache-airflow-providers-amazon/operators/databrew.rst
index aa55d1be3ca50..2611c09279b48 100644
--- a/docs/apache-airflow-providers-amazon/operators/databrew.rst
+++ b/docs/apache-airflow-providers-amazon/operators/databrew.rst
@@ -39,7 +39,7 @@ Operators
 Start an AWS Glue DataBrew job
 ==============================
 
-To submit a new AWS Glue DataBrew job you can use :class:`~airflow.providers.amazon.aws.operators.glue.GlueDataBrewStartJobOperator`.
+To submit a new AWS Glue DataBrew job you can use :class:`~airflow.providers.amazon.aws.operators.databrew.GlueDataBrewStartJobOperator`.
 
 .. exampleinclude:: /../../tests/system/providers/amazon/aws/example_databrew.py
     :language: python
diff --git a/tests/providers/amazon/aws/triggers/test_glue_databrew.py b/tests/providers/amazon/aws/triggers/test_databrew.py
similarity index 94%
rename from tests/providers/amazon/aws/triggers/test_glue_databrew.py
rename to tests/providers/amazon/aws/triggers/test_databrew.py
index 87577c9f809e1..6df4a6319f71d 100644
--- a/tests/providers/amazon/aws/triggers/test_glue_databrew.py
+++ b/tests/providers/amazon/aws/triggers/test_databrew.py
@@ -18,7 +18,7 @@
 
 import pytest
 
-from airflow.providers.amazon.aws.triggers.glue import GlueDataBrewJobCompleteTrigger
+from airflow.providers.amazon.aws.triggers.databrew import GlueDataBrewJobCompleteTrigger
 
 TEST_JOB_NAME = "test_job_name"
 TEST_JOB_RUN_ID = "a1234"

From c7410fd4992017c41d8ba3a15d87560e6f7d25a6 Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Thu, 12 Oct 2023 14:37:22 -0400
Subject: [PATCH 20/21] provider.yaml fix and added DataBrew logo

---
 airflow/providers/amazon/provider.yaml          |  12 +++++++++---
 .../aws/AWS-Glue-DataBrew_64.png                | Bin 0 -> 14575 bytes
 2 files changed, 9 insertions(+), 3 deletions(-)
 create mode 100644 docs/integration-logos/aws/AWS-Glue-DataBrew_64.png

diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml
index bfb32c1e5b22a..7362b18d3f9db 100644
--- a/airflow/providers/amazon/provider.yaml
+++ b/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/databrew.rst
+    logo: /integration-logos/aws/AWS-Glue-DataBrew_64.png
+    tags: [aws]
 
 operators:
   - integration-name: Amazon Athena
@@ -364,7 +370,7 @@ operators:
   - integration-name: Amazon Appflow
     python-modules:
       - airflow.providers.amazon.aws.operators.appflow
-  - integration-name: Amazon Glue DataBrew
+  - integration-name: AWS Glue DataBrew
     python-modules:
       - airflow.providers.amazon.aws.operators.databrew
 
@@ -543,7 +549,7 @@ hooks:
   - integration-name: Amazon Appflow
     python-modules:
       - airflow.providers.amazon.aws.hooks.appflow
-  - integration-name: Amazon Glue DataBrew
+  - integration-name: AWS Glue DataBrew
     python-modules:
       - airflow.providers.amazon.aws.hooks.databrew
 
@@ -594,7 +600,7 @@ triggers:
   - integration-name: Amazon Simple Queue Service (SQS)
     python-modules:
       - airflow.providers.amazon.aws.triggers.sqs
-  - integration-name: Amazon Glue DataBrew
+  - integration-name: AWS Glue DataBrew
     python-modules:
       - airflow.providers.amazon.aws.triggers.databrew
 
diff --git a/docs/integration-logos/aws/AWS-Glue-DataBrew_64.png b/docs/integration-logos/aws/AWS-Glue-DataBrew_64.png
new file mode 100644
index 0000000000000000000000000000000000000000..9821fcdf26e94832189fe5595dbad1fa8fa5fb67
GIT binary patch

From 357164bf21307881ace099cefb9b22e91f541407 Mon Sep 17 00:00:00 2001
From: Mike Ellis <ellisms@amazon.com>
Date: Fri, 13 Oct 2023 09:50:02 -0400
Subject: [PATCH 21/21] change databrew filenames to glue_databrew

---
 .../amazon/aws/hooks/{databrew.py => glue_databrew.py}    | 0
 .../aws/operators/{databrew.py => glue_databrew.py}       | 4 ++--
 .../amazon/aws/triggers/{databrew.py => glue_databrew.py} | 2 +-
 airflow/providers/amazon/provider.yaml                    | 8 ++++----
 .../operators/{databrew.rst => glue_databrew.rst}         | 4 ++--
 .../aws/hooks/{test_databrew.py => test_glue_databrew.py} | 2 +-
 .../operators/{test_databrew.py => test_glue_databrew.py} | 4 ++--
 .../triggers/{test_databrew.py => test_glue_databrew.py}  | 2 +-
 .../waiters/{test_databrew.py => test_glue_databrew.py}   | 2 +-
 .../aws/{example_databrew.py => example_glue_databrew.py} | 2 +-
 10 files changed, 15 insertions(+), 15 deletions(-)
 rename airflow/providers/amazon/aws/hooks/{databrew.py => glue_databrew.py} (100%)
 rename airflow/providers/amazon/aws/operators/{databrew.py => glue_databrew.py} (95%)
 rename airflow/providers/amazon/aws/triggers/{databrew.py => glue_databrew.py} (96%)
 rename docs/apache-airflow-providers-amazon/operators/{databrew.rst => glue_databrew.rst} (95%)
 rename tests/providers/amazon/aws/hooks/{test_databrew.py => test_glue_databrew.py} (94%)
 rename tests/providers/amazon/aws/operators/{test_databrew.py => test_glue_databrew.py} (92%)
 rename tests/providers/amazon/aws/triggers/{test_databrew.py => test_glue_databrew.py} (93%)
 rename tests/providers/amazon/aws/waiters/{test_databrew.py => test_glue_databrew.py} (96%)
 rename tests/system/providers/amazon/aws/{example_databrew.py => example_glue_databrew.py} (96%)

diff --git a/airflow/providers/amazon/aws/hooks/databrew.py b/airflow/providers/amazon/aws/hooks/glue_databrew.py
similarity index 100%
rename from airflow/providers/amazon/aws/hooks/databrew.py
rename to airflow/providers/amazon/aws/hooks/glue_databrew.py
diff --git a/airflow/providers/amazon/aws/operators/databrew.py b/airflow/providers/amazon/aws/operators/glue_databrew.py
similarity index 95%
rename from airflow/providers/amazon/aws/operators/databrew.py
rename to airflow/providers/amazon/aws/operators/glue_databrew.py
index 5cd0ee4b5dd93..596a507397f5f 100644
--- a/airflow/providers/amazon/aws/operators/databrew.py
+++ b/airflow/providers/amazon/aws/operators/glue_databrew.py
@@ -22,8 +22,8 @@
 
 from airflow.configuration import conf
 from airflow.models import BaseOperator
-from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
-from airflow.providers.amazon.aws.triggers.databrew import GlueDataBrewJobCompleteTrigger
+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
diff --git a/airflow/providers/amazon/aws/triggers/databrew.py b/airflow/providers/amazon/aws/triggers/glue_databrew.py
similarity index 96%
rename from airflow/providers/amazon/aws/triggers/databrew.py
rename to airflow/providers/amazon/aws/triggers/glue_databrew.py
index 697df713a3235..595b653e2fbf8 100644
--- a/airflow/providers/amazon/aws/triggers/databrew.py
+++ b/airflow/providers/amazon/aws/triggers/glue_databrew.py
@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 
 
diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml
index 7362b18d3f9db..83b43ee588783 100644
--- a/airflow/providers/amazon/provider.yaml
+++ b/airflow/providers/amazon/provider.yaml
@@ -291,7 +291,7 @@ integrations:
   - 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/databrew.rst
+      - /docs/apache-airflow-providers-amazon/operators/glue_databrew.rst
     logo: /integration-logos/aws/AWS-Glue-DataBrew_64.png
     tags: [aws]
 
@@ -372,7 +372,7 @@ operators:
       - airflow.providers.amazon.aws.operators.appflow
   - integration-name: AWS Glue DataBrew
     python-modules:
-      - airflow.providers.amazon.aws.operators.databrew
+      - airflow.providers.amazon.aws.operators.glue_databrew
 
 sensors:
   - integration-name: Amazon Athena
@@ -551,7 +551,7 @@ hooks:
       - airflow.providers.amazon.aws.hooks.appflow
   - integration-name: AWS Glue DataBrew
     python-modules:
-      - airflow.providers.amazon.aws.hooks.databrew
+      - airflow.providers.amazon.aws.hooks.glue_databrew
 
 triggers:
   - integration-name: Amazon Web Services
@@ -602,7 +602,7 @@ triggers:
       - airflow.providers.amazon.aws.triggers.sqs
   - integration-name: AWS Glue DataBrew
     python-modules:
-      - airflow.providers.amazon.aws.triggers.databrew
+      - airflow.providers.amazon.aws.triggers.glue_databrew
 
 transfers:
   - source-integration-name: Amazon DynamoDB
diff --git a/docs/apache-airflow-providers-amazon/operators/databrew.rst b/docs/apache-airflow-providers-amazon/operators/glue_databrew.rst
similarity index 95%
rename from docs/apache-airflow-providers-amazon/operators/databrew.rst
rename to docs/apache-airflow-providers-amazon/operators/glue_databrew.rst
index 2611c09279b48..d14f89848653d 100644
--- a/docs/apache-airflow-providers-amazon/operators/databrew.rst
+++ b/docs/apache-airflow-providers-amazon/operators/glue_databrew.rst
@@ -39,9 +39,9 @@ Operators
 Start an AWS Glue DataBrew job
 ==============================
 
-To submit a new AWS Glue DataBrew job you can use :class:`~airflow.providers.amazon.aws.operators.databrew.GlueDataBrewStartJobOperator`.
+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_databrew.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_glue_databrew.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_glue_databrew_start]
diff --git a/tests/providers/amazon/aws/hooks/test_databrew.py b/tests/providers/amazon/aws/hooks/test_glue_databrew.py
similarity index 94%
rename from tests/providers/amazon/aws/hooks/test_databrew.py
rename to tests/providers/amazon/aws/hooks/test_glue_databrew.py
index 9b5153edd5fad..e3ea047f90064 100644
--- a/tests/providers/amazon/aws/hooks/test_databrew.py
+++ b/tests/providers/amazon/aws/hooks/test_glue_databrew.py
@@ -20,7 +20,7 @@
 from typing import TYPE_CHECKING
 from unittest import mock
 
-from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
 
 if TYPE_CHECKING:
     from unittest.mock import MagicMock
diff --git a/tests/providers/amazon/aws/operators/test_databrew.py b/tests/providers/amazon/aws/operators/test_glue_databrew.py
similarity index 92%
rename from tests/providers/amazon/aws/operators/test_databrew.py
rename to tests/providers/amazon/aws/operators/test_glue_databrew.py
index 07071ddc04070..571e4816b557d 100644
--- a/tests/providers/amazon/aws/operators/test_databrew.py
+++ b/tests/providers/amazon/aws/operators/test_glue_databrew.py
@@ -22,8 +22,8 @@
 import pytest
 from moto import mock_databrew
 
-from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
-from airflow.providers.amazon.aws.operators.databrew import GlueDataBrewStartJobOperator
+from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.operators.glue_databrew import GlueDataBrewStartJobOperator
 
 JOB_NAME = "test_job"
 
diff --git a/tests/providers/amazon/aws/triggers/test_databrew.py b/tests/providers/amazon/aws/triggers/test_glue_databrew.py
similarity index 93%
rename from tests/providers/amazon/aws/triggers/test_databrew.py
rename to tests/providers/amazon/aws/triggers/test_glue_databrew.py
index 6df4a6319f71d..7352fffcd83f3 100644
--- a/tests/providers/amazon/aws/triggers/test_databrew.py
+++ b/tests/providers/amazon/aws/triggers/test_glue_databrew.py
@@ -18,7 +18,7 @@
 
 import pytest
 
-from airflow.providers.amazon.aws.triggers.databrew import GlueDataBrewJobCompleteTrigger
+from airflow.providers.amazon.aws.triggers.glue_databrew import GlueDataBrewJobCompleteTrigger
 
 TEST_JOB_NAME = "test_job_name"
 TEST_JOB_RUN_ID = "a1234"
diff --git a/tests/providers/amazon/aws/waiters/test_databrew.py b/tests/providers/amazon/aws/waiters/test_glue_databrew.py
similarity index 96%
rename from tests/providers/amazon/aws/waiters/test_databrew.py
rename to tests/providers/amazon/aws/waiters/test_glue_databrew.py
index 2a0fad503edad..2393898102923 100644
--- a/tests/providers/amazon/aws/waiters/test_databrew.py
+++ b/tests/providers/amazon/aws/waiters/test_glue_databrew.py
@@ -22,7 +22,7 @@
 import boto3
 import pytest
 
-from airflow.providers.amazon.aws.hooks.databrew import GlueDataBrewHook
+from airflow.providers.amazon.aws.hooks.glue_databrew import GlueDataBrewHook
 
 RUNNING_STATES = ["STARTING", "RUNNING", "STOPPING"]
 TERMINAL_STATES = ["STOPPED", "SUCCEEDED", "FAILED"]
diff --git a/tests/system/providers/amazon/aws/example_databrew.py b/tests/system/providers/amazon/aws/example_glue_databrew.py
similarity index 96%
rename from tests/system/providers/amazon/aws/example_databrew.py
rename to tests/system/providers/amazon/aws/example_glue_databrew.py
index 7d35df98498ad..08625b5611927 100644
--- a/tests/system/providers/amazon/aws/example_databrew.py
+++ b/tests/system/providers/amazon/aws/example_glue_databrew.py
@@ -20,7 +20,7 @@
 
 from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
-from airflow.providers.amazon.aws.operators.databrew import (
+from airflow.providers.amazon.aws.operators.glue_databrew import (
     GlueDataBrewStartJobOperator,
 )
 from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder
