apache-airflow-providers-amazon 9.5.0rc1__py3-none-any.whl → 9.5.0rc3__py3-none-any.whl

airflow/providers/amazon/aws/operators/sagemaker.py

@@ -44,7 +44,6 @@ from airflow.providers.amazon.aws.utils import trim_none_values, validate_execut
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.utils.helpers import prune_dict
-from airflow.utils.json import AirflowJsonEncoder
 
 if TYPE_CHECKING:
     from airflow.providers.common.compat.openlineage.facet import Dataset
@@ -56,7 +55,7 @@ CHECK_INTERVAL_SECOND: int = 30
 
 
 def serialize(result: dict) -> dict:
-    return json.loads(json.dumps(result, cls=AirflowJsonEncoder))
+    return json.loads(json.dumps(result, default=repr))
 
 
 class SageMakerBaseOperator(BaseOperator):
@@ -171,7 +170,7 @@ class SageMakerBaseOperator(BaseOperator):
                 timestamp = str(
                     time.time_ns() // 1000000000
                 )  # only keep the relevant datetime (first 10 digits)
-                name = f"{proposed_name[:max_name_len - len(timestamp) - 1]}-{timestamp}"  # we subtract one to make provision for the dash between the truncated name and timestamp
+                name = f"{proposed_name[: max_name_len - len(timestamp) - 1]}-{timestamp}"  # we subtract one to make provision for the dash between the truncated name and timestamp
                 self.log.info("Changed %s name to '%s' to avoid collision.", resource_type, name)
         return name
 
@@ -179,8 +178,7 @@ class SageMakerBaseOperator(BaseOperator):
         """Raise exception if resource type is not 'model' or 'job'."""
         if resource_type not in ("model", "job"):
             raise AirflowException(
-                "Argument resource_type accepts only 'model' and 'job'. "
-                f"Provided value: '{resource_type}'."
+                f"Argument resource_type accepts only 'model' and 'job'. Provided value: '{resource_type}'."
             )
 
     def _check_if_job_exists(self, job_name: str, describe_func: Callable[[str], Any]) -> bool:
@@ -560,8 +558,7 @@ class SageMakerEndpointOperator(SageMakerBaseOperator):
                 self.operation = "update"
                 sagemaker_operation = self.hook.update_endpoint
                 self.log.warning(
-                    "cannot create already existing endpoint %s, "
-                    "updating it with the given config instead",
+                    "cannot create already existing endpoint %s, updating it with the given config instead",
                     endpoint_info["EndpointName"],
                 )
                 if "Tags" in endpoint_info:

airflow/providers/amazon/aws/sensors/ec2.py

@@ -18,21 +18,21 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.ec2 import EC2StateSensorTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class EC2InstanceStateSensor(BaseSensorOperator):
+class EC2InstanceStateSensor(AwsBaseSensor[EC2Hook]):
     """
     Poll the state of the AWS EC2 instance until the instance reaches the target state.
 
@@ -46,7 +46,8 @@ class EC2InstanceStateSensor(BaseSensorOperator):
     :param deferrable: if True, the sensor will run in deferrable mode
     """
 
-    template_fields: Sequence[str] = ("target_state", "instance_id", "region_name")
+    aws_hook_class = EC2Hook
+    template_fields: Sequence[str] = aws_template_fields("target_state", "instance_id", "region_name")
     ui_color = "#cc8811"
     ui_fgcolor = "#ffffff"
     valid_states = ["running", "stopped", "terminated"]
@@ -56,8 +57,6 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         *,
         target_state: str,
         instance_id: str,
-        aws_conn_id: str | None = "aws_default",
-        region_name: str | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
@@ -66,8 +65,6 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         super().__init__(**kwargs)
         self.target_state = target_state
         self.instance_id = instance_id
-        self.aws_conn_id = aws_conn_id
-        self.region_name = region_name
         self.deferrable = deferrable
 
     def execute(self, context: Context) -> Any:
@@ -85,10 +82,6 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         else:
             super().execute(context=context)
 
-    @cached_property
-    def hook(self):
-        return EC2Hook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-
     def poke(self, context: Context):
         instance_state = self.hook.get_instance_state(instance_id=self.instance_id)
         self.log.info("instance state: %s", instance_state)

airflow/providers/amazon/aws/sensors/glacier.py

@@ -95,5 +95,5 @@ class GlacierJobOperationSensor(AwsBaseSensor[GlacierHook]):
             return False
         else:
             raise AirflowException(
-                f'Sensor failed. Job status: {response["Action"]}, code status: {response["StatusCode"]}'
+                f"Sensor failed. Job status: {response['Action']}, code status: {response['StatusCode']}"
             )

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

@@ -18,13 +18,16 @@
 from __future__ import annotations
 
 from collections.abc import Collection, 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.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.utils.mixins import aws_template_fields
-from airflow.utils.state import State
+from airflow.utils.state import DagRunState
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -46,9 +49,24 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         (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)
+        ``{airflow.utils.state.DagRunState.SUCCESS}`` (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)
+        AirflowException, default is ``{airflow.utils.state.DagRunState.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
@@ -58,6 +76,9 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         "external_dag_run_id",
         "success_states",
         "failure_states",
+        "deferrable",
+        "max_retries",
+        "poke_interval",
     )
 
     def __init__(
@@ -68,19 +89,25 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         external_dag_run_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 State.success_states)
-        self.failure_states = set(failure_states if failure_states else State.failed_states)
+        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}
 
         if len(self.success_states & self.failure_states):
-            raise AirflowException("allowed_states and failed_states must not have any values in common")
+            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.deferrable = deferrable
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
 
     def poke(self, context: Context) -> bool:
         self.log.info(
@@ -102,12 +129,33 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         # 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}."
+                f"failed with state: {state}"
             )
-        return False
+
+        return state in self.success_states
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        validate_execute_complete_event(event)
+
+    def execute(self, context: Context):
+        if self.deferrable:
+            self.defer(
+                trigger=MwaaDagRunCompletedTrigger(
+                    external_env_name=self.external_env_name,
+                    external_dag_id=self.external_dag_id,
+                    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_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        else:
+            super().execute(context=context)

airflow/providers/amazon/aws/sensors/s3.py

@@ -192,7 +192,7 @@ class S3KeySensor(AwsBaseSensor[S3Hook]):
         self.defer(
             timeout=timedelta(seconds=self.timeout),
             trigger=S3KeyTrigger(
-                bucket_name=cast(str, self.bucket_name),
+                bucket_name=cast("str", self.bucket_name),
                 bucket_key=self.bucket_key,
                 wildcard_match=self.wildcard_match,
                 aws_conn_id=self.aws_conn_id,

airflow/providers/amazon/aws/sensors/step_function.py

@@ -81,5 +81,6 @@ class StepFunctionExecutionSensor(AwsBaseSensor[StepFunctionHook]):
             return False
 
         self.log.info("Doing xcom_push of output")
-        self.xcom_push(context, "output", output)
+
+        context["ti"].xcom_push(key="output", value=output)
         return True

airflow/providers/amazon/aws/transfers/mongo_to_s3.py

@@ -103,7 +103,7 @@ class MongoToS3Operator(BaseOperator):
         if self.is_pipeline:
             results: CommandCursor[Any] | Cursor = MongoHook(self.mongo_conn_id).aggregate(
                 mongo_collection=self.mongo_collection,
-                aggregate_query=cast(list, self.mongo_query),
+                aggregate_query=cast("list", self.mongo_query),
                 mongo_db=self.mongo_db,
                 allowDiskUse=self.allow_disk_use,
             )
@@ -111,7 +111,7 @@ class MongoToS3Operator(BaseOperator):
         else:
             results = MongoHook(self.mongo_conn_id).find(
                 mongo_collection=self.mongo_collection,
-                query=cast(dict, self.mongo_query),
+                query=cast("dict", self.mongo_query),
                 projection=self.mongo_projection,
                 mongo_db=self.mongo_db,
                 find_one=False,

airflow/providers/amazon/aws/transfers/redshift_to_s3.py

@@ -29,6 +29,7 @@ from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -102,7 +103,7 @@ class RedshiftToS3Operator(BaseOperator):
         table: str | None = None,
         select_query: str | None = None,
         redshift_conn_id: str = "redshift_default",
-        aws_conn_id: str | None = "aws_default",
+        aws_conn_id: str | None | ArgNotSet = NOTSET,
         verify: bool | str | None = None,
         unload_options: list | None = None,
         autocommit: bool = False,
@@ -118,7 +119,6 @@ class RedshiftToS3Operator(BaseOperator):
         self.schema = schema
         self.table = table
         self.redshift_conn_id = redshift_conn_id
-        self.aws_conn_id = aws_conn_id
         self.verify = verify
         self.unload_options = unload_options or []
         self.autocommit = autocommit
@@ -127,6 +127,16 @@ class RedshiftToS3Operator(BaseOperator):
         self.table_as_file_name = table_as_file_name
         self.redshift_data_api_kwargs = redshift_data_api_kwargs or {}
         self.select_query = select_query
+        # In execute() we attempt to fetch this aws connection to check for extras. If the user didn't
+        # actually provide a connection note that, because we don't want to let the exception bubble up in
+        # that case (since we're silently injecting a connection on their behalf).
+        self._aws_conn_id: str | None
+        if isinstance(aws_conn_id, ArgNotSet):
+            self.conn_set = False
+            self._aws_conn_id = "aws_default"
+        else:
+            self.conn_set = True
+            self._aws_conn_id = aws_conn_id
 
     def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str
@@ -176,11 +186,16 @@ class RedshiftToS3Operator(BaseOperator):
                     raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")
         else:
             redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
-        conn = S3Hook.get_connection(conn_id=self.aws_conn_id) if self.aws_conn_id else None
+        conn = (
+            S3Hook.get_connection(conn_id=self._aws_conn_id)
+            # Only fetch the connection if it was set by the user and it is not None
+            if self.conn_set and self._aws_conn_id
+            else None
+        )
         if conn and conn.extra_dejson.get("role_arn", False):
             credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
         else:
-            s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
+            s3_hook = S3Hook(aws_conn_id=self._aws_conn_id, verify=self.verify)
             credentials = s3_hook.get_credentials()
             credentials_block = build_credentials_block(credentials)
 

airflow/providers/amazon/aws/transfers/s3_to_redshift.py

@@ -25,6 +25,7 @@ from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
+from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -93,7 +94,7 @@ class S3ToRedshiftOperator(BaseOperator):
         s3_key: str,
         schema: str | None = None,
         redshift_conn_id: str = "redshift_default",
-        aws_conn_id: str | None = "aws_default",
+        aws_conn_id: str | None | ArgNotSet = NOTSET,
         verify: bool | str | None = None,
         column_list: list[str] | None = None,
         copy_options: list | None = None,
@@ -117,6 +118,16 @@ class S3ToRedshiftOperator(BaseOperator):
         self.method = method
         self.upsert_keys = upsert_keys
         self.redshift_data_api_kwargs = redshift_data_api_kwargs or {}
+        # In execute() we attempt to fetch this aws connection to check for extras. If the user didn't
+        # actually provide a connection note that, because we don't want to let the exception bubble up in
+        # that case (since we're silently injecting a connection on their behalf).
+        self._aws_conn_id: str | None
+        if isinstance(aws_conn_id, ArgNotSet):
+            self.conn_set = False
+            self._aws_conn_id = "aws_default"
+        else:
+            self.conn_set = True
+            self._aws_conn_id = aws_conn_id
 
         if self.redshift_data_api_kwargs:
             for arg in ["sql", "parameters"]:
@@ -149,14 +160,19 @@ class S3ToRedshiftOperator(BaseOperator):
         else:
             redshift_sql_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
 
-        conn = S3Hook.get_connection(conn_id=self.aws_conn_id) if self.aws_conn_id else None
+        conn = (
+            S3Hook.get_connection(conn_id=self._aws_conn_id)
+            # Only fetch the connection if it was set by the user and it is not None
+            if self.conn_set and self._aws_conn_id
+            else None
+        )
         region_info = ""
         if conn and conn.extra_dejson.get("region", False):
             region_info = f"region '{conn.extra_dejson['region']}'"
         if conn and conn.extra_dejson.get("role_arn", False):
             credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
         else:
-            s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
+            s3_hook = S3Hook(aws_conn_id=self._aws_conn_id, verify=self.verify)
             credentials = s3_hook.get_credentials()
             credentials_block = build_credentials_block(credentials)
 

airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -223,7 +223,7 @@ class SqlToS3Operator(BaseOperator):
             return
         for group_label in (grouped_df := df.groupby(**self.groupby_kwargs)).groups:
             yield (
-                cast(str, group_label),
+                cast("str", group_label),
                 grouped_df.get_group(group_label)
                 .drop(random_column_name, axis=1, errors="ignore")
                 .reset_index(drop=True),

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

@@ -55,6 +55,8 @@ class AwsBaseWaiterTrigger(BaseTrigger):
 
     :param waiter_delay: The amount of time in seconds to wait between attempts.
     :param waiter_max_attempts: The maximum number of attempts to be made.
+    :param waiter_config_overrides: A dict to update waiter's default configuration. Only specified keys will
+        be updated.
     :param aws_conn_id: The Airflow connection used for AWS credentials. To be used to build the hook.
     :param region_name: The AWS region where the resources to watch are. To be used to build the hook.
     :param verify: Whether or not to verify SSL certificates. To be used to build the hook.
@@ -77,6 +79,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
         return_value: Any,
         waiter_delay: int,
         waiter_max_attempts: int,
+        waiter_config_overrides: dict[str, Any] | None = None,
         aws_conn_id: str | None,
         region_name: str | None = None,
         verify: bool | str | None = None,
@@ -91,6 +94,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
         self.failure_message = failure_message
         self.status_message = status_message
         self.status_queries = status_queries
+        self.waiter_config_overrides = waiter_config_overrides
 
         self.return_key = return_key
         self.return_value = return_value
@@ -140,7 +144,12 @@ class AwsBaseWaiterTrigger(BaseTrigger):
     async def run(self) -> AsyncIterator[TriggerEvent]:
         hook = self.hook()
         async with await hook.get_async_conn() as client:
-            waiter = hook.get_waiter(self.waiter_name, deferrable=True, client=client)
+            waiter = hook.get_waiter(
+                self.waiter_name,
+                deferrable=True,
+                client=client,
+                config_overrides=self.waiter_config_overrides,
+            )
             await async_wait(
                 waiter,
                 self.waiter_delay,

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

@@ -0,0 +1,128 @@
+# 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
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.utils.state import DagRunState
+
+if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class MwaaDagRunCompletedTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when an MWAA Dag Run is complete.
+
+    :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.DagRunState.SUCCESS}`` (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.DagRunState.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,
+        *,
+        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,
+    ) -> 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}
+
+        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 DagRunState} - 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,
+                "success_states": success_states,
+                "failure_states": failure_states,
+            },
+            waiter_name="mwaa_dag_run_complete",
+            waiter_args={
+                "Name": external_env_name,
+                "Path": f"/dags/{external_dag_id}/dagRuns/{external_dag_run_id}",
+                "Method": "GET",
+            },
+            failure_message=f"The 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="dag_run_id",
+            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,
+                )
+            },
+        )
+
+    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,
+        )
+
+
+def _build_waiter_acceptors(
+    success_states: set[str], failure_states: set[str], in_progress_states: set[str]
+) -> list:
+    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(
+                {
+                    "matcher": "path",
+                    "argument": "RestApiResponse.state",
+                    "expected": dag_run_state,
+                    "state": state_waiter_category,
+                }
+            )
+
+    return acceptors

airflow/providers/amazon/aws/utils/waiter_with_logging.py

@@ -136,15 +136,16 @@ async def async_wait(
             last_response = error.last_response
 
             if "terminal failure" in error_reason:
-                log.error("%s: %s", failure_message, _LazyStatusFormatter(status_args, last_response))
-                raise AirflowException(f"{failure_message}: {error}")
+                raise AirflowException(
+                    f"{failure_message}: {_LazyStatusFormatter(status_args, last_response)}\n{error}"
+                )
 
             if (
                 "An error occurred" in error_reason
                 and isinstance(last_response.get("Error"), dict)
                 and "Code" in last_response.get("Error")
             ):
-                raise AirflowException(f"{failure_message}: {error}")
+                raise AirflowException(f"{failure_message}\n{last_response}\n{error}")
 
             log.info("%s: %s", status_message, _LazyStatusFormatter(status_args, last_response))
         else:

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

@@ -0,0 +1,36 @@
+{
+  "version": 2,
+  "waiters": {
+    "mwaa_dag_run_complete": {
+      "delay": 60,
+      "maxAttempts": 720,
+      "operation": "InvokeRestApi",
+      "acceptors": [
+        {
+          "matcher": "path",
+          "argument": "RestApiResponse.state",
+          "expected": "queued",
+          "state": "retry"
+        },
+        {
+          "matcher": "path",
+          "argument": "RestApiResponse.state",
+          "expected": "running",
+          "state": "retry"
+        },
+        {
+          "matcher": "path",
+          "argument": "RestApiResponse.state",
+          "expected": "success",
+          "state": "success"
+        },
+        {
+          "matcher": "path",
+          "argument": "RestApiResponse.state",
+          "expected": "failed",
+          "state": "failure"
+        }
+      ]
+    }
+  }
+}