apache-airflow-providers-amazon 8.3.1rc1__py3-none-any.whl → 8.4.0__py3-none-any.whl

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

@@ -20,15 +20,20 @@ from __future__ import annotations
 import json
 import warnings
 from datetime import timedelta
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
+from astroid.decorators import cachedproperty
 from mypy_boto3_rds.type_defs import TagTypeDef
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
-from airflow.providers.amazon.aws.triggers.rds import RdsDbInstanceTrigger
+from airflow.providers.amazon.aws.triggers.rds import (
+    RdsDbAvailableTrigger,
+    RdsDbDeletedTrigger,
+    RdsDbStoppedTrigger,
+)
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
@@ -62,13 +67,17 @@ class RdsBaseOperator(BaseOperator):
                 AirflowProviderDeprecationWarning,
                 stacklevel=3,  # 2 is in the operator's init, 3 is in the user code creating the operator
             )
-        hook_params = hook_params or {}
-        self.region_name = region_name or hook_params.pop("region_name", None)
-        self.hook = RdsHook(aws_conn_id=aws_conn_id, region_name=self.region_name, **(hook_params))
+        self.hook_params = hook_params or {}
+        self.aws_conn_id = aws_conn_id
+        self.region_name = region_name or self.hook_params.pop("region_name", None)
         super().__init__(*args, **kwargs)
 
         self._await_interval = 60  # seconds
 
+    @cachedproperty
+    def hook(self) -> RdsHook:
+        return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name, **self.hook_params)
+
     def execute(self, context: Context) -> str:
         """Different implementations for snapshots, tasks and events."""
         raise NotImplementedError
@@ -106,10 +115,9 @@ class RdsCreateDbSnapshotOperator(RdsBaseOperator):
         db_snapshot_identifier: str,
         tags: Sequence[TagTypeDef] | dict | None = None,
         wait_for_completion: bool = True,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_type = RdsDbType(db_type)
         self.db_identifier = db_identifier
         self.db_snapshot_identifier = db_snapshot_identifier
@@ -194,10 +202,9 @@ class RdsCopyDbSnapshotOperator(RdsBaseOperator):
         target_custom_availability_zone: str = "",
         source_region: str = "",
         wait_for_completion: bool = True,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.db_type = RdsDbType(db_type)
         self.source_db_snapshot_identifier = source_db_snapshot_identifier
@@ -274,10 +281,9 @@ class RdsDeleteDbSnapshotOperator(RdsBaseOperator):
         db_type: str,
         db_snapshot_identifier: str,
         wait_for_completion: bool = True,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.db_type = RdsDbType(db_type)
         self.db_snapshot_identifier = db_snapshot_identifier
@@ -345,10 +351,9 @@ class RdsStartExportTaskOperator(RdsBaseOperator):
         s3_prefix: str = "",
         export_only: list[str] | None = None,
         wait_for_completion: bool = True,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.export_task_identifier = export_task_identifier
         self.source_arn = source_arn
@@ -397,10 +402,9 @@ class RdsCancelExportTaskOperator(RdsBaseOperator):
         export_task_identifier: str,
         wait_for_completion: bool = True,
         check_interval: int = 30,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.export_task_identifier = export_task_identifier
         self.wait_for_completion = wait_for_completion
@@ -461,10 +465,9 @@ class RdsCreateEventSubscriptionOperator(RdsBaseOperator):
         enabled: bool = True,
         tags: Sequence[TagTypeDef] | dict | None = None,
         wait_for_completion: bool = True,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.subscription_name = subscription_name
         self.sns_topic_arn = sns_topic_arn
@@ -511,10 +514,9 @@ class RdsDeleteEventSubscriptionOperator(RdsBaseOperator):
         self,
         *,
         subscription_name: str,
-        aws_conn_id: str = "aws_default",
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.subscription_name = subscription_name
 
@@ -545,7 +547,6 @@ class RdsCreateDbInstanceOperator(RdsBaseOperator):
     :param engine: The name of the database engine to be used for this instance
     :param rds_kwargs: Named arguments to pass to boto3 RDS client function ``create_db_instance``
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rds.html#RDS.Client.create_db_instance
-    :param aws_conn_id: The Airflow connection used for AWS credentials.
     :param wait_for_completion:  If True, waits for creation of the DB instance to complete. (default: True)
     :param waiter_delay: Time (in seconds) to wait between two consecutive calls to check DB instance state
     :param waiter_max_attempts: The maximum number of attempts to check DB instance state
@@ -563,14 +564,13 @@ class RdsCreateDbInstanceOperator(RdsBaseOperator):
         db_instance_class: str,
         engine: str,
         rds_kwargs: dict | None = None,
-        aws_conn_id: str = "aws_default",
         wait_for_completion: bool = True,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         waiter_delay: int = 30,
         waiter_max_attempts: int = 60,
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
 
         self.db_instance_identifier = db_instance_identifier
         self.db_instance_class = db_instance_class
@@ -580,7 +580,6 @@ class RdsCreateDbInstanceOperator(RdsBaseOperator):
         self.deferrable = deferrable
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context) -> str:
         self.log.info("Creating new DB instance %s", self.db_instance_identifier)
@@ -593,15 +592,15 @@ class RdsCreateDbInstanceOperator(RdsBaseOperator):
         )
         if self.deferrable:
             self.defer(
-                trigger=RdsDbInstanceTrigger(
-                    db_instance_identifier=self.db_instance_identifier,
+                trigger=RdsDbAvailableTrigger(
+                    db_identifier=self.db_instance_identifier,
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
-                    waiter_name="db_instance_available",
                     # ignoring type because create_db_instance is a dict
                     response=create_db_instance,  # type: ignore[arg-type]
+                    db_type=RdsDbType.INSTANCE,
                 ),
                 method_name="execute_complete",
                 timeout=timedelta(seconds=self.waiter_delay * self.waiter_max_attempts),
@@ -638,7 +637,6 @@ class RdsDeleteDbInstanceOperator(RdsBaseOperator):
     :param db_instance_identifier: The DB instance identifier for the DB instance to be deleted
     :param rds_kwargs: Named arguments to pass to boto3 RDS client function ``delete_db_instance``
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rds.html#RDS.Client.delete_db_instance
-    :param aws_conn_id: The Airflow connection used for AWS credentials.
     :param wait_for_completion:  If True, waits for deletion of the DB instance to complete. (default: True)
     :param waiter_delay: Time (in seconds) to wait between two consecutive calls to check DB instance state
     :param waiter_max_attempts: The maximum number of attempts to check DB instance state
@@ -654,21 +652,19 @@ class RdsDeleteDbInstanceOperator(RdsBaseOperator):
         *,
         db_instance_identifier: str,
         rds_kwargs: dict | None = None,
-        aws_conn_id: str = "aws_default",
         wait_for_completion: bool = True,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         waiter_delay: int = 30,
         waiter_max_attempts: int = 60,
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_instance_identifier = db_instance_identifier
         self.rds_kwargs = rds_kwargs or {}
         self.wait_for_completion = False if deferrable else wait_for_completion
         self.deferrable = deferrable
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context) -> str:
         self.log.info("Deleting DB instance %s", self.db_instance_identifier)
@@ -679,15 +675,15 @@ class RdsDeleteDbInstanceOperator(RdsBaseOperator):
         )
         if self.deferrable:
             self.defer(
-                trigger=RdsDbInstanceTrigger(
-                    db_instance_identifier=self.db_instance_identifier,
+                trigger=RdsDbDeletedTrigger(
+                    db_identifier=self.db_instance_identifier,
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
-                    waiter_name="db_instance_deleted",
                     # ignoring type because delete_db_instance is a dict
                     response=delete_db_instance,  # type: ignore[arg-type]
+                    db_type=RdsDbType.INSTANCE,
                 ),
                 method_name="execute_complete",
                 timeout=timedelta(seconds=self.waiter_delay * self.waiter_max_attempts),
@@ -723,8 +719,11 @@ class RdsStartDbOperator(RdsBaseOperator):
 
     :param db_identifier: The AWS identifier of the DB to start
     :param db_type: Type of the DB - either "instance" or "cluster" (default: "instance")
-    :param aws_conn_id: The Airflow connection used for AWS credentials. (default: "aws_default")
     :param wait_for_completion:  If True, waits for DB to start. (default: True)
+    :param waiter_delay: Time (in seconds) to wait between two consecutive calls to check DB instance state
+    :param waiter_max_attempts: The maximum number of attempts to check DB instance state
+    :param deferrable: If True, the operator will wait asynchronously for the DB instance to be created.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
     """
 
     template_fields = ("db_identifier", "db_type")
@@ -734,26 +733,52 @@ class RdsStartDbOperator(RdsBaseOperator):
         *,
         db_identifier: str,
         db_type: RdsDbType | str = RdsDbType.INSTANCE,
-        aws_conn_id: str = "aws_default",
         wait_for_completion: bool = True,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 40,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_identifier = db_identifier
         self.db_type = db_type
         self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
 
     def execute(self, context: Context) -> str:
         self.db_type = RdsDbType(self.db_type)
-        start_db_response = self._start_db()
-        if self.wait_for_completion:
+        start_db_response: dict[str, Any] = self._start_db()
+        if self.deferrable:
+            self.defer(
+                trigger=RdsDbAvailableTrigger(
+                    db_identifier=self.db_identifier,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                    region_name=self.region_name,
+                    response=start_db_response,
+                    db_type=RdsDbType.INSTANCE,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
             self._wait_until_db_available()
         return json.dumps(start_db_response, default=str)
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        if event is None or event["status"] != "success":
+            raise AirflowException(f"Failed to start DB: {event}")
+        else:
+            return json.dumps(event["response"], default=str)
+
     def _start_db(self):
         self.log.info("Starting DB %s '%s'", self.db_type.value, self.db_identifier)
         if self.db_type == RdsDbType.INSTANCE:
-            response = self.hook.conn.start_db_instance(DBInstanceIdentifier=self.db_identifier)
+            response = self.hook.conn.start_db_instance(
+                DBInstanceIdentifier=self.db_identifier,
+            )
         else:
             response = self.hook.conn.start_db_cluster(DBClusterIdentifier=self.db_identifier)
         return response
@@ -761,9 +786,19 @@ class RdsStartDbOperator(RdsBaseOperator):
     def _wait_until_db_available(self):
         self.log.info("Waiting for DB %s to reach 'available' state", self.db_type.value)
         if self.db_type == RdsDbType.INSTANCE:
-            self.hook.wait_for_db_instance_state(self.db_identifier, target_state="available")
+            self.hook.wait_for_db_instance_state(
+                self.db_identifier,
+                target_state="available",
+                check_interval=self.waiter_delay,
+                max_attempts=self.waiter_max_attempts,
+            )
         else:
-            self.hook.wait_for_db_cluster_state(self.db_identifier, target_state="available")
+            self.hook.wait_for_db_cluster_state(
+                self.db_identifier,
+                target_state="available",
+                check_interval=self.waiter_delay,
+                max_attempts=self.waiter_max_attempts,
+            )
 
 
 class RdsStopDbOperator(RdsBaseOperator):
@@ -779,8 +814,11 @@ class RdsStopDbOperator(RdsBaseOperator):
     :param db_snapshot_identifier: The instance identifier of the DB Snapshot to create before
         stopping the DB instance. The default value (None) skips snapshot creation. This
         parameter is ignored when ``db_type`` is "cluster"
-    :param aws_conn_id: The Airflow connection used for AWS credentials. (default: "aws_default")
     :param wait_for_completion:  If True, waits for DB to stop. (default: True)
+    :param waiter_delay: Time (in seconds) to wait between two consecutive calls to check DB instance state
+    :param waiter_max_attempts: The maximum number of attempts to check DB instance state
+    :param deferrable: If True, the operator will wait asynchronously for the DB instance to be created.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
     """
 
     template_fields = ("db_identifier", "db_snapshot_identifier", "db_type")
@@ -791,23 +829,47 @@ class RdsStopDbOperator(RdsBaseOperator):
         db_identifier: str,
         db_type: RdsDbType | str = RdsDbType.INSTANCE,
         db_snapshot_identifier: str | None = None,
-        aws_conn_id: str = "aws_default",
         wait_for_completion: bool = True,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 40,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
-        super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+        super().__init__(**kwargs)
         self.db_identifier = db_identifier
         self.db_type = db_type
         self.db_snapshot_identifier = db_snapshot_identifier
         self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
 
     def execute(self, context: Context) -> str:
         self.db_type = RdsDbType(self.db_type)
-        stop_db_response = self._stop_db()
-        if self.wait_for_completion:
+        stop_db_response: dict[str, Any] = self._stop_db()
+        if self.deferrable:
+            self.defer(
+                trigger=RdsDbStoppedTrigger(
+                    db_identifier=self.db_identifier,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                    region_name=self.region_name,
+                    response=stop_db_response,
+                    db_type=RdsDbType.INSTANCE,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
             self._wait_until_db_stopped()
         return json.dumps(stop_db_response, default=str)
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        if event is None or event["status"] != "success":
+            raise AirflowException(f"Failed to start DB: {event}")
+        else:
+            return json.dumps(event["response"], default=str)
+
     def _stop_db(self):
         self.log.info("Stopping DB %s '%s'", self.db_type.value, self.db_identifier)
         if self.db_type == RdsDbType.INSTANCE:
@@ -829,9 +891,19 @@ class RdsStopDbOperator(RdsBaseOperator):
     def _wait_until_db_stopped(self):
         self.log.info("Waiting for DB %s to reach 'stopped' state", self.db_type.value)
         if self.db_type == RdsDbType.INSTANCE:
-            self.hook.wait_for_db_instance_state(self.db_identifier, target_state="stopped")
+            self.hook.wait_for_db_instance_state(
+                self.db_identifier,
+                target_state="stopped",
+                check_interval=self.waiter_delay,
+                max_attempts=self.waiter_max_attempts,
+            )
         else:
-            self.hook.wait_for_db_cluster_state(self.db_identifier, target_state="stopped")
+            self.hook.wait_for_db_cluster_state(
+                self.db_identifier,
+                target_state="stopped",
+                check_interval=self.waiter_delay,
+                max_attempts=self.waiter_max_attempts,
+            )
 
 
 __all__ = [

airflow/providers/amazon/aws/operators/redshift_data.py

@@ -51,6 +51,9 @@ class RedshiftDataOperator(BaseOperator):
         if False (default) will return statement ID
     :param aws_conn_id: aws connection to use
     :param region: aws region to use
+    :param workgroup_name: name of the Redshift Serverless workgroup. Mutually exclusive with
+        `cluster_identifier`. Specify this parameter to query Redshift Serverless. More info
+        https://docs.aws.amazon.com/redshift/latest/mgmt/working-with-serverless.html
     """
 
     template_fields = (
@@ -62,6 +65,7 @@ class RedshiftDataOperator(BaseOperator):
         "statement_name",
         "aws_conn_id",
         "region",
+        "workgroup_name",
     )
     template_ext = (".sql",)
     template_fields_renderers = {"sql": "sql"}
@@ -82,12 +86,14 @@ class RedshiftDataOperator(BaseOperator):
         return_sql_result: bool = False,
         aws_conn_id: str = "aws_default",
         region: str | None = None,
+        workgroup_name: str | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         self.database = database
         self.sql = sql
         self.cluster_identifier = cluster_identifier
+        self.workgroup_name = workgroup_name
         self.db_user = db_user
         self.parameters = parameters
         self.secret_arn = secret_arn
@@ -119,6 +125,7 @@ class RedshiftDataOperator(BaseOperator):
             database=self.database,
             sql=self.sql,
             cluster_identifier=self.cluster_identifier,
+            workgroup_name=self.workgroup_name,
             db_user=self.db_user,
             parameters=self.parameters,
             secret_arn=self.secret_arn,

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

@@ -30,7 +30,10 @@ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarni
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
-from airflow.providers.amazon.aws.triggers.sagemaker import SageMakerTrigger
+from airflow.providers.amazon.aws.triggers.sagemaker import (
+    SageMakerPipelineTrigger,
+    SageMakerTrigger,
+)
 from airflow.providers.amazon.aws.utils import trim_none_values
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
@@ -998,8 +1001,10 @@ class SageMakerStartPipelineOperator(SageMakerBaseOperator):
         All parameters supplied need to already be present in the pipeline definition.
     :param wait_for_completion: If true, this operator will only complete once the pipeline is complete.
     :param check_interval: How long to wait between checks for pipeline status when waiting for completion.
+    :param waiter_max_attempts: How many times to check the status before failing.
     :param verbose: Whether to print steps details when waiting for completion.
         Defaults to true, consider turning off for pipelines that have thousands of steps.
+    :param deferrable: Run operator in the deferrable mode.
 
     :return str: Returns The ARN of the pipeline execution created in Amazon SageMaker.
     """
@@ -1015,7 +1020,9 @@ class SageMakerStartPipelineOperator(SageMakerBaseOperator):
         pipeline_params: dict | None = None,
         wait_for_completion: bool = False,
         check_interval: int = CHECK_INTERVAL_SECOND,
+        waiter_max_attempts: int = 9999,
         verbose: bool = True,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
         super().__init__(config={}, aws_conn_id=aws_conn_id, **kwargs)
@@ -1024,22 +1031,46 @@ class SageMakerStartPipelineOperator(SageMakerBaseOperator):
         self.pipeline_params = pipeline_params
         self.wait_for_completion = wait_for_completion
         self.check_interval = check_interval
+        self.waiter_max_attempts = waiter_max_attempts
         self.verbose = verbose
+        self.deferrable = deferrable
 
     def execute(self, context: Context) -> str:
         arn = self.hook.start_pipeline(
             pipeline_name=self.pipeline_name,
             display_name=self.display_name,
             pipeline_params=self.pipeline_params,
-            wait_for_completion=self.wait_for_completion,
-            check_interval=self.check_interval,
-            verbose=self.verbose,
         )
         self.log.info(
             "Starting a new execution for pipeline %s, running with ARN %s", self.pipeline_name, arn
         )
+        if self.deferrable:
+            self.defer(
+                trigger=SageMakerPipelineTrigger(
+                    waiter_type=SageMakerPipelineTrigger.Type.COMPLETE,
+                    pipeline_execution_arn=arn,
+                    waiter_delay=self.check_interval,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
+            self.hook.check_status(
+                arn,
+                "PipelineExecutionStatus",
+                lambda p: self.hook.describe_pipeline_exec(p, self.verbose),
+                self.check_interval,
+                non_terminal_states=self.hook.pipeline_non_terminal_states,
+                max_ingestion_time=self.waiter_max_attempts * self.check_interval,
+            )
         return arn
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        if event is None or event["status"] != "success":
+            raise AirflowException(f"Failure during pipeline execution: {event}")
+        return event["value"]
+
 
 class SageMakerStopPipelineOperator(SageMakerBaseOperator):
     """
@@ -1057,6 +1088,7 @@ class SageMakerStopPipelineOperator(SageMakerBaseOperator):
     :param verbose: Whether to print steps details when waiting for completion.
         Defaults to true, consider turning off for pipelines that have thousands of steps.
     :param fail_if_not_running: raises an exception if the pipeline stopped or succeeded before this was run
+    :param deferrable: Run operator in the deferrable mode.
 
     :return str: Returns the status of the pipeline execution after the operation has been done.
     """
@@ -1073,23 +1105,24 @@ class SageMakerStopPipelineOperator(SageMakerBaseOperator):
         pipeline_exec_arn: str,
         wait_for_completion: bool = False,
         check_interval: int = CHECK_INTERVAL_SECOND,
+        waiter_max_attempts: int = 9999,
         verbose: bool = True,
         fail_if_not_running: bool = False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
         super().__init__(config={}, aws_conn_id=aws_conn_id, **kwargs)
         self.pipeline_exec_arn = pipeline_exec_arn
         self.wait_for_completion = wait_for_completion
         self.check_interval = check_interval
+        self.waiter_max_attempts = waiter_max_attempts
         self.verbose = verbose
         self.fail_if_not_running = fail_if_not_running
+        self.deferrable = deferrable
 
     def execute(self, context: Context) -> str:
         status = self.hook.stop_pipeline(
             pipeline_exec_arn=self.pipeline_exec_arn,
-            wait_for_completion=self.wait_for_completion,
-            check_interval=self.check_interval,
-            verbose=self.verbose,
             fail_if_not_running=self.fail_if_not_running,
         )
         self.log.info(
@@ -1097,8 +1130,43 @@ class SageMakerStopPipelineOperator(SageMakerBaseOperator):
             self.pipeline_exec_arn,
             status,
         )
+
+        if status not in self.hook.pipeline_non_terminal_states:
+            # pipeline already stopped
+            return status
+
+        # else, eventually wait for completion
+        if self.deferrable:
+            self.defer(
+                trigger=SageMakerPipelineTrigger(
+                    waiter_type=SageMakerPipelineTrigger.Type.STOPPED,
+                    pipeline_execution_arn=self.pipeline_exec_arn,
+                    waiter_delay=self.check_interval,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+        elif self.wait_for_completion:
+            status = self.hook.check_status(
+                self.pipeline_exec_arn,
+                "PipelineExecutionStatus",
+                lambda p: self.hook.describe_pipeline_exec(p, self.verbose),
+                self.check_interval,
+                non_terminal_states=self.hook.pipeline_non_terminal_states,
+                max_ingestion_time=self.waiter_max_attempts * self.check_interval,
+            )["PipelineExecutionStatus"]
+
         return status
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        if event is None or event["status"] != "success":
+            raise AirflowException(f"Failure during pipeline execution: {event}")
+        else:
+            # theoretically we should do a `describe` call to know this,
+            # but if we reach this point, this is the only possible status
+            return "Stopped"
+
 
 class SageMakerRegisterModelVersionOperator(SageMakerBaseOperator):
     """

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

@@ -17,11 +17,14 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Sequence
 
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
+from airflow.providers.amazon.aws.triggers.step_function import StepFunctionsExecutionCompleteTrigger
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -42,6 +45,11 @@ class StepFunctionStartExecutionOperator(BaseOperator):
     :param state_machine_input: JSON data input to pass to the State Machine
     :param aws_conn_id: aws connection to uses
     :param do_xcom_push: if True, execution_arn is pushed to XCom with key execution_arn.
+    :param waiter_max_attempts: Maximum number of attempts to poll the execution.
+    :param waiter_delay: Number of seconds between polling the state of the execution.
+    :param deferrable: If True, the operator will wait asynchronously for the job to complete.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False, but can be overridden in config file by setting default_deferrable to True)
     """
 
     template_fields: Sequence[str] = ("state_machine_arn", "name", "input")
@@ -56,6 +64,9 @@ class StepFunctionStartExecutionOperator(BaseOperator):
         state_machine_input: dict | str | None = None,
         aws_conn_id: str = "aws_default",
         region_name: str | None = None,
+        waiter_max_attempts: int = 30,
+        waiter_delay: int = 60,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -64,6 +75,9 @@ class StepFunctionStartExecutionOperator(BaseOperator):
         self.input = state_machine_input
         self.aws_conn_id = aws_conn_id
         self.region_name = region_name
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
 
     def execute(self, context: Context):
         hook = StepFunctionHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
@@ -74,9 +88,27 @@ class StepFunctionStartExecutionOperator(BaseOperator):
             raise AirflowException(f"Failed to start State Machine execution for: {self.state_machine_arn}")
 
         self.log.info("Started State Machine execution for %s: %s", self.state_machine_arn, execution_arn)
-
+        if self.deferrable:
+            self.defer(
+                trigger=StepFunctionsExecutionCompleteTrigger(
+                    execution_arn=execution_arn,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                    region_name=self.region_name,
+                ),
+                method_name="execute_complete",
+                timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
+            )
         return execution_arn
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        if event is None or event["status"] != "success":
+            raise AirflowException(f"Trigger error: event is {event}")
+
+        self.log.info("State Machine execution completed successfully")
+        return event["execution_arn"]
+
 
 class StepFunctionGetExecutionOutputOperator(BaseOperator):
     """