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

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

@@ -22,7 +22,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.ftp.hooks.ftp import FTPHook
 
 if TYPE_CHECKING:

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

@@ -25,9 +25,8 @@ from typing import TYPE_CHECKING
 
 from packaging.version import Version
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.amazon.version_compat import BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 
 if TYPE_CHECKING:
@@ -39,6 +38,11 @@ class GCSToS3Operator(BaseOperator):
     """
     Synchronizes a Google Cloud Storage bucket with an S3 bucket.
 
+    .. note::
+        When flatten_structure=True, it takes precedence over keep_directory_structure.
+        For example, with flatten_structure=True, "folder/subfolder/file.txt" becomes "file.txt"
+        regardless of the keep_directory_structure setting.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:GCSToS3Operator`
@@ -79,6 +83,9 @@ class GCSToS3Operator(BaseOperator):
         object to be uploaded in S3
     :param keep_directory_structure: (Optional) When set to False the path of the file
          on the bucket is recreated within path passed in dest_s3_key.
+    :param flatten_structure: (Optional) When set to True, places all files directly
+        in the dest_s3_key directory without preserving subdirectory structure.
+        Takes precedence over keep_directory_structure when enabled.
     :param match_glob: (Optional) filters objects based on the glob pattern given by the string
         (e.g, ``'**/*/.json'``)
     :param gcp_user_project: (Optional) The identifier of the Google Cloud project to bill for this request.
@@ -108,6 +115,7 @@ class GCSToS3Operator(BaseOperator):
         dest_s3_extra_args: dict | None = None,
         s3_acl_policy: str | None = None,
         keep_directory_structure: bool = True,
+        flatten_structure: bool = False,
         match_glob: str | None = None,
         gcp_user_project: str | None = None,
         **kwargs,
@@ -124,6 +132,10 @@ class GCSToS3Operator(BaseOperator):
         self.dest_s3_extra_args = dest_s3_extra_args or {}
         self.s3_acl_policy = s3_acl_policy
         self.keep_directory_structure = keep_directory_structure
+        self.flatten_structure = flatten_structure
+
+        if self.flatten_structure and self.keep_directory_structure:
+            self.log.warning("flatten_structure=True takes precedence over keep_directory_structure=True")
         try:
             from airflow.providers.google import __version__ as _GOOGLE_PROVIDER_VERSION
 
@@ -140,6 +152,17 @@ class GCSToS3Operator(BaseOperator):
         self.match_glob = match_glob
         self.gcp_user_project = gcp_user_project
 
+    def _transform_file_path(self, file_path: str) -> str:
+        """
+        Transform the GCS file path according to the specified options.
+
+        :param file_path: The original GCS file path
+        :return: The transformed file path for S3 destination
+        """
+        if self.flatten_structure:
+            return os.path.basename(file_path)
+        return file_path
+
     def execute(self, context: Context) -> list[str]:
         # list all files in an Google Cloud Storage bucket
         gcs_hook = GCSHook(
@@ -167,7 +190,7 @@ class GCSToS3Operator(BaseOperator):
             aws_conn_id=self.dest_aws_conn_id, verify=self.dest_verify, extra_args=self.dest_s3_extra_args
         )
 
-        if not self.keep_directory_structure and self.prefix:
+        if not self.keep_directory_structure and self.prefix and not self.flatten_structure:
             self.dest_s3_key = os.path.join(self.dest_s3_key, self.prefix)
 
         if not self.replace:
@@ -187,15 +210,34 @@ class GCSToS3Operator(BaseOperator):
             existing_files = existing_files or []
             # remove the prefix for the existing files to allow the match
             existing_files = [file.replace(prefix, "", 1) for file in existing_files]
-            gcs_files = list(set(gcs_files) - set(existing_files))
+
+            # Transform GCS files for comparison and filter out existing ones
+            existing_files_set = set(existing_files)
+            filtered_files = []
+            seen_transformed = set()
+
+            for file in gcs_files:
+                transformed = self._transform_file_path(file)
+                if transformed not in existing_files_set and transformed not in seen_transformed:
+                    filtered_files.append(file)
+                    seen_transformed.add(transformed)
+                elif transformed in seen_transformed:
+                    self.log.warning(
+                        "Skipping duplicate file %s (transforms to %s)",
+                        file,
+                        transformed,
+                    )
+
+            gcs_files = filtered_files
 
         if gcs_files:
             for file in gcs_files:
                 with gcs_hook.provide_file(
                     object_name=file, bucket_name=str(self.gcs_bucket), user_project=self.gcp_user_project
                 ) as local_tmp_file:
-                    dest_key = os.path.join(self.dest_s3_key, file)
-                    self.log.info("Saving file to %s", dest_key)
+                    transformed_path = self._transform_file_path(file)
+                    dest_key = os.path.join(self.dest_s3_key, transformed_path)
+                    self.log.info("Saving file from %s to %s", file, dest_key)
                     s3_hook.load_file(
                         filename=local_tmp_file.name,
                         key=dest_key,

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

@@ -22,7 +22,7 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook
-from airflow.providers.amazon.version_compat import BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 
 if TYPE_CHECKING:

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

@@ -26,14 +26,11 @@ from typing import TYPE_CHECKING
 
 from airflow.models.xcom import XCOM_RETURN_KEY
 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.google.common.hooks.discovery_api import GoogleDiscoveryApiHook
 
 if TYPE_CHECKING:
-    try:
-        from airflow.sdk.types import RuntimeTaskInstanceProtocol
-    except ImportError:
-        from airflow.models import TaskInstance as RuntimeTaskInstanceProtocol  # type: ignore[assignment]
+    from airflow.providers.common.compat.sdk import RuntimeTaskInstanceProtocol
     from airflow.utils.context import Context
 
 # MAX XCOM Size is 48KB

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

@@ -24,8 +24,8 @@ from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING, Literal
 
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
-from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook
+from airflow.providers.common.compat.sdk import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -23,7 +23,7 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 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.http.hooks.http import HttpHook
 
 if TYPE_CHECKING:

airflow/providers/amazon/aws/transfers/imap_attachment_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
 from airflow.providers.imap.hooks.imap import ImapHook
 
 if TYPE_CHECKING:

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

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

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

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, cast
 from bson import json_util
 
 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.mongo.hooks.mongo import MongoHook
 
 if TYPE_CHECKING:

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

@@ -23,13 +23,12 @@ import re
 from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
-from airflow.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 AirflowException, BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -131,12 +130,12 @@ class RedshiftToS3Operator(BaseOperator):
         # actually provide a connection note that, because we don't want to let the exception bubble up in
         # that case (since we're silently injecting a connection on their behalf).
         self._aws_conn_id: str | None
-        if isinstance(aws_conn_id, ArgNotSet):
-            self.conn_set = False
-            self._aws_conn_id = "aws_default"
-        else:
+        if is_arg_set(aws_conn_id):
             self.conn_set = True
             self._aws_conn_id = aws_conn_id
+        else:
+            self.conn_set = False
+            self._aws_conn_id = "aws_default"
 
     def _build_unload_query(
         self, credentials_block: str, select_query: str, s3_key: str, unload_options: str

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

@@ -22,9 +22,8 @@ from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
 from botocore.exceptions import ClientError, WaiterError
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
-from airflow.providers.amazon.version_compat import BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,7 +22,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.ftp.hooks.ftp import FTPHook
 
 if TYPE_CHECKING:

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

@@ -19,13 +19,12 @@ 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.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
-from airflow.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 AirflowException, BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -122,12 +121,12 @@ class S3ToRedshiftOperator(BaseOperator):
         # actually provide a connection note that, because we don't want to let the exception bubble up in
         # that case (since we're silently injecting a connection on their behalf).
         self._aws_conn_id: str | None
-        if isinstance(aws_conn_id, ArgNotSet):
-            self.conn_set = False
-            self._aws_conn_id = "aws_default"
-        else:
+        if is_arg_set(aws_conn_id):
             self.conn_set = True
             self._aws_conn_id = aws_conn_id
+        else:
+            self.conn_set = False
+            self._aws_conn_id = "aws_default"
 
         if self.redshift_data_api_kwargs:
             for arg in ["sql", "parameters"]:

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

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 
 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.ssh.hooks.ssh import SSHHook
 
 if TYPE_CHECKING:

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

@@ -21,9 +21,8 @@ from functools import cached_property
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.amazon.version_compat import BaseHook, BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook, BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,7 +22,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
 from airflow.providers.salesforce.hooks.salesforce import SalesforceHook
 
 if TYPE_CHECKING:

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

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 
 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.ssh.hooks.ssh import SSHHook
 
 if TYPE_CHECKING:

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

@@ -25,9 +25,9 @@ from collections import namedtuple
 from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Literal, cast
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.amazon.version_compat import BaseHook, BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook, BaseOperator
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -191,7 +191,7 @@ class SqlToS3Operator(BaseOperator):
             if df[col].dtype.name == "object" and file_format == FILE_FORMAT.PARQUET:
                 # if the type wasn't identified or converted, change it to a string so if can still be
                 # processed.
-                df[col] = df[col].astype(str)
+                df[col] = cast("pd.Series", df[col].astype(str))  # type: ignore[call-overload]
 
             if "float" in df[col].dtype.name and df[col].hasnans:
                 # inspect values to determine if dtype of non-null values is int or float
@@ -201,13 +201,13 @@ class SqlToS3Operator(BaseOperator):
                     # The type ignore can be removed here if https://github.com/numpy/numpy/pull/23690
                     # is merged and released as currently NumPy does not consider None as valid for x/y.
                     df[col] = np.where(df[col].isnull(), None, df[col])  # type: ignore[call-overload]
-                    df[col] = df[col].astype(pd.Int64Dtype())
+                    df[col] = cast("pd.Series", df[col].astype(pd.Int64Dtype()))  # type: ignore[call-overload]
                 elif np.isclose(notna_series, notna_series.astype(int)).all():
                     # set to float dtype that retains floats and supports NaNs
                     # The type ignore can be removed here if https://github.com/numpy/numpy/pull/23690
                     # is merged and released
                     df[col] = np.where(df[col].isnull(), None, df[col])  # type: ignore[call-overload]
-                    df[col] = df[col].astype(pd.Float64Dtype())
+                    df[col] = cast("pd.Series", df[col].astype(pd.Float64Dtype()))  # type: ignore[call-overload]
 
     @staticmethod
     def _strip_suffixes(
@@ -304,12 +304,11 @@ class SqlToS3Operator(BaseOperator):
                     group_df.reset_index(drop=True),
                 )
         elif isinstance(df, pl.DataFrame):
-            for group_label, group_df in df.group_by(**self.groupby_kwargs):  # type: ignore[assignment]
-                if random_column_name:
-                    group_df = group_df.drop(random_column_name)
+            for group_label, group_df_in in df.group_by(**self.groupby_kwargs):  # type: ignore[assignment]
+                group_df2 = group_df_in.drop(random_column_name) if random_column_name else group_df_in
                 yield (
                     cast("str", group_label[0] if isinstance(group_label, tuple) else group_label),
-                    group_df,
+                    group_df2,
                 )
 
     def _get_hook(self) -> DbApiHook:

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

@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
-from airflow.utils.types import NOTSET, ArgNotSet
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook

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

@@ -23,11 +23,11 @@ from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError, WaiterError
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:

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

@@ -20,10 +20,10 @@ from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.eks import EksHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.triggers.base import TriggerEvent
 
 if TYPE_CHECKING:

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

@@ -41,6 +41,11 @@ class S3KeyTrigger(BaseTrigger):
         Unix wildcard pattern
     :param aws_conn_id: reference to the s3 connection
     :param use_regex: whether to use regex to check bucket
+    :param metadata_keys: List of head_object attributes to gather and send to ``check_fn``.
+        Acceptable values: Any top level attribute returned by s3.head_object. Specify * to return
+        all available attributes.
+        Default value: "Size".
+        If the requested attribute is not found, the key is still included and the value is None.
     :param hook_params: params for hook its optional
     """
 
@@ -56,6 +61,7 @@ class S3KeyTrigger(BaseTrigger):
         region_name: str | None = None,
         verify: bool | str | None = None,
         botocore_config: dict | None = None,
+        metadata_keys: list[str] | None = None,
         **hook_params: Any,
     ):
         super().__init__()
@@ -70,6 +76,7 @@ class S3KeyTrigger(BaseTrigger):
         self.region_name = region_name
         self.verify = verify
         self.botocore_config = botocore_config
+        self.metadata_keys = metadata_keys if metadata_keys else ["Size", "Key"]
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize S3KeyTrigger arguments and classpath."""
@@ -87,6 +94,7 @@ class S3KeyTrigger(BaseTrigger):
                 "region_name": self.region_name,
                 "verify": self.verify,
                 "botocore_config": self.botocore_config,
+                "metadata_keys": self.metadata_keys,
             },
         )
 
@@ -108,11 +116,30 @@ class S3KeyTrigger(BaseTrigger):
                         client, self.bucket_name, self.bucket_key, self.wildcard_match, self.use_regex
                     ):
                         if self.should_check_fn:
-                            s3_objects = await self.hook.get_files_async(
+                            raw_objects = await self.hook.get_files_async(
                                 client, self.bucket_name, self.bucket_key, self.wildcard_match
                             )
+                            files = []
+                            for f in raw_objects:
+                                metadata = {}
+                                obj = await self.hook.get_head_object_async(
+                                    client=client, key=f, bucket_name=self.bucket_name
+                                )
+                                if obj is None:
+                                    return
+
+                                if "*" in self.metadata_keys:
+                                    metadata = obj
+                                else:
+                                    for mk in self.metadata_keys:
+                                        if mk == "Size":
+                                            metadata[mk] = obj.get("ContentLength")
+                                        else:
+                                            metadata[mk] = obj.get(mk, None)
+                                metadata["Key"] = f
+                                files.append(metadata)
                             await asyncio.sleep(self.poke_interval)
-                            yield TriggerEvent({"status": "running", "files": s3_objects})
+                            yield TriggerEvent({"status": "running", "files": files})
                         else:
                             yield TriggerEvent({"status": "success"})
                         return

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

@@ -26,9 +26,9 @@ from typing import Any
 
 from botocore.exceptions import WaiterError
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 

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

@@ -20,10 +20,10 @@ import asyncio
 from collections.abc import AsyncIterator, Collection
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.utils.sqs import process_response
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.common.compat.sdk import AirflowException
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.triggers.base import BaseEventTrigger, TriggerEvent

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

@@ -36,6 +36,11 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
     :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
     :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
     :param aws_conn_id: The Airflow connection used for AWS credentials.
+    :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
     """
 
     def __init__(
@@ -45,6 +50,9 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
         waiter_delay: int = 120,
         waiter_max_attempts: int = 75,
         aws_conn_id: str | None = None,
+        region_name: str | None = None,
+        verify: bool | str | None = None,
+        botocore_config: dict | None = None,
     ) -> None:
         super().__init__(
             serialized_fields={"command_id": command_id},
@@ -58,11 +66,19 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
             waiter_delay=waiter_delay,
             waiter_max_attempts=waiter_max_attempts,
             aws_conn_id=aws_conn_id,
+            region_name=region_name,
+            verify=verify,
+            botocore_config=botocore_config,
         )
         self.command_id = command_id
 
     def hook(self) -> AwsGenericHook:
-        return SsmHook(aws_conn_id=self.aws_conn_id)
+        return SsmHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         hook = self.hook()

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

@@ -23,7 +23,7 @@ from enum import Enum
 from importlib import metadata
 from typing import Any
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.helpers import prune_dict
 from airflow.version import version
 

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

@@ -26,16 +26,13 @@ from typing import TYPE_CHECKING, Any
 from botocore import UNSIGNED
 from botocore.config import Config
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.utils import trim_none_values
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
-    try:
-        from airflow.sdk import Connection
-    except ImportError:
-        from airflow.models.connection import Connection  # type: ignore[assignment]
+    from airflow.providers.common.compat.sdk import Connection
 
 
 @dataclass

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

@@ -158,4 +158,4 @@ def aws_template_fields(*template_fields: str) -> tuple[str, ...]:
             f"{', '.join(map(repr, template_fields))}."
         )
         raise TypeError(msg)
-    return tuple(sorted(list({"aws_conn_id", "region_name", "verify"} | set(template_fields))))
+    return tuple(sorted({"aws_conn_id", "region_name", "verify"} | set(template_fields)))

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

@@ -22,7 +22,7 @@ import time
 from collections.abc import Callable
 from enum import Enum
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 
 log = logging.getLogger(__name__)
 
@@ -51,8 +51,8 @@ def waiter(
     :param desired_state: Wait until the getter returns this value
     :param failure_states: A set of states which indicate failure and should throw an
         exception if any are reached before the desired_state
-    :param object_type: Used for the reporting string. What are you waiting for? (application, job, etc)
-    :param action: Used for the reporting string. What action are you waiting for? (created, deleted, etc)
+    :param object_type: Used for the reporting string. What are you waiting for? (application, job, etc.)
+    :param action: Used for the reporting string. What action are you waiting for? (created, deleted, etc.)
     :param countdown: Number of seconds the waiter should wait for the desired state before timing out.
         Defaults to 25 * 60 seconds. None = infinite.
     :param check_interval_seconds: Number of seconds waiter should wait before attempting

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

@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Any
 import jmespath
 from botocore.exceptions import NoCredentialsError, WaiterError
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from botocore.waiter import Waiter