apache-airflow-providers-amazon 8.17.0rc2__py3-none-any.whl → 8.18.0rc2__py3-none-any.whl

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

@@ -0,0 +1,218 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.hooks.neptune import NeptuneHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune import (
+    NeptuneClusterAvailableTrigger,
+    NeptuneClusterStoppedTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
+    """Starts an Amazon Neptune DB cluster.
+
+    Amazon Neptune Database is a serverless graph database designed for superior scalability
+    and availability. Neptune Database provides built-in security, continuous backups, and
+    integrations with other AWS services
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:NeptuneStartDbClusterOperator`
+
+    :param db_cluster_id: The DB cluster identifier of the Neptune DB cluster to be started.
+    :param wait_for_completion: Whether to wait for the cluster to start. (default: True)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to start.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param waiter_delay: Time in seconds to wait between status checks.
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    :return: dictionary with Neptune cluster id
+    """
+
+    aws_hook_class = NeptuneHook
+    template_fields: Sequence[str] = aws_template_fields("cluster_id")
+
+    def __init__(
+        self,
+        db_cluster_id: str,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_id = db_cluster_id
+        self.wait_for_completion = wait_for_completion
+        self.deferrable = deferrable
+        self.delay = waiter_delay
+        self.max_attempts = waiter_max_attempts
+
+    def execute(self, context: Context) -> dict[str, str]:
+        self.log.info("Starting Neptune cluster: %s", self.cluster_id)
+
+        # Check to make sure the cluster is not already available.
+        status = self.hook.get_cluster_status(self.cluster_id)
+        if status.lower() in NeptuneHook.AVAILABLE_STATES:
+            self.log.info("Neptune cluster %s is already available.", self.cluster_id)
+            return {"db_cluster_id": self.cluster_id}
+
+        resp = self.hook.conn.start_db_cluster(DBClusterIdentifier=self.cluster_id)
+        status = resp.get("DBClusters", {}).get("Status", "Unknown")
+
+        if self.deferrable:
+            self.log.info("Deferring for cluster start: %s", self.cluster_id)
+
+            self.defer(
+                trigger=NeptuneClusterAvailableTrigger(
+                    aws_conn_id=self.aws_conn_id,
+                    db_cluster_id=self.cluster_id,
+                    waiter_delay=self.delay,
+                    waiter_max_attempts=self.max_attempts,
+                ),
+                method_name="execute_complete",
+            )
+
+        elif self.wait_for_completion:
+            self.log.info("Waiting for Neptune cluster %s to start.", self.cluster_id)
+            self.hook.wait_for_cluster_availability(self.cluster_id, self.delay, self.max_attempts)
+
+        return {"db_cluster_id": self.cluster_id}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, str]:
+        status = ""
+        cluster_id = ""
+
+        if event:
+            status = event.get("status", "")
+            cluster_id = event.get("cluster_id", "")
+
+        self.log.info("Neptune cluster %s available with status: %s", cluster_id, status)
+
+        return {"db_cluster_id": cluster_id}
+
+
+class NeptuneStopDbClusterOperator(AwsBaseOperator[NeptuneHook]):
+    """
+    Stops an Amazon Neptune DB cluster.
+
+    Amazon Neptune Database is a serverless graph database designed for superior scalability
+    and availability. Neptune Database provides built-in security, continuous backups, and
+    integrations with other AWS services
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:NeptuneStartDbClusterOperator`
+
+    :param db_cluster_id: The DB cluster identifier of the Neptune DB cluster to be stopped.
+    :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param waiter_delay: Time in seconds to wait between status checks.
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    :return: dictionary with Neptune cluster id
+    """
+
+    aws_hook_class = NeptuneHook
+    template_fields: Sequence[str] = aws_template_fields("cluster_id")
+
+    def __init__(
+        self,
+        db_cluster_id: str,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_id = db_cluster_id
+        self.wait_for_completion = wait_for_completion
+        self.deferrable = deferrable
+        self.delay = waiter_delay
+        self.max_attempts = waiter_max_attempts
+
+    def execute(self, context: Context) -> dict[str, str]:
+        self.log.info("Stopping Neptune cluster: %s", self.cluster_id)
+
+        # Check to make sure the cluster is not already stopped.
+        status = self.hook.get_cluster_status(self.cluster_id)
+        if status.lower() in NeptuneHook.STOPPED_STATES:
+            self.log.info("Neptune cluster %s is already stopped.", self.cluster_id)
+            return {"db_cluster_id": self.cluster_id}
+
+        resp = self.hook.conn.stop_db_cluster(DBClusterIdentifier=self.cluster_id)
+        status = resp.get("DBClusters", {}).get("Status", "Unknown")
+
+        if self.deferrable:
+            self.log.info("Deferring for cluster stop: %s", self.cluster_id)
+
+            self.defer(
+                trigger=NeptuneClusterStoppedTrigger(
+                    aws_conn_id=self.aws_conn_id,
+                    db_cluster_id=self.cluster_id,
+                    waiter_delay=self.delay,
+                    waiter_max_attempts=self.max_attempts,
+                ),
+                method_name="execute_complete",
+            )
+
+        elif self.wait_for_completion:
+            self.log.info("Waiting for Neptune cluster %s to start.", self.cluster_id)
+            self.hook.wait_for_cluster_stopped(self.cluster_id, self.delay, self.max_attempts)
+
+        return {"db_cluster_id": self.cluster_id}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, str]:
+        status = ""
+        cluster_id = ""
+
+        if event:
+            status = event.get("status", "")
+            cluster_id = event.get("cluster_id", "")
+
+        self.log.info("Neptune cluster %s stopped with status: %s", cluster_id, status)
+
+        return {"db_cluster_id": cluster_id}

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

@@ -32,6 +32,7 @@ from airflow.providers.amazon.aws.triggers.rds import (
     RdsDbDeletedTrigger,
     RdsDbStoppedTrigger,
 )
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 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
@@ -637,11 +638,13 @@ class RdsCreateDbInstanceOperator(RdsBaseOperator):
             )
         return json.dumps(create_db_instance, default=str)
 
-    def execute_complete(self, context, event=None) -> str:
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"DB instance creation failed: {event}")
-        else:
-            return json.dumps(event["response"], default=str)
+
+        return json.dumps(event["response"], default=str)
 
 
 class RdsDeleteDbInstanceOperator(RdsBaseOperator):
@@ -720,11 +723,13 @@ class RdsDeleteDbInstanceOperator(RdsBaseOperator):
             )
         return json.dumps(delete_db_instance, default=str)
 
-    def execute_complete(self, context, event=None) -> str:
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"DB instance deletion failed: {event}")
-        else:
-            return json.dumps(event["response"], default=str)
+
+        return json.dumps(event["response"], default=str)
 
 
 class RdsStartDbOperator(RdsBaseOperator):
@@ -786,10 +791,12 @@ class RdsStartDbOperator(RdsBaseOperator):
         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":
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
             raise AirflowException(f"Failed to start DB: {event}")
-        else:
-            return json.dumps(event["response"], default=str)
+
+        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)
@@ -883,10 +890,12 @@ class RdsStopDbOperator(RdsBaseOperator):
         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":
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
             raise AirflowException(f"Failed to start DB: {event}")
-        else:
-            return json.dumps(event["response"], default=str)
+
+        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)

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

@@ -31,6 +31,7 @@ from airflow.providers.amazon.aws.triggers.redshift_cluster import (
     RedshiftPauseClusterTrigger,
     RedshiftResumeClusterTrigger,
 )
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -314,10 +315,11 @@ class RedshiftCreateClusterOperator(BaseOperator):
         self.log.info("Created Redshift cluster %s", self.cluster_identifier)
         self.log.info(cluster)
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"Error creating cluster: {event}")
-        return
 
 
 class RedshiftCreateClusterSnapshotOperator(BaseOperator):
@@ -409,12 +411,13 @@ class RedshiftCreateClusterSnapshotOperator(BaseOperator):
                 },
             )
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"Error creating snapshot: {event}")
-        else:
-            self.log.info("Cluster snapshot created.")
-        return
+
+        self.log.info("Cluster snapshot created.")
 
 
 class RedshiftDeleteClusterSnapshotOperator(BaseOperator):
@@ -569,10 +572,7 @@ class RedshiftResumeClusterOperator(BaseOperator):
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        if event is None:
-            err_msg = "Trigger error: event is None"
-            self.log.info(err_msg)
-            raise AirflowException(err_msg)
+        event = validate_execute_complete_event(event)
 
         if event["status"] != "success":
             raise AirflowException(f"Error resuming cluster: {event}")
@@ -659,10 +659,7 @@ class RedshiftPauseClusterOperator(BaseOperator):
                 )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        if event is None:
-            err_msg = "Trigger error: event is None"
-            self.log.info(err_msg)
-            raise AirflowException(err_msg)
+        event = validate_execute_complete_event(event)
 
         if event["status"] != "success":
             raise AirflowException(f"Error pausing cluster: {event}")
@@ -767,10 +764,7 @@ class RedshiftDeleteClusterOperator(BaseOperator):
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        if event is None:
-            err_msg = "Trigger error: event is None"
-            self.log.info(err_msg)
-            raise AirflowException(err_msg)
+        event = validate_execute_complete_event(event)
 
         if event["status"] != "success":
             raise AirflowException(f"Error deleting cluster: {event}")

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

@@ -24,6 +24,7 @@ from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.redshift_data import RedshiftDataTrigger
+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:
@@ -170,10 +171,7 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
     def execute_complete(
         self, context: Context, event: dict[str, Any] | None = None
     ) -> GetStatementResultResponseTypeDef | str:
-        if event is None:
-            err_msg = "Trigger error: event is None"
-            self.log.info(err_msg)
-            raise AirflowException(err_msg)
+        event = validate_execute_complete_event(event)
 
         if event["status"] == "error":
             msg = f"context: {context}, error message: {event['message']}"

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

@@ -29,12 +29,17 @@ from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 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.hooks.sagemaker import (
+    LogState,
+    SageMakerHook,
+    secondary_training_status_message,
+)
 from airflow.providers.amazon.aws.triggers.sagemaker import (
     SageMakerPipelineTrigger,
+    SageMakerTrainingPrintLogTrigger,
     SageMakerTrigger,
 )
-from airflow.providers.amazon.aws.utils import trim_none_values
+from airflow.providers.amazon.aws.utils import trim_none_values, validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.utils.helpers import prune_dict
@@ -310,11 +315,13 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
         self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
         return {"Processing": self.serialized_job}
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
-        else:
-            self.log.info(event["message"])
+
+        self.log.info(event["message"])
         self.serialized_job = serialize(self.hook.describe_processing_job(self.config["ProcessingJobName"]))
         self.log.info("%s completed successfully.", self.task_id)
         return {"Processing": self.serialized_job}
@@ -561,7 +568,9 @@ class SageMakerEndpointOperator(SageMakerBaseOperator):
             "Endpoint": serialize(self.hook.describe_endpoint(endpoint_info["EndpointName"])),
         }
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
         endpoint_info = self.config.get("Endpoint", self.config)
@@ -744,10 +753,7 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
         return self.serialize_result()
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
-        if event is None:
-            err_msg = "Trigger error: event is None"
-            self.log.error(err_msg)
-            raise AirflowException(err_msg)
+        event = validate_execute_complete_event(event)
 
         self.log.info(event["message"])
         return self.serialize_result()
@@ -899,9 +905,11 @@ class SageMakerTuningOperator(SageMakerBaseOperator):
                     aws_conn_id=self.aws_conn_id,
                 ),
                 method_name="execute_complete",
-                timeout=datetime.timedelta(seconds=self.max_ingestion_time)
-                if self.max_ingestion_time is not None
-                else None,
+                timeout=(
+                    datetime.timedelta(seconds=self.max_ingestion_time)
+                    if self.max_ingestion_time is not None
+                    else None
+                ),
             )
             description = {}  # never executed but makes static checkers happy
         elif self.wait_for_completion:
@@ -917,7 +925,9 @@ class SageMakerTuningOperator(SageMakerBaseOperator):
 
         return {"Tuning": serialize(description)}
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
         return {
@@ -1085,28 +1095,77 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
             raise AirflowException(f"Sagemaker Training Job creation failed: {response}")
 
         if self.deferrable and self.wait_for_completion:
-            self.defer(
-                timeout=self.execution_timeout,
-                trigger=SageMakerTrigger(
+            description = self.hook.describe_training_job(self.config["TrainingJobName"])
+            status = description["TrainingJobStatus"]
+
+            if self.print_log:
+                instance_count = description["ResourceConfig"]["InstanceCount"]
+                last_describe_job_call = time.monotonic()
+                job_already_completed = status not in self.hook.non_terminal_states
+                _, description, last_describe_job_call = self.hook.describe_training_job_with_log(
+                    self.config["TrainingJobName"],
+                    {},
+                    [],
+                    instance_count,
+                    LogState.COMPLETE if job_already_completed else LogState.TAILING,
+                    description,
+                    last_describe_job_call,
+                )
+                self.log.info(secondary_training_status_message(description, None))
+
+            if status in self.hook.failed_states:
+                reason = description.get("FailureReason", "(No reason provided)")
+                raise AirflowException(f"SageMaker job failed because {reason}")
+            elif status == "Completed":
+                log_message = f"{self.task_id} completed successfully."
+                if self.print_log:
+                    billable_seconds = SageMakerHook.count_billable_seconds(
+                        training_start_time=description["TrainingStartTime"],
+                        training_end_time=description["TrainingEndTime"],
+                        instance_count=instance_count,
+                    )
+                    log_message = f"Billable seconds: {billable_seconds}\n{log_message}"
+                self.log.info(log_message)
+                return {"Training": serialize(description)}
+
+            timeout = self.execution_timeout
+            if self.max_ingestion_time:
+                timeout = datetime.timedelta(seconds=self.max_ingestion_time)
+
+            trigger: SageMakerTrainingPrintLogTrigger | SageMakerTrigger
+            if self.print_log:
+                trigger = SageMakerTrainingPrintLogTrigger(
+                    job_name=self.config["TrainingJobName"],
+                    poke_interval=self.check_interval,
+                    aws_conn_id=self.aws_conn_id,
+                )
+            else:
+                trigger = SageMakerTrigger(
                     job_name=self.config["TrainingJobName"],
                     job_type="Training",
                     poke_interval=self.check_interval,
                     max_attempts=self.max_attempts,
                     aws_conn_id=self.aws_conn_id,
-                ),
+                )
+
+            self.defer(
+                timeout=timeout,
+                trigger=trigger,
                 method_name="execute_complete",
             )
 
-        self.serialized_training_data = serialize(
-            self.hook.describe_training_job(self.config["TrainingJobName"])
-        )
-        return {"Training": self.serialized_training_data}
+        return self.serialize_result()
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        event = validate_execute_complete_event(event)
 
-    def execute_complete(self, context, event=None):
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
-        else:
-            self.log.info(event["message"])
+
+        self.log.info(event["message"])
+        return self.serialize_result()
+
+    def serialize_result(self) -> dict[str, dict]:
         self.serialized_training_data = serialize(
             self.hook.describe_training_job(self.config["TrainingJobName"])
         )
@@ -1237,7 +1296,9 @@ class SageMakerStartPipelineOperator(SageMakerBaseOperator):
         return arn
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        if event is None or event["status"] != "success":
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
             raise AirflowException(f"Failure during pipeline execution: {event}")
         return event["value"]
 
@@ -1330,12 +1391,14 @@ class SageMakerStopPipelineOperator(SageMakerBaseOperator):
         return status
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
-        if event is None or event["status"] != "success":
+        event = validate_execute_complete_event(event)
+
+        if 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"
+
+        # 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

@@ -29,6 +29,7 @@ from airflow.providers.amazon.aws.links.step_function import (
 )
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.step_function import StepFunctionsExecutionCompleteTrigger
+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:
@@ -129,7 +130,9 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
         return execution_arn
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
-        if event is None or event["status"] != "success":
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
             raise AirflowException(f"Trigger error: event is {event}")
 
         self.log.info("State Machine execution completed successfully")

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

@@ -60,8 +60,8 @@ class BatchSensor(BaseSensorOperator):
         aws_conn_id: str = "aws_default",
         region_name: str | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        poke_interval: float = 5,
-        max_retries: int = 5,
+        poke_interval: float = 30,
+        max_retries: int = 4200,
         **kwargs,
     ):
         super().__init__(**kwargs)

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

@@ -24,6 +24,7 @@ from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook
 from airflow.providers.amazon.aws.triggers.ec2 import EC2StateSensorTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
@@ -92,11 +93,12 @@ class EC2InstanceStateSensor(BaseSensorOperator):
         self.log.info("instance state: %s", instance_state)
         return instance_state == self.target_state
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        event = validate_execute_complete_event(event)
+
         if event["status"] != "success":
             # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
             message = f"Error: {event}"
             if self.soft_fail:
                 raise AirflowSkipException(message)
             raise AirflowException(message)
-        return