apache-airflow-providers-amazon 8.26.0rc2__py3-none-any.whl → 8.27.0__py3-none-any.whl

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

@@ -0,0 +1,348 @@
+# 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 botocore.exceptions import ClientError
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.kinesis_analytics import KinesisAnalyticsV2Hook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.kinesis_analytics import (
+    KinesisAnalyticsV2ApplicationOperationCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class KinesisAnalyticsV2CreateApplicationOperator(AwsBaseOperator[KinesisAnalyticsV2Hook]):
+    """
+    Creates an AWS Managed Service for Apache Flink application.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KinesisAnalyticsV2CreateApplicationOperator`
+
+    :param application_name: The name of application. (templated)
+    :param runtime_environment: The runtime environment for the application. (templated)
+    :param service_execution_role: The IAM role used by the application to access services. (templated)
+    :param create_application_kwargs: Create application extra properties. (templated)
+    :param application_description: A summary description of the application. (templated)
+
+    :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
+    """
+
+    aws_hook_class = KinesisAnalyticsV2Hook
+    ui_color = "#44b5e2"
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "application_name",
+        "runtime_environment",
+        "service_execution_role",
+        "create_application_kwargs",
+        "application_description",
+    )
+    template_fields_renderers: dict = {
+        "create_application_kwargs": "json",
+    }
+
+    def __init__(
+        self,
+        application_name: str,
+        runtime_environment: str,
+        service_execution_role: str,
+        create_application_kwargs: dict[str, Any] | None = None,
+        application_description: str = "Managed Service for Apache Flink application created from Airflow",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.application_name = application_name
+        self.runtime_environment = runtime_environment
+        self.service_execution_role = service_execution_role
+        self.create_application_kwargs = create_application_kwargs or {}
+        self.application_description = application_description
+
+    def execute(self, context: Context) -> dict[str, str]:
+        self.log.info("Creating AWS Managed Service for Apache Flink application %s.", self.application_name)
+        try:
+            response = self.hook.conn.create_application(
+                ApplicationName=self.application_name,
+                ApplicationDescription=self.application_description,
+                RuntimeEnvironment=self.runtime_environment,
+                ServiceExecutionRole=self.service_execution_role,
+                **self.create_application_kwargs,
+            )
+        except ClientError as error:
+            raise AirflowException(
+                f"AWS Managed Service for Apache Flink application creation failed: {error.response['Error']['Message']}"
+            )
+
+        self.log.info(
+            "AWS Managed Service for Apache Flink application created successfully %s.",
+            self.application_name,
+        )
+
+        return {"ApplicationARN": response["ApplicationDetail"]["ApplicationARN"]}
+
+
+class KinesisAnalyticsV2StartApplicationOperator(AwsBaseOperator[KinesisAnalyticsV2Hook]):
+    """
+    Starts an AWS Managed Service for Apache Flink application.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KinesisAnalyticsV2StartApplicationOperator`
+
+    :param application_name: The name of application. (templated)
+    :param run_configuration: Application properties to start Apache Flink Job. (templated)
+
+    :param wait_for_completion: Whether to wait for job to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
+    :param deferrable: If True, the operator will wait asynchronously for the job to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :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
+    """
+
+    aws_hook_class = KinesisAnalyticsV2Hook
+    ui_color = "#44b5e2"
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "application_name",
+        "run_configuration",
+    )
+    template_fields_renderers: dict = {
+        "run_configuration": "json",
+    }
+
+    def __init__(
+        self,
+        application_name: str,
+        run_configuration: dict[str, Any] | None = None,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 20,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.application_name = application_name
+        self.run_configuration = run_configuration or {}
+        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) -> dict[str, Any]:
+        msg = "AWS Managed Service for Apache Flink application"
+
+        try:
+            self.log.info("Starting %s %s.", msg, self.application_name)
+            self.hook.conn.start_application(
+                ApplicationName=self.application_name, RunConfiguration=self.run_configuration
+            )
+        except ClientError as error:
+            raise AirflowException(
+                f"Failed to start {msg} {self.application_name}: {error.response['Error']['Message']}"
+            )
+
+        describe_response = self.hook.conn.describe_application(ApplicationName=self.application_name)
+
+        if self.deferrable:
+            self.log.info("Deferring for %s to start: %s.", msg, self.application_name)
+            self.defer(
+                trigger=KinesisAnalyticsV2ApplicationOperationCompleteTrigger(
+                    application_name=self.application_name,
+                    waiter_name="application_start_complete",
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    region_name=self.region_name,
+                    verify=self.verify,
+                    botocore_config=self.botocore_config,
+                ),
+                method_name="execute_complete",
+            )
+        if self.wait_for_completion:
+            self.log.info("Waiting for %s to start: %s.", msg, self.application_name)
+
+            self.hook.get_waiter("application_start_complete").wait(
+                ApplicationName=self.application_name,
+                WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+            )
+
+        self.log.info("%s started successfully %s.", msg, self.application_name)
+
+        return {"ApplicationARN": describe_response["ApplicationDetail"]["ApplicationARN"]}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException(
+                "Error while starting AWS Managed Service for Apache Flink application: %s", event
+            )
+
+        response = self.hook.conn.describe_application(
+            ApplicationName=event["application_name"],
+        )
+
+        self.log.info(
+            "AWS Managed Service for Apache Flink application %s started successfully.",
+            event["application_name"],
+        )
+
+        return {"ApplicationARN": response["ApplicationDetail"]["ApplicationARN"]}
+
+
+class KinesisAnalyticsV2StopApplicationOperator(AwsBaseOperator[KinesisAnalyticsV2Hook]):
+    """
+    Stop an AWS Managed Service for Apache Flink application.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KinesisAnalyticsV2StopApplicationOperator`
+
+    :param application_name: The name of your application. (templated)
+    :param force: Set to true to force the application to stop. If you set Force to true, Managed Service for
+        Apache Flink stops the application without taking a snapshot. (templated)
+
+    :param wait_for_completion: Whether to wait for job to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
+    :param deferrable: If True, the operator will wait asynchronously for the job to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :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
+    """
+
+    aws_hook_class = KinesisAnalyticsV2Hook
+    ui_color = "#44b5e2"
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "application_name",
+        "force",
+    )
+
+    def __init__(
+        self,
+        application_name: str,
+        force: bool = False,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 20,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.application_name = application_name
+        self.force = force
+        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) -> dict[str, Any]:
+        msg = "AWS Managed Service for Apache Flink application"
+
+        try:
+            self.log.info("Stopping %s %s.", msg, self.application_name)
+
+            self.hook.conn.stop_application(ApplicationName=self.application_name, Force=self.force)
+        except ClientError as error:
+            raise AirflowException(
+                f"Failed to stop {msg} {self.application_name}: {error.response['Error']['Message']}"
+            )
+
+        describe_response = self.hook.conn.describe_application(ApplicationName=self.application_name)
+
+        if self.deferrable:
+            self.log.info("Deferring for %s to stop: %s.", msg, self.application_name)
+            self.defer(
+                trigger=KinesisAnalyticsV2ApplicationOperationCompleteTrigger(
+                    application_name=self.application_name,
+                    waiter_name="application_stop_complete",
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    region_name=self.region_name,
+                    verify=self.verify,
+                    botocore_config=self.botocore_config,
+                ),
+                method_name="execute_complete",
+            )
+        if self.wait_for_completion:
+            self.log.info("Waiting for %s to stop: %s.", msg, self.application_name)
+
+            self.hook.get_waiter("application_stop_complete").wait(
+                ApplicationName=self.application_name,
+                WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+            )
+
+        self.log.info("%s stopped successfully %s.", msg, self.application_name)
+
+        return {"ApplicationARN": describe_response["ApplicationDetail"]["ApplicationARN"]}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException("Error while stopping AWS Managed Service for Apache Flink application")
+
+        response = self.hook.conn.describe_application(
+            ApplicationName=event["application_name"],
+        )
+
+        self.log.info(
+            "AWS Managed Service for Apache Flink application %s stopped successfully.",
+            event["application_name"],
+        )
+
+        return {"ApplicationARN": response["ApplicationDetail"]["ApplicationARN"]}

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

@@ -36,6 +36,7 @@ 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
+from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
     from mypy_boto3_rds.type_defs import TagTypeDef
@@ -782,7 +783,7 @@ class RdsStartDbOperator(RdsBaseOperator):
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
                     response=start_db_response,
-                    db_type=RdsDbType.INSTANCE,
+                    db_type=self.db_type,
                 ),
                 method_name="execute_complete",
             )
@@ -881,12 +882,25 @@ class RdsStopDbOperator(RdsBaseOperator):
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
                     response=stop_db_response,
-                    db_type=RdsDbType.INSTANCE,
+                    db_type=self.db_type,
                 ),
                 method_name="execute_complete",
             )
         elif self.wait_for_completion:
-            self._wait_until_db_stopped()
+            waiter = self.hook.get_waiter(f"db_{self.db_type.value}_stopped")
+            waiter_key = (
+                "DBInstanceIdentifier" if self.db_type == RdsDbType.INSTANCE else "DBClusterIdentifier"
+            )
+            kwargs = {waiter_key: self.db_identifier}
+            waiter.wait(
+                WaiterConfig=prune_dict(
+                    {
+                        "Delay": self.waiter_delay,
+                        "MaxAttempts": self.waiter_max_attempts,
+                    }
+                ),
+                **kwargs,
+            )
         return json.dumps(stop_db_response, default=str)
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
@@ -915,23 +929,6 @@ class RdsStopDbOperator(RdsBaseOperator):
             response = self.hook.conn.stop_db_cluster(DBClusterIdentifier=self.db_identifier)
         return response
 
-    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",
-                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",
-                check_interval=self.waiter_delay,
-                max_attempts=self.waiter_max_attempts,
-            )
-
 
 __all__ = [
     "RdsCreateDbSnapshotOperator",

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

@@ -32,6 +32,7 @@ from airflow.providers.amazon.aws.triggers.redshift_cluster import (
     RedshiftResumeClusterTrigger,
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -256,8 +257,6 @@ class RedshiftCreateClusterOperator(BaseOperator):
             params["ClusterVersion"] = self.cluster_version
         if self.allow_version_upgrade:
             params["AllowVersionUpgrade"] = self.allow_version_upgrade
-        if self.publicly_accessible:
-            params["PubliclyAccessible"] = self.publicly_accessible
         if self.encrypted:
             params["Encrypted"] = self.encrypted
         if self.hsm_client_certificate_identifier:
@@ -287,6 +286,10 @@ class RedshiftCreateClusterOperator(BaseOperator):
         if self.default_iam_role_arn:
             params["DefaultIamRoleArn"] = self.default_iam_role_arn
 
+        # PubliclyAccessible is True by default on Redshift side, hence, we should always set it regardless
+        # of its value
+        params["PubliclyAccessible"] = self.publicly_accessible
+
         cluster = redshift_hook.create_cluster(
             self.cluster_identifier,
             self.node_type,
@@ -505,8 +508,8 @@ class RedshiftResumeClusterOperator(BaseOperator):
         aws_conn_id: str | None = "aws_default",
         wait_for_completion: bool = False,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        poll_interval: int = 10,
-        max_attempts: int = 10,
+        poll_interval: int = 30,
+        max_attempts: int = 30,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -540,38 +543,38 @@ class RedshiftResumeClusterOperator(BaseOperator):
                 else:
                     raise error
 
-        if self.deferrable:
-            cluster_state = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
-            if cluster_state == "available":
-                self.log.info("Resumed cluster successfully")
-            elif cluster_state == "deleting":
-                raise AirflowException(
-                    "Unable to resume cluster since cluster is currently in status: %s", cluster_state
-                )
+        if self.wait_for_completion:
+            if self.deferrable:
+                cluster_state = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
+                if cluster_state == "available":
+                    self.log.info("Resumed cluster successfully")
+                elif cluster_state == "deleting":
+                    raise AirflowException(
+                        "Unable to resume cluster since cluster is currently in status: %s", cluster_state
+                    )
+                else:
+                    self.defer(
+                        trigger=RedshiftResumeClusterTrigger(
+                            cluster_identifier=self.cluster_identifier,
+                            waiter_delay=self.poll_interval,
+                            waiter_max_attempts=self.max_attempts,
+                            aws_conn_id=self.aws_conn_id,
+                        ),
+                        method_name="execute_complete",
+                        # timeout is set to ensure that if a trigger dies, the timeout does not restart
+                        # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
+                        timeout=timedelta(seconds=self.max_attempts * self.poll_interval + 60),
+                    )
             else:
-                self.defer(
-                    trigger=RedshiftResumeClusterTrigger(
-                        cluster_identifier=self.cluster_identifier,
-                        waiter_delay=self.poll_interval,
-                        waiter_max_attempts=self.max_attempts,
-                        aws_conn_id=self.aws_conn_id,
-                    ),
-                    method_name="execute_complete",
-                    # timeout is set to ensure that if a trigger dies, the timeout does not restart
-                    # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
-                    timeout=timedelta(seconds=self.max_attempts * self.poll_interval + 60),
+                waiter = redshift_hook.get_waiter("cluster_resumed")
+                waiter.wait(
+                    ClusterIdentifier=self.cluster_identifier,
+                    WaiterConfig={
+                        "Delay": self.poll_interval,
+                        "MaxAttempts": self.max_attempts,
+                    },
                 )
 
-        if self.wait_for_completion:
-            waiter = redshift_hook.get_waiter("cluster_resumed")
-            waiter.wait(
-                ClusterIdentifier=self.cluster_identifier,
-                WaiterConfig={
-                    "Delay": self.poll_interval,
-                    "MaxAttempts": self.max_attempts,
-                },
-            )
-
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
         event = validate_execute_complete_event(event)
 
@@ -594,6 +597,7 @@ class RedshiftPauseClusterOperator(BaseOperator):
         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 wait_for_completion: If True, waits for the cluster to be paused. (default: False)
     :param deferrable: Run operator in the deferrable mode
     :param poll_interval: Time (in seconds) to wait between two consecutive calls to check cluster state
     :param max_attempts: Maximum number of attempts to poll the cluster
@@ -608,14 +612,16 @@ class RedshiftPauseClusterOperator(BaseOperator):
         *,
         cluster_identifier: str,
         aws_conn_id: str | None = "aws_default",
+        wait_for_completion: bool = False,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        poll_interval: int = 10,
-        max_attempts: int = 15,
+        poll_interval: int = 30,
+        max_attempts: int = 30,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.cluster_identifier = cluster_identifier
         self.aws_conn_id = aws_conn_id
+        self.wait_for_completion = wait_for_completion
         self.deferrable = deferrable
         self.max_attempts = max_attempts
         self.poll_interval = poll_interval
@@ -641,26 +647,38 @@ class RedshiftPauseClusterOperator(BaseOperator):
                     time.sleep(self._attempt_interval)
                 else:
                     raise error
-        if self.deferrable:
-            cluster_state = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
-            if cluster_state == "paused":
-                self.log.info("Paused cluster successfully")
-            elif cluster_state == "deleting":
-                raise AirflowException(
-                    f"Unable to pause cluster since cluster is currently in status: {cluster_state}"
-                )
+        if self.wait_for_completion:
+            if self.deferrable:
+                cluster_state = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
+                if cluster_state == "paused":
+                    self.log.info("Paused cluster successfully")
+                elif cluster_state == "deleting":
+                    raise AirflowException(
+                        f"Unable to pause cluster since cluster is currently in status: {cluster_state}"
+                    )
+                else:
+                    self.defer(
+                        trigger=RedshiftPauseClusterTrigger(
+                            cluster_identifier=self.cluster_identifier,
+                            waiter_delay=self.poll_interval,
+                            waiter_max_attempts=self.max_attempts,
+                            aws_conn_id=self.aws_conn_id,
+                        ),
+                        method_name="execute_complete",
+                        # timeout is set to ensure that if a trigger dies, the timeout does not restart
+                        # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
+                        timeout=timedelta(seconds=self.max_attempts * self.poll_interval + 60),
+                    )
             else:
-                self.defer(
-                    trigger=RedshiftPauseClusterTrigger(
-                        cluster_identifier=self.cluster_identifier,
-                        waiter_delay=self.poll_interval,
-                        waiter_max_attempts=self.max_attempts,
-                        aws_conn_id=self.aws_conn_id,
+                waiter = redshift_hook.get_waiter("cluster_paused")
+                waiter.wait(
+                    ClusterIdentifier=self.cluster_identifier,
+                    WaiterConfig=prune_dict(
+                        {
+                            "Delay": self.poll_interval,
+                            "MaxAttempts": self.max_attempts,
+                        }
                     ),
-                    method_name="execute_complete",
-                    # timeout is set to ensure that if a trigger dies, the timeout does not restart
-                    # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
-                    timeout=timedelta(seconds=self.max_attempts * self.poll_interval + 60),
                 )
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:

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

@@ -324,8 +324,7 @@ class S3CopyObjectOperator(BaseOperator):
         )
 
     def get_openlineage_facets_on_start(self):
-        from openlineage.client.run import Dataset
-
+        from airflow.providers.common.compat.openlineage.facet import Dataset
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         dest_bucket_name, dest_bucket_key = S3Hook.get_s3_bucket_key(
@@ -439,8 +438,7 @@ class S3CreateObjectOperator(BaseOperator):
             s3_hook.load_bytes(self.data, s3_key, s3_bucket, self.replace, self.encrypt, self.acl_policy)
 
     def get_openlineage_facets_on_start(self):
-        from openlineage.client.run import Dataset
-
+        from airflow.providers.common.compat.openlineage.facet import Dataset
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         bucket, key = S3Hook.get_s3_bucket_key(self.s3_bucket, self.s3_key, "dest_bucket", "dest_key")
@@ -546,13 +544,12 @@ class S3DeleteObjectsOperator(BaseOperator):
 
     def get_openlineage_facets_on_complete(self, task_instance):
         """Implement _on_complete because object keys are resolved in execute()."""
-        from openlineage.client.facet import (
+        from airflow.providers.common.compat.openlineage.facet import (
+            Dataset,
             LifecycleStateChange,
             LifecycleStateChangeDatasetFacet,
-            LifecycleStateChangeDatasetFacetPreviousIdentifier,
+            PreviousIdentifier,
         )
-        from openlineage.client.run import Dataset
-
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         if not self._keys:
@@ -570,7 +567,7 @@ class S3DeleteObjectsOperator(BaseOperator):
                 facets={
                     "lifecycleStateChange": LifecycleStateChangeDatasetFacet(
                         lifecycleStateChange=LifecycleStateChange.DROP.value,
-                        previousIdentifier=LifecycleStateChangeDatasetFacetPreviousIdentifier(
+                        previousIdentifier=PreviousIdentifier(
                             namespace=bucket_url,
                             name=key,
                         ),
@@ -725,8 +722,7 @@ class S3FileTransformOperator(BaseOperator):
             self.log.info("Upload successful")
 
     def get_openlineage_facets_on_start(self):
-        from openlineage.client.run import Dataset
-
+        from airflow.providers.common.compat.openlineage.facet import Dataset
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         dest_bucket_name, dest_bucket_key = S3Hook.get_s3_bucket_key(