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

airflow/providers/amazon/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.17.0"
+__version__ = "8.18.0"
 
 try:
     from airflow import __version__ as airflow_version

airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py

@@ -80,7 +80,7 @@ def update_schema(args):
 
 
 def _get_client():
-    """Returns Amazon Verified Permissions client."""
+    """Return Amazon Verified Permissions client."""
     region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY)
     return boto3.client("verifiedpermissions", region_name=region_name)
 

airflow/providers/amazon/aws/auth_manager/views/auth.py

@@ -71,7 +71,7 @@ class AwsAuthManagerAuthenticationViews(AirflowBaseView):
     @expose("/login_callback", methods=("GET", "POST"))
     def login_callback(self):
         """
-        Callback where the user is redirected to after successful login.
+        Redirect the user to this callback after successful login.
 
         CSRF protection needs to be disabled otherwise the callback won't work.
         """

airflow/providers/amazon/aws/executors/ecs/Dockerfile

@@ -37,7 +37,7 @@ COPY <<"EOF" /install_dags_entrypoint.sh
 #!/bin/bash
 
 echo "Downloading DAGs from S3 bucket"
-aws s3 sync "$S3_URL" "$CONTAINER_DAG_PATH"
+aws s3 sync "$S3_URI" "$CONTAINER_DAG_PATH"
 
 /home/airflow/entrypoint.sh "$@"
 EOF
@@ -98,8 +98,8 @@ ENV CONTAINER_DAG_PATH=$container_dag_path
 
 
 # Use these arguments to load DAGs onto the container from S3
-ARG s3_url
-ENV S3_URL=$s3_url
+ARG s3_uri
+ENV S3_URI=$s3_uri
 # If using S3 bucket as source of DAGs, uncommenting the next ENTRYPOINT command will overwrite this one.
 ENTRYPOINT ["/usr/bin/dumb-init", "--", "/home/airflow/entrypoint.sh"]
 

airflow/providers/amazon/aws/executors/ecs/boto_schema.py

@@ -61,7 +61,7 @@ class BotoTaskSchema(Schema):
 
     @post_load
     def make_task(self, data, **kwargs):
-        """Overwrites marshmallow load() to return an instance of EcsExecutorTask instead of a dictionary."""
+        """Overwrite marshmallow load() to return an EcsExecutorTask instance instead of a dictionary."""
         # Imported here to avoid circular import.
         from airflow.providers.amazon.aws.executors.ecs.utils import EcsExecutorTask
 

airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -42,9 +42,13 @@ from airflow.providers.amazon.aws.executors.ecs.utils import (
     EcsQueuedTask,
     EcsTaskCollection,
 )
-from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry import exponential_backoff_retry
+from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry import (
+    calculate_next_attempt_delay,
+    exponential_backoff_retry,
+)
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.utils import timezone
+from airflow.utils.helpers import merge_dicts
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
@@ -108,7 +112,7 @@ class AwsEcsExecutor(BaseExecutor):
         self.run_task_kwargs = self._load_run_kwargs()
 
     def start(self):
-        """This is called by the scheduler when the Executor is being run for the first time."""
+        """Call this when the Executor is run for the first time by the scheduler."""
         check_health = conf.getboolean(
             CONFIG_GROUP_NAME, AllEcsConfigKeys.CHECK_HEALTH_ON_STARTUP, fallback=False
         )
@@ -213,7 +217,7 @@ class AwsEcsExecutor(BaseExecutor):
             self.log.exception("Failed to sync %s", self.__class__.__name__)
 
     def sync_running_tasks(self):
-        """Checks and update state on all running tasks."""
+        """Check and update state on all running tasks."""
         all_task_arns = self.active_workers.get_all_arns()
         if not all_task_arns:
             self.log.debug("No active Airflow tasks, skipping sync.")
@@ -300,7 +304,14 @@ class AwsEcsExecutor(BaseExecutor):
             )
             self.active_workers.increment_failure_count(task_key)
             self.pending_tasks.appendleft(
-                EcsQueuedTask(task_key, task_cmd, queue, exec_info, failure_count + 1)
+                EcsQueuedTask(
+                    task_key,
+                    task_cmd,
+                    queue,
+                    exec_info,
+                    failure_count + 1,
+                    timezone.utcnow() + calculate_next_attempt_delay(failure_count),
+                )
             )
         else:
             self.log.error(
@@ -313,7 +324,7 @@ class AwsEcsExecutor(BaseExecutor):
 
     def attempt_task_runs(self):
         """
-        Takes tasks from the pending_tasks queue, and attempts to find an instance to run it on.
+        Take tasks from the pending_tasks queue, and attempts to find an instance to run it on.
 
         If the launch type is EC2, this will attempt to place tasks on empty EC2 instances.  If
             there are no EC2 instances available, no task is placed and this function will be
@@ -331,6 +342,8 @@ class AwsEcsExecutor(BaseExecutor):
             exec_config = ecs_task.executor_config
             attempt_number = ecs_task.attempt_number
             _failure_reasons = []
+            if timezone.utcnow() < ecs_task.next_attempt_time:
+                continue
             try:
                 run_task_response = self._run_task(task_key, cmd, queue, exec_config)
             except NoCredentialsError:
@@ -361,6 +374,9 @@ class AwsEcsExecutor(BaseExecutor):
                 # Make sure the number of attempts does not exceed MAX_RUN_TASK_ATTEMPTS
                 if int(attempt_number) <= int(self.__class__.MAX_RUN_TASK_ATTEMPTS):
                     ecs_task.attempt_number += 1
+                    ecs_task.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
+                        attempt_number
+                    )
                     self.pending_tasks.appendleft(ecs_task)
                 else:
                     self.log.error(
@@ -393,8 +409,8 @@ class AwsEcsExecutor(BaseExecutor):
         The command and executor config will be placed in the container-override
         section of the JSON request before calling Boto3's "run_task" function.
         """
-        run_task_api = self._run_task_kwargs(task_id, cmd, queue, exec_config)
-        boto_run_task = self.ecs.run_task(**run_task_api)
+        run_task_kwargs = self._run_task_kwargs(task_id, cmd, queue, exec_config)
+        boto_run_task = self.ecs.run_task(**run_task_kwargs)
         run_task_response = BotoRunTaskSchema().load(boto_run_task)
         return run_task_response
 
@@ -402,30 +418,32 @@ class AwsEcsExecutor(BaseExecutor):
         self, task_id: TaskInstanceKey, cmd: CommandType, queue: str, exec_config: ExecutorConfigType
     ) -> dict:
         """
-        Overrides the Airflow command to update the container overrides so kwargs are specific to this task.
+        Update the Airflow command by modifying container overrides for task-specific kwargs.
 
         One last chance to modify Boto3's "run_task" kwarg params before it gets passed into the Boto3 client.
         """
-        run_task_api = deepcopy(self.run_task_kwargs)
-        container_override = self.get_container(run_task_api["overrides"]["containerOverrides"])
+        run_task_kwargs = deepcopy(self.run_task_kwargs)
+        run_task_kwargs = merge_dicts(run_task_kwargs, exec_config)
+        container_override = self.get_container(run_task_kwargs["overrides"]["containerOverrides"])
         container_override["command"] = cmd
-        container_override.update(exec_config)
 
         # Inject the env variable to configure logging for containerized execution environment
         if "environment" not in container_override:
             container_override["environment"] = []
         container_override["environment"].append({"name": "AIRFLOW_IS_EXECUTOR_CONTAINER", "value": "true"})
 
-        return run_task_api
+        return run_task_kwargs
 
     def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
         """Save the task to be executed in the next sync by inserting the commands into a queue."""
         if executor_config and ("name" in executor_config or "command" in executor_config):
             raise ValueError('Executor Config should never override "name" or "command"')
-        self.pending_tasks.append(EcsQueuedTask(key, command, queue, executor_config or {}, 1))
+        self.pending_tasks.append(
+            EcsQueuedTask(key, command, queue, executor_config or {}, 1, timezone.utcnow())
+        )
 
     def end(self, heartbeat_interval=10):
-        """Waits for all currently running tasks to end, and doesn't launch any tasks."""
+        """Wait for all currently running tasks to end, and don't launch any tasks."""
         try:
             while True:
                 self.sync()
@@ -465,8 +483,13 @@ class AwsEcsExecutor(BaseExecutor):
         return ecs_executor_run_task_kwargs
 
     def get_container(self, container_list):
-        """Searches task list for core Airflow container."""
+        """Search task list for core Airflow container."""
         for container in container_list:
-            if container["name"] == self.container_name:
-                return container
+            try:
+                if container["name"] == self.container_name:
+                    return container
+            except KeyError:
+                raise EcsExecutorException(
+                    'container "name" must be provided in "containerOverrides" configuration'
+                )
         raise KeyError(f"No such container found by container name: {self.container_name}")

airflow/providers/amazon/aws/executors/ecs/utils.py

@@ -23,6 +23,7 @@ Data classes and utility functions used by the ECS executor.
 
 from __future__ import annotations
 
+import datetime
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Dict, List
@@ -58,6 +59,7 @@ class EcsQueuedTask:
     queue: str
     executor_config: ExecutorConfigType
     attempt_number: int
+    next_attempt_time: datetime.datetime
 
 
 @dataclass
@@ -125,9 +127,9 @@ class EcsExecutorTask:
 
     def get_task_state(self) -> str:
         """
-        This is the primary logic that handles state in an ECS task.
+        Determine the state of an ECS task based on its status and other relevant attributes.
 
-        It will determine if a status is:
+        It can return one of the following statuses:
             QUEUED - Task is being provisioned.
             RUNNING - Task is launched on ECS.
             REMOVED - Task provisioning has failed for some reason. See `stopped_reason`.
@@ -171,7 +173,7 @@ class EcsTaskCollection:
         exec_config: ExecutorConfigType,
         attempt_number: int,
     ):
-        """Adds a task to the collection."""
+        """Add a task to the collection."""
         arn = task.task_arn
         self.tasks[arn] = task
         self.key_to_arn[airflow_task_key] = arn
@@ -180,7 +182,7 @@ class EcsTaskCollection:
         self.key_to_failure_counts[airflow_task_key] = attempt_number
 
     def update_task(self, task: EcsExecutorTask):
-        """Updates the state of the given task based on task ARN."""
+        """Update the state of the given task based on task ARN."""
         self.tasks[task.task_arn] = task
 
     def task_by_key(self, task_key: TaskInstanceKey) -> EcsExecutorTask:
@@ -193,7 +195,7 @@ class EcsTaskCollection:
         return self.tasks[arn]
 
     def pop_by_key(self, task_key: TaskInstanceKey) -> EcsExecutorTask:
-        """Deletes task from collection based off of Airflow Task Instance Key."""
+        """Delete task from collection based off of Airflow Task Instance Key."""
         arn = self.key_to_arn[task_key]
         task = self.tasks[arn]
         del self.key_to_arn[task_key]
@@ -225,11 +227,11 @@ class EcsTaskCollection:
         return self.key_to_task_info[task_key]
 
     def __getitem__(self, value):
-        """Gets a task by AWS ARN."""
+        """Get a task by AWS ARN."""
         return self.task_by_arn(value)
 
     def __len__(self):
-        """Determines the number of tasks in collection."""
+        """Determine the number of tasks in collection."""
         return len(self.tasks)
 
 

airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py

@@ -25,6 +25,21 @@ from airflow.utils import timezone
 log = logging.getLogger(__name__)
 
 
+def calculate_next_attempt_delay(
+    attempt_number: int,
+    max_delay: int = 60 * 2,
+    exponent_base: int = 4,
+) -> timedelta:
+    """
+    Calculate the exponential backoff (in seconds) until the next attempt.
+
+    :param attempt_number: Number of attempts since last success.
+    :param max_delay: Maximum delay in seconds between retries. Default 120.
+    :param exponent_base: Exponent base to calculate delay. Default 4.
+    """
+    return timedelta(seconds=min((exponent_base**attempt_number), max_delay))
+
+
 def exponential_backoff_retry(
     last_attempt_time: datetime,
     attempts_since_last_successful: int,
@@ -34,7 +49,7 @@ def exponential_backoff_retry(
     exponent_base: int = 4,
 ) -> None:
     """
-    Retries a callable function with exponential backoff between attempts if it raises an exception.
+    Retry a callable function with exponential backoff between attempts if it raises an exception.
 
     :param last_attempt_time: Timestamp of last attempt call.
     :param attempts_since_last_successful: Number of attempts since last success.
@@ -47,8 +62,10 @@ def exponential_backoff_retry(
         log.error("Max attempts reached. Exiting.")
         return
 
-    delay = min((exponent_base**attempts_since_last_successful), max_delay)
-    next_retry_time = last_attempt_time + timedelta(seconds=delay)
+    next_retry_time = last_attempt_time + calculate_next_attempt_delay(
+        attempt_number=attempts_since_last_successful, max_delay=max_delay, exponent_base=exponent_base
+    )
+
     current_time = timezone.utcnow()
 
     if current_time >= next_retry_time:
@@ -56,5 +73,7 @@ def exponential_backoff_retry(
             callable_function()
         except Exception:
             log.exception("Error calling %r", callable_function.__name__)
-            next_delay = min((exponent_base ** (attempts_since_last_successful + 1)), max_delay)
+            next_delay = calculate_next_attempt_delay(
+                attempts_since_last_successful + 1, max_delay, exponent_base
+            )
             log.info("Waiting for %s seconds before retrying.", next_delay)

airflow/providers/amazon/aws/hooks/athena.py

@@ -25,7 +25,7 @@ This module contains AWS Athena hook.
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Collection
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -34,6 +34,19 @@ from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 if TYPE_CHECKING:
     from botocore.paginate import PageIterator
 
+MULTI_LINE_QUERY_LOG_PREFIX = "\n\t\t"
+
+
+def query_params_to_string(params: dict[str, str | Collection[str]]) -> str:
+    result = ""
+    for key, value in params.items():
+        if key == "QueryString":
+            value = (
+                MULTI_LINE_QUERY_LOG_PREFIX + str(value).replace("\n", MULTI_LINE_QUERY_LOG_PREFIX).rstrip()
+            )
+        result += f"\t{key}: {value}\n"
+    return result.rstrip()
+
 
 class AthenaHook(AwsBaseHook):
     """Interact with Amazon Athena.
@@ -115,7 +128,7 @@ class AthenaHook(AwsBaseHook):
         if client_request_token:
             params["ClientRequestToken"] = client_request_token
         if self.log_query:
-            self.log.info("Running Query with params: %s", params)
+            self.log.info("Running Query with params:\n%s", query_params_to_string(params))
         response = self.get_conn().start_query_execution(**params)
         query_execution_id = response["QueryExecutionId"]
         self.log.info("Query execution id: %s", query_execution_id)

airflow/providers/amazon/aws/hooks/base_aws.py

@@ -29,7 +29,6 @@ import inspect
 import json
 import logging
 import os
-import warnings
 from copy import deepcopy
 from functools import cached_property, wraps
 from pathlib import Path
@@ -44,6 +43,7 @@ import tenacity
 from botocore.config import Config
 from botocore.waiter import Waiter, WaiterModel
 from dateutil.tz import tzlocal
+from deprecated import deprecated
 from slugify import slugify
 
 from airflow.configuration import conf
@@ -314,7 +314,7 @@ class BaseSessionFactory(LoggingMixin):
             idp_request_retry_kwargs = saml_config["idp_request_retry_kwargs"]
             self.log.info("idp_request_retry_kwargs= %s", idp_request_retry_kwargs)
             from requests.adapters import HTTPAdapter
-            from requests.packages.urllib3.util.retry import Retry
+            from urllib3.util.retry import Retry
 
             retry_strategy = Retry(**idp_request_retry_kwargs)
             adapter = HTTPAdapter(max_retries=retry_strategy)
@@ -1020,6 +1020,13 @@ except ImportError:
     pass
 
 
+@deprecated(
+    reason=(
+        "airflow.providers.amazon.aws.hook.base_aws.BaseAsyncSessionFactory "
+        "has been deprecated and will be removed in future"
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class BaseAsyncSessionFactory(BaseSessionFactory):
     """
     Base AWS Session Factory class to handle aiobotocore session creation.
@@ -1029,12 +1036,6 @@ class BaseAsyncSessionFactory(BaseSessionFactory):
     """
 
     def __init__(self, *args, **kwargs):
-        warnings.warn(
-            "airflow.providers.amazon.aws.hook.base_aws.BaseAsyncSessionFactory has been deprecated and "
-            "will be removed in future",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
         super().__init__(*args, **kwargs)
 
     async def get_role_credentials(self) -> dict:
@@ -1113,6 +1114,13 @@ class BaseAsyncSessionFactory(BaseSessionFactory):
         return self._get_session_with_assume_role()
 
 
+@deprecated(
+    reason=(
+        "airflow.providers.amazon.aws.hook.base_aws.AwsBaseAsyncHook "
+        "has been deprecated and will be removed in future"
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class AwsBaseAsyncHook(AwsBaseHook):
     """Interacts with AWS using aiobotocore asynchronously.
 
@@ -1129,12 +1137,6 @@ class AwsBaseAsyncHook(AwsBaseHook):
     """
 
     def __init__(self, *args, **kwargs):
-        warnings.warn(
-            "airflow.providers.amazon.aws.hook.base_aws.AwsBaseAsyncHook has been deprecated and "
-            "will be removed in future",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
         super().__init__(*args, **kwargs)
 
     def get_async_session(self) -> AioSession:

airflow/providers/amazon/aws/hooks/logs.py

@@ -17,8 +17,11 @@
 # under the License.
 from __future__ import annotations
 
+import asyncio
 import warnings
-from typing import Generator
+from typing import Any, AsyncGenerator, Generator
+
+from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -151,3 +154,84 @@ class AwsLogsHook(AwsBaseHook):
                 num_consecutive_empty_response = 0
 
             continuation_token.value = response["nextForwardToken"]
+
+    async def describe_log_streams_async(
+        self, log_group: str, stream_prefix: str, order_by: str, count: int
+    ) -> dict[str, Any] | None:
+        """Async function to get the list of log streams for the specified log group.
+
+        You can list all the log streams or filter the results by prefix. You can also control
+        how the results are ordered.
+
+        :param log_group: The name of the log group.
+        :param stream_prefix: The prefix to match.
+        :param order_by: If the value is LogStreamName , the results are ordered by log stream name.
+         If the value is LastEventTime , the results are ordered by the event time. The default value is LogStreamName.
+        :param count: The maximum number of items returned
+        """
+        async with self.async_conn as client:
+            try:
+                response: dict[str, Any] = await client.describe_log_streams(
+                    logGroupName=log_group,
+                    logStreamNamePrefix=stream_prefix,
+                    orderBy=order_by,
+                    limit=count,
+                )
+                return response
+            except ClientError as error:
+                # On the very first training job run on an account, there's no log group until
+                # the container starts logging, so ignore any errors thrown about that
+                if error.response["Error"]["Code"] == "ResourceNotFoundException":
+                    return None
+                raise error
+
+    async def get_log_events_async(
+        self,
+        log_group: str,
+        log_stream_name: str,
+        start_time: int = 0,
+        skip: int = 0,
+        start_from_head: bool = True,
+    ) -> AsyncGenerator[Any, dict[str, Any]]:
+        """Yield all the available items in a single log stream.
+
+        :param log_group: The name of the log group.
+        :param log_stream_name: The name of the specific stream.
+        :param start_time: The time stamp value to start reading the logs from (default: 0).
+        :param skip: The number of log entries to skip at the start (default: 0).
+            This is for when there are multiple entries at the same timestamp.
+        :param start_from_head: whether to start from the beginning (True) of the log or
+            at the end of the log (False).
+        """
+        next_token = None
+        while True:
+            if next_token is not None:
+                token_arg: dict[str, str] = {"nextToken": next_token}
+            else:
+                token_arg = {}
+
+            async with self.async_conn as client:
+                response = await client.get_log_events(
+                    logGroupName=log_group,
+                    logStreamName=log_stream_name,
+                    startTime=start_time,
+                    startFromHead=start_from_head,
+                    **token_arg,
+                )
+
+                events = response["events"]
+                event_count = len(events)
+
+                if event_count > skip:
+                    events = events[skip:]
+                    skip = 0
+                else:
+                    skip -= event_count
+                    events = []
+
+                for event in events:
+                    await asyncio.sleep(1)
+                    yield event
+
+                if next_token != response["nextForwardToken"]:
+                    next_token = response["nextForwardToken"]

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

@@ -0,0 +1,85 @@
+#
+# 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 airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class NeptuneHook(AwsBaseHook):
+    """
+    Interact with Amazon Neptune.
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    .. seealso::
+        - :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    AVAILABLE_STATES = ["available"]
+    STOPPED_STATES = ["stopped"]
+
+    def __init__(self, *args, **kwargs):
+        kwargs["client_type"] = "neptune"
+        super().__init__(*args, **kwargs)
+
+    def wait_for_cluster_availability(self, cluster_id: str, delay: int = 30, max_attempts: int = 60) -> str:
+        """
+        Wait for Neptune cluster to start.
+
+        :param cluster_id: The ID of the cluster to wait for.
+        :param delay: Time in seconds to delay between polls.
+        :param max_attempts: Maximum number of attempts to poll for completion.
+        :return: The status of the cluster.
+        """
+        self.get_waiter("cluster_available").wait(
+            DBClusterIdentifier=cluster_id, WaiterConfig={"Delay": delay, "MaxAttempts": max_attempts}
+        )
+
+        status = self.get_cluster_status(cluster_id)
+        self.log.info("Finished waiting for cluster %s. Status is now %s", cluster_id, status)
+
+        return status
+
+    def wait_for_cluster_stopped(self, cluster_id: str, delay: int = 30, max_attempts: int = 60) -> str:
+        """
+        Wait for Neptune cluster to stop.
+
+        :param cluster_id: The ID of the cluster to wait for.
+        :param delay: Time in seconds to delay between polls.
+        :param max_attempts: Maximum number of attempts to poll for completion.
+        :return: The status of the cluster.
+        """
+        self.get_waiter("cluster_stopped").wait(
+            DBClusterIdentifier=cluster_id, WaiterConfig={"Delay": delay, "MaxAttempts": max_attempts}
+        )
+
+        status = self.get_cluster_status(cluster_id)
+        self.log.info("Finished waiting for cluster %s. Status is now %s", cluster_id, status)
+
+        return status
+
+    def get_cluster_status(self, cluster_id: str) -> str:
+        """
+        Get the status of a Neptune cluster.
+
+        :param cluster_id: The ID of the cluster to get the status of.
+        :return: The status of the cluster.
+        """
+        return self.get_conn().describe_db_clusters(DBClusterIdentifier=cluster_id)["DBClusters"][0]["Status"]

airflow/providers/amazon/aws/hooks/quicksight.py

@@ -18,10 +18,10 @@
 from __future__ import annotations
 
 import time
-import warnings
 from functools import cached_property
 
 from botocore.exceptions import ClientError
+from deprecated import deprecated
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -172,14 +172,15 @@ class QuickSightHook(AwsBaseHook):
         return status
 
     @cached_property
-    def sts_hook(self):
-        warnings.warn(
-            f"`{type(self).__name__}.sts_hook` property is deprecated and will be removed in the future. "
+    @deprecated(
+        reason=(
+            "`QuickSightHook.sts_hook` property is deprecated and will be removed in the future. "
             "This property used for obtain AWS Account ID, "
-            f"please consider to use `{type(self).__name__}.account_id` instead",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
+            "please consider to use `QuickSightHook.account_id` instead"
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
+    def sts_hook(self):
         from airflow.providers.amazon.aws.hooks.sts import StsHook
 
         return StsHook(aws_conn_id=self.aws_conn_id)