From 369a8ff63b55c47fc65fe72d65d6dc5b6230cc8d Mon Sep 17 00:00:00 2001 From: Mike Ellis 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 `__ -* `Glue IAM Role creation `__ +* `AWS boto3 library documentation for Glue DataBrew `__ 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 --- 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 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 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 `__ +* `Glue IAM Role creation `__ * `AWS boto3 library documentation for Glue DataBrew `__ 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 --- 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 --- 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 --- .../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 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 `__ 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 `__ 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 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 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 literal 14575 zcmeHu-K97bC%6<1E=7wME7IaFg#tHy z&pGS<1NX!IaI)63GV?rpW{>XK@|!3v4aFB&6j(?|NH3I?9f<>x1|lIxyhj53BSAbAhzAJ?wE!6j74eV!?_B}P z|9nMGEI|Fg>3@Wje9jX{NMJ){c^O?Vk(mAP)Q<78HMjBI zM=_hJ_2>H=jEAHB=cC7sEFt@a@5Yi1`fTz+Ks;nxJjhgpN<;jAGNDMAF{jA*bSUd| z((sG2D##lKX@C;M4m>ESrC~VO4A6sWUI;i)YUNswn853tW zDuUVz!F@cZ)7jaV?A@rM&?9HHny>V>=P>Ux_WL z3zE#>SFx56RN<%1{GWwqtJFerf z3HRgx{NYGXiO!uB;$HQBkEOTA7p8Ne^~B3#4~vN}l*o|(iVWL?Bn*0<=fQXhcb3L( zTNAJAkWlW_&g$>D0|-f)j=mfHdBSRVB9=p<9WRgr>Az(_^Hzhv%zLxx!V$Z`6Ywcu zDpT_TL<3rbI`3Na!7M5!sCe92$;Hf*ZA)VRauNFl(uOyBrVO8Z{a&o^{6Yj1rPS{aE0@C+Vsd{v9%MLD`bm}N~+cWEnqm5Dhm9OJ4PWm9J39qg5mF-;v0PSQv2574S_aX&m zHz6sbZns-6HGV|e#D>F&k|mE+Ct(`k|LlI+m;MG~mvo*C+YFW2n2*h(Tjpp!r2?f3 zqImr|$7Bnc!k|q`_{j|2?+4i~7q%AR1&cEVX(G{vFx)GaFjYGxW`uT{MMMnLTAMU zk(tJ-0~i>ztg2)scdDhml3^%)gCp{E1SoM!m5HPx0t>hXkk>p0I zBxa}x%AoB(A}k}~3uc@4O;xy?Oq6S+G!tO+_2ZctcGFQuV&extJQa9@>;O=njspc; zj?QXd^VVr9BDd619zJ^xWy6oe2># zwB2$rQ=X3fUkgTRPsw4(u;a$LbE9&YagCXQy@>(d>CL{H@+p z!LKmsV5S1^7b{K<*c`>ODXA{vEmIbq*}oZ$-E$pq|7gXlP%gOBHYB;BT=xzBZ6lVU za`4#LNvMB?>#$%f+&e#+JeNGUk+_Jl_ik(TVxB>j0NnsBc!5&o7gxy~$!zJ#M!W>T zBv%;OhWebceLLeKU+8_++lFg;D07T6pXXo0p?UX*pQwNS%4NNBJZ=u#X~xau_~F(IgYo?KbFXe;Kxy<*L+O!EYNv~41?)(Y5vo_?{OW+E?f%rBBZx19zOG_#gDU>A zY8InrWlUded{0a%;keZQSEVk!UEwMYBUE89alfAqc3u`@-y6Vkj3Hn(c9580T-1arARyHP^b_)pd@L%i7-@IR_xR|u+H+DChT%^p> z3T$#lSJ|&xb?_I!>ut8k>J*YGlLB3SiTPgE)bhS| zp=>p?(0)T=!p(FD5PIH;i3vKESZ`>aao>9TH(0G>Bewmbg_o!{qM&HyV_RWmd`|c{ z-EoC}U8-wqSyyuXVnP_PSc$TAWZDtxWb#@t_`WPV`|t)3KIUo0HJ zO9$sYBeAZ0@AlR8y;-Y4xG^hwkh1qPD%sBFPmc96YcHN0nH`1su>-@g124CUIXzy8 z$)MA+*GVy?Jad@^k~`l#?Jg8X(4RZrD7`K;1(^kK#QYj1#g3c7bO<`Q&Zk;^5&M}eH7b`jM`9J9Kt*})aa$)}(3(0CYGcgp@L>z0P z$1H+(Rv^#cN{|ysctVuKC)z0?0~tHJlFM}a|-&kC(m*4=(s^It2IKTjD$a^dB634|6ftzqg-P|k1LDdoMUf_6n1TbG4$&}_F9 zzmZ+h2q`V45Z={7=82x_h(3KHD>NL5tAgn+8kRw9uqgSd9j1eWm0b|R^zb~euM=%`P{r~6eCp) zUTeWM{!!+?dk3c*xC| z2zTsoVG0RE#Z&{~-iNy?oK2|b8D%Yy9aa_&hZ|x5B!Ww%_4W

5R~P;wTkZph752 zaZ(%k3~juVD2~@F%$YW-)wh;S*QeutpJ`gwm2mH~93**NddyjSDHF!WUzUxVyjQ42g19&Hs<)-_mD7NRG+!wWK4pYJ-+vL}zwIdXO7$J@#m9}GF{N^C++Jt!CQ7@$ zMA(V#|1?->VSEQ@)E$7O=`LKxEZC&|OF&-NOD z$~6w(_2Rz@k{S6wC1r7LJV&H95Es-wCQN*`iV*+x060d9N`x6}sW?;Rn)2Oix@|b4 zdQtSz9WF_IYEoCH06OT3o@j=Ao`^^Bn7{NK;dA=0#4!0JiPcq6TsXSLKXb{ zOW)lxPE{gjh<()1s;W#n%~uxA5fPg)RUdv9?`|w*X?y5O#yV+T-qFkUAz|CZ{(Rqe z4{3(<^m<PP0DR^znWOaoua+`Nmh9mQG=* z-qMe(F=&zUB6h?aKl=r;O`jaQ`Bcs6w9DfNzCQSexKg;x0xBt%AI4iW)4Z2rW}?o} zI5PV%isjScT8Ha-hW3ZQM=(8P#k&?3BQ*a~FdQoyf^)t#Q*n#8T$0gr47-snJ0mh} z+oZ9dj6pFFm_ZZ6JMEH=|68-B{v!7y!^-ELWFDy~(`p@D-#pvSk`7?C6>WEFz@bUO zSO67){_Z^)w!0V&r~fCQ@U2i<$egEm1kblER&3qzMzhpKnyCB2kf#$-&G`GzO{cR> z6X=>S>(h~q(Y0TOf9GB)5#9_JBsGk0X++ud&?Ok}$8Ck^q>!0lUlx%;Q%=MiO;RT6v%1WC79IvcN{!6O zT{ZY4tmo6Zb~T39Rj677H9L)lOf%++gR0k3SAk5`R$pnhi<5q#eN3D!*T}41o%<4q?h8JACA%SNMtc}{apNY>?Bg)bST&)%3C0M8oFW3#C+>X%` zK8Xw(p^T1dM&B-2B&&lJH`a_FMYd{cILDb+Y$32A{6~Y-&mR@3m`)D3=0em zF%Jm}-QK3iOsz%!*z|W=nOs-Ydrb#o3Gc@AUH;HV9|U<7%iv<|-m7^WwAt2c^Ne3g zRrU=HMUX14JqLsOpozsc{>*2FYU_H2%5oDoc;0t${oQW>W=h z-rcIz)TzC90Sn`#@JVKk&d6;0_mQV(W-g-)Tu#Yw50J3bJ?JdxZwyAIjz5^kSg=l1 z!Jt%^KtC7Ox;{M`L@9p#LkbY@e~3LH=DKzE;STr-74@*>W%^HuO8=XoH#@s;8Vm!M zW#_s3Mz?I7ff43N zI-LO|$!PA2{jT4KM*NiljL}OY1dm$gugq`zA?%sk9AriG4ctrVWG3nU^a-Zc^opq$ z#{wZ)J(lSoHZ_mG`kB)mSd3G2WlhjiB4*FxP?3y$D^*8O-#*qJNHnsLj^R8po!`r0=b8R`=x-{jI z9iJ@%3A2idKRCaiTL$0Xja0PIv^?2yjtsLUPXs0JRDr^O>^8wE>&<(_IWZb+Imqtx&T<8`cY9v$WP zNZ8Jko>$$zT0)_8RX3{LY?4jy_@%u3eW;-JHCTqeOq`;D7g*U#ss z{HLRWGeHAIM8H`kWB;oij71>l3L6>#439ompB~mKjF5%6>d5oxDwhNL-swr%V_24d zLrT$I+Jk@)W(6=^*zPJ;J+%BE(_! zk?opxVYVQ^0%4f&j%Wi+TYG(8<+JMy1Iw}$vb++^8RHc!&OL3}8TZUT>sWAeZY(Du z)=1!Iv1(t`@p(DTa?-clnunK*uZ^fuoFu2Ibjq-C3H>gfchVadCTU5hNUc;cV}m9) zoBFYE*a*MVBW0jQ_reLkW`~ZwUW>g%SCGX0dRM*^MOSWO+F_Q!a8Fm2NhThygLVp5 zqFX;$p{oB-wA*Z!`$P`nRHN9uDkUc?VBcj_p zb%LDQr_~e)9)HWim?|xu@2nl6hr%!ECWB-n)DKs}S39B7e92`MGIH;SH#q1f{b5wp zG7*imAbL}7jLJAEZDTlSM^RrNuvU6!7E7dky5-g!H-GTMA~XsE)ez)rqj~l z?n$veMRfjmg!NqEJ9M0}*PG?txtl-!yw^;_-nk^hn|^lMb?ipi-&<*w`GY$V)|+Zvo1**~(E+Vk+4+DW1=FOKbw?ApXx zksjy0%XA$M3o5I9+CWs4Qq=|f%>ADQU@_H5)5o~S)Cz`yJzQbF7MacC`sNymY$9ZB11!!-R`}N17#}4v=UkP)dj8$NPGAE<=iHxg@+%P=z&* zVM?nhtPS(H-emNHq=cGnH@T3JzgM0>Tz7_!Y`$bjM+Y2$yj{$$Ino;%`@Su}4_fGv zsS2-RJrV|)Qy+x)81AVfJ9no+qt*pfV& z#u)raq+tnbLw{8gj{T>vQEw~c1#;|W%IZ8-?fJ@3>8n@ z`h%`krfMJaGkpc5vF$%|&171Z3Z$QPia22N_srStZ~0M85DG;*!=RTm(gGy9BrvQr z|EyB5CR#?%>WrhM!+3_N+HK^_8x!zBNI!yRw$d?5cU>H-s&7GLFlX;UYbF-7#&ziG z9_z1zwYaVQu*d0Vvyy=?TV@@3M z(T6s2PwLS=p>9V~cYm6*Jl8E6s_+JQ>YjliCq#O>?%PgDn5wL4u+V~Lyvi)a2{ECr z)ZJ*SIM^m7;riYtQ{y0d?Hamvf4J`17G|9n@uk4-+n&_H2Zvx4h$(Glp32bduN-e9 zwh?-pbopqKI%C5|1+S=0Tsm3M%5) zj&6Uh49TEdN9hQsMz?ROvvoltg}`_CM=7tokQg!*MgT9fE|Vy9KeXU1w@5cz=7}yb z9A7`z;aYtv_$iUAlpO`4J4jx&3lKP5p!_Eic!7z1hsApmQ0P6$>Qbd7uSLV|4;Veu z9_iA-c(^)ko$oCDppWO{w%H>3dMPn-4<`28o5EohAQ%7Q%NR}A5s^(zRsD*GhU@%% zC9}}RL{_%)AlR;8I1CLNy&99{>vGW$Q6Uv01#y7X+#w^UXJ|4_=%I63x!&oyEIx9cjb z$8E;qtjv3e@Oyd_8VUIY7qZM10eID=9H&SUT@7u&KmWU2z^r$_#RJ_BVisclu4-+I z2k`q$VDT&NDT9hfp-sGB@U;5UO-{g247Rso|CF`kT- zzZkACWs#N&GIGL~unx&c>RcZ{CllxYRda}Zm3GQVDVo4_+rDZL(Gw`~(tWuWPfB4h z-dQrtWj&SwLg-LpZNNYpnh3xop25~5c*nyw!wxNS`7^@l^SPp5}@N=i6>!=<7Nf?3(<#cN{-`lKPC7RO2elr*KQ< zN^qWI9=EE!k(FOF$86uuAfXSmrX6v?uLLPf7@8wYhDF!uG%!|M@L~Qy2xI`3py8h5 zLONb-xI-6Hbr@d7*s8I^c-iU*AL+Qh7;EQpIm9%K(m$+pby$H-_`GB*OMw=kP}P>_ znXI#8d#%CJoRNM)hQ%I`!AmsjFPtW)z45y^p=WG&|H#ew%DLM08eu_}l z4;=%+GD)~f^H%xOwE>X=2?;tfk0M0W^TRPUI4L%egaoCPJKRW?u#!CLNxrZe$nud7r1B<~ z$bjnRS{pu5mX{L8=eF+ag>S``g3U0oS;r3pcg|ZsOr0X6l#oBkmHxloLTie!Fp}|7 zu|3PSN;=+puen%lpKPrR=$7J5v*wCW68pT52seGaSKXFI#{eFLMA4}66}Kt9eqsX- zGSiov4)iI>fD%xvHhlUZ&i;a8QFiGgbFz>(lYx^zZnFeb*w0%A?hrd21&~53jA+6_ zUL?BS{es*#VvY2GcUUS)=A0BUL`@f36H|e6(}$}m^|VasvE%P7=&QfyU6nRMm-qAq zq-mdiIZ_-Rrs6r{=`ReMW}Iah+DA`koXPc!FU5p{Oe-?=O<6hr5}&(;YS^aVCSq*+ zpL3+=AOu&exgGCTEuda@Z`4p=K&aDh-prSU}J(bSCiGbOr ztOf@K5nbbjBO5p7X#zl2IY0JI@A5q4u#VXLwoa5f1aRE3S}rz3h478M{S)WutJh6v ztHWqe+dCP&b4=UKlUu;+gd|A*e-G|NupPo3QV2P54+4+Pyp>qixnAYk?Um#Ql65mo zkDUry)YSN8)WCnu3GV)UnJsB_TVk}~aoZUa67-a}6K5r%-$81I;xIt%-|lxGuvUM4 zPy5?~Vz+ETO(2ERsL;xqN2AbmTry*SD%(i#g(8E7H{5{|PUJQR|9+U(-6YH8 zFVHF)hV)XB9;cZ%wHLih@~dGy4`**wJ5x7ksUd?9vnX z+RnCYe3Hfs7QL2@^QFk2SEq#xLK2d<$|?=9R*hk#W1t**6JF;1yT5h17+!|vea>Vb zribLeOv-=G?Y?!TV;krqOZp>`gj+tyrV4(2ZfpstbB>8=B|X1@x|BqUH3pd$rsZD+ zxvC4?DZSq?9MyZLp2#CL!w`m@Bc;wgTO zUpoISnK{LqKvDN=t)aA@+hnyxvm9kADH!6&agTE#@k=ay8FF`?p&&+!%v3#g-MAc6 zyuxhGY96^Zu0LhQ!z$2;;({iW9%8-$RWTNj_$8RIJVppW%~CTc(;+G62u9t5Wm|%I_zZ# zWD6~T8)O`OqY@Y$-;MOkl_^jADoI&s^%=xk6q>7jZ*^;%B`vu-|w zc7eG)?AEnEg0*67IrnH683u@jEpxMV;;vS8;u5oar)hg?Nj23U8no~U^nZ5QuEm5> zn3tYTyE7+7pjz^6XjF;c{H%slDnx>C(r@p7zU$d?B_8z4ro+p&&e_=+L`9TZv-($R zQ&UiwfZDUd?>1Rf%CIA^nY;TgaLA{|$q&ZLy=5wJ`iL&D*R zrOrR=hi{E@O}=|$bD{J{24N#$GNh^S;pKq*v-i;EAomH4ugcy2uS4gEznN0aTS&Hd zfmJNaR`o^rMpT;RvtF*0ohb{6t9(wW)M@yW3`-uo-N+ctw3okSsa`~?3I=~W!dR@G z8(ZbtQ`|RXA-G?VAFn5Z(G89^f2>#-cCI70)TKq{@qo-nd}~M7lI?SF8zbA7SOQ~v zEK0oMC@e^euz2b>Os|=YS=ao?qfs1999U~;ZS1r~cI<}of|NvVYZUbP=eGIau&CmZ zKY_kH(pvH&g>!k^ram$)m+oHSBK&HU&MtHgM9{F*B24_wEc3Px_BGeMk4Uof>auV1 z7f5fMWp_pUxjn(7)+{jLM1;B8)&>z-5%;V3;aArAlv({2KX8V4DW+JL!l}$GP#2Tb zf$PUL+keI@9tl%*m^Ep7N-FtFZ3d-g8m1@K!lFg$((O;2r49A-D|JdO&A$uDQ|Z??PFa>~@IvHm=`R9?pqaw7ZALqjTAzvt^{SQK^U zYbK|yj{SWei8@(PkP%=eG$c?|Ias( z2JebmBFFevNk^d%2Sl5THJw<=DvuZBBD``c!`ND1(s~~J+AodWXVrk3Zkv$?wDFBx zV)3@Tp8Y^XaA9}nIi=edqVnfpuK7zZ$2lx>si>J?l8!-(0^tOHZ(Mmzq>sJTbio@M zeQrzz4kPz*z{GLEnu#T-eF|WdWk3~!K>|V^nGl* zgkQ17FV@%9+|6}*YCJ!YV+gz*X=dM@&|}QTvt}bhIN2%c=oHm&&M*x`6l0Xuz9M|< zjiYDWba9RlMCh-nR;sZxbr&=Gq$I?XI8crso}}s12zEV#V=1lcnKvb*t>b^(h^*Ga0L%B9%U~H zERibH>KC#8IvS#Vx$bQZNkYcarZkqH(=IXl~os?wg5>eG0<#-r_daNpTz zl^Z?QX;Oo*-?htLPi%0lcxI|o-RZnvoe`278W%xqC>}GU@(6QKV1EVwQ_IriU@&3M zH;lK_cXmm!U*`$F1i)f$s98Ty`L#Gg45W0{e~}@2GW_{v zZbVzlSvDq&AauB`EA?<;0}^V1h33&%3=)jDn)o5yw0Ac-4{q(f%~M7#>5PG(DHRaF*Cx0iQjfK?oizVy$<0jyg6I+4=K2)8LSbo5D% z%+UfoXMyUgn%`x(9xzDrjRJMW2WResBJK!)6`gutx_IFy8GQ<(V*cft^w-xvOylpn zQhESh5&_p6=|6siZb@NCV%Ib%JBwV2pRCmXOdUqVR(@QI{_UswBNhxF;Q-rBcXSP9 zD;EP{{aywVeK5Shm|(ZfsUub7IR0N#6aN4?ex$E#>5R|gz^kJ#H;TMN1jecb+R)p- z{tbG4gk>ncLF17F9lBT_v@Vd~85+Stj7#-ZNewlOG6+KwEH>msR@IyQWY&{JJYJ38 zFCRo2{Se5bO0qoXB=MaSdwM>fy5h$G9dx&-s=n&anHeSRF!#mau-j=JSA=W(PVqo+ zos3I*%*@|+s5_Z4CD*T2-rR)-UPs3nT~YW+|63)wrwu2mS+# z)rH?3R2>e?!%uumTM~AoI9LC*ZhmQB2l$dF47oTSF{qMIw!q_e$ABa+6wOYo5~2l8 z`V7svRVERR8|xdiphwa8uHeS)Ya?*;%mjc3;lJ`2V^b<=Pnlr1;SMp};v>awk$5M! zIe}r61cOOqh6VN4s|R0L?tGv_`l_Wx9Lr0MrXgRbJNq)pbArQG5WIv!e(c&QK|Nut zn8yRI^7z4Zepe17%~Fd+F$(s3z~MuipvqU{Gn_~ZOKUbA9Om5-8~TInM1{t^rD(zc zN--q-P=TsZ9rQSm$0}N5ej*6ZT^6|43@H}%OJy)Byl`t9c(YOv~dmF z1?sL8%UBBSkXg2EO2TAQT)<>hJz?NhAxvZ272inp5c!Y2O0na2b=;ob@}gBwTrD}m8&T8G`Zgb*u+Vx8|@@DNxw+XMCXKJH)7pu1;7pvQw>};kY-Q_^9si21R zDkgw(=bO8{T(_Dd7ir!WxC&M0S8*?B<#Hv`tN34iVEY?vTU%k zU>MSdjc%PuPvfR>@S~q{vLi8fcRi)l(oF?Iks-jktzMc>QiA zvAW#s)6TV`bsr8Ql8VF#5yW(&ap`+=Q1XCA8KljCHW@(g#WY?vNyR+IQ} zBJvQ{5&_wy08#8>{#}){^KKodpQ4kL(T#{sK{*za(J^P+o@R#fZ7CW8R8~OaK#>(T zPn4o79?V5B1e~tuuN2Ge?XVjoL_|txY^$qZ82@P&>GxPCT7LMjSd%m;?y6#+umOae z?qbAhJA6B^Q1WmZF?{2~3RRHdNoqP%M#_UHjt$1Rwo@QHnM1}PFdnaJh)D$YdNS%1 z0oeZ~he=Cn=#(&tsk3@2`BWA(<2(8{#+JLC>S?WHn5-PByX@Kbg|Eu_Yo&6xN+Omi zFVa!We7}!6s+J^ICV9;~xDAWU21=2C4Cw_ywPcc!AF{A%05Q@|Pr zeL&o86}_3*t36VFRkP^EG9uU-wKlQ#E2(}G2%Wj3=X?dqi#FA>gt6$q+Pi?cTPgO7 z+572Z-zdHaKKN4vV@-vnRlWdx0xyQGj~c`cOxGfvNcEnxdTJU=GN7A;^U2a3=014P zL5a=QN}PCO;hc4s>7Tw7CkVU`X%%U$MywPWz zh~iCuIYm<41p|rEV-6UeDWqJwl3;=ke1Veiuebt}*XPcv#q1H^h~7~hi!MH{j*{Vt z-2cUSOAI0FOCb`*`DKVDFwn{Z)ZnG`(jptw*SR`KF`y?}@yTTZJcT@RO+NHh?jn$Z zSkDV4bky$QLhjfDhGaG}Zep({#Gs!gD)$ew?^&DqZIXL~iOP6WQBdfJ*iM5m9qgek zBoD7`Z4b8whz;h^giF2msu6g=`HdcF;-D8Zu{OKAZ!|EH(8ok5llazyBkC4ul8$fE zO>n1-W;p0KOA^(tR2{gl2Rx6OKn3n_32B#z!<~d+Yh)xhe2qceTWZJ z!aWU#^z_C zLZQSlMhwR8f1o@C3EY~22OR)Q(sTuDYkXc2x&e&i3=RF)i?+>^A_U&!BabCkL{!=$ zL7jgrsa7Jaazh>5N0h#B+6BcS+5|)=v!Q5BY&T&WQ#J`Wnff^JNVH~5=0r4QYBp6a zOw?NnI?KW~%Wx6Z_he2G>{bHwKq7zNnXpv!C7{?Qv4&F;%;?)u^KJ zheTZ-5Fb8;XaB9z{1#rQv$a>HS$X>Dy-A$+7Y3vYpIae!Vl+i1hwmO-;mwCuABCjx z7mCsZhkHCo^D?Wf`4oP>f*l1O+z+TKLy`Z32vmiBfvN1vB7HJh(#yLVa)Wf7{PF+t zDYS;By_fl~?bTARBaj&r#BdFjWnMg&I5>NS8v0YycS1uT@G*$hfT$*-e51 z@goU?`xHlDgrE5*E+;8L(oCEn`B-TL3nWcnv3Qp;s|A3!VBG`21nFn$G#c{|JxUTD zq+a)=OFEsi#VP~=OII0=3@hF4*ZPOX1>Y+D;+zDS1hGKd3kJaQbm#yg1cXWCC!&*p z=$S}q9AN4WXG5A;0~4obHCYO9AV`{VG9+E1b^T7ZKpBpOOp$^YtOO(i+wa@Jh58D> zv(@x$7O|QD3u&ZWk6{s%Nq80#ktaf3Jc01BAX<-1r6r3Qjr1i<2H(;Tfvv9mgwW&d zrW?%>=@7(?5v*vkZ9-%QEg?h^z0f=~L>liOdPQph5v|7yu0%xADk-iwGZEh74-}Zun-v8qq18*$QiVV5)?;U`s?^kHTh#~?# zU!uRKrae=egc|sIQk&kw{53w3(bjGTMUD5~WyL&0iSUFboWT1Szv40Blig*y+wVH7 zJNEB3@7ee`Re{pZ;iNSG0Qm}MTbq~Xx_&|p#^PpW@oC5p6|lvT0y(RPF6%w7EHUK* z>#*afFE?&>V&s`-0CQ?AUl4@XTpjQ2M-ykgAWT|ZNm|W8gah~o;c0S!ui;ZB9Yu#r z1nR^;sJ;0-*5%e6s9OUw82m|<@Z^ltdm$I7i;V2YpSRLR+Ix`+n+_U%IeukRIeo~P z47`g(k|gotH;vf{L;IcI`C2k}*Mk&xAzwfM+!W{={`u8Q6ESzAp`!VB_XJE}47+q}*; zpVNOA2gBkznwI-WLErT;sf9X6I7>59Ivu}PQpUj{z@uq+p-|=4Dh0M6-1fTTMz{%0 zY*U!J#1X@iu=*iRFVO-IUx^-NsXdsFsPnw9hxo=VsB)rqPwiVaML`}uwSK5Dap3nc z#l~;{Mv#sYs5O@o7hna7l;%cw3wDcgmX#P1Um*6v@^8hK5a139Vx4n$WW_O|EK3K4 lrTs7I|Bk@_s}Tr(#-PtIkxe7)?nXlVlod4OYh}$t{ttT=+T8#E literal 0 HcmV?d00001 From 357164bf21307881ace099cefb9b22e91f541407 Mon Sep 17 00:00:00 2001 From: Mike Ellis 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