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

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,8 +757,8 @@ 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.wait_policy:
-            waiter_name = WAITER_POLICY_NAME_MAPPING[self.wait_policy]
+        if self.wait_for_completion:
+            waiter_name = WAITER_POLICY_NAME_MAPPING[WaitPolicy.WAIT_FOR_COMPLETION]
 
             if self.deferrable:
                 self.defer(

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"]
 

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

@@ -36,7 +36,7 @@ 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 poke_mode_only
 
 
 class S3KeySensor(AwsBaseSensor[S3Hook]):
@@ -215,6 +215,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 +227,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_unified_studio.py

@@ -25,7 +25,7 @@ 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 BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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/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