Add MwaaDagRunSensor to Amazon Provider Package (apache#46945)

Includes the doc page, unit tests and system test.

Support for deferrable mode will be added soon.

Author: ramitkataria

providers/amazon/docs/operators/mwaa.rst

@@ -24,6 +24,9 @@ is a managed service for Apache Airflow that lets you use your current, familiar
 your workflows. You gain improved scalability, availability, and security without the operational burden of managing
 underlying infrastructure.
 
+Note: Unlike Airflow's built-in operators, these operators are meant for interaction with external Airflow environments
+hosted on AWS MWAA.
+
 Prerequisite Tasks
 ------------------
 
@@ -45,19 +48,34 @@ Trigger a DAG run in an Amazon MWAA environment
 To trigger a DAG run in an Amazon MWAA environment you can use the
 :class:`~airflow.providers.amazon.aws.operators.mwaa.MwaaTriggerDagRunOperator`
 
-Note: Unlike :class:`~airflow.providers.standard.operators.trigger_dagrun.TriggerDagRunOperator`, this operator is capable of
-triggering a DAG in a separate Airflow environment as long as the environment with the DAG being triggered is running on
-AWS MWAA.
-
-In the following example, the task ``trigger_dag_run`` triggers a dag run for a DAG with with the ID ``hello_world`` in
-the environment ``MyAirflowEnvironment``.
+In the following example, the task ``trigger_dag_run`` triggers a DAG run for the DAG ``hello_world`` in the environment
+``MyAirflowEnvironment``.
 
 .. exampleinclude:: /../../providers/amazon/tests/system/amazon/aws/example_mwaa.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_mwaa_trigger_dag_run]
     :end-before: [END howto_operator_mwaa_trigger_dag_run]
 
+Sensors
+-------
+
+.. _howto/sensor:MwaaDagRunSensor:
+
+Wait on the state of an AWS MWAA DAG Run
+========================================
+
+To wait for a DAG Run running on Amazon MWAA until it reaches one of the given states, you can use the
+:class:`~airflow.providers.amazon.aws.sensors.mwaa.MwaaDagRunSensor`
+
+In the following example, the task ``wait_for_dag_run`` waits for the DAG run created in the above task to complete.
+
+.. exampleinclude:: /../../providers/amazon/tests/system/amazon/aws/example_mwaa.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_mwaa_dag_run]
+    :end-before: [END howto_sensor_mwaa_dag_run]
+
 References
 ----------
 

providers/amazon/provider.yaml

@@ -485,6 +485,9 @@ sensors:
   - integration-name: Amazon Managed Service for Apache Flink
     python-modules:
       - airflow.providers.amazon.aws.sensors.kinesis_analytics
+  - integration-name: Amazon Managed Workflows for Apache Airflow (MWAA)
+    python-modules:
+      - airflow.providers.amazon.aws.sensors.mwaa
   - integration-name: Amazon OpenSearch Serverless
     python-modules:
       - airflow.providers.amazon.aws.sensors.opensearch_serverless

providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py

@@ -0,0 +1,113 @@
+#
+# 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 collections.abc import Collection, Sequence
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.utils.state import State
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
+    """
+    Waits for a DAG Run in an MWAA Environment to complete.
+
+    If the DAG Run fails, an AirflowException is thrown.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:MwaaDagRunSensor`
+
+    :param external_env_name: The external MWAA environment name that contains the DAG Run you want to wait for
+        (templated)
+    :param external_dag_id: The DAG ID in the external MWAA environment that contains the DAG Run you want to wait for
+        (templated)
+    :param external_dag_run_id: The DAG Run ID in the external MWAA environment that you want to wait for (templated)
+    :param success_states: Collection of DAG Run states that would make this task marked as successful, default is
+        ``airflow.utils.state.State.success_states`` (templated)
+    :param failure_states: Collection of DAG Run states that would make this task marked as failed and raise an
+        AirflowException, default is ``airflow.utils.state.State.failed_states`` (templated)
+    """
+
+    aws_hook_class = MwaaHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "external_env_name",
+        "external_dag_id",
+        "external_dag_run_id",
+        "success_states",
+        "failure_states",
+    )
+
+    def __init__(
+        self,
+        *,
+        external_env_name: str,
+        external_dag_id: str,
+        external_dag_run_id: str,
+        success_states: Collection[str] | None = None,
+        failure_states: Collection[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.success_states = set(success_states if success_states else State.success_states)
+        self.failure_states = set(failure_states if failure_states else State.failed_states)
+
+        if len(self.success_states & self.failure_states):
+            raise AirflowException("allowed_states and failed_states must not have any values in common")
+
+        self.external_env_name = external_env_name
+        self.external_dag_id = external_dag_id
+        self.external_dag_run_id = external_dag_run_id
+
+    def poke(self, context: Context) -> bool:
+        self.log.info(
+            "Poking for DAG run %s of DAG %s in MWAA environment %s",
+            self.external_dag_run_id,
+            self.external_dag_id,
+            self.external_env_name,
+        )
+        response = self.hook.invoke_rest_api(
+            env_name=self.external_env_name,
+            path=f"/dags/{self.external_dag_id}/dagRuns/{self.external_dag_run_id}",
+            method="GET",
+        )
+
+        # If RestApiStatusCode == 200, the RestApiResponse must have the "state" key, otherwise something terrible has
+        # happened in the API and KeyError would be raised
+        # If RestApiStatusCode >= 300, a botocore exception would've already been raised during the
+        # self.hook.invoke_rest_api call
+        # The scope of this sensor is going to only be raising AirflowException due to failure of the DAGRun
+
+        state = response["RestApiResponse"]["state"]
+        if state in self.success_states:
+            return True
+
+        if state in self.failure_states:
+            raise AirflowException(
+                f"The DAG run {self.external_dag_run_id} of DAG {self.external_dag_id} in MWAA environment {self.external_env_name} "
+                f"failed with state {state}."
+            )
+        return False

providers/amazon/src/airflow/providers/amazon/get_provider_info.py

@@ -604,6 +604,10 @@ def get_provider_info():
                 "integration-name": "Amazon Managed Service for Apache Flink",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.kinesis_analytics"],
             },
+            {
+                "integration-name": "Amazon Managed Workflows for Apache Airflow (MWAA)",
+                "python-modules": ["airflow.providers.amazon.aws.sensors.mwaa"],
+            },
             {
                 "integration-name": "Amazon OpenSearch Serverless",
                 "python-modules": ["airflow.providers.amazon.aws.sensors.opensearch_serverless"],

providers/amazon/tests/system/amazon/aws/example_mwaa.py

@@ -21,6 +21,7 @@
 from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.mwaa import MwaaTriggerDagRunOperator
+from airflow.providers.amazon.aws.sensors.mwaa import MwaaDagRunSensor
 from system.amazon.aws.utils import SystemTestContextBuilder
 
 DAG_ID = "example_mwaa"
@@ -29,7 +30,6 @@
 EXISTING_ENVIRONMENT_NAME_KEY = "ENVIRONMENT_NAME"
 EXISTING_DAG_ID_KEY = "DAG_ID"
 
-
 sys_test_context_task = (
     SystemTestContextBuilder()
     # NOTE: Creating a functional MWAA environment is time-consuming and requires
@@ -67,11 +67,22 @@
     )
     # [END howto_operator_mwaa_trigger_dag_run]
 
+    # [START howto_sensor_mwaa_dag_run]
+    wait_for_dag_run = MwaaDagRunSensor(
+        task_id="wait_for_dag_run",
+        external_env_name=env_name,
+        external_dag_id=trigger_dag_id,
+        external_dag_run_id="{{ task_instance.xcom_pull(task_ids='trigger_dag_run')['RestApiResponse']['dag_run_id'] }}",
+        poke_interval=5,
+    )
+    # [END howto_sensor_mwaa_dag_run]
+
     chain(
         # TEST SETUP
         test_context,
         # TEST BODY
         trigger_dag_run,
+        wait_for_dag_run,
     )
 
     from tests_common.test_utils.watcher import watcher

providers/amazon/tests/unit/amazon/aws/sensors/test_mwaa.py

@@ -0,0 +1,75 @@
+# 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 airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
+from airflow.providers.amazon.aws.sensors.mwaa import MwaaDagRunSensor
+from airflow.utils.state import State
+
+SENSOR_KWARGS = {
+    "task_id": "test_mwaa_sensor",
+    "external_env_name": "test_env",
+    "external_dag_id": "test_dag",
+    "external_dag_run_id": "test_run_id",
+}
+
+
+@pytest.fixture
+def mock_invoke_rest_api():
+    with mock.patch.object(MwaaHook, "invoke_rest_api") as m:
+        yield m
+
+
+class TestMwaaDagRunSuccessSensor:
+    def test_init_success(self):
+        success_states = {"state1", "state2"}
+        failure_states = {"state3", "state4"}
+        sensor = MwaaDagRunSensor(
+            **SENSOR_KWARGS, success_states=success_states, failure_states=failure_states
+        )
+        assert sensor.external_env_name == SENSOR_KWARGS["external_env_name"]
+        assert sensor.external_dag_id == SENSOR_KWARGS["external_dag_id"]
+        assert sensor.external_dag_run_id == SENSOR_KWARGS["external_dag_run_id"]
+        assert set(sensor.success_states) == success_states
+        assert set(sensor.failure_states) == failure_states
+
+    def test_init_failure(self):
+        with pytest.raises(AirflowException):
+            MwaaDagRunSensor(
+                **SENSOR_KWARGS, success_states={"state1", "state2"}, failure_states={"state2", "state3"}
+            )
+
+    @pytest.mark.parametrize("status", sorted(State.success_states))
+    def test_poke_completed(self, mock_invoke_rest_api, status):
+        mock_invoke_rest_api.return_value = {"RestApiResponse": {"state": status}}
+        assert MwaaDagRunSensor(**SENSOR_KWARGS).poke({})
+
+    @pytest.mark.parametrize("status", ["running", "queued"])
+    def test_poke_not_completed(self, mock_invoke_rest_api, status):
+        mock_invoke_rest_api.return_value = {"RestApiResponse": {"state": status}}
+        assert not MwaaDagRunSensor(**SENSOR_KWARGS).poke({})
+
+    @pytest.mark.parametrize("status", sorted(State.failed_states))
+    def test_poke_terminated(self, mock_invoke_rest_api, status):
+        mock_invoke_rest_api.return_value = {"RestApiResponse": {"state": status}}
+        with pytest.raises(AirflowException):
+            MwaaDagRunSensor(**SENSOR_KWARGS).poke({})