Add wait/defer support - MwaaTriggerDagRunOperator (#47528)

* Add wait/defer support - MwaaTriggerDagRunOperator

Depends on #47527 because of the MwaaDagRunCompletedTrigger

Author: ramitkataria

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

@@ -19,10 +19,14 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.mwaa import MwaaDagRunCompletedTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
@@ -48,6 +52,23 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
     :param conf: Additional configuration parameters. The value of this field can be set only when creating
         the object. (templated)
     :param note: Contains manually entered notes by the user about the DagRun. (templated)
+
+    :param wait_for_completion: Whether to wait for DAG run to stop. (default: False)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
+    :param waiter_max_attempts: Maximum number of attempts to check for DAG run completion. (default: 720)
+    :param deferrable: If True, the operator will wait asynchronously for the DAG run to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
     """
 
     aws_hook_class = MwaaHook
@@ -74,6 +95,10 @@ def __init__(
         data_interval_end: str | None = None,
         conf: dict | None = None,
         note: str | None = None,
+        wait_for_completion: bool = False,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 720,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -85,6 +110,21 @@ def __init__(
         self.data_interval_end = data_interval_end
         self.conf = conf if conf else {}
         self.note = note
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict:
+        validated_event = validate_execute_complete_event(event)
+        if validated_event["status"] != "success":
+            raise AirflowException(f"DAG run failed: {validated_event}")
+
+        dag_run_id = validated_event["dag_run_id"]
+        self.log.info("DAG run %s of DAG %s completed", dag_run_id, self.trigger_dag_id)
+        return self.hook.invoke_rest_api(
+            env_name=self.env_name, path=f"/dags/{self.trigger_dag_id}/dagRuns/{dag_run_id}", method="GET"
+        )
 
     def execute(self, context: Context) -> dict:
         """
@@ -94,7 +134,7 @@ def execute(self, context: Context) -> dict:
         :return: dict with information about the Dag run
             For details of the returned dict, see :py:meth:`botocore.client.MWAA.invoke_rest_api`
         """
-        return self.hook.invoke_rest_api(
+        response = self.hook.invoke_rest_api(
             env_name=self.env_name,
             path=f"/dags/{self.trigger_dag_id}/dagRuns",
             method="POST",
@@ -107,3 +147,34 @@ def execute(self, context: Context) -> dict:
                 "note": self.note,
             },
         )
+
+        dag_run_id = response["RestApiResponse"]["dag_run_id"]
+        self.log.info("DAG run %s of DAG %s created", dag_run_id, self.trigger_dag_id)
+
+        task_description = f"DAG run {dag_run_id} of DAG {self.trigger_dag_id} to complete"
+        if self.deferrable:
+            self.log.info("Deferring for %s", task_description)
+            self.defer(
+                trigger=MwaaDagRunCompletedTrigger(
+                    external_env_name=self.env_name,
+                    external_dag_id=self.trigger_dag_id,
+                    external_dag_run_id=dag_run_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
+            self.log.info("Waiting for %s", task_description)
+            api_kwargs = {
+                "Name": self.env_name,
+                "Path": f"/dags/{self.trigger_dag_id}/dagRuns/{dag_run_id}",
+                "Method": "GET",
+            }
+            self.hook.get_waiter("mwaa_dag_run_complete").wait(
+                **api_kwargs,
+                WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+            )
+
+        return response

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

@@ -109,21 +109,20 @@ def hook(self) -> AwsGenericHook:
 def _build_waiter_acceptors(
     success_states: set[str], failure_states: set[str], in_progress_states: set[str]
 ) -> list:
-    def build_acceptor(dag_run_state: str, state_waiter_category: str):
-        return {
-            "matcher": "path",
-            "argument": "RestApiResponse.state",
-            "expected": dag_run_state,
-            "state": state_waiter_category,
-        }
-
     acceptors = []
     for state_set, state_waiter_category in (
         (success_states, "success"),
         (failure_states, "failure"),
         (in_progress_states, "retry"),
     ):
         for dag_run_state in state_set:
-            acceptors.append(build_acceptor(dag_run_state, state_waiter_category))
+            acceptors.append(
+                {
+                    "matcher": "path",
+                    "argument": "RestApiResponse.state",
+                    "expected": dag_run_state,
+                    "state": state_waiter_category,
+                }
+            )
 
     return acceptors

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

@@ -17,7 +17,11 @@
 from __future__ import annotations
 
 from unittest import mock
+from unittest.mock import MagicMock
 
+import pytest
+
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
 from airflow.providers.amazon.aws.operators.mwaa import MwaaTriggerDagRunOperator
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
@@ -31,6 +35,10 @@
     "data_interval_end": "2025-01-03T00:00:01Z",
     "conf": {"key": "value"},
     "note": "test note",
+    "wait_for_completion": False,
+    "waiter_delay": 5,
+    "waiter_max_attempts": 20,
+    "deferrable": False,
 }
 HOOK_RETURN_VALUE = {
     "ResponseMetadata": {},
@@ -53,6 +61,10 @@ def test_init(self):
         assert op.data_interval_end == OP_KWARGS["data_interval_end"]
         assert op.conf == OP_KWARGS["conf"]
         assert op.note == OP_KWARGS["note"]
+        assert op.wait_for_completion == OP_KWARGS["wait_for_completion"]
+        assert op.waiter_delay == OP_KWARGS["waiter_delay"]
+        assert op.waiter_max_attempts == OP_KWARGS["waiter_max_attempts"]
+        assert op.deferrable == OP_KWARGS["deferrable"]
 
     @mock.patch.object(MwaaTriggerDagRunOperator, "hook")
     def test_execute(self, mock_hook):
@@ -78,3 +90,26 @@ def test_execute(self, mock_hook):
     def test_template_fields(self):
         operator = MwaaTriggerDagRunOperator(**OP_KWARGS)
         validate_template_fields(operator)
+
+    @pytest.mark.parametrize(
+        "wait_for_completion, deferrable",
+        [
+            pytest.param(False, False, id="no_wait"),
+            pytest.param(True, False, id="wait"),
+            pytest.param(False, True, id="defer"),
+        ],
+    )
+    @mock.patch.object(MwaaHook, "get_waiter")
+    @mock.patch.object(MwaaTriggerDagRunOperator, "hook")
+    def test_execute_wait_combinations(self, mock_hook, _, wait_for_completion, deferrable):
+        kwargs = OP_KWARGS
+        kwargs["wait_for_completion"] = wait_for_completion
+        kwargs["deferrable"] = deferrable
+        op = MwaaTriggerDagRunOperator(**OP_KWARGS)
+        mock_hook.invoke_rest_api.return_value = HOOK_RETURN_VALUE
+        op.defer = MagicMock()
+        response = op.execute({})
+
+        assert response == HOOK_RETURN_VALUE
+        assert mock_hook.get_waiter.call_count == wait_for_completion
+        assert op.defer.call_count == deferrable