apache-airflow-providers-amazon 8.22.0__py3-none-any.whl → 8.23.0__py3-none-any.whl

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

@@ -617,6 +617,14 @@ class EmrContainerOperator(BaseOperator):
                     job_id=self.job_id,
                     aws_conn_id=self.aws_conn_id,
                     waiter_delay=self.poll_interval,
+                    waiter_max_attempts=self.max_polling_attempts,
+                )
+                if self.max_polling_attempts
+                else EmrContainerTrigger(
+                    virtual_cluster_id=self.virtual_cluster_id,
+                    job_id=self.job_id,
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.poll_interval,
                 ),
                 method_name="execute_complete",
             )
@@ -734,10 +742,20 @@ class EmrCreateJobFlowOperator(BaseOperator):
         waiter_max_attempts: int | None = None,
         waiter_delay: int | None = None,
         waiter_countdown: int | None = None,
-        waiter_check_interval_seconds: int = 60,
+        waiter_check_interval_seconds: int | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs: Any,
     ):
+        if waiter_check_interval_seconds:
+            warnings.warn(
+                "The parameter `waiter_check_interval_seconds` has been deprecated to "
+                "standardize naming conventions.  Please `use waiter_delay instead`.  In the "
+                "future this will default to None and defer to the waiter's default value.",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+        else:
+            waiter_check_interval_seconds = 60
         if waiter_countdown:
             warnings.warn(
                 "The parameter waiter_countdown has been deprecated to standardize "
@@ -749,15 +767,7 @@ class EmrCreateJobFlowOperator(BaseOperator):
             # waiter_countdown defaults to never timing out, which is not supported
             # by boto waiters, so we will set it here to "a very long time" for now.
             waiter_max_attempts = (waiter_countdown or 999) // waiter_check_interval_seconds
-        if waiter_check_interval_seconds:
-            warnings.warn(
-                "The parameter waiter_check_interval_seconds has been deprecated to "
-                "standardize naming conventions.  Please use waiter_delay instead.  In the "
-                "future this will default to None and defer to the waiter's default value.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-            waiter_delay = waiter_check_interval_seconds
+
         super().__init__(**kwargs)
         self.aws_conn_id = aws_conn_id
         self.emr_conn_id = emr_conn_id
@@ -765,7 +775,7 @@ class EmrCreateJobFlowOperator(BaseOperator):
         self.region_name = region_name
         self.wait_for_completion = wait_for_completion
         self.waiter_max_attempts = waiter_max_attempts or 60
-        self.waiter_delay = waiter_delay or 30
+        self.waiter_delay = waiter_delay or waiter_check_interval_seconds or 60
         self.deferrable = deferrable
 
     @cached_property

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

@@ -19,11 +19,15 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
 
+from botocore.exceptions import ClientError
+
 from airflow.configuration import conf
+from airflow.exceptions import AirflowException
 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,
+    NeptuneClusterInstancesAvailableTrigger,
     NeptuneClusterStoppedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
@@ -32,6 +36,50 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
+def handle_waitable_exception(
+    operator: NeptuneStartDbClusterOperator | NeptuneStopDbClusterOperator, err: str
+):
+    """
+    Handle client exceptions for invalid cluster or invalid instance status that are temporary.
+
+    After status change, it's possible to retry. Waiter will handle terminal status.
+    """
+    code = err
+
+    if code in ("InvalidDBInstanceStateFault", "InvalidDBInstanceState"):
+        if operator.deferrable:
+            operator.log.info("Deferring until instances become available: %s", operator.cluster_id)
+            operator.defer(
+                trigger=NeptuneClusterInstancesAvailableTrigger(
+                    aws_conn_id=operator.aws_conn_id,
+                    db_cluster_id=operator.cluster_id,
+                    region_name=operator.region_name,
+                    botocore_config=operator.botocore_config,
+                    verify=operator.verify,
+                ),
+                method_name="execute",
+            )
+        else:
+            operator.log.info("Need to wait for instances to become available: %s", operator.cluster_id)
+            operator.hook.wait_for_cluster_instance_availability(cluster_id=operator.cluster_id)
+    if code in ["InvalidClusterState", "InvalidDBClusterStateFault"]:
+        if operator.deferrable:
+            operator.log.info("Deferring until cluster becomes available: %s", operator.cluster_id)
+            operator.defer(
+                trigger=NeptuneClusterAvailableTrigger(
+                    aws_conn_id=operator.aws_conn_id,
+                    db_cluster_id=operator.cluster_id,
+                    region_name=operator.region_name,
+                    botocore_config=operator.botocore_config,
+                    verify=operator.verify,
+                ),
+                method_name="execute",
+            )
+        else:
+            operator.log.info("Need to wait for cluster to become available: %s", operator.cluster_id)
+            operator.hook.wait_for_cluster_availability(operator.cluster_id)
+
+
 class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
     """Starts an Amazon Neptune DB cluster.
 
@@ -78,10 +126,10 @@ class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
         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
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
 
-    def execute(self, context: Context) -> dict[str, str]:
+    def execute(self, context: Context, event: dict[str, Any] | None = None, **kwargs) -> dict[str, str]:
         self.log.info("Starting Neptune cluster: %s", self.cluster_id)
 
         # Check to make sure the cluster is not already available.
@@ -89,9 +137,32 @@ class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
         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")
+        elif status.lower() in NeptuneHook.ERROR_STATES:
+            # some states will not allow you to start the cluster
+            self.log.error(
+                "Neptune cluster %s is in error state %s and cannot be started", self.cluster_id, status
+            )
+            raise AirflowException(f"Neptune cluster {self.cluster_id} is in error state {status}")
+
+        """
+        A cluster and its instances must be in a valid state to send the start request.
+        This loop covers the case where the cluster is not available and also the case where
+        the cluster is available, but one or more of the instances are in an invalid state.
+        If either are in an invalid state, wait for the availability and retry.
+        Let the waiters handle retries and detecting the error states.
+        """
+        try:
+            self.hook.conn.start_db_cluster(DBClusterIdentifier=self.cluster_id)
+        except ClientError as ex:
+            code = ex.response["Error"]["Code"]
+            self.log.warning("Received client error when attempting to start the cluster: %s", code)
+
+            if code in ["InvalidDBInstanceState", "InvalidClusterState", "InvalidDBClusterStateFault"]:
+                handle_waitable_exception(operator=self, err=code)
+
+            else:
+                # re raise for any other type of client error
+                raise
 
         if self.deferrable:
             self.log.info("Deferring for cluster start: %s", self.cluster_id)
@@ -100,15 +171,17 @@ class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
                 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,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_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)
+            self.hook.wait_for_cluster_availability(
+                self.cluster_id, self.waiter_delay, self.waiter_max_attempts
+            )
 
         return {"db_cluster_id": self.cluster_id}
 
@@ -171,20 +244,53 @@ class NeptuneStopDbClusterOperator(AwsBaseOperator[NeptuneHook]):
         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
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
 
-    def execute(self, context: Context) -> dict[str, str]:
+    def execute(self, context: Context, event: dict[str, Any] | None = None, **kwargs) -> dict[str, str]:
         self.log.info("Stopping Neptune cluster: %s", self.cluster_id)
 
-        # Check to make sure the cluster is not already stopped.
+        # Check to make sure the cluster is not already stopped or that its not in a bad state
         status = self.hook.get_cluster_status(self.cluster_id)
+        self.log.info("Current status: %s", status)
+
         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")
+        elif status.lower() in NeptuneHook.ERROR_STATES:
+            # some states will not allow you to stop the cluster
+            self.log.error(
+                "Neptune cluster %s is in error state %s and cannot be stopped", self.cluster_id, status
+            )
+            raise AirflowException(f"Neptune cluster {self.cluster_id} is in error state {status}")
+
+        """
+        A cluster and its instances must be in a valid state to send the stop request.
+        This loop covers the case where the cluster is not available and also the case where
+        the cluster is available, but one or more of the instances are in an invalid state.
+        If either are in an invalid state, wait for the availability and retry.
+        Let the waiters handle retries and detecting the error states.
+        """
+
+        try:
+            self.hook.conn.stop_db_cluster(DBClusterIdentifier=self.cluster_id)
+
+        # cluster must be in available state to stop it
+        except ClientError as ex:
+            code = ex.response["Error"]["Code"]
+            self.log.warning("Received client error when attempting to stop the cluster: %s", code)
+
+            # these can be handled by a waiter
+            if code in [
+                "InvalidDBInstanceState",
+                "InvalidDBInstanceStateFault",
+                "InvalidClusterState",
+                "InvalidDBClusterStateFault",
+            ]:
+                handle_waitable_exception(self, code)
+            else:
+                # re raise for any other type of client error
+                raise
 
         if self.deferrable:
             self.log.info("Deferring for cluster stop: %s", self.cluster_id)
@@ -193,22 +299,23 @@ class NeptuneStopDbClusterOperator(AwsBaseOperator[NeptuneHook]):
                 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,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_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)
+            self.log.info("Waiting for Neptune cluster %s to stop.", self.cluster_id)
+
+            self.hook.wait_for_cluster_stopped(self.cluster_id, self.waiter_delay, self.waiter_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 = ""
-
+        self.log.info(event)
         if event:
             status = event.get("status", "")
             cluster_id = event.get("cluster_id", "")

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

@@ -750,20 +750,18 @@ class SageMakerTransformOperator(SageMakerBaseOperator):
                 method_name="execute_complete",
             )
 
-        return self.serialize_result()
+        return self.serialize_result(transform_config["TransformJobName"])
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
         event = validate_execute_complete_event(event)
 
         self.log.info(event["message"])
-        return self.serialize_result()
+        return self.serialize_result(event["job_name"])
 
-    def serialize_result(self) -> dict[str, dict]:
-        transform_config = self.config.get("Transform", self.config)
-        self.serialized_model = serialize(self.hook.describe_model(transform_config["ModelName"]))
-        self.serialized_transform = serialize(
-            self.hook.describe_transform_job(transform_config["TransformJobName"])
-        )
+    def serialize_result(self, job_name: str) -> dict[str, dict]:
+        job_description = self.hook.describe_transform_job(job_name)
+        self.serialized_model = serialize(self.hook.describe_model(job_description["ModelName"]))
+        self.serialized_transform = serialize(job_description)
         return {"Model": self.serialized_model, "Transform": self.serialized_transform}
 
     def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:
@@ -1154,7 +1152,7 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
                 method_name="execute_complete",
             )
 
-        return self.serialize_result()
+        return self.serialize_result(self.config["TrainingJobName"])
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
         event = validate_execute_complete_event(event)
@@ -1163,12 +1161,10 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
             raise AirflowException(f"Error while running job: {event}")
 
         self.log.info(event["message"])
-        return self.serialize_result()
+        return self.serialize_result(event["job_name"])
 
-    def serialize_result(self) -> dict[str, dict]:
-        self.serialized_training_data = serialize(
-            self.hook.describe_training_job(self.config["TrainingJobName"])
-        )
+    def serialize_result(self, job_name: str) -> dict[str, dict]:
+        self.serialized_training_data = serialize(self.hook.describe_training_job(job_name))
         return {"Training": self.serialized_training_data}
 
     def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage:

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

@@ -0,0 +1,147 @@
+# 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
+
+import abc
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException, AirflowSkipException
+from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.triggers.comprehend import ComprehendPiiEntitiesDetectionJobCompletedTrigger
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class ComprehendBaseSensor(AwsBaseSensor[ComprehendHook]):
+    """
+    General sensor behavior for Amazon Comprehend.
+
+    Subclasses must implement following methods:
+        - ``get_state()``
+
+    Subclasses must set the following fields:
+        - ``INTERMEDIATE_STATES``
+        - ``FAILURE_STATES``
+        - ``SUCCESS_STATES``
+        - ``FAILURE_MESSAGE``
+
+    :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)
+    """
+
+    aws_hook_class = ComprehendHook
+
+    INTERMEDIATE_STATES: tuple[str, ...] = ()
+    FAILURE_STATES: tuple[str, ...] = ()
+    SUCCESS_STATES: tuple[str, ...] = ()
+    FAILURE_MESSAGE = ""
+
+    ui_color = "#66c3ff"
+
+    def __init__(
+        self,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.deferrable = deferrable
+
+    def poke(self, context: Context, **kwargs) -> bool:
+        state = self.get_state()
+        if state in self.FAILURE_STATES:
+            # TODO: remove this if block when min_airflow_version is set to higher than 2.7.1
+            if self.soft_fail:
+                raise AirflowSkipException(self.FAILURE_MESSAGE)
+            raise AirflowException(self.FAILURE_MESSAGE)
+
+        return state not in self.INTERMEDIATE_STATES
+
+    @abc.abstractmethod
+    def get_state(self) -> str:
+        """Implement in subclasses."""
+
+
+class ComprehendStartPiiEntitiesDetectionJobCompletedSensor(ComprehendBaseSensor):
+    """
+    Poll the state of the pii entities detection job until it reaches a completed state; fails if the job fails.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:ComprehendStartPiiEntitiesDetectionJobCompletedSensor`
+
+    :param job_id: The id of the Comprehend pii entities detection job.
+
+    :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
+        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 verify: Whether to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    INTERMEDIATE_STATES: tuple[str, ...] = ("IN_PROGRESS",)
+    FAILURE_STATES: tuple[str, ...] = ("FAILED", "STOP_REQUESTED", "STOPPED")
+    SUCCESS_STATES: tuple[str, ...] = ("COMPLETED",)
+    FAILURE_MESSAGE = "Comprehend start pii entities detection job sensor failed."
+
+    template_fields: Sequence[str] = aws_template_fields("job_id")
+
+    def __init__(
+        self,
+        *,
+        job_id: str,
+        max_retries: int = 75,
+        poke_interval: int = 120,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.job_id = job_id
+        self.max_retries = max_retries
+        self.poke_interval = poke_interval
+
+    def execute(self, context: Context) -> Any:
+        if self.deferrable:
+            self.defer(
+                trigger=ComprehendPiiEntitiesDetectionJobCompletedTrigger(
+                    job_id=self.job_id,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="poke",
+            )
+        else:
+            super().execute(context=context)
+
+    def get_state(self) -> str:
+        return self.hook.conn.describe_pii_entities_detection_job(JobId=self.job_id)[
+            "PiiEntitiesDetectionJobProperties"
+        ]["JobStatus"]

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

@@ -354,6 +354,14 @@ class EmrContainerSensor(BaseSensorOperator):
                     job_id=self.job_id,
                     aws_conn_id=self.aws_conn_id,
                     waiter_delay=self.poll_interval,
+                    waiter_max_attempts=self.max_retries,
+                )
+                if self.max_retries
+                else EmrContainerTrigger(
+                    virtual_cluster_id=self.virtual_cluster_id,
+                    job_id=self.job_id,
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.poll_interval,
                 ),
                 method_name="execute_complete",
             )

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

@@ -0,0 +1,61 @@
+# 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
+
+if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+
+class ComprehendPiiEntitiesDetectionJobCompletedTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when a Comprehend pii entities detection job is complete.
+
+    :param job_id: The id of the Comprehend pii entities detection job.
+    :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
+    :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        *,
+        job_id: str,
+        waiter_delay: int = 120,
+        waiter_max_attempts: int = 75,
+        aws_conn_id: str | None = "aws_default",
+    ) -> None:
+        super().__init__(
+            serialized_fields={"job_id": job_id},
+            waiter_name="pii_entities_detection_job_complete",
+            waiter_args={"JobId": job_id},
+            failure_message="Comprehend start pii entities detection job failed.",
+            status_message="Status of Comprehend start pii entities detection job is",
+            status_queries=["PiiEntitiesDetectionJobProperties.JobStatus"],
+            return_key="job_id",
+            return_value=job_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return ComprehendHook(aws_conn_id=self.aws_conn_id)

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

@@ -113,3 +113,48 @@ class NeptuneClusterStoppedTrigger(AwsBaseWaiterTrigger):
             verify=self.verify,
             config=self.botocore_config,
         )
+
+
+class NeptuneClusterInstancesAvailableTrigger(AwsBaseWaiterTrigger):
+    """
+    Triggers when a Neptune Cluster Instance is available.
+
+    :param db_cluster_id: Cluster ID to wait on instances from
+    :param waiter_delay: The amount of time in seconds to wait between attempts.
+    :param waiter_max_attempts: The maximum number of attempts to be made.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    :param region_name: AWS region name (example: us-east-1)
+    """
+
+    def __init__(
+        self,
+        *,
+        db_cluster_id: str,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str | None = None,
+        region_name: str | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            serialized_fields={"db_cluster_id": db_cluster_id},
+            waiter_name="db_instance_available",
+            waiter_args={"Filters": [{"Name": "db-cluster-id", "Values": [db_cluster_id]}]},
+            failure_message="Failed to start Neptune instances",
+            status_message="Status of Neptune instances are",
+            status_queries=["DBInstances[].Status"],
+            return_key="db_cluster_id",
+            return_value=db_cluster_id,
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+            **kwargs,
+        )
+
+    def hook(self) -> AwsGenericHook:
+        return NeptuneHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region_name,
+            verify=self.verify,
+            config=self.botocore_config,
+        )

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

@@ -121,7 +121,7 @@ class SageMakerTrigger(BaseTrigger):
                 status_message=f"{self.job_type} job not done yet",
                 status_args=[self._get_response_status_key(self.job_type)],
             )
-            yield TriggerEvent({"status": "success", "message": "Job completed."})
+            yield TriggerEvent({"status": "success", "message": "Job completed.", "job_name": self.job_name})
 
 
 class SageMakerPipelineTrigger(BaseTrigger):

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

@@ -22,6 +22,8 @@ from datetime import datetime, timezone
 from enum import Enum
 from typing import Any
 
+import importlib_metadata
+
 from airflow.exceptions import AirflowException
 from airflow.utils.helpers import prune_dict
 from airflow.version import version
@@ -74,6 +76,11 @@ def get_airflow_version() -> tuple[int, ...]:
     return tuple(int(x) for x in match.groups())
 
 
+def get_botocore_version() -> tuple[int, ...]:
+    """Return the version number of the installed botocore package in the form of a tuple[int,...]."""
+    return tuple(map(int, importlib_metadata.version("botocore").split(".")[:3]))
+
+
 def validate_execute_complete_event(event: dict[str, Any] | None = None) -> dict[str, Any]:
     if event is None:
         err_msg = "Trigger error: event is None"