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

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]):
@@ -122,8 +122,19 @@ class S3KeySensor(AwsBaseSensor[S3Hook]):
         """
         if self.wildcard_match:
             prefix = re.split(r"[\[*?]", key, 1)[0]
-            keys = self.hook.get_file_metadata(prefix, bucket_name)
-            key_matches = [k for k in keys if fnmatch.fnmatch(k["Key"], key)]
+
+            key_matches: list[str] = []
+
+            # Is check_fn is None, then we can return True without having to iterate through each value in
+            # yielded by iter_file_metadata. Otherwise, we'll check for a match, and add all matches to the
+            # key_matches list
+            for k in self.hook.iter_file_metadata(prefix, bucket_name):
+                if fnmatch.fnmatch(k["Key"], key):
+                    if self.check_fn is None:
+                        # This will only wait for a single match, and will immediately return
+                        return True
+                    key_matches.append(k)
+
             if not key_matches:
                 return False
 
@@ -132,21 +143,23 @@ class S3KeySensor(AwsBaseSensor[S3Hook]):
             for f in key_matches:
                 metadata = {}
                 if "*" in self.metadata_keys:
-                    metadata = self.hook.head_object(f["Key"], bucket_name)
+                    metadata = self.hook.head_object(f["Key"], bucket_name)  # type: ignore[index]
                 else:
-                    for key in self.metadata_keys:
+                    for mk in self.metadata_keys:
                         try:
-                            metadata[key] = f[key]
+                            metadata[mk] = f[mk]  # type: ignore[index]
                         except KeyError:
                             # supplied key might be from head_object response
-                            self.log.info("Key %s not found in response, performing head_object", key)
-                            metadata[key] = self.hook.head_object(f["Key"], bucket_name).get(key, None)
+                            self.log.info("Key %s not found in response, performing head_object", mk)
+                            metadata[mk] = self.hook.head_object(f["Key"], bucket_name).get(mk, None)  # type: ignore[index]
                 files.append(metadata)
+
         elif self.use_regex:
-            keys = self.hook.get_file_metadata("", bucket_name)
-            key_matches = [k for k in keys if re.match(pattern=key, string=k["Key"])]
-            if not key_matches:
-                return False
+            for k in self.hook.iter_file_metadata("", bucket_name):
+                if re.match(pattern=key, string=k["Key"]):
+                    return True
+            return False
+
         else:
             obj = self.hook.head_object(key, bucket_name)
             if obj is None:
@@ -202,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",
         )
@@ -213,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

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:

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

@@ -27,7 +27,7 @@ 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 BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 
 if TYPE_CHECKING:
@@ -39,6 +39,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 +84,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 +116,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 +133,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 +153,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 +191,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 +211,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

@@ -28,8 +28,8 @@ 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 BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -131,12 +131,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

@@ -24,7 +24,7 @@ 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 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

@@ -24,8 +24,8 @@ 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 BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -122,12 +122,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

@@ -23,7 +23,7 @@ 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 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

@@ -27,7 +27,7 @@ from typing import TYPE_CHECKING, Any, Literal, cast
 
 from airflow.exceptions import AirflowException, 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 BaseHook, BaseOperator
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -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