apache-airflow-providers-amazon 9.15.0__py3-none-any.whl → 9.18.1rc1__py3-none-any.whl

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

@@ -21,11 +21,11 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.batch import BatchJobTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,7 +22,6 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, TypeVar
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.bedrock import (
@@ -34,6 +33,7 @@ from airflow.providers.amazon.aws.triggers.bedrock import (
     BedrockProvisionModelThroughputCompletedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.triggers.bedrock import BedrockBaseBatchInferenceTrigger

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

@@ -21,7 +21,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.comprehend import ComprehendHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.comprehend import (
@@ -29,6 +28,7 @@ from airflow.providers.amazon.aws.triggers.comprehend import (
     ComprehendPiiEntitiesDetectionJobCompletedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -20,10 +20,10 @@ from __future__ import annotations
 from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.dms import DmsHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,12 +21,12 @@ 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.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.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -20,7 +20,6 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ecs import (
     EcsClusterStates,
     EcsHook,
@@ -29,6 +28,7 @@ from airflow.providers.amazon.aws.hooks.ecs import (
 )
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     import boto3

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

@@ -23,7 +23,7 @@ from abc import abstractmethod
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.eks import (
     ClusterStates,
     EksHook,
@@ -32,6 +32,7 @@ from airflow.providers.amazon.aws.hooks.eks import (
 )
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,9 +22,6 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import (
-    AirflowException,
-)
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.links.emr import EmrClusterLink, EmrLogsLink, get_log_uri
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
@@ -35,6 +32,7 @@ from airflow.providers.amazon.aws.triggers.emr import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,10 +21,10 @@ from collections.abc import Sequence
 from enum import Enum
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,7 +21,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.glue import GlueDataQualityHook, GlueJobHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.glue import (
@@ -31,6 +30,7 @@ from airflow.providers.amazon.aws.triggers.glue import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,12 +22,12 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.glue_catalog import GlueCatalogHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.glue import GlueCatalogPartitionTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -20,10 +20,10 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.glue_crawler import GlueCrawlerHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -20,13 +20,13 @@ 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.kinesis_analytics import KinesisAnalyticsV2Hook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.kinesis_analytics import (
     KinesisAnalyticsV2ApplicationOperationCompleteTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -20,11 +20,11 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils import trim_none_values
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -18,14 +18,14 @@
 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
 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, MwaaTaskCompletedTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.state import DagRunState, TaskInstanceState
 
 if TYPE_CHECKING:
@@ -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"]
 

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

@@ -20,13 +20,13 @@ 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.opensearch_serverless import OpenSearchServerlessHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.opensearch_serverless import (
     OpenSearchServerlessCollectionActiveTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.helpers import exactly_one
 
 if TYPE_CHECKING:

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

@@ -20,9 +20,9 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -19,11 +19,11 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,12 +21,12 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftClusterTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -31,12 +31,11 @@ from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.s3 import S3KeysUnchangedTrigger, S3KeyTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.sensors.base import poke_mode_only
+from airflow.providers.common.compat.sdk import AirflowException, poke_mode_only
 
 
 class S3KeySensor(AwsBaseSensor[S3Hook]):
@@ -215,6 +214,7 @@ class S3KeySensor(AwsBaseSensor[S3Hook]):
                 poke_interval=self.poke_interval,
                 should_check_fn=bool(self.check_fn),
                 use_regex=self.use_regex,
+                metadata_keys=self.metadata_keys,
             ),
             method_name="execute_complete",
         )
@@ -226,7 +226,7 @@ class S3KeySensor(AwsBaseSensor[S3Hook]):
         Relies on trigger to throw an exception, otherwise it assumes execution was successful.
         """
         if event["status"] == "running":
-            found_keys = self.check_fn(event["files"])  # type: ignore[misc]
+            found_keys = self.check_fn(event["files"], **context)  # type: ignore[misc]
             if not found_keys:
                 self._defer()
         elif event["status"] == "error":

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

@@ -20,10 +20,10 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker import LogState, SageMakerHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,11 +21,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
     SageMakerNotebookHook,
 )
-from airflow.providers.amazon.version_compat import BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -24,13 +24,13 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.sqs import process_response
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection

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

@@ -21,7 +21,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.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.ssm import SsmRunCommandTrigger
@@ -34,32 +33,45 @@ if TYPE_CHECKING:
 
 class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
     """
-    Poll the state of an AWS SSM Run Command until all instance jobs reach a terminal state. Fails if any instance job ends in a failed state.
+    Poll the state of an AWS SSM Run Command until completion.
+
+    Waits until all instance jobs reach a terminal state. Fails if any
+    instance job ends in a failed state.
 
     .. seealso::
-        For more information on how to use this sensor, take a look at the guide:
+        For more information on how to use this sensor, take a look at the
+        guide:
         :ref:`howto/sensor:SsmRunCommandCompletedSensor`
 
     :param command_id: The ID of the AWS SSM Run Command.
-
-    :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: 120)
-    :param max_retries: Number of times before returning the current state. (default: 75)
+    :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: 120)
+    :param max_retries: Number of times before returning the current state.
+        (default: 75)
     :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
     """
 
-    INTERMEDIATE_STATES: tuple[str, ...] = ("Pending", "Delayed", "InProgress", "Cancelling")
+    INTERMEDIATE_STATES: tuple[str, ...] = (
+        "Pending",
+        "Delayed",
+        "InProgress",
+        "Cancelling",
+    )
     FAILURE_STATES: tuple[str, ...] = ("Cancelled", "TimedOut", "Failed")
     SUCCESS_STATES: tuple[str, ...] = ("Success",)
     FAILURE_MESSAGE = "SSM run command sensor failed."
@@ -89,14 +101,18 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
         command_invocations = response.get("CommandInvocations", [])
 
         if not command_invocations:
-            self.log.info("No command invocations found for command_id=%s yet, waiting...", self.command_id)
+            self.log.info(
+                "No command invocations found",
+                "command_id=%s yet, waiting...",
+                self.command_id,
+            )
             return False
 
         for invocation in command_invocations:
             state = invocation["Status"]
 
             if state in self.FAILURE_STATES:
-                raise AirflowException(self.FAILURE_MESSAGE)
+                raise RuntimeError(self.FAILURE_MESSAGE)
 
             if state in self.INTERMEDIATE_STATES:
                 return False
@@ -122,6 +138,6 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[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"])

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

@@ -20,10 +20,10 @@ import json
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -23,7 +23,7 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.amazon.version_compat import BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperator
 
 try:
     from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
@@ -49,12 +49,12 @@ class AzureBlobStorageToS3Operator(BaseOperator):
     :param prefix: Prefix string which filters objects whose name begin with
         this prefix. (templated)
     :param delimiter: The delimiter by which you want to filter the objects. (templated)
-        For e.g to lists the CSV files from in a directory in GCS you would use
+        For e.g. to lists the CSV files from in a directory in GCS you would use
         delimiter='.csv'.
     :param aws_conn_id: Connection id of the S3 connection to use
     :param dest_s3_key: The base S3 key to be used to store the files. (templated)
     :param dest_verify: Whether or not to verify SSL certificates for S3 connection.
-        By default SSL certificates are verified.
+        By default, SSL certificates are verified.
         You can provide the following values:
 
         - ``False``: do not validate SSL certificates. SSL will still be used

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

@@ -22,8 +22,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.version_compat import BaseOperator
-from airflow.utils.types import NOTSET, ArgNotSet
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet, is_arg_set
+from airflow.providers.common.compat.sdk import BaseOperator
 
 
 class AwsToAwsBaseOperator(BaseOperator):
@@ -55,7 +55,7 @@ class AwsToAwsBaseOperator(BaseOperator):
         self.source_aws_conn_id = source_aws_conn_id
         self.dest_aws_conn_id = dest_aws_conn_id
         self.source_aws_conn_id = source_aws_conn_id
-        if isinstance(dest_aws_conn_id, ArgNotSet):
-            self.dest_aws_conn_id = self.source_aws_conn_id
-        else:
+        if is_arg_set(dest_aws_conn_id):
             self.dest_aws_conn_id = dest_aws_conn_id
+        else:
+            self.dest_aws_conn_id = self.source_aws_conn_id

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

@@ -36,8 +36,8 @@ from airflow.providers.amazon.aws.transfers.base import AwsToAwsBaseOperator
 from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
-    from airflow.utils.types import ArgNotSet
+    from airflow.providers.amazon.version_compat import ArgNotSet
+    from airflow.sdk import Context
 
 
 class JSONEncoder(json.JSONEncoder):
@@ -216,9 +216,9 @@ class DynamoDBToS3Operator(AwsToAwsBaseOperator):
         scan_kwargs = copy(self.dynamodb_scan_kwargs) if self.dynamodb_scan_kwargs else {}
         err = None
         f: IO[Any]
-        with NamedTemporaryFile() as f:
+        with NamedTemporaryFile() as f_tmp:
             try:
-                f = self._scan_dynamodb_and_upload_to_s3(f, scan_kwargs, table)
+                f = self._scan_dynamodb_and_upload_to_s3(f_tmp, scan_kwargs, table)
             except Exception as e:
                 err = e
                 raise e

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

@@ -24,7 +24,7 @@ from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.amazon.version_compat import BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperator
 from airflow.providers.exasol.hooks.exasol import ExasolHook
 
 if TYPE_CHECKING: