apache-airflow-providers-amazon 9.14.0__py3-none-any.whl → 9.18.0rc2__py3-none-any.whl

airflow/providers/amazon/aws/notifications/sns.py

@@ -21,6 +21,7 @@ from collections.abc import Sequence
 from functools import cached_property
 
 from airflow.providers.amazon.aws.hooks.sns import SnsHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_1_PLUS
 from airflow.providers.common.compat.notifier import BaseNotifier
 
 
@@ -60,8 +61,13 @@ class SnsNotifier(BaseNotifier):
         subject: str | None = None,
         message_attributes: dict | None = None,
         region_name: str | None = None,
+        **kwargs,
     ):
-        super().__init__()
+        if AIRFLOW_V_3_1_PLUS:
+            #  Support for passing context was added in 3.1.0
+            super().__init__(**kwargs)
+        else:
+            super().__init__()
         self.aws_conn_id = aws_conn_id
         self.region_name = region_name
         self.target_arn = target_arn
@@ -83,5 +89,14 @@ class SnsNotifier(BaseNotifier):
             message_attributes=self.message_attributes,
         )
 
+    async def async_notify(self, context):
+        """Publish the notification message to Amazon SNS (async)."""
+        await self.hook.apublish_to_target(
+            target_arn=self.target_arn,
+            message=self.message,
+            subject=self.subject,
+            message_attributes=self.message_attributes,
+        )
+
 
 send_sns_notification = SnsNotifier

airflow/providers/amazon/aws/notifications/sqs.py

@@ -21,6 +21,7 @@ from collections.abc import Sequence
 from functools import cached_property
 
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_1_PLUS
 from airflow.providers.common.compat.notifier import BaseNotifier
 
 
@@ -64,8 +65,13 @@ class SqsNotifier(BaseNotifier):
         message_group_id: str | None = None,
         delay_seconds: int = 0,
         region_name: str | None = None,
+        **kwargs,
     ):
-        super().__init__()
+        if AIRFLOW_V_3_1_PLUS:
+            #  Support for passing context was added in 3.1.0
+            super().__init__(**kwargs)
+        else:
+            super().__init__()
         self.aws_conn_id = aws_conn_id
         self.region_name = region_name
         self.queue_url = queue_url
@@ -89,5 +95,15 @@ class SqsNotifier(BaseNotifier):
             message_group_id=self.message_group_id,
         )
 
+    async def async_notify(self, context):
+        """Publish the notification message to Amazon SQS queue (async)."""
+        await self.hook.asend_message(
+            queue_url=self.queue_url,
+            message_body=self.message_body,
+            delay_seconds=self.delay_seconds,
+            message_attributes=self.message_attributes,
+            message_group_id=self.message_group_id,
+        )
+
 
 send_sqs_notification = SqsNotifier

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

@@ -25,8 +25,8 @@ from airflow.providers.amazon.aws.utils.mixins import (
     AwsHookType,
     aws_template_fields,
 )
-from airflow.providers.amazon.version_compat import BaseOperator
-from airflow.utils.types import NOTSET, ArgNotSet
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
+from airflow.providers.common.compat.sdk import BaseOperator
 
 
 class AwsBaseOperator(BaseOperator, AwsBaseHookMixin[AwsHookType]):

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

@@ -482,6 +482,8 @@ class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
                     # It may also be that permissions haven't even propagated yet to check for the index
                     or "server returned 401" in error_message
                     or "user does not have permissions" in error_message
+                    or "status code: 403" in error_message
+                    or "bad authorization" in error_message
                 )
                 if all(
                     [

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

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
 
 class CloudFormationCreateStackOperator(AwsBaseOperator[CloudFormationHook]):
     """
-    An operator that creates a AWS CloudFormation stack.
+    An operator that creates an AWS CloudFormation stack.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -68,7 +68,7 @@ class CloudFormationCreateStackOperator(AwsBaseOperator[CloudFormationHook]):
 
 class CloudFormationDeleteStackOperator(AwsBaseOperator[CloudFormationHook]):
     """
-    An operator that deletes a AWS CloudFormation stack.
+    An operator that deletes an AWS CloudFormation stack.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:

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

@@ -23,11 +23,12 @@ import random
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException, AirflowTaskTimeout
+from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook
 from airflow.providers.amazon.aws.links.datasync import DataSyncTaskExecutionLink, DataSyncTaskLink
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowTaskTimeout
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -57,8 +57,8 @@ from airflow.providers.amazon.aws.utils.waiter import (
     waiter,
 )
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
 from airflow.utils.helpers import exactly_one, prune_dict
-from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -654,11 +654,10 @@ class EmrCreateJobFlowOperator(AwsBaseOperator[EmrHook]):
     :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 wait_for_completion: Deprecated - use `wait_policy` instead.
-        Whether to finish task immediately after creation (False) or wait for jobflow
+    :param wait_for_completion: Whether to finish task immediately after creation (False) or wait for jobflow
         completion (True)
         (default: None)
-    :param wait_policy: Whether to finish the task immediately after creation (None) or:
+    :param wait_policy: Deprecated. Use `wait_for_completion` instead. Whether to finish the task immediately after creation (None) or:
         - wait for the jobflow completion (WaitPolicy.WAIT_FOR_COMPLETION)
         - wait for the jobflow completion and cluster to terminate (WaitPolicy.WAIT_FOR_STEPS_COMPLETION)
         (default: None)
@@ -698,19 +697,29 @@ class EmrCreateJobFlowOperator(AwsBaseOperator[EmrHook]):
         super().__init__(**kwargs)
         self.emr_conn_id = emr_conn_id
         self.job_flow_overrides = job_flow_overrides or {}
-        self.wait_policy = wait_policy
+        self.wait_for_completion = wait_for_completion
         self.waiter_max_attempts = waiter_max_attempts or 60
         self.waiter_delay = waiter_delay or 60
         self.deferrable = deferrable
 
-        if wait_for_completion is not None:
+        if wait_policy is not None:
             warnings.warn(
-                "`wait_for_completion` parameter is deprecated, please use `wait_policy` instead.",
+                "`wait_policy` parameter is deprecated and will be removed in a future release; "
+                "please use `wait_for_completion` (bool) instead.",
                 AirflowProviderDeprecationWarning,
                 stacklevel=2,
             )
-            # preserve previous behaviour
-            self.wait_policy = WaitPolicy.WAIT_FOR_COMPLETION if wait_for_completion else None
+
+            if wait_for_completion is not None:
+                raise ValueError(
+                    "Cannot specify both `wait_for_completion` and deprecated `wait_policy`. "
+                    "Please use `wait_for_completion` (bool)."
+                )
+
+            self.wait_for_completion = wait_policy in (
+                WaitPolicy.WAIT_FOR_COMPLETION,
+                WaitPolicy.WAIT_FOR_STEPS_COMPLETION,
+            )
 
     @property
     def _hook_parameters(self):
@@ -748,30 +757,32 @@ class EmrCreateJobFlowOperator(AwsBaseOperator[EmrHook]):
                 job_flow_id=self._job_flow_id,
                 log_uri=get_log_uri(emr_client=self.hook.conn, job_flow_id=self._job_flow_id),
             )
-        if self.deferrable:
-            self.defer(
-                trigger=EmrCreateJobFlowTrigger(
-                    job_flow_id=self._job_flow_id,
-                    aws_conn_id=self.aws_conn_id,
-                    waiter_delay=self.waiter_delay,
-                    waiter_max_attempts=self.waiter_max_attempts,
-                ),
-                method_name="execute_complete",
-                # timeout is set to ensure that if a trigger dies, the timeout does not restart
-                # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
-                timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay + 60),
-            )
-        if self.wait_policy:
-            waiter_name = WAITER_POLICY_NAME_MAPPING[self.wait_policy]
-            self.hook.get_waiter(waiter_name).wait(
-                ClusterId=self._job_flow_id,
-                WaiterConfig=prune_dict(
-                    {
-                        "Delay": self.waiter_delay,
-                        "MaxAttempts": self.waiter_max_attempts,
-                    }
-                ),
-            )
+        if self.wait_for_completion:
+            waiter_name = WAITER_POLICY_NAME_MAPPING[WaitPolicy.WAIT_FOR_COMPLETION]
+
+            if self.deferrable:
+                self.defer(
+                    trigger=EmrCreateJobFlowTrigger(
+                        job_flow_id=self._job_flow_id,
+                        aws_conn_id=self.aws_conn_id,
+                        waiter_delay=self.waiter_delay,
+                        waiter_max_attempts=self.waiter_max_attempts,
+                    ),
+                    method_name="execute_complete",
+                    # timeout is set to ensure that if a trigger dies, the timeout does not restart
+                    # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
+                    timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay + 60),
+                )
+            else:
+                self.hook.get_waiter(waiter_name).wait(
+                    ClusterId=self._job_flow_id,
+                    WaiterConfig=prune_dict(
+                        {
+                            "Delay": self.waiter_delay,
+                            "MaxAttempts": self.waiter_max_attempts,
+                        }
+                    ),
+                )
         return self._job_flow_id
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:

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

@@ -19,7 +19,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -46,12 +46,14 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
     :param trigger_run_id: The Run ID. This together with trigger_dag_id are a unique key. (templated)
     :param logical_date: The logical date (previously called execution date). This is the time or interval
         covered by this DAG run, according to the DAG definition. This together with trigger_dag_id are a
-        unique key. (templated)
+        unique key. This field is required if your environment is running with Airflow 3. (templated)
     :param data_interval_start: The beginning of the interval the DAG run covers
     :param data_interval_end: The end of the interval the DAG run covers
     :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 airflow_version: The Airflow major version the MWAA environment runs.
+            This parameter is only used if the local web token method is used to call Airflow API. (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)
@@ -81,6 +83,7 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
         "data_interval_end",
         "conf",
         "note",
+        "airflow_version",
     )
     template_fields_renderers = {"conf": "json"}
 
@@ -95,6 +98,7 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
         data_interval_end: str | None = None,
         conf: dict | None = None,
         note: str | None = None,
+        airflow_version: Literal[2, 3] | None = None,
         wait_for_completion: bool = False,
         waiter_delay: int = 60,
         waiter_max_attempts: int = 20,
@@ -110,6 +114,7 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
         self.data_interval_end = data_interval_end
         self.conf = conf if conf else {}
         self.note = note
+        self.airflow_version = airflow_version
         self.wait_for_completion = wait_for_completion
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
@@ -123,7 +128,10 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
         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"
+            env_name=self.env_name,
+            path=f"/dags/{self.trigger_dag_id}/dagRuns/{dag_run_id}",
+            method="GET",
+            airflow_version=self.airflow_version,
         )
 
     def execute(self, context: Context) -> dict:
@@ -146,6 +154,7 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
                 "conf": self.conf,
                 "note": self.note,
             },
+            airflow_version=self.airflow_version,
         )
 
         dag_run_id = response["RestApiResponse"]["dag_run_id"]

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

@@ -33,7 +33,7 @@ from airflow.providers.amazon.aws.links.sagemaker_unified_studio import (
 from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import (
     SageMakerNotebookJobTrigger,
 )
-from airflow.providers.amazon.version_compat import BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -20,7 +20,6 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ssm import SsmHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.ssm import SsmRunCommandTrigger
@@ -36,27 +35,35 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
     Executes the SSM Run Command to perform actions on managed instances.
 
     .. seealso::
-        For more information on how to use this operator, take a look at the guide:
+        For more information on how to use this operator, take a look at the
+        guide:
         :ref:`howto/operator:SsmRunCommandOperator`
 
-    :param document_name: The name of the Amazon Web Services Systems Manager document (SSM document) to run.
-    :param run_command_kwargs: Optional parameters to pass to the send_command API.
-
-    :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
-    :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
-    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 75)
-    :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
-        This implies waiting for completion. This mode requires aiobotocore module to be installed.
-        (default: False)
+    :param document_name: The name of the Amazon Web Services Systems Manager
+        document (SSM document) to run.
+    :param run_command_kwargs: Optional parameters to pass to the send_command
+        API.
+
+    :param wait_for_completion: Whether to wait for cluster to stop.
+        (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks.
+        (default: 120)
+    :param waiter_max_attempts: Maximum number of attempts to check for job
+        completion. (default: 75)
+    :param deferrable: If True, the operator will wait asynchronously for the
+        cluster 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
+        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 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:
+    :param botocore_config: Configuration dictionary (key-values) for botocore
+        client. See:
         https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
     """
 
@@ -90,7 +97,7 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
         event = validate_execute_complete_event(event)
 
         if event["status"] != "success":
-            raise AirflowException(f"Error while running run command: {event}")
+            raise RuntimeError(f"Error while running run command: {event}")
 
         self.log.info("SSM run command `%s` completed.", event["command_id"])
         return event["command_id"]
@@ -112,6 +119,9 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
                     aws_conn_id=self.aws_conn_id,
+                    region_name=self.region_name,
+                    verify=self.verify,
+                    botocore_config=self.botocore_config,
                 ),
                 method_name="execute_complete",
             )
@@ -125,7 +135,102 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
                 waiter.wait(
                     CommandId=command_id,
                     InstanceId=instance_id,
-                    WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+                    WaiterConfig={
+                        "Delay": self.waiter_delay,
+                        "MaxAttempts": self.waiter_max_attempts,
+                    },
                 )
 
         return command_id
+
+
+class SsmGetCommandInvocationOperator(AwsBaseOperator[SsmHook]):
+    """
+    Retrieves the output and execution details of an SSM command invocation.
+
+    This operator allows you to fetch the standard output, standard error,
+    execution status, and other details from SSM commands. It can be used to
+    retrieve output from commands executed by SsmRunCommandOperator in previous
+    tasks, or from commands executed outside of Airflow entirely.
+
+    The operator returns structured data including stdout, stderr, execution
+    times, and status information for each instance that executed the command.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the
+        guide:
+        :ref:`howto/operator:SsmGetCommandInvocationOperator`
+
+    :param command_id: The ID of the SSM command to retrieve output for.
+    :param instance_id: The ID of the specific instance to retrieve output
+        for. If not provided, retrieves output from all instances that
+        executed the command.
+    :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 = SsmHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "command_id",
+        "instance_id",
+    )
+
+    def __init__(
+        self,
+        *,
+        command_id: str,
+        instance_id: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.command_id = command_id
+        self.instance_id = instance_id
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        """Execute the operator to retrieve command invocation output."""
+        if self.instance_id:
+            self.log.info(
+                "Retrieving output for command %s on instance %s",
+                self.command_id,
+                self.instance_id,
+            )
+            invocations = [{"InstanceId": self.instance_id}]
+        else:
+            self.log.info("Retrieving output for command %s from all instances", self.command_id)
+            response = self.hook.list_command_invocations(self.command_id)
+            invocations = response.get("CommandInvocations", [])
+
+        output_data: dict[str, Any] = {"command_id": self.command_id, "invocations": []}
+
+        for invocation in invocations:
+            instance_id = invocation["InstanceId"]
+            try:
+                invocation_details = self.hook.get_command_invocation(self.command_id, instance_id)
+                output_data["invocations"].append(
+                    {
+                        "instance_id": instance_id,
+                        "status": invocation_details.get("Status", ""),
+                        "response_code": invocation_details.get("ResponseCode", ""),
+                        "standard_output": invocation_details.get("StandardOutputContent", ""),
+                        "standard_error": invocation_details.get("StandardErrorContent", ""),
+                        "execution_start_time": invocation_details.get("ExecutionStartDateTime", ""),
+                        "execution_end_time": invocation_details.get("ExecutionEndDateTime", ""),
+                        "document_name": invocation_details.get("DocumentName", ""),
+                        "comment": invocation_details.get("Comment", ""),
+                    }
+                )
+            except Exception as e:
+                self.log.warning("Failed to get output for instance %s: %s", instance_id, e)
+                output_data["invocations"].append({"instance_id": instance_id, "error": str(e)})
+
+        return output_data

airflow/providers/amazon/aws/secrets/secrets_manager.py

@@ -187,10 +187,9 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
         }
 
         for conn_field, extra_words in self.extra_conn_words.items():
-            if conn_field == "user":
-                # Support `user` for backwards compatibility.
-                conn_field = "login"
-            possible_words_for_conn_fields[conn_field].extend(extra_words)
+            # Support `user` for backwards compatibility.
+            conn_field_backcompat = "login" if conn_field == "user" else conn_field
+            possible_words_for_conn_fields[conn_field_backcompat].extend(extra_words)
 
         conn_d: dict[str, Any] = {}
         for conn_field, possible_words in possible_words_for_conn_fields.items():

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

@@ -25,8 +25,8 @@ from airflow.providers.amazon.aws.utils.mixins import (
     AwsHookType,
     aws_template_fields,
 )
-from airflow.providers.amazon.version_compat import BaseSensorOperator
-from airflow.utils.types import NOTSET, ArgNotSet
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 
 
 class AwsBaseSensor(BaseSensorOperator, AwsBaseHookMixin[AwsHookType]):

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

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 from collections.abc import Collection, Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -51,6 +51,8 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         ``{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 airflow_version: The Airflow major version the MWAA environment runs.
+            This parameter is only used if the local web token method is used to call Airflow API. (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)
@@ -75,6 +77,7 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         "external_dag_run_id",
         "success_states",
         "failure_states",
+        "airflow_version",
         "deferrable",
         "max_retries",
         "poke_interval",
@@ -88,6 +91,7 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         external_dag_run_id: str,
         success_states: Collection[str] | None = None,
         failure_states: Collection[str] | None = None,
+        airflow_version: Literal[2, 3] | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poke_interval: int = 60,
         max_retries: int = 720,
@@ -104,6 +108,7 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
         self.external_env_name = external_env_name
         self.external_dag_id = external_dag_id
         self.external_dag_run_id = external_dag_run_id
+        self.airflow_version = airflow_version
         self.deferrable = deferrable
         self.poke_interval = poke_interval
         self.max_retries = max_retries
@@ -119,6 +124,7 @@ class MwaaDagRunSensor(AwsBaseSensor[MwaaHook]):
             env_name=self.external_env_name,
             path=f"/dags/{self.external_dag_id}/dagRuns/{self.external_dag_run_id}",
             method="GET",
+            airflow_version=self.airflow_version,
         )
 
         # If RestApiStatusCode == 200, the RestApiResponse must have the "state" key, otherwise something terrible has
@@ -179,6 +185,8 @@ class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
         ``{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 airflow_version: The Airflow major version the MWAA environment runs.
+            This parameter is only used if the local web token method is used to call Airflow API. (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)
@@ -204,6 +212,7 @@ class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
         "external_task_id",
         "success_states",
         "failure_states",
+        "airflow_version",
         "deferrable",
         "max_retries",
         "poke_interval",
@@ -218,6 +227,7 @@ class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
         external_task_id: str,
         success_states: Collection[str] | None = None,
         failure_states: Collection[str] | None = None,
+        airflow_version: Literal[2, 3] | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poke_interval: int = 60,
         max_retries: int = 720,
@@ -235,6 +245,7 @@ class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
         self.external_dag_id = external_dag_id
         self.external_dag_run_id = external_dag_run_id
         self.external_task_id = external_task_id
+        self.airflow_version = airflow_version
         self.deferrable = deferrable
         self.poke_interval = poke_interval
         self.max_retries = max_retries
@@ -252,6 +263,7 @@ class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
             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",
+            airflow_version=self.airflow_version,
         )
         # If RestApiStatusCode == 200, the RestApiResponse must have the "state" key, otherwise something terrible has
         # happened in the API and KeyError would be raised
@@ -278,6 +290,7 @@ class MwaaTaskSensor(AwsBaseSensor[MwaaHook]):
                 env_name=self.external_env_name,
                 path=f"/dags/{self.external_dag_id}/dagRuns",
                 method="GET",
+                airflow_version=self.airflow_version,
             )
             self.external_dag_run_id = response["RestApiResponse"]["dag_runs"][-1]["dag_run_id"]