Source code for airflow.providers.amazon.aws.hooks.quicksight

#
# 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 time

from botocore.exceptions import ClientError

from airflow import AirflowException
from airflow.compat.functools import cached_property
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.sts import StsHook


[docs]class QuickSightHook(AwsBaseHook):
    """
    Interact with Amazon QuickSight.

    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`
    """

[docs]    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
[docs]    FAILED_STATES = {"FAILED"}

    def __init__(self, *args, **kwargs):
        super().__init__(client_type="quicksight", *args, **kwargs)

    @cached_property
[docs]    def sts_hook(self):
        return StsHook(aws_conn_id=self.aws_conn_id)

[docs]    def create_ingestion(
        self,
        data_set_id: str,
        ingestion_id: str,
        ingestion_type: str,
        wait_for_completion: bool = True,
        check_interval: int = 30,
    ) -> dict:
        """
        Creates and starts a new SPICE ingestion for a dataset. Refreshes the SPICE datasets

        :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 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
        :return: Returns descriptive information about the created data ingestion
            having Ingestion ARN, HTTP status, ingestion ID and ingestion status.
        """
        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(
                DataSetId=data_set_id,
                IngestionId=ingestion_id,
                IngestionType=ingestion_type,
                AwsAccountId=aws_account_id,
            )

            if wait_for_completion:
                self.wait_for_state(
                    aws_account_id=aws_account_id,
                    data_set_id=data_set_id,
                    ingestion_id=ingestion_id,
                    target_state={"COMPLETED"},
                    check_interval=check_interval,
                )
            return create_ingestion_response
        except Exception as general_error:
            self.log.error("Failed to run Amazon QuickSight create_ingestion API, error: %s", general_error)
            raise

[docs]    def get_status(self, aws_account_id: str, data_set_id: str, ingestion_id: str) -> str:
        """
        Get the current status of QuickSight Create Ingestion API.

        :param aws_account_id: An AWS Account ID
        :param data_set_id: QuickSight Data Set ID
        :param ingestion_id: QuickSight Ingestion ID
        :return: An QuickSight Ingestion Status
        """
        try:
            describe_ingestion_response = self.get_conn().describe_ingestion(
                AwsAccountId=aws_account_id, DataSetId=data_set_id, IngestionId=ingestion_id
            )
            return describe_ingestion_response["Ingestion"]["IngestionStatus"]
        except KeyError:
            raise AirflowException("Could not get status of the Amazon QuickSight Ingestion")
        except ClientError:
            raise AirflowException("AWS request failed, check logs for more info")

[docs]    def wait_for_state(
        self,
        aws_account_id: str,
        data_set_id: str,
        ingestion_id: str,
        target_state: set,
        check_interval: int,
    ):
        """
        Check status of a QuickSight Create Ingestion API

        :param aws_account_id: An 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
        :param check_interval: the time interval in seconds which the operator
            will check the status of QuickSight Ingestion
        :return: response of describe_ingestion call after Ingestion is is done
        """
        sec = 0
        status = self.get_status(aws_account_id, data_set_id, ingestion_id)
        while status in self.NON_TERMINAL_STATES and status != target_state:
            self.log.info("Current status is %s", status)
            time.sleep(check_interval)
            sec += check_interval
            if status in self.FAILED_STATES:
                raise AirflowException("The Amazon QuickSight Ingestion failed!")
            if status == "CANCELLED":
                raise AirflowException("The Amazon QuickSight SPICE ingestion cancelled!")
            status = self.get_status(aws_account_id, data_set_id, ingestion_id)

        self.log.info("QuickSight Ingestion completed")
        return status