Add MwaaTaskSensor to Amazon Provider Package

providers/amazon/docs/operators/mwaa.rst

@@ -76,6 +76,22 @@ In the following example, the task ``wait_for_dag_run`` waits for the DAG run cr
     :start-after: [START howto_sensor_mwaa_dag_run]
     :end-before: [END howto_sensor_mwaa_dag_run]
 
+.. _howto/sensor:MwaaTaskSensor:
+
+Wait on the state of an AWS MWAA Task
+========================================
+
+To wait for a DAG task instance across MWAA environments until it reaches one of the given states, you can use the
+:class:`~airflow.providers.amazon.aws.sensors.mwaa.MwaaTaskSensor`
+
+In the following example, the task ``wait_for_task`` waits for the DAG run created in the above task to complete.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_mwaa.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_mwaa_task]
+    :end-before: [END howto_sensor_mwaa_task]
+
 References
 ----------
 

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

@@ -24,10 +24,9 @@
 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.triggers.mwaa import MwaaDagRunCompletedTrigger
-from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.triggers.mwaa import MwaaDagRunCompletedTrigger, MwaaTaskCompletedTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.utils.state import DagRunState
+from airflow.utils.state import DagRunState, TaskInstanceState
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -139,7 +138,7 @@ def poke(self, context: Context) -> bool:
         return state in self.success_states
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        validate_execute_complete_event(event)
+        return None
 
     def execute(self, context: Context):
         if self.deferrable:
@@ -150,10 +149,152 @@ def execute(self, context: Context):
                     external_dag_run_id=self.external_dag_run_id,
                     success_states=self.success_states,
                     failure_states=self.failure_states,
-                    # somehow the type of poke_interval is derived as float ??
-                    waiter_delay=self.poke_interval,  # type: ignore[arg-type]
+                    waiter_delay=int(self.poke_interval),
                     waiter_max_attempts=self.max_retries,
                     aws_conn_id=self.aws_conn_id,
+                    end_from_trigger=True,
+                ),
+                method_name="execute_complete",
+            )
+        else:
+            super().execute(context=context)
+
+
+class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
+    """
+    Waits for a task in an MWAA Environment to complete.
+
+    If the task fails, an AirflowException is thrown.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:MwaaTaskSensor`
+
+    :param external_env_name: The external MWAA environment name that contains the Task Instance you want to wait for
+        (templated)
+    :param external_dag_id: The DAG ID in the external MWAA environment that contains the Task Instance 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 external_task_id: The Task ID in the external MWAA environment that you want to wait for (templated)
+    :param success_states: Collection of task instance states that would make this task marked as successful, default is
+        ``{airflow.utils.state.TaskInstanceState.SUCCESS}`` (templated)
+    :param failure_states: Collection of task instance states that would make this task marked as failed and raise an
+        AirflowException, default is ``{airflow.utils.state.TaskInstanceState.FAILED}`` (templated)
+    :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
+        module to be installed.
+        (default: False, but can be overridden in config file by setting default_deferrable to True)
+    :param poke_interval: Polling period in seconds to check for the status of the job. (default: 60)
+    :param max_retries: Number of times before returning the current state. (default: 720)
+    :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
+    template_fields: Sequence[str] = aws_template_fields(
+        "external_env_name",
+        "external_dag_id",
+        "external_dag_run_id",
+        "external_task_id",
+        "success_states",
+        "failure_states",
+        "deferrable",
+        "max_retries",
+        "poke_interval",
+    )
+
+    def __init__(
+        self,
+        *,
+        external_env_name: str,
+        external_dag_id: str,
+        external_dag_run_id: str | None = None,
+        external_task_id: str,
+        success_states: Collection[str] | None = None,
+        failure_states: Collection[str] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poke_interval: int = 60,
+        max_retries: int = 720,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.success_states = set(success_states) if success_states else {TaskInstanceState.SUCCESS.value}
+        self.failure_states = set(failure_states) if failure_states else {TaskInstanceState.FAILED.value}
+
+        if len(self.success_states & self.failure_states):
+            raise ValueError("success_states and failure_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
+        self.external_task_id = external_task_id
+        self.deferrable = deferrable
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
+
+    def poke(self, context: Context) -> bool:
+        self.log.info(
+            "Poking for task %s of DAG run %s of DAG %s in MWAA environment %s",
+            self.external_task_id,
+            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}/taskInstances/{self.external_task_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 task
+
+        state = response["RestApiResponse"]["state"]
+
+        if state in self.failure_states:
+            raise AirflowException(
+                f"The task {self.external_task_id} of 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 state in self.success_states
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        return None
+
+    def execute(self, context: Context):
+        if self.external_dag_run_id is None:
+            response = self.hook.invoke_rest_api(
+                env_name=self.external_env_name,
+                path=f"/dags/{self.external_dag_id}/dagRuns",
+                method="GET",
+            )
+            self.external_dag_run_id = response["RestApiResponse"]["dag_runs"][-1]["dag_run_id"]
+
+        if self.deferrable:
+            self.defer(
+                trigger=MwaaTaskCompletedTrigger(
+                    external_env_name=self.external_env_name,
+                    external_dag_id=self.external_dag_id,
+                    external_dag_run_id=self.external_dag_run_id,
+                    external_task_id=self.external_task_id,
+                    success_states=self.success_states,
+                    failure_states=self.failure_states,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                    end_from_trigger=True,
                 ),
                 method_name="execute_complete",
             )

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

@@ -80,7 +80,7 @@ def __init__(
         waiter_delay: int,
         waiter_max_attempts: int,
         waiter_config_overrides: dict[str, Any] | None = None,
-        aws_conn_id: str | None,
+        aws_conn_id: str | None = "aws_default",
         region_name: str | None = None,
         verify: bool | str | None = None,
         botocore_config: dict | None = None,

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

@@ -22,7 +22,7 @@
 
 from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
-from airflow.utils.state import DagRunState
+from airflow.utils.state import DagRunState, State, TaskInstanceState
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
@@ -48,15 +48,15 @@ class MwaaDagRunCompletedTrigger(AwsBaseWaiterTrigger):
 
     def __init__(
         self,
-        *,
+        *args,
         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,
         waiter_delay: int = 60,
         waiter_max_attempts: int = 720,
-        aws_conn_id: str | None = None,
+        **kwargs,
     ) -> None:
         self.success_states = set(success_states) if success_states else {DagRunState.SUCCESS.value}
         self.failure_states = set(failure_states) if failure_states else {DagRunState.FAILED.value}
@@ -87,14 +87,100 @@ def __init__(
             return_value=external_dag_run_id,
             waiter_delay=waiter_delay,
             waiter_max_attempts=waiter_max_attempts,
-            aws_conn_id=aws_conn_id,
             waiter_config_overrides={
                 "acceptors": _build_waiter_acceptors(
                     success_states=self.success_states,
                     failure_states=self.failure_states,
                     in_progress_states=in_progress_states,
                 )
             },
+            **kwargs,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return MwaaHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
+
+
+class MwaaTaskCompletedTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when an MWAA Task is complete.
+
+    :param external_env_name: The external MWAA environment name that contains the Task Instance you want to wait for
+        (templated)
+    :param external_dag_id: The DAG ID in the external MWAA environment that contains the Task Instance 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).
+        If not provided, the latest DAG run is used by default.
+    :param external_task_id: The Task ID in the external MWAA environment that you want to wait for (templated)
+    :param success_states: Collection of task instance states that would make this task marked as successful, default is
+        ``{airflow.utils.state.TaskInstanceState.SUCCESS}`` (templated)
+    :param failure_states: Collection of task instance states that would make this task marked as failed and raise an
+        AirflowException, default is ``{airflow.utils.state.TaskInstanceState.FAILED}`` (templated)
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 720)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *args,
+        external_env_name: str,
+        external_dag_id: str,
+        external_dag_run_id: str | None = None,
+        external_task_id: str,
+        success_states: Collection[str] | None = None,
+        failure_states: Collection[str] | None = None,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 720,
+        **kwargs,
+    ) -> None:
+        self.success_states = (
+            set(success_states) if success_states else {state.value for state in State.success_states}
+        )
+        self.failure_states = (
+            set(failure_states) if failure_states else {state.value for state in State.failed_states}
+        )
+
+        if len(self.success_states & self.failure_states):
+            raise ValueError("success_states and failure_states must not have any values in common")
+
+        in_progress_states = {s.value for s in TaskInstanceState} - self.success_states - self.failure_states
+
+        super().__init__(
+            serialized_fields={
+                "external_env_name": external_env_name,
+                "external_dag_id": external_dag_id,
+                "external_dag_run_id": external_dag_run_id,
+                "external_task_id": external_task_id,
+                "success_states": success_states,
+                "failure_states": failure_states,
+            },
+            waiter_name="mwaa_task_complete",
+            waiter_args={
+                "Name": external_env_name,
+                "Path": f"/dags/{external_dag_id}/dagRuns/{external_dag_run_id}/taskInstances/{external_task_id}",
+                "Method": "GET",
+            },
+            failure_message=f"The task {external_task_id} of DAG run {external_dag_run_id} of DAG {external_dag_id} in MWAA environment {external_env_name} failed with state",
+            status_message="State of DAG run",
+            status_queries=["RestApiResponse.state"],
+            return_key="task_id",
+            return_value=external_task_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            waiter_config_overrides={
+                "acceptors": _build_waiter_acceptors(
+                    success_states=self.success_states,
+                    failure_states=self.failure_states,
+                    in_progress_states=in_progress_states,
+                )
+            },
+            **kwargs,
         )
 
     def hook(self) -> AwsGenericHook:

providers/amazon/src/airflow/providers/amazon/aws/waiters/mwaa.json

@@ -31,6 +31,11 @@
           "state": "failure"
         }
       ]
+    },
+    "mwaa_task_complete": {
+      "delay": 60,
+      "maxAttempts": 20,
+      "operation": "InvokeRestApi"
     }
   }
 }