apache-airflow-providers-amazon 8.16.0rc1__py3-none-any.whl → 8.17.0__py3-none-any.whl

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

@@ -28,7 +28,7 @@ from collections import defaultdict, deque
 from copy import deepcopy
 from typing import TYPE_CHECKING
 
-from botocore.exceptions import ClientError
+from botocore.exceptions import ClientError, NoCredentialsError
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -42,6 +42,9 @@ 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.hooks.ecs import EcsHook
+from airflow.utils import timezone
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
@@ -51,6 +54,12 @@ if TYPE_CHECKING:
         ExecutorConfigType,
     )
 
+INVALID_CREDENTIALS_EXCEPTIONS = [
+    "ExpiredTokenException",
+    "InvalidClientTokenId",
+    "UnrecognizedClientException",
+]
+
 
 class AwsEcsExecutor(BaseExecutor):
     """
@@ -91,30 +100,15 @@ class AwsEcsExecutor(BaseExecutor):
 
         self.cluster = conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.CLUSTER)
         self.container_name = conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.CONTAINER_NAME)
-        aws_conn_id = conf.get(
-            CONFIG_GROUP_NAME,
-            AllEcsConfigKeys.AWS_CONN_ID,
-            fallback=CONFIG_DEFAULTS[AllEcsConfigKeys.AWS_CONN_ID],
-        )
-        region_name = conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.REGION_NAME)
-        from airflow.providers.amazon.aws.hooks.ecs import EcsHook
+        self.attempts_since_last_successful_connection = 0
+
+        self.load_ecs_connection(check_connection=False)
+        self.IS_BOTO_CONNECTION_HEALTHY = False
 
-        self.ecs = EcsHook(aws_conn_id=aws_conn_id, region_name=region_name).conn
         self.run_task_kwargs = self._load_run_kwargs()
 
     def start(self):
-        """
-        Make a test API call to check the health of the ECS Executor.
-
-        Deliberately use an invalid task ID, some potential outcomes in order:
-          1. "AccessDeniedException" is raised if there are insufficient permissions.
-          2. "ClusterNotFoundException" is raised if permissions exist but the cluster does not.
-          3. The API responds with a failure message if the cluster is found and there
-             are permissions, but the cluster itself has issues.
-          4. "InvalidParameterException" is raised if the permissions and cluster exist but the task does not.
-
-        The last one is considered a success state for the purposes of this check.
-        """
+        """This is called by the scheduler when the Executor is being run for the first time."""
         check_health = conf.getboolean(
             CONFIG_GROUP_NAME, AllEcsConfigKeys.CHECK_HEALTH_ON_STARTUP, fallback=False
         )
@@ -123,7 +117,25 @@ class AwsEcsExecutor(BaseExecutor):
             return
 
         self.log.info("Starting ECS Executor and determining health...")
+        try:
+            self.check_health()
+        except AirflowException:
+            self.log.error("Stopping the Airflow Scheduler from starting until the issue is resolved.")
+            raise
+
+    def check_health(self):
+        """
+        Make a test API call to check the health of the ECS Executor.
 
+        Deliberately use an invalid task ID, some potential outcomes in order:
+          1. `AccessDeniedException` is raised if there are insufficient permissions.
+          2. `ClusterNotFoundException` is raised if permissions exist but the cluster does not.
+          3. The API responds with a failure message if the cluster is found and there
+             are permissions, but the cluster itself has issues.
+          4. `InvalidParameterException` is raised if the permissions and cluster exist but the task does not.
+
+        The last one is considered a success state for the purposes of this check.
+        """
         success_status = "succeeded."
         status = success_status
 
@@ -151,18 +163,50 @@ class AwsEcsExecutor(BaseExecutor):
         finally:
             msg_prefix = "ECS Executor health check has %s"
             if status == success_status:
+                self.IS_BOTO_CONNECTION_HEALTHY = True
                 self.log.info(msg_prefix, status)
             else:
                 msg_error_suffix = (
-                    "The ECS executor will not be able to run Airflow tasks until the issue is addressed. "
-                    "Stopping the Airflow Scheduler from starting until the issue is resolved."
+                    "The ECS executor will not be able to run Airflow tasks until the issue is addressed."
                 )
                 raise AirflowException(msg_prefix % status + msg_error_suffix)
 
+    def load_ecs_connection(self, check_connection: bool = True):
+        self.log.info("Loading Connection information")
+        aws_conn_id = conf.get(
+            CONFIG_GROUP_NAME,
+            AllEcsConfigKeys.AWS_CONN_ID,
+            fallback=CONFIG_DEFAULTS[AllEcsConfigKeys.AWS_CONN_ID],
+        )
+        region_name = conf.get(CONFIG_GROUP_NAME, AllEcsConfigKeys.REGION_NAME)
+        self.ecs = EcsHook(aws_conn_id=aws_conn_id, region_name=region_name).conn
+        self.attempts_since_last_successful_connection += 1
+        self.last_connection_reload = timezone.utcnow()
+
+        if check_connection:
+            self.check_health()
+            self.attempts_since_last_successful_connection = 0
+
     def sync(self):
+        if not self.IS_BOTO_CONNECTION_HEALTHY:
+            exponential_backoff_retry(
+                self.last_connection_reload,
+                self.attempts_since_last_successful_connection,
+                self.load_ecs_connection,
+            )
+            if not self.IS_BOTO_CONNECTION_HEALTHY:
+                return
         try:
             self.sync_running_tasks()
             self.attempt_task_runs()
+        except (ClientError, NoCredentialsError) as error:
+            error_code = error.response["Error"]["Code"]
+            if error_code in INVALID_CREDENTIALS_EXCEPTIONS:
+                self.IS_BOTO_CONNECTION_HEALTHY = False
+                self.log.warning(
+                    f"AWS credentials are either missing or expired: {error}.\nRetrying connection"
+                )
+
         except Exception:
             # We catch any and all exceptions because otherwise they would bubble
             # up and kill the scheduler process
@@ -176,6 +220,7 @@ class AwsEcsExecutor(BaseExecutor):
             return
 
         describe_tasks_response = self.__describe_tasks(all_task_arns)
+
         self.log.debug("Active Workers: %s", describe_tasks_response)
 
         if describe_tasks_response["failures"]:
@@ -288,6 +333,15 @@ class AwsEcsExecutor(BaseExecutor):
             _failure_reasons = []
             try:
                 run_task_response = self._run_task(task_key, cmd, queue, exec_config)
+            except NoCredentialsError:
+                self.pending_tasks.appendleft(ecs_task)
+                raise
+            except ClientError as e:
+                error_code = e.response["Error"]["Code"]
+                if error_code in INVALID_CREDENTIALS_EXCEPTIONS:
+                    self.pending_tasks.appendleft(ecs_task)
+                    raise
+                _failure_reasons.append(str(e))
             except Exception as e:
                 # Failed to even get a response back from the Boto3 API or something else went
                 # wrong.  For any possible failure we want to add the exception reasons to the

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

@@ -40,6 +40,7 @@ from airflow.providers.amazon.aws.executors.ecs.utils import (
     camelize_dict_keys,
     parse_assign_public_ip,
 )
+from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.utils.helpers import prune_dict
 
 
@@ -60,6 +61,22 @@ def build_task_kwargs() -> dict:
     task_kwargs = _fetch_config_values()
     task_kwargs.update(_fetch_templated_kwargs())
 
+    has_launch_type: bool = "launch_type" in task_kwargs
+    has_capacity_provider: bool = "capacity_provider_strategy" in task_kwargs
+
+    if has_capacity_provider and has_launch_type:
+        raise ValueError(
+            "capacity_provider_strategy and launch_type are mutually exclusive, you can not provide both."
+        )
+    elif "cluster" in task_kwargs and not (has_capacity_provider or has_launch_type):
+        # Default API behavior if neither is provided is to fall back on the default capacity
+        # provider if it exists. Since it is not a required value, check if there is one
+        # before using it, and if there is not then use the FARGATE launch_type as
+        # the final fallback.
+        cluster = EcsHook().conn.describe_clusters(clusters=[task_kwargs["cluster"]])["clusters"][0]
+        if not cluster.get("defaultCapacityProviderStrategy"):
+            task_kwargs["launch_type"] = "FARGATE"
+
     # There can only be 1 count of these containers
     task_kwargs["count"] = 1  # type: ignore
     # There could be a generic approach to the below, but likely more convoluted then just manually ensuring

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

@@ -44,7 +44,6 @@ CONFIG_DEFAULTS = {
     "conn_id": "aws_default",
     "max_run_task_attempts": "3",
     "assign_public_ip": "False",
-    "launch_type": "FARGATE",
     "platform_version": "LATEST",
     "check_health_on_startup": "True",
 }
@@ -81,6 +80,7 @@ class RunTaskKwargsConfigKeys(BaseConfigKeys):
     """Keys loaded into the config which are valid ECS run_task kwargs."""
 
     ASSIGN_PUBLIC_IP = "assign_public_ip"
+    CAPACITY_PROVIDER_STRATEGY = "capacity_provider_strategy"
     CLUSTER = "cluster"
     LAUNCH_TYPE = "launch_type"
     PLATFORM_VERSION = "platform_version"

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

@@ -0,0 +1,16 @@
+# 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.

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

@@ -0,0 +1,60 @@
+# 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 logging
+from datetime import datetime, timedelta
+from typing import Callable
+
+from airflow.utils import timezone
+
+log = logging.getLogger(__name__)
+
+
+def exponential_backoff_retry(
+    last_attempt_time: datetime,
+    attempts_since_last_successful: int,
+    callable_function: Callable,
+    max_delay: int = 60 * 2,
+    max_attempts: int = -1,
+    exponent_base: int = 4,
+) -> None:
+    """
+    Retries 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.
+    :param callable_function: Callable function that will be called if enough time has passed.
+    :param max_delay: Maximum delay in seconds between retries. Default 120.
+    :param max_attempts: Maximum number of attempts before giving up. Default -1 (no limit).
+    :param exponent_base: Exponent base to calculate delay. Default 4.
+    """
+    if max_attempts != -1 and attempts_since_last_successful >= max_attempts:
+        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)
+    current_time = timezone.utcnow()
+
+    if current_time >= next_retry_time:
+        try:
+            callable_function()
+        except Exception:
+            log.exception("Error calling %r", callable_function.__name__)
+            next_delay = min((exponent_base ** (attempts_since_last_successful + 1)), max_delay)
+            log.info("Waiting for %s seconds before retrying.", next_delay)

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

@@ -0,0 +1,168 @@
+# 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 json
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+import pyathena
+from sqlalchemy.engine.url import URL
+
+from airflow.exceptions import AirflowException, AirflowNotFoundException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
+from airflow.providers.common.sql.hooks.sql import DbApiHook
+
+if TYPE_CHECKING:
+    from pyathena.connection import Connection as AthenaConnection
+
+
+class AthenaSQLHook(AwsBaseHook, DbApiHook):
+    """Interact with Amazon Athena.
+
+    Provide wrapper around PyAthena library.
+
+    :param athena_conn_id: :ref:`Amazon Athena Connection <howto/connection:athena>`.
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    You can specify ``driver`` in ``extra`` of your connection in order to use
+    a different driver than the default ``rest``.
+
+    Also, aws_domain could be specified in ``extra`` of your connection.
+
+    PyAthena and AWS Authentication parameters could be passed in extra field of ``athena_conn_id`` connection.
+
+    Passing authentication parameters in ``athena_conn_id`` will override those in ``aws_conn_id``.
+
+    .. seealso::
+        :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+
+    .. note::
+        get_uri() depends on SQLAlchemy and PyAthena.
+    """
+
+    conn_name_attr = "athena_conn_id"
+    default_conn_name = "athena_default"
+    conn_type = "athena"
+    hook_name = "Amazon Athena"
+    supports_autocommit = True
+
+    def __init__(self, athena_conn_id: str = default_conn_name, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.athena_conn_id = athena_conn_id
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Return custom UI field behaviour for AWS Athena Connection."""
+        return {
+            "hidden_fields": ["host", "port"],
+            "relabeling": {
+                "login": "AWS Access Key ID",
+                "password": "AWS Secret Access Key",
+            },
+            "placeholders": {
+                "login": "AKIAIOSFODNN7EXAMPLE",
+                "password": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+                "extra": json.dumps(
+                    {
+                        "aws_domain": "amazonaws.com",
+                        "driver": "rest",
+                        "s3_staging_dir": "s3://bucket_name/staging/",
+                        "work_group": "primary",
+                        "region_name": "us-east-1",
+                        "session_kwargs": {"profile_name": "default"},
+                        "config_kwargs": {"retries": {"mode": "standard", "max_attempts": 10}},
+                        "role_arn": "arn:aws:iam::123456789098:role/role-name",
+                        "assume_role_method": "assume_role",
+                        "assume_role_kwargs": {"RoleSessionName": "airflow"},
+                        "aws_session_token": "AQoDYXdzEJr...EXAMPLETOKEN",
+                        "endpoint_url": "http://localhost:4566",
+                    },
+                    indent=2,
+                ),
+            },
+        }
+
+    @cached_property
+    def conn_config(self) -> AwsConnectionWrapper:
+        """Get the Airflow Connection object and wrap it in helper (cached)."""
+        athena_conn = self.get_connection(self.athena_conn_id)
+        if self.aws_conn_id:
+            try:
+                connection = self.get_connection(self.aws_conn_id)
+                connection.login = athena_conn.login
+                connection.password = athena_conn.password
+                connection.schema = athena_conn.schema
+                connection.set_extra(json.dumps({**athena_conn.extra_dejson, **connection.extra_dejson}))
+            except AirflowNotFoundException:
+                connection = athena_conn
+                connection.conn_type = "aws"
+                self.log.warning(
+                    "Unable to find AWS Connection ID '%s', switching to empty.", self.aws_conn_id
+                )
+
+        return AwsConnectionWrapper(
+            conn=connection, region_name=self._region_name, botocore_config=self._config, verify=self._verify
+        )
+
+    @property
+    def conn(self) -> AwsConnectionWrapper:
+        """Get Aws Connection Wrapper object."""
+        return self.conn_config
+
+    def _get_conn_params(self) -> dict[str, str | None]:
+        """Retrieve connection parameters."""
+        if not self.conn.region_name:
+            raise AirflowException("region_name must be specified in the connection's extra")
+
+        return dict(
+            driver=self.conn.extra_dejson.get("driver", "rest"),
+            schema_name=self.conn.schema,
+            region_name=self.conn.region_name,
+            aws_domain=self.conn.extra_dejson.get("aws_domain", "amazonaws.com"),
+        )
+
+    def get_uri(self) -> str:
+        """Overridden to use the Athena dialect as driver name."""
+        conn_params = self._get_conn_params()
+        creds = self.get_credentials(region_name=conn_params["region_name"])
+
+        return URL.create(
+            f'awsathena+{conn_params["driver"]}',
+            username=creds.access_key,
+            password=creds.secret_key,
+            host=f'athena.{conn_params["region_name"]}.{conn_params["aws_domain"]}',
+            port=443,
+            database=conn_params["schema_name"],
+            query={"aws_session_token": creds.token, **self.conn.extra_dejson},
+        )
+
+    def get_conn(self) -> AthenaConnection:
+        """Get a ``pyathena.Connection`` object."""
+        conn_params = self._get_conn_params()
+
+        conn_kwargs: dict = {
+            "schema_name": conn_params["schema_name"],
+            "region_name": conn_params["region_name"],
+            "session": self.get_session(region_name=conn_params["region_name"]),
+            **self.conn.extra_dejson,
+        }
+
+        return pyathena.connect(**conn_kwargs)

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

@@ -629,6 +629,20 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         """Verify or not SSL certificates boto3 client/resource read-only property."""
         return self.conn_config.verify
 
+    @cached_property
+    def account_id(self) -> str:
+        """Return associated AWS Account ID."""
+        return (
+            self.get_session(region_name=self.region_name)
+            .client(
+                service_name="sts",
+                endpoint_url=self.conn_config.get_service_endpoint_url("sts"),
+                config=self.config,
+                verify=self.verify,
+            )
+            .get_caller_identity()["Account"]
+        )
+
     def get_session(self, region_name: str | None = None, deferrable: bool = False) -> boto3.session.Session:
         """Get the underlying boto3.session.Session(region_name=region_name)."""
         return SessionFactory(

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

@@ -18,13 +18,13 @@
 from __future__ import annotations
 
 import time
+import warnings
 from functools import cached_property
 
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.hooks.sts import StsHook
 
 
 class QuickSightHook(AwsBaseHook):
@@ -46,10 +46,6 @@ class QuickSightHook(AwsBaseHook):
     def __init__(self, *args, **kwargs):
         super().__init__(client_type="quicksight", *args, **kwargs)
 
-    @cached_property
-    def sts_hook(self):
-        return StsHook(aws_conn_id=self.aws_conn_id)
-
     def create_ingestion(
         self,
         data_set_id: str,
@@ -57,6 +53,7 @@ class QuickSightHook(AwsBaseHook):
         ingestion_type: str,
         wait_for_completion: bool = True,
         check_interval: int = 30,
+        aws_account_id: str | None = None,
     ) -> dict:
         """
         Create and start a new SPICE ingestion for a dataset; refresh the SPICE datasets.
@@ -66,18 +63,18 @@ class QuickSightHook(AwsBaseHook):
 
         :param data_set_id:  ID of the dataset used in the ingestion.
         :param ingestion_id: ID for the ingestion.
-        :param ingestion_type: Type of ingestion . "INCREMENTAL_REFRESH"|"FULL_REFRESH"
+        :param ingestion_type: Type of ingestion: "INCREMENTAL_REFRESH"|"FULL_REFRESH"
         :param wait_for_completion: if the program should keep running until job finishes
         :param check_interval: the time interval in seconds which the operator
             will check the status of QuickSight Ingestion
+        :param aws_account_id: An AWS Account ID, if set to ``None`` then use associated AWS Account ID.
         :return: Returns descriptive information about the created data ingestion
             having Ingestion ARN, HTTP status, ingestion ID and ingestion status.
         """
+        aws_account_id = aws_account_id or self.account_id
         self.log.info("Creating QuickSight Ingestion for data set id %s.", data_set_id)
-        quicksight_client = self.get_conn()
         try:
-            aws_account_id = self.sts_hook.get_account_number()
-            create_ingestion_response = quicksight_client.create_ingestion(
+            create_ingestion_response = self.conn.create_ingestion(
                 DataSetId=data_set_id,
                 IngestionId=ingestion_id,
                 IngestionType=ingestion_type,
@@ -97,20 +94,21 @@ class QuickSightHook(AwsBaseHook):
             self.log.error("Failed to run Amazon QuickSight create_ingestion API, error: %s", general_error)
             raise
 
-    def get_status(self, aws_account_id: str, data_set_id: str, ingestion_id: str) -> str:
+    def get_status(self, aws_account_id: str | None, data_set_id: str, ingestion_id: str) -> str:
         """
         Get the current status of QuickSight Create Ingestion API.
 
         .. seealso::
             - :external+boto3:py:meth:`QuickSight.Client.describe_ingestion`
 
-        :param aws_account_id: An AWS Account ID
+        :param aws_account_id: An AWS Account ID, if set to ``None`` then use associated AWS Account ID.
         :param data_set_id: QuickSight Data Set ID
         :param ingestion_id: QuickSight Ingestion ID
         :return: An QuickSight Ingestion Status
         """
+        aws_account_id = aws_account_id or self.account_id
         try:
-            describe_ingestion_response = self.get_conn().describe_ingestion(
+            describe_ingestion_response = self.conn.describe_ingestion(
                 AwsAccountId=aws_account_id, DataSetId=data_set_id, IngestionId=ingestion_id
             )
             return describe_ingestion_response["Ingestion"]["IngestionStatus"]
@@ -119,17 +117,19 @@ class QuickSightHook(AwsBaseHook):
         except ClientError as e:
             raise AirflowException(f"AWS request failed: {e}")
 
-    def get_error_info(self, aws_account_id: str, data_set_id: str, ingestion_id: str) -> dict | None:
+    def get_error_info(self, aws_account_id: str | None, data_set_id: str, ingestion_id: str) -> dict | None:
         """
         Get info about the error if any.
 
-        :param aws_account_id: An AWS Account ID
+        :param aws_account_id: An AWS Account ID, if set to ``None`` then use associated AWS Account ID.
         :param data_set_id: QuickSight Data Set ID
         :param ingestion_id: QuickSight Ingestion ID
         :return: Error info dict containing the error type (key 'Type') and message (key 'Message')
             if available. Else, returns None.
         """
-        describe_ingestion_response = self.get_conn().describe_ingestion(
+        aws_account_id = aws_account_id or self.account_id
+
+        describe_ingestion_response = self.conn.describe_ingestion(
             AwsAccountId=aws_account_id, DataSetId=data_set_id, IngestionId=ingestion_id
         )
         # using .get() to get None if the key is not present, instead of an exception.
@@ -137,7 +137,7 @@ class QuickSightHook(AwsBaseHook):
 
     def wait_for_state(
         self,
-        aws_account_id: str,
+        aws_account_id: str | None,
         data_set_id: str,
         ingestion_id: str,
         target_state: set,
@@ -146,7 +146,7 @@ class QuickSightHook(AwsBaseHook):
         """
         Check status of a QuickSight Create Ingestion API.
 
-        :param aws_account_id: An AWS Account ID
+        :param aws_account_id: An AWS Account ID, if set to ``None`` then use associated AWS Account ID.
         :param data_set_id: QuickSight Data Set ID
         :param ingestion_id: QuickSight Ingestion ID
         :param target_state: Describes the QuickSight Job's Target State
@@ -154,6 +154,8 @@ class QuickSightHook(AwsBaseHook):
             will check the status of QuickSight Ingestion
         :return: response of describe_ingestion call after Ingestion is done
         """
+        aws_account_id = aws_account_id or self.account_id
+
         while True:
             status = self.get_status(aws_account_id, data_set_id, ingestion_id)
             self.log.info("Current status is %s", status)
@@ -168,3 +170,16 @@ class QuickSightHook(AwsBaseHook):
 
         self.log.info("QuickSight Ingestion completed")
         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. "
+            "This property used for obtain AWS Account ID, "
+            f"please consider to use `{type(self).__name__}.account_id` instead",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+        from airflow.providers.amazon.aws.hooks.sts import StsHook
+
+        return StsHook(aws_conn_id=self.aws_conn_id)