apache-airflow-providers-amazon 9.0.0rc1__tar.gz → 9.1.0rc1__tar.gz

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-amazon
-Version: 9.0.0rc1
+Version: 9.1.0rc1
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: PyAthena>=3.0.10
-Requires-Dist: apache-airflow-providers-common-compat>=1.1.0rc0
+Requires-Dist: apache-airflow-providers-common-compat>=1.2.1rc0
 Requires-Dist: apache-airflow-providers-common-sql>=1.3.1rc0
 Requires-Dist: apache-airflow-providers-http
 Requires-Dist: apache-airflow>=2.8.0rc0
@@ -33,7 +33,6 @@ Requires-Dist: jmespath>=0.7.0
 Requires-Dist: jsonpath_ng>=1.5.3
 Requires-Dist: python3-saml>=1.16.0
 Requires-Dist: redshift_connector>=2.0.918
-Requires-Dist: sqlalchemy_redshift>=0.8.6
 Requires-Dist: watchtower>=3.0.0,!=3.3.0,<4
 Requires-Dist: aiobotocore[boto3]>=2.13.0 ; extra == "aiobotocore"
 Requires-Dist: apache-airflow-providers-apache-hive ; extra == "apache.hive"
@@ -55,8 +54,8 @@ Requires-Dist: s3fs>=2023.10.0 ; extra == "s3fs"
 Requires-Dist: apache-airflow-providers-salesforce ; extra == "salesforce"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -124,7 +123,7 @@ Provides-Extra: ssh
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.0.0.rc1``
+Release: ``9.1.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -137,7 +136,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/>`_.
 
 Installation
 ------------
@@ -155,7 +154,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ======================
 ``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.1.0``
+``apache-airflow-providers-common-compat``  ``>=1.2.1``
 ``apache-airflow-providers-common-sql``     ``>=1.3.1``
 ``apache-airflow-providers-http``
 ``boto3``                                   ``>=1.34.90``
@@ -164,7 +163,6 @@ PIP package                                 Version required
 ``watchtower``                              ``>=3.0.0,!=3.3.0,<4``
 ``jsonpath_ng``                             ``>=1.5.3``
 ``redshift_connector``                      ``>=2.0.918``
-``sqlalchemy_redshift``                     ``>=0.8.6``
 ``asgiref``                                 ``>=2.3.0``
 ``PyAthena``                                ``>=3.0.10``
 ``jmespath``                                ``>=0.7.0``
@@ -204,4 +202,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/changelog.html>`_.

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-amazon``
 
-Release: ``9.0.0.rc1``
+Release: ``9.1.0.rc1``
 
 
 Amazon integration (including `Amazon Web Services (AWS) <https://aws.amazon.com/>`__).
@@ -55,7 +55,7 @@ This is a provider package for ``amazon`` provider. All classes for this provide
 are in ``airflow.providers.amazon`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/>`_.
 
 Installation
 ------------
@@ -73,7 +73,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ======================
 ``apache-airflow``                          ``>=2.8.0``
-``apache-airflow-providers-common-compat``  ``>=1.1.0``
+``apache-airflow-providers-common-compat``  ``>=1.2.1``
 ``apache-airflow-providers-common-sql``     ``>=1.3.1``
 ``apache-airflow-providers-http``
 ``boto3``                                   ``>=1.34.90``
@@ -82,7 +82,6 @@ PIP package                                 Version required
 ``watchtower``                              ``>=3.0.0,!=3.3.0,<4``
 ``jsonpath_ng``                             ``>=1.5.3``
 ``redshift_connector``                      ``>=2.0.918``
-``sqlalchemy_redshift``                     ``>=0.8.6``
 ``asgiref``                                 ``>=2.3.0``
 ``PyAthena``                                ``>=3.0.10``
 ``jmespath``                                ``>=0.7.0``
@@ -122,4 +121,4 @@ Dependent package
 ======================================================================================================================  ===================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-amazon/9.1.0/changelog.html>`_.

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.0.0"
+__version__ = "9.1.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.8.0"

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/assets/s3.py

@@ -19,16 +19,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-
-try:
-    from airflow.assets import Asset
-except ModuleNotFoundError:
-    from airflow.datasets import Dataset as Asset  # type: ignore[no-redef]
+from airflow.providers.common.compat.assets import Asset
 
 if TYPE_CHECKING:
     from urllib.parse import SplitResult
 
-    from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
+    from airflow.providers.common.compat.openlineage.facet import (
+        Dataset as OpenLineageDataset,
+    )
 
 
 def create_asset(*, bucket: str, key: str, extra=None) -> Asset:
@@ -43,7 +41,9 @@ def sanitize_uri(uri: SplitResult) -> SplitResult:
 
 def convert_asset_to_openlineage(asset: Asset, lineage_context) -> OpenLineageDataset:
     """Translate Asset with valid AIP-60 uri to OpenLineage with assistance from the hook."""
-    from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
+    from airflow.providers.common.compat.openlineage.facet import (
+        Dataset as OpenLineageDataset,
+    )
 
     bucket, key = S3Hook.parse_s3_url(asset.uri)
     return OpenLineageDataset(namespace=f"s3://{bucket}", name=key if key else "/")

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/hooks/appflow.py

@@ -117,9 +117,9 @@ class AppflowHook(AwsGenericHook["AppflowClient"]):
 
         self.conn.update_flow(
             flowName=response["flowName"],
-            destinationFlowConfigList=response["destinationFlowConfigList"],
-            sourceFlowConfig=response["sourceFlowConfig"],
-            triggerConfig=response["triggerConfig"],
+            destinationFlowConfigList=response["destinationFlowConfigList"],  # type: ignore[arg-type]
+            sourceFlowConfig=response["sourceFlowConfig"],  # type: ignore[arg-type]
+            triggerConfig=response["triggerConfig"],  # type: ignore[arg-type]
             description=response.get("description", "Flow description."),
-            tasks=tasks,
+            tasks=tasks,  # type: ignore[arg-type]
         )

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -163,7 +163,7 @@ class RedshiftSQLHook(DbApiHook):
         # Compatibility: The 'create' factory method was added in SQLAlchemy 1.4
         # to replace calling the default URL constructor directly.
         create_url = getattr(URL, "create", URL)
-        return str(create_url(drivername="redshift+redshift_connector", **conn_params))
+        return str(create_url(drivername="postgresql", **conn_params))
 
     def get_sqlalchemy_engine(self, engine_kwargs=None):
         """Overridden to pass Redshift-specific arguments."""

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/hooks/s3.py

@@ -41,14 +41,16 @@ from urllib.parse import urlsplit
 from uuid import uuid4
 
 if TYPE_CHECKING:
-    from mypy_boto3_s3.service_resource import Bucket as S3Bucket, Object as S3ResourceObject
+    from mypy_boto3_s3.service_resource import (
+        Bucket as S3Bucket,
+        Object as S3ResourceObject,
+    )
 
     from airflow.utils.types import ArgNotSet
 
     with suppress(ImportError):
         from aiobotocore.client import AioBaseClient
 
-from importlib.util import find_spec
 
 from asgiref.sync import sync_to_async
 from boto3.s3.transfer import S3Transfer, TransferConfig
@@ -58,15 +60,9 @@ from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.tags import format_tags
+from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
 from airflow.utils.helpers import chunks
 
-if find_spec("airflow.assets"):
-    from airflow.lineage.hook import get_hook_lineage_collector
-else:
-    # TODO: import from common.compat directly after common.compat providers with
-    # asset_compat_lineage_collector released
-    from airflow.providers.amazon.aws.utils.asset_compat_lineage_collector import get_hook_lineage_collector
-
 logger = logging.getLogger(__name__)
 
 
@@ -148,9 +144,10 @@ def unify_bucket_name_and_key(func: Callable) -> Callable:
 
         if "bucket_name" not in bound_args.arguments:
             with suppress(S3HookUriParseFailure):
-                bound_args.arguments["bucket_name"], bound_args.arguments[key_name] = S3Hook.parse_s3_url(
-                    bound_args.arguments[key_name]
-                )
+                (
+                    bound_args.arguments["bucket_name"],
+                    bound_args.arguments[key_name],
+                ) = S3Hook.parse_s3_url(bound_args.arguments[key_name])
 
         return func(*bound_args.args, **bound_args.kwargs)
 
@@ -318,7 +315,8 @@ class S3Hook(AwsBaseHook):
                 self.log.info('Bucket "%s" does not exist', bucket_name)
             elif return_code == 403:
                 self.log.error(
-                    'Access to bucket "%s" is forbidden or there was an error with the request', bucket_name
+                    'Access to bucket "%s" is forbidden or there was an error with the request',
+                    bucket_name,
                 )
                 self.log.error(e)
             return False
@@ -359,7 +357,8 @@ class S3Hook(AwsBaseHook):
             self.get_conn().create_bucket(Bucket=bucket_name)
         else:
             self.get_conn().create_bucket(
-                Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region_name}
+                Bucket=bucket_name,
+                CreateBucketConfiguration={"LocationConstraint": region_name},
             )
 
     @provide_bucket_name
@@ -410,7 +409,10 @@ class S3Hook(AwsBaseHook):
 
         paginator = self.get_conn().get_paginator("list_objects_v2")
         response = paginator.paginate(
-            Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+            Bucket=bucket_name,
+            Prefix=prefix,
+            Delimiter=delimiter,
+            PaginationConfig=config,
         )
 
         prefixes: list[str] = []
@@ -471,7 +473,10 @@ class S3Hook(AwsBaseHook):
 
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(
-            Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+            Bucket=bucket_name,
+            Prefix=prefix,
+            Delimiter=delimiter,
+            PaginationConfig=config,
         )
 
         prefixes = []
@@ -569,7 +574,11 @@ class S3Hook(AwsBaseHook):
         return await self._check_key_async(client, bucket, wildcard_match, bucket_keys, use_regex)
 
     async def check_for_prefix_async(
-        self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None
+        self,
+        client: AioBaseClient,
+        prefix: str,
+        delimiter: str,
+        bucket_name: str | None = None,
     ) -> bool:
         """
         Check that a prefix exists in a bucket.
@@ -587,7 +596,11 @@ class S3Hook(AwsBaseHook):
         return prefix in plist
 
     async def _check_for_prefix_async(
-        self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None
+        self,
+        client: AioBaseClient,
+        prefix: str,
+        delimiter: str,
+        bucket_name: str | None = None,
     ) -> bool:
         return await self.check_for_prefix_async(
             client, prefix=prefix, delimiter=delimiter, bucket_name=bucket_name
@@ -643,7 +656,10 @@ class S3Hook(AwsBaseHook):
 
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(
-            Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter, PaginationConfig=config
+            Bucket=bucket_name,
+            Prefix=prefix,
+            Delimiter=delimiter,
+            PaginationConfig=config,
         )
 
         keys = []
@@ -655,7 +671,10 @@ class S3Hook(AwsBaseHook):
         return keys
 
     def _list_key_object_filter(
-        self, keys: list, from_datetime: datetime | None = None, to_datetime: datetime | None = None
+        self,
+        keys: list,
+        from_datetime: datetime | None = None,
+        to_datetime: datetime | None = None,
     ) -> list:
         def _is_in_period(input_date: datetime) -> bool:
             if from_datetime is not None and input_date <= from_datetime:
@@ -766,7 +785,10 @@ class S3Hook(AwsBaseHook):
                     "message": success_message,
                 }
 
-            self.log.error("FAILURE: Inactivity Period passed, not enough objects found in %s", path)
+            self.log.error(
+                "FAILURE: Inactivity Period passed, not enough objects found in %s",
+                path,
+            )
             return {
                 "status": "error",
                 "message": f"FAILURE: Inactivity Period passed, not enough objects found in {path}",
@@ -1109,7 +1131,13 @@ class S3Hook(AwsBaseHook):
             extra_args["ACL"] = acl_policy
 
         client = self.get_conn()
-        client.upload_file(filename, bucket_name, key, ExtraArgs=extra_args, Config=self.transfer_config)
+        client.upload_file(
+            filename,
+            bucket_name,
+            key,
+            ExtraArgs=extra_args,
+            Config=self.transfer_config,
+        )
         get_hook_lineage_collector().add_input_asset(
             context=self, scheme="file", asset_kwargs={"path": filename}
         )
@@ -1308,18 +1336,32 @@ class S3Hook(AwsBaseHook):
         )
 
         source_bucket_name, source_bucket_key = self.get_s3_bucket_key(
-            source_bucket_name, source_bucket_key, "source_bucket_name", "source_bucket_key"
+            source_bucket_name,
+            source_bucket_key,
+            "source_bucket_name",
+            "source_bucket_key",
         )
 
-        copy_source = {"Bucket": source_bucket_name, "Key": source_bucket_key, "VersionId": source_version_id}
+        copy_source = {
+            "Bucket": source_bucket_name,
+            "Key": source_bucket_key,
+            "VersionId": source_version_id,
+        }
         response = self.get_conn().copy_object(
-            Bucket=dest_bucket_name, Key=dest_bucket_key, CopySource=copy_source, **kwargs
+            Bucket=dest_bucket_name,
+            Key=dest_bucket_key,
+            CopySource=copy_source,
+            **kwargs,
         )
         get_hook_lineage_collector().add_input_asset(
-            context=self, scheme="s3", asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key}
+            context=self,
+            scheme="s3",
+            asset_kwargs={"bucket": source_bucket_name, "key": source_bucket_key},
         )
         get_hook_lineage_collector().add_output_asset(
-            context=self, scheme="s3", asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key}
+            context=self,
+            scheme="s3",
+            asset_kwargs={"bucket": dest_bucket_name, "key": dest_bucket_key},
         )
         return response
 
@@ -1435,7 +1477,10 @@ class S3Hook(AwsBaseHook):
             file_path = Path(local_dir, subdir, filename_in_s3)
 
             if file_path.is_file():
-                self.log.error("file '%s' already exists. Failing the task and not overwriting it", file_path)
+                self.log.error(
+                    "file '%s' already exists. Failing the task and not overwriting it",
+                    file_path,
+                )
                 raise FileExistsError
 
             file_path.parent.mkdir(exist_ok=True, parents=True)
@@ -1484,7 +1529,10 @@ class S3Hook(AwsBaseHook):
         s3_client = self.get_conn()
         try:
             return s3_client.generate_presigned_url(
-                ClientMethod=client_method, Params=params, ExpiresIn=expires_in, HttpMethod=http_method
+                ClientMethod=client_method,
+                Params=params,
+                ExpiresIn=expires_in,
+                HttpMethod=http_method,
             )
 
         except ClientError as e:

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/hooks/sagemaker.py

@@ -153,7 +153,9 @@ class SageMakerHook(AwsBaseHook):
     non_terminal_states = {"InProgress", "Stopping"}
     endpoint_non_terminal_states = {"Creating", "Updating", "SystemUpdating", "RollingBack", "Deleting"}
     pipeline_non_terminal_states = {"Executing", "Stopping"}
+    processing_job_non_terminal_states = {"InProgress", "Stopping"}
     failed_states = {"Failed"}
+    processing_job_failed_states = {*failed_states, "Stopped"}
     training_failed_states = {*failed_states, "Stopped"}
 
     def __init__(self, *args, **kwargs):

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/operators/redshift_data.py

@@ -28,7 +28,10 @@ from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
-    from mypy_boto3_redshift_data.type_defs import GetStatementResultResponseTypeDef
+    from mypy_boto3_redshift_data.type_defs import (
+        DescribeStatementResponseTypeDef,
+        GetStatementResultResponseTypeDef,
+    )
 
     from airflow.utils.context import Context
 
@@ -37,7 +40,7 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
     """
     Executes SQL Statements against an Amazon Redshift cluster using Redshift Data.
 
-    .. seealso::
+    ... see also::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:RedshiftDataOperator`
 
@@ -84,7 +87,6 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
     )
     template_ext = (".sql",)
     template_fields_renderers = {"sql": "sql"}
-    statement_id: str | None
 
     def __init__(
         self,
@@ -124,12 +126,11 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
                 poll_interval,
             )
         self.return_sql_result = return_sql_result
-        self.statement_id: str | None = None
         self.deferrable = deferrable
         self.session_id = session_id
         self.session_keep_alive_seconds = session_keep_alive_seconds
 
-    def execute(self, context: Context) -> GetStatementResultResponseTypeDef | str:
+    def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | list[str]:
         """Execute a statement against Amazon Redshift."""
         self.log.info("Executing statement: %s", self.sql)
 
@@ -154,13 +155,14 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
             session_keep_alive_seconds=self.session_keep_alive_seconds,
         )
 
-        self.statement_id = query_execution_output.statement_id
+        # Pull the statement ID, session ID
+        self.statement_id: str = query_execution_output.statement_id
 
         if query_execution_output.session_id:
             self.xcom_push(context, key="session_id", value=query_execution_output.session_id)
 
         if self.deferrable and self.wait_for_completion:
-            is_finished = self.hook.check_query_is_finished(self.statement_id)
+            is_finished: bool = self.hook.check_query_is_finished(self.statement_id)
             if not is_finished:
                 self.defer(
                     timeout=self.execution_timeout,
@@ -176,16 +178,13 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
                     method_name="execute_complete",
                 )
 
-        if self.return_sql_result:
-            result = self.hook.conn.get_statement_result(Id=self.statement_id)
-            self.log.debug("Statement result: %s", result)
-            return result
-        else:
-            return self.statement_id
+        # Use the get_sql_results method to return the results of the SQL query, or the statement_ids,
+        # depending on the value of self.return_sql_result
+        return self.get_sql_results(statement_id=self.statement_id, return_sql_result=self.return_sql_result)
 
     def execute_complete(
         self, context: Context, event: dict[str, Any] | None = None
-    ) -> GetStatementResultResponseTypeDef | str:
+    ) -> list[GetStatementResultResponseTypeDef] | list[str]:
         event = validate_execute_complete_event(event)
 
         if event["status"] == "error":
@@ -197,16 +196,40 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
             raise AirflowException("statement_id should not be empty.")
 
         self.log.info("%s completed successfully.", self.task_id)
-        if self.return_sql_result:
-            result = self.hook.conn.get_statement_result(Id=statement_id)
-            self.log.debug("Statement result: %s", result)
-            return result
 
-        return statement_id
+        # Use the get_sql_results method to return the results of the SQL query, or the statement_ids,
+        # depending on the value of self.return_sql_result
+        return self.get_sql_results(statement_id=statement_id, return_sql_result=self.return_sql_result)
+
+    def get_sql_results(
+        self, statement_id: str, return_sql_result: bool
+    ) -> list[GetStatementResultResponseTypeDef] | list[str]:
+        """
+        Retrieve either the result of the SQL query, or the statement ID(s).
+
+        :param statement_id: Statement ID of the running queries
+        :param return_sql_result: Boolean, true if results should be returned
+        """
+        # ISSUE-40427: Pull the statement, and check to see if there are sub-statements. If that is the
+        # case, pull each of the sub-statement ID's, and grab the results. Otherwise, just use statement_id
+        statement: DescribeStatementResponseTypeDef = self.hook.conn.describe_statement(Id=statement_id)
+        statement_ids: list[str] = (
+            [sub_statement["Id"] for sub_statement in statement["SubStatements"]]
+            if len(statement.get("SubStatements", [])) > 0
+            else [statement_id]
+        )
+
+        # If returning the SQL result, use get_statement_result to return the records for each query
+        if return_sql_result:
+            results: list = [self.hook.conn.get_statement_result(Id=sid) for sid in statement_ids]
+            self.log.debug("Statement result(s): %s", results)
+            return results
+        else:
+            return statement_ids
 
     def on_kill(self) -> None:
         """Cancel the submitted redshift query."""
-        if self.statement_id:
+        if hasattr(self, "statement_id"):
             self.log.info("Received a kill signal.")
             self.log.info("Stopping Query with statementId - %s", self.statement_id)
 

{apache_airflow_providers_amazon-9.0.0rc1 → apache_airflow_providers_amazon-9.1.0rc1}/airflow/providers/amazon/aws/sensors/sagemaker.py

@@ -330,3 +330,35 @@ class SageMakerAutoMLSensor(SageMakerBaseSensor):
 
     def state_from_response(self, response: dict) -> str:
         return response["AutoMLJobStatus"]
+
+
+class SageMakerProcessingSensor(SageMakerBaseSensor):
+    """
+    Poll the processing job until it reaches a terminal state; raise AirflowException with the failure reason.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:SageMakerProcessingSensor`
+
+    :param job_name: Name of the processing job to watch.
+    """
+
+    template_fields: Sequence[str] = ("job_name",)
+    template_ext: Sequence[str] = ()
+
+    def __init__(self, *, job_name: str, **kwargs):
+        super().__init__(**kwargs)
+        self.job_name = job_name
+
+    def non_terminal_states(self) -> set[str]:
+        return SageMakerHook.processing_job_non_terminal_states
+
+    def failed_states(self) -> set[str]:
+        return SageMakerHook.processing_job_failed_states
+
+    def get_sagemaker_response(self) -> dict:
+        self.log.info("Poking Sagemaker ProcessingJob %s", self.job_name)
+        return self.hook.describe_processing_job(self.job_name)
+
+    def state_from_response(self, response: dict) -> str:
+        return response["ProcessingJobStatus"]