You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2022/02/28 09:02:47 UTC

[GitHub] [airflow] hsrocks opened a new pull request #21863: Add Quicksight create ingestion Hook and Operator

hsrocks opened a new pull request #21863:
URL: https://github.com/apache/airflow/pull/21863


   This PR creates the Hook and Operator for Quicksight Create ingestion API which is helpful for event driven Quicksight Spice data refresh. 
   
   
   <!--
   Thank you for contributing! Please make sure that your code changes
   are covered with tests. And in case of new features or big changes
   remember to adjust the documentation.
   
   Feel free to ping committers for the review!
   
   In case of existing issue, reference it using one of the following:
   
   closes: #ISSUE
   related: #ISSUE
   
   How to write a good git commit message:
   http://chris.beams.io/posts/git-commit/
   -->
   
   ---
   **^ Add meaningful description above**
   
   Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
   In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
   In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
   In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/main/UPDATING.md).
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] eladkal commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r826376515



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,141 @@
+#
+# 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.
+
+import time
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,

Review comment:
       This is a parameter we extract from AWS connection object are we not? (it's used to construct role_arn)




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] eladkal commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821767350



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,143 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)

Review comment:
       Ah I see now.
   Other hooks just did it via kwargs
   https://github.com/apache/airflow/blob/602abe8394fafe7de54df7e73af56de848cdf617/airflow/providers/amazon/aws/hooks/kinesis.py#L41
   
   https://github.com/apache/airflow/blob/037865970ba628265afd44fe2bddbc6b15996fa6/airflow/providers/amazon/aws/hooks/rds.py#L48




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] potiuk commented on pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
potiuk commented on pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#issuecomment-1061233328


   You need to rebase.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] eladkal commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821011801



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,152 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    non_terminal_states = {"INITIALIZED", "QUEUED", "RUNNING"}
+    failed_states = {"FAILED"}

Review comment:
       Constants should be capitalized see: https://www.python.org/dev/peps/pep-0008/#constants
   
   You can also see example for this in EmrHook:
   https://github.com/apache/airflow/blob/6126c4e40f36f76d294cfd72f84cb83d97d622a5/airflow/providers/amazon/aws/hooks/emr.py#L106-L116

##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,152 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    non_terminal_states = {"INITIALIZED", "QUEUED", "RUNNING"}
+    failed_states = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,
+        ingestion_type: str,
+        wait_for_completion: bool = True,
+        check_interval: int = 30,
+        max_ingestion_time: Optional[int] = None,
+    ):
+        """
+        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 aws_account_id: Amazon Web Services account ID.
+        :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 max_ingestion_time: the maximum ingestion time in seconds. If QuickSight ingestion runs
+         longer than this will fail. Setting this to None implies no timeout for Ingestion.
+         :return: Returns descriptive information about the created data ingestion
+         having Ingestion ARN, HTTP status,
+         ingestion ID and ingestion status.
+        :rtype: Dict
+        """
+        self.log.info("Creating QuickSight Ingestion for data set id %s.", data_set_id)
+        quicksight_client = self.get_conn()
+        try:
+            create_ingestion_response = quicksight_client.create_ingestion(
+                DataSetId=data_set_id,
+                IngestionId=ingestion_id,
+                AwsAccountId=aws_account_id,
+                IngestionType=ingestion_type,
+            )
+            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,
+                    max_ingestion_time=max_ingestion_time,
+                )
+            return create_ingestion_response
+        except Exception as general_error:
+            self.log.error("Failed to run QuickSight create_ingestion API, error: %s", general_error)
+            raise
+
+    def get_status(self, aws_account_id: str, data_set_id: str, ingestion_id: 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
+        :rtype: str
+        """
+        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 QuickSight Ingestion")
+        except ClientError:
+            raise AirflowException("AWS request failed, check logs for more info")
+
+    def wait_for_state(
+        self,
+        aws_account_id: str,
+        data_set_id: str,
+        ingestion_id: str,
+        target_state: set,
+        check_interval: int,
+        max_ingestion_time: Optional[int] = None,
+    ):
+        """
+        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
+        :param max_ingestion_time: the maximum ingestion time in seconds. QuickSight API if
+         run longer than this will fail. Setting this to None implies no timeout.
+        :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("QuickSight SPICE ingestion failed")
+            if status == "CANCELLED":
+                raise AirflowException("QuickSight SPICE ingestion cancelled")
+            if max_ingestion_time and sec > max_ingestion_time:
+                raise AirflowException(f"QuickSight Ingestion took more than {max_ingestion_time} seconds")

Review comment:
       I don't think we need `max_ingestion_time`
   Operators have `execution_timeout` :
   https://github.com/apache/airflow/blob/6126c4e40f36f76d294cfd72f84cb83d97d622a5/airflow/models/baseoperator.py#L722
   
   which takes care of the need to fail a task if it didn't finish on specified timedelta.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] potiuk commented on pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
potiuk commented on pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#issuecomment-1061583458


   Reopened to rebuild. it does happen sometimes.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r823794724



##########
File path: airflow/providers/amazon/aws/example_dags/example_quicksight.py
##########
@@ -0,0 +1,42 @@
+# 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 datetime import datetime
+
+from airflow import DAG
+from airflow.providers.amazon.aws.operators.quicksight import QuickSightCreateIngestionOperator
+
+DATA_SET_ID = "DemoDataSet_Test"
+INGESTION_ID = "DemoDataSet_Ingestion_Test"
+AWS_ACCOUNT_ID = "123456789012"
+
+# [START howto_operator_quicksight]
+with DAG(
+    "sample_quicksight_dag",
+    schedule_interval=None,
+    start_date=datetime(2022, 2, 21),
+    catchup=False,
+) as dag:
+    quicksight_create_ingestion = QuickSightCreateIngestionOperator(
+        data_set_id=DATA_SET_ID,
+        ingestion_id=INGESTION_ID,
+        aws_account_id=AWS_ACCOUNT_ID,
+        task_id="sample_quicksight_dag",
+    )
+
+    quicksight_create_ingestion

Review comment:
       Done :) . Please check. Thanks!




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] ferruzzi commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r827405626



##########
File path: airflow/providers/amazon/aws/example_dags/example_quicksight.py
##########
@@ -0,0 +1,67 @@
+# 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 datetime import datetime
+
+from airflow import DAG
+from airflow.providers.amazon.aws.operators.quicksight import QuickSightCreateIngestionOperator
+from airflow.providers.amazon.aws.sensors.quicksight import QuickSightSensor
+
+DATA_SET_ID = "DemoDataSet_Test"
+INGESTION_WAITING_ID = "DemoDataSet_Ingestion_Waiting_Test"
+INGESTION_NO_WAITING_ID = "DemoDataSet_Ingestion_No_Waiting_Test"
+AWS_ACCOUNT_ID = "123456789012"

Review comment:
       ```suggestion
   
   DATA_SET_ID = os.getenv("DATA_SET_ID", "DemoDataSet_Test")
   INGESTION_WAITING_ID = os.getenv("INGESTION_WAITING_ID", "DemoDataSet_Ingestion_Waiting_Test")
   INGESTION_NO_WAITING_ID = os.getenv("INGESTION_NO_WAITING_ID", "DemoDataSet_Ingestion_No_Waiting_Test")
   AWS_ACCOUNT_ID = os.getenv("AWS_ACCOUNT_ID", "123456789012")
   ```
   
   @o-nikolas is looking to add System Tests in the near future and this should make his job easier.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821034236



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,152 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    non_terminal_states = {"INITIALIZED", "QUEUED", "RUNNING"}
+    failed_states = {"FAILED"}

Review comment:
       Thanks I agree! I have made constant as caps in other file but missed it over here




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#issuecomment-1061669791


   > Static/docs failures.
   
   Let me fix them and check on my local again before pushing . Appreciate the support  @potiuk 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821032876



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,145 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,
+        ingestion_type: str,
+        wait_for_completion: bool = True,
+        check_interval: int = 30,
+        max_ingestion_time: Optional[int] = None,
+    ):
+        """
+        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 aws_account_id: Amazon Web Services account ID.
+        :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 max_ingestion_time: the maximum ingestion time in seconds. If QuickSight ingestion runs
+         longer than this will fail. Setting this to None implies no timeout for Ingestion.
+         :return: Returns descriptive information about the created data ingestion
+         having Ingestion ARN, HTTP status,
+         ingestion ID and ingestion status.
+        :rtype: Dict
+        """
+        self.log.info('Creating QuickSight Ingestion for data set id %s.', data_set_id)
+        quicksight_client = self.get_conn()
+        try:
+            create_ingestion_response = quicksight_client.create_ingestion(
+                DataSetId=data_set_id,
+                IngestionId=ingestion_id,
+                AwsAccountId=aws_account_id,
+                IngestionType=ingestion_type,
+            )
+            if wait_for_completion:
+                self.check_status(
+                    aws_account_id=aws_account_id,
+                    ingestion_id=ingestion_id,
+                    data_set_id=data_set_id,
+                    check_interval=check_interval,
+                    max_ingestion_time=max_ingestion_time,
+                )
+            return create_ingestion_response
+        except Exception as general_error:
+            self.log.error("Failed to run QuickSight create_ingestion API, error: %s", general_error)
+            raise
+
+    def check_status(
+        self,
+        aws_account_id: str,
+        ingestion_id: str,
+        data_set_id: str,
+        check_interval: int,
+        max_ingestion_time: Optional[int] = None,

Review comment:
       Thanks ! Have done the changes . Appreciate the suggestion




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r826550067



##########
File path: airflow/providers/amazon/aws/sensors/quicksight.py
##########
@@ -0,0 +1,80 @@
+#
+# 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 typing import TYPE_CHECKING, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class QuickSightSensor(BaseSensorOperator):
+    """
+    Watches for the status of a QuickSight Ingestion.
+
+    :param aws_account_id: An AWS Account ID
+    :param data_set_id: QuickSight Data Set ID
+    :param ingestion_id: QuickSight Ingestion ID
+    :param aws_conn_id: aws connection to use, defaults to "aws_default"
+    """
+
+    def __init__(
+        self,
+        *,
+        aws_account_id: str,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.aws_account_id = aws_account_id
+        self.data_set_id = data_set_id
+        self.aws_account_id = aws_account_id
+        self.ingestion_id = ingestion_id
+        self.aws_conn_id = aws_conn_id
+        self.success_status = "COMPLETED"
+        self.errored_statuses = ("FAILED", "CANCELLED")
+        self.hook: Optional[QuickSightHook] = None
+
+    def poke(self, context: "Context"):
+        """
+        Pokes until the QuickSight Ingestion has successfully finished.
+
+        :param context: The task context during execution.
+        :return: True if it COMPLETED and False if not.
+        :rtype: bool
+        """
+        hook = self.get_hook()
+        self.log.info("Poking for Amazon QuickSight Ingestion ID: %s", self.ingestion_id)
+        quicksight_ingestion_state = hook.get_status(self.aws_account_id, self.data_set_id, self.ingestion_id)
+        self.log.info("QuickSight Status: %s", quicksight_ingestion_state)
+        if quicksight_ingestion_state in self.errored_statuses:
+            raise AirflowException("The QuickSight Ingestion failed!")

Review comment:
       Sure @ferruzzi :) Thanks a lot !! Will standardise and ensure that message or doc comments are consistent




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] ferruzzi commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r826433470



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,141 @@
+#
+# 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.
+
+import time
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,

Review comment:
       It seems odd that all of their client calls for Quicksight take the account ID as a parameter.  If you don't want to take it as an input, you get get it from [STS.get_caller_identity(): ]( https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.get_caller_identity) and I'd suggest fetching it once programmatically then storing it as a field in the QuicksightHook rather than passing it all over the place.
   
   This should do the trick inline:
   
   ```
   import boto3
   boto3.client('sts').get_caller_identity()['Account']
   ``` 
   
   But if you want to follow Airflow structure and add an STS hook it would look like this:
   
   ```
   class StsHook(AwsBaseHook):
       def __init__(self) -> None:
           super().__init__(client_type='sts')
   
       def get_account_number(self) -> Dict:
           return self.conn.get_caller_identity()['Account']
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] eladkal commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821586952



##########
File path: airflow/providers/amazon/aws/example_dags/example_quicksight.py
##########
@@ -0,0 +1,42 @@
+# 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 datetime import datetime
+
+from airflow import DAG
+from airflow.providers.amazon.aws.operators.quicksight import QuickSightCreateIngestionOperator
+
+DATA_SET_ID = "DemoDataSet_Test"
+INGESTION_ID = "DemoDataSet_Ingestion_Test"
+AWS_ACCOUNT_ID = "123456789012"
+
+# [START howto_operator_quicksight]
+with DAG(
+    "sample_quicksight_dag",
+    schedule_interval=None,
+    start_date=datetime(2022, 2, 21),
+    catchup=False,
+) as dag:
+    quicksight_create_ingestion = QuickSightCreateIngestionOperator(
+        data_set_id=DATA_SET_ID,
+        ingestion_id=INGESTION_ID,
+        aws_account_id=AWS_ACCOUNT_ID,
+        task_id="sample_quicksight_dag",
+    )
+
+    quicksight_create_ingestion

Review comment:
       probably can also be useful to add the sensor to the example dag

##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,143 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)

Review comment:
       Why do you need to specify `client_type` here?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] eladkal commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
eladkal commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r820485967



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,145 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,
+        ingestion_type: str,
+        wait_for_completion: bool = True,
+        check_interval: int = 30,
+        max_ingestion_time: Optional[int] = None,
+    ):
+        """
+        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 aws_account_id: Amazon Web Services account ID.
+        :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 max_ingestion_time: the maximum ingestion time in seconds. If QuickSight ingestion runs
+         longer than this will fail. Setting this to None implies no timeout for Ingestion.
+         :return: Returns descriptive information about the created data ingestion
+         having Ingestion ARN, HTTP status,
+         ingestion ID and ingestion status.
+        :rtype: Dict
+        """
+        self.log.info('Creating QuickSight Ingestion for data set id %s.', data_set_id)
+        quicksight_client = self.get_conn()
+        try:
+            create_ingestion_response = quicksight_client.create_ingestion(
+                DataSetId=data_set_id,
+                IngestionId=ingestion_id,
+                AwsAccountId=aws_account_id,
+                IngestionType=ingestion_type,
+            )
+            if wait_for_completion:
+                self.check_status(
+                    aws_account_id=aws_account_id,
+                    ingestion_id=ingestion_id,
+                    data_set_id=data_set_id,
+                    check_interval=check_interval,
+                    max_ingestion_time=max_ingestion_time,
+                )
+            return create_ingestion_response
+        except Exception as general_error:
+            self.log.error("Failed to run QuickSight create_ingestion API, error: %s", general_error)
+            raise
+
+    def check_status(
+        self,
+        aws_account_id: str,
+        ingestion_id: str,
+        data_set_id: str,
+        check_interval: int,
+        max_ingestion_time: Optional[int] = None,

Review comment:
       this function should be broken into two functions: `get_status` and `wait_for_state`
   see similar functionality in TableauHook:
   https://github.com/apache/airflow/blob/4e05bbc92557f5eafaa37509387b2fb2b0ab3d4a/airflow/providers/tableau/hooks/tableau.py#L150
   https://github.com/apache/airflow/blob/4e05bbc92557f5eafaa37509387b2fb2b0ab3d4a/airflow/providers/tableau/hooks/tableau.py#L161
   
   
   You then can also even add Sensor very easily:
   https://github.com/apache/airflow/blob/4e05bbc92557f5eafaa37509387b2fb2b0ab3d4a/airflow/providers/tableau/sensors/tableau_job_status.py#L30




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks closed pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks closed pull request #21863:
URL: https://github.com/apache/airflow/pull/21863


   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] ferruzzi commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r826462628



##########
File path: airflow/providers/amazon/aws/sensors/quicksight.py
##########
@@ -0,0 +1,80 @@
+#
+# 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 typing import TYPE_CHECKING, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class QuickSightSensor(BaseSensorOperator):
+    """
+    Watches for the status of a QuickSight Ingestion.

Review comment:
       ```suggestion
       Watches for the status of an Amazon QuickSight Ingestion.
   ```

##########
File path: airflow/providers/amazon/aws/example_dags/example_quicksight.py
##########
@@ -0,0 +1,67 @@
+# 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 datetime import datetime
+
+from airflow import DAG
+from airflow.providers.amazon.aws.operators.quicksight import QuickSightCreateIngestionOperator
+from airflow.providers.amazon.aws.sensors.quicksight import QuickSightSensor
+
+DATA_SET_ID = "DemoDataSet_Test"
+INGESTION_WAITING_ID = "DemoDataSet_Ingestion_Waiting_Test"
+INGESTION_NO_WAITING_ID = "DemoDataSet_Ingestion_No_Waiting_Test"
+AWS_ACCOUNT_ID = "123456789012"
+
+with DAG(
+    "sample_quicksight_dag",
+    schedule_interval=None,
+    start_date=datetime(2022, 2, 21),
+    catchup=False,

Review comment:
       ```suggestion
       dag_id='example_quicksight',
       schedule_interval=None,
       start_date=datetime(2021, 1, 1),
       tags=['example'],
       catchup=False,
   ```
   
   Nitpick:  We're working on organizing and standardizing the AWS examples and docs.  If you don't mind applying these changes to match what we're doing, I'd appreciate it.

##########
File path: airflow/providers/amazon/aws/operators/quicksight.py
##########
@@ -0,0 +1,100 @@
+# 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 typing import TYPE_CHECKING, Optional, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+DEFAULT_CONN_ID = "aws_default"
+
+
+class QuickSightCreateIngestionOperator(BaseOperator):
+    """
+    Creates and starts a new SPICE ingestion for a dataset.
+    Also, helps to Refresh existing SPICE datasets
+
+    :param data_set_id:  ID of the dataset used in the ingestion.
+    :param ingestion_id: ID for the ingestion.
+    :param aws_account_id: Amazon Web Services account ID.
+    :param ingestion_type: Type of ingestion. Values Can be  INCREMENTAL_REFRESH or FULL_REFRESH.
+        Default FULL_REFRESH.
+    :param wait_for_completion: If wait is set to True, the time interval, in seconds,
+        that the operation waits to check the status of the QuickSight Ingestion.

Review comment:
       ```suggestion
           that the operation waits to check the status of the Amazon QuickSight Ingestion.
   ```
   
   Nitpicking, if you don't mind.

##########
File path: airflow/providers/amazon/aws/sensors/quicksight.py
##########
@@ -0,0 +1,80 @@
+#
+# 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 typing import TYPE_CHECKING, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class QuickSightSensor(BaseSensorOperator):
+    """
+    Watches for the status of a QuickSight Ingestion.
+
+    :param aws_account_id: An AWS Account ID
+    :param data_set_id: QuickSight Data Set ID
+    :param ingestion_id: QuickSight Ingestion ID
+    :param aws_conn_id: aws connection to use, defaults to "aws_default"
+    """
+
+    def __init__(
+        self,
+        *,
+        aws_account_id: str,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.aws_account_id = aws_account_id
+        self.data_set_id = data_set_id
+        self.aws_account_id = aws_account_id
+        self.ingestion_id = ingestion_id
+        self.aws_conn_id = aws_conn_id
+        self.success_status = "COMPLETED"
+        self.errored_statuses = ("FAILED", "CANCELLED")
+        self.hook: Optional[QuickSightHook] = None
+
+    def poke(self, context: "Context"):
+        """
+        Pokes until the QuickSight Ingestion has successfully finished.
+
+        :param context: The task context during execution.
+        :return: True if it COMPLETED and False if not.
+        :rtype: bool
+        """
+        hook = self.get_hook()
+        self.log.info("Poking for Amazon QuickSight Ingestion ID: %s", self.ingestion_id)
+        quicksight_ingestion_state = hook.get_status(self.aws_account_id, self.data_set_id, self.ingestion_id)
+        self.log.info("QuickSight Status: %s", quicksight_ingestion_state)
+        if quicksight_ingestion_state in self.errored_statuses:
+            raise AirflowException("The QuickSight Ingestion failed!")

Review comment:
       ```suggestion
               raise AirflowException("The Amazon QuickSight Ingestion failed!")
   ```
   
   Nitpick for consistency.

##########
File path: airflow/providers/amazon/aws/operators/quicksight.py
##########
@@ -0,0 +1,100 @@
+# 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 typing import TYPE_CHECKING, Optional, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+DEFAULT_CONN_ID = "aws_default"
+
+
+class QuickSightCreateIngestionOperator(BaseOperator):
+    """
+    Creates and starts a new SPICE ingestion for a dataset.
+    Also, helps to Refresh existing SPICE datasets
+
+    :param data_set_id:  ID of the dataset used in the ingestion.
+    :param ingestion_id: ID for the ingestion.
+    :param aws_account_id: Amazon Web Services account ID.
+    :param ingestion_type: Type of ingestion. Values Can be  INCREMENTAL_REFRESH or FULL_REFRESH.
+        Default FULL_REFRESH.
+    :param wait_for_completion: If wait is set to True, the time interval, in seconds,
+        that the operation waits to check the status of the QuickSight Ingestion.
+    :param check_interval: if wait is set to be true, this is the time interval
+        in seconds which the operator will check the status of the QuickSight Ingestion
+    :param aws_conn_id: The Airflow connection used for AWS credentials. (templated)
+         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 the default boto3 configuration would be used (and must be
+         maintained on each worker node).
+    :param region: Which AWS region the connection should use. (templated)
+         If this is None or empty then the default boto3 behaviour is used.
+    """
+
+    template_fields: Sequence[str] = (
+        "data_set_id",
+        "ingestion_id",
+        "aws_account_id",
+        "ingestion_type",
+        "wait_for_completion",
+        "check_interval",
+        "aws_conn_id",
+        "region",
+    )
+    ui_color = "#ffd700"
+
+    def __init__(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,
+        ingestion_type: str = "FULL_REFRESH",
+        wait_for_completion: bool = True,
+        check_interval: int = 30,
+        aws_conn_id: str = DEFAULT_CONN_ID,
+        region: Optional[str] = None,
+        **kwargs,
+    ):
+        self.data_set_id = data_set_id
+        self.ingestion_id = ingestion_id
+        self.aws_account_id = aws_account_id
+        self.ingestion_type = ingestion_type
+        self.wait_for_completion = wait_for_completion
+        self.check_interval = check_interval
+        self.aws_conn_id = aws_conn_id
+        self.region = region
+        super().__init__(**kwargs)
+
+    def execute(self, context: "Context"):
+        hook = QuickSightHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region,
+        )
+        self.log.info("Running the QuickSight Spice Ingestion on Dataset ID: %s)", self.data_set_id)

Review comment:
       ```suggestion
           self.log.info("Running the Amazon QuickSight SPICE Ingestion on Dataset ID: %s)", self.data_set_id)
   ```
   
   Based on other uses, it looks like this is an acronym.

##########
File path: airflow/providers/amazon/aws/sensors/quicksight.py
##########
@@ -0,0 +1,80 @@
+#
+# 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 typing import TYPE_CHECKING, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class QuickSightSensor(BaseSensorOperator):
+    """
+    Watches for the status of a QuickSight Ingestion.
+
+    :param aws_account_id: An AWS Account ID
+    :param data_set_id: QuickSight Data Set ID
+    :param ingestion_id: QuickSight Ingestion ID

Review comment:
       Nitpick: Same parameter names in the Operator above has a different descriptions, consider standardizing?

##########
File path: airflow/providers/amazon/aws/operators/quicksight.py
##########
@@ -0,0 +1,100 @@
+# 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 typing import TYPE_CHECKING, Optional, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+DEFAULT_CONN_ID = "aws_default"
+
+
+class QuickSightCreateIngestionOperator(BaseOperator):
+    """
+    Creates and starts a new SPICE ingestion for a dataset.
+    Also, helps to Refresh existing SPICE datasets
+
+    :param data_set_id:  ID of the dataset used in the ingestion.
+    :param ingestion_id: ID for the ingestion.
+    :param aws_account_id: Amazon Web Services account ID.
+    :param ingestion_type: Type of ingestion. Values Can be  INCREMENTAL_REFRESH or FULL_REFRESH.
+        Default FULL_REFRESH.
+    :param wait_for_completion: If wait is set to True, the time interval, in seconds,
+        that the operation waits to check the status of the QuickSight Ingestion.
+    :param check_interval: if wait is set to be true, this is the time interval
+        in seconds which the operator will check the status of the QuickSight Ingestion
+    :param aws_conn_id: The Airflow connection used for AWS credentials. (templated)
+         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 the default boto3 configuration would be used (and must be
+         maintained on each worker node).
+    :param region: Which AWS region the connection should use. (templated)
+         If this is None or empty then the default boto3 behaviour is used.
+    """
+
+    template_fields: Sequence[str] = (
+        "data_set_id",
+        "ingestion_id",
+        "aws_account_id",
+        "ingestion_type",
+        "wait_for_completion",
+        "check_interval",
+        "aws_conn_id",
+        "region",
+    )
+    ui_color = "#ffd700"
+
+    def __init__(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,
+        ingestion_type: str = "FULL_REFRESH",
+        wait_for_completion: bool = True,
+        check_interval: int = 30,
+        aws_conn_id: str = DEFAULT_CONN_ID,
+        region: Optional[str] = None,
+        **kwargs,
+    ):
+        self.data_set_id = data_set_id
+        self.ingestion_id = ingestion_id
+        self.aws_account_id = aws_account_id
+        self.ingestion_type = ingestion_type
+        self.wait_for_completion = wait_for_completion
+        self.check_interval = check_interval
+        self.aws_conn_id = aws_conn_id
+        self.region = region
+        super().__init__(**kwargs)
+
+    def execute(self, context: "Context"):
+        hook = QuickSightHook(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.region,
+        )
+        self.log.info("Running the QuickSight Spice Ingestion on Dataset ID: %s)", self.data_set_id)
+        create_ingestion_response = hook.create_ingestion(
+            data_set_id=self.data_set_id,
+            ingestion_id=self.ingestion_id,
+            aws_account_id=self.aws_account_id,
+            ingestion_type=self.ingestion_type,
+            wait_for_completion=self.wait_for_completion,
+            check_interval=self.check_interval,
+        )
+        return create_ingestion_response

Review comment:
       Is there a particular reason to store it and then return it instead of just returning it on L92?

##########
File path: airflow/providers/amazon/aws/operators/quicksight.py
##########
@@ -0,0 +1,100 @@
+# 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 typing import TYPE_CHECKING, Optional, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+DEFAULT_CONN_ID = "aws_default"
+
+
+class QuickSightCreateIngestionOperator(BaseOperator):
+    """
+    Creates and starts a new SPICE ingestion for a dataset.
+    Also, helps to Refresh existing SPICE datasets
+
+    :param data_set_id:  ID of the dataset used in the ingestion.
+    :param ingestion_id: ID for the ingestion.
+    :param aws_account_id: Amazon Web Services account ID.
+    :param ingestion_type: Type of ingestion. Values Can be  INCREMENTAL_REFRESH or FULL_REFRESH.
+        Default FULL_REFRESH.
+    :param wait_for_completion: If wait is set to True, the time interval, in seconds,
+        that the operation waits to check the status of the QuickSight Ingestion.
+    :param check_interval: if wait is set to be true, this is the time interval
+        in seconds which the operator will check the status of the QuickSight Ingestion

Review comment:
       ```suggestion
           in seconds which the operator will check the status of the Amazon QuickSight Ingestion
   ```

##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,141 @@
+#
+# 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.
+
+import time
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,

Review comment:
       Yeah, it seems odd that all of their client calls for Quicksight take the account ID as a parameter.  If you don't want to take it as an input, you get get it from [STS.get_caller_identity(): ]( https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.get_caller_identity) and I'd suggest storing it as a field in the QuicksightHook perhaps rather than passing it all over the place.
   
   This should do the trick inline:
   
   ```
   import boto3
   boto3.client('sts').get_caller_identity()['Account']
   ``` 
   
   But if you want to follow Airflow structure and add an STS hook it would look like this:
   
   ```
   class StsHook(AwsBaseHook):
       def __init__(self) -> None:
           super().__init__(client_type='sts')
   
       def get_account_number(self) -> Dict:
           return self.conn.get_caller_identity()['Account']
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821034985



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,152 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    non_terminal_states = {"INITIALIZED", "QUEUED", "RUNNING"}
+    failed_states = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,
+        ingestion_type: str,
+        wait_for_completion: bool = True,
+        check_interval: int = 30,
+        max_ingestion_time: Optional[int] = None,
+    ):
+        """
+        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 aws_account_id: Amazon Web Services account ID.
+        :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 max_ingestion_time: the maximum ingestion time in seconds. If QuickSight ingestion runs
+         longer than this will fail. Setting this to None implies no timeout for Ingestion.
+         :return: Returns descriptive information about the created data ingestion
+         having Ingestion ARN, HTTP status,
+         ingestion ID and ingestion status.
+        :rtype: Dict
+        """
+        self.log.info("Creating QuickSight Ingestion for data set id %s.", data_set_id)
+        quicksight_client = self.get_conn()
+        try:
+            create_ingestion_response = quicksight_client.create_ingestion(
+                DataSetId=data_set_id,
+                IngestionId=ingestion_id,
+                AwsAccountId=aws_account_id,
+                IngestionType=ingestion_type,
+            )
+            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,
+                    max_ingestion_time=max_ingestion_time,
+                )
+            return create_ingestion_response
+        except Exception as general_error:
+            self.log.error("Failed to run QuickSight create_ingestion API, error: %s", general_error)
+            raise
+
+    def get_status(self, aws_account_id: str, data_set_id: str, ingestion_id: 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
+        :rtype: str
+        """
+        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 QuickSight Ingestion")
+        except ClientError:
+            raise AirflowException("AWS request failed, check logs for more info")
+
+    def wait_for_state(
+        self,
+        aws_account_id: str,
+        data_set_id: str,
+        ingestion_id: str,
+        target_state: set,
+        check_interval: int,
+        max_ingestion_time: Optional[int] = None,
+    ):
+        """
+        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
+        :param max_ingestion_time: the maximum ingestion time in seconds. QuickSight API if
+         run longer than this will fail. Setting this to None implies no timeout.
+        :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("QuickSight SPICE ingestion failed")
+            if status == "CANCELLED":
+                raise AirflowException("QuickSight SPICE ingestion cancelled")
+            if max_ingestion_time and sec > max_ingestion_time:
+                raise AirflowException(f"QuickSight Ingestion took more than {max_ingestion_time} seconds")

Review comment:
       Thanks for the details ! I have made the changes. Please let me know if it looks fine.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#issuecomment-1061454981


   > You need to rebase.
   
   Done :) Seems like build got queued and eventually failed.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] potiuk commented on pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
potiuk commented on pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#issuecomment-1061665998


   Static/docs failures.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] hsrocks commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
hsrocks commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r821714920



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,143 @@
+#
+# 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.
+
+import time
+from typing import Optional
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)

Review comment:
       > Why do you need to specify `client_type` here?
   
   I was initialising the boto3.client as quicksight . Please let me know if I am doing it wrong way . 
   
   As per basehook : https://airflow.apache.org/docs/apache-airflow-providers-amazon/1.0.0/_modules/airflow/providers/amazon/aws/hooks/base_aws.html
   
   :param client_type: boto3.client client_type. Eg 's3', 'emr' etc
   
   So for quicksight I put it like this . Is there better way or suggested approach for this ?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] potiuk closed pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
potiuk closed pull request #21863:
URL: https://github.com/apache/airflow/pull/21863


   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] eladkal commented on pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
eladkal commented on pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#issuecomment-1067284270


   @ferruzzi can you take a look?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] ferruzzi commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r826433470



##########
File path: airflow/providers/amazon/aws/hooks/quicksight.py
##########
@@ -0,0 +1,141 @@
+#
+# 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.
+
+import time
+
+from botocore.exceptions import ClientError
+
+from airflow import AirflowException
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+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`
+    """
+
+    NON_TERMINAL_STATES = {"INITIALIZED", "QUEUED", "RUNNING"}
+    FAILED_STATES = {"FAILED"}
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="quicksight", *args, **kwargs)
+
+    def create_ingestion(
+        self,
+        data_set_id: str,
+        ingestion_id: str,
+        aws_account_id: str,

Review comment:
       Yeah, it seems odd that all of their client calls for Quicksight take the account ID as a parameter.  If you don't want to take it as an input, you get get it from [STS.get_caller_identity(): ]( https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.get_caller_identity) and I'd suggest fetching it once programmatically then storing it as a field in the QuicksightHook rather than passing it all over the place.
   
   This should do the trick inline:
   
   ```
   import boto3
   boto3.client('sts').get_caller_identity()['Account']
   ``` 
   
   But if you want to follow Airflow structure and add an STS hook it would look like this:
   
   ```
   class StsHook(AwsBaseHook):
       def __init__(self) -> None:
           super().__init__(client_type='sts')
   
       def get_account_number(self) -> Dict:
           return self.conn.get_caller_identity()['Account']
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [airflow] ferruzzi commented on a change in pull request #21863: Add Quicksight create ingestion Hook and Operator

Posted by GitBox <gi...@apache.org>.
ferruzzi commented on a change in pull request #21863:
URL: https://github.com/apache/airflow/pull/21863#discussion_r827383898



##########
File path: docs/apache-airflow-providers-amazon/operators/quicksight.rst
##########
@@ -0,0 +1,68 @@
+ .. 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.
+
+
+Amazon QuickSight Operators
+========================================
+
+Prerequisite Tasks
+------------------
+
+.. include:: _partials/prerequisite_tasks.rst
+
+Overview
+--------
+
+Airflow to Amazon QuickSight integration allows to create and start the SPICE ingestion for dataset.
+
+  - :class:`~airflow.providers.amazon.aws.operators.quicksight.QuickSightCreateIngestionOperator`
+  - :class:`~airflow.providers.amazon.aws.sensor.quicksight.QuickSightSensor`
+
+Purpose
+"""""""
+
+This example DAG ``example_quicksight.py`` uses ``QuickSightCreateIngestionOperator`` for
+creating and starting the SPICE ingestion for the dataset configured to use SPICE. In the example,
+we created two ingestion. One of the ingestion waits for the SPICE ingestion to complete while
+other ingestion does not wait for completion and uses ``QuickSightSensor`` to check for ingestion
+status until it completes
+
+Defining tasks
+""""""""""""""
+
+In the following code we create and start a QuickSight SPICE ingestion for the dataset and wait
+for its completion.
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_quicksight.py
+    :language: python

Review comment:
       ```suggestion
       :language: python
       :dedent: 4
   ```
   
   Something someone showed me recently,. un-indents the code sample for a cleaner look.

##########
File path: docs/apache-airflow-providers-amazon/operators/quicksight.rst
##########
@@ -0,0 +1,68 @@
+ .. 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.
+
+
+Amazon QuickSight Operators
+========================================
+
+Prerequisite Tasks
+------------------
+
+.. include:: _partials/prerequisite_tasks.rst
+
+Overview
+--------
+
+Airflow to Amazon QuickSight integration allows to create and start the SPICE ingestion for dataset.

Review comment:
       Might just be me, but "allows to create" feels awkward?  Maybe "allows the user to create" or "allows users" or something?

##########
File path: docs/apache-airflow-providers-amazon/operators/quicksight.rst
##########
@@ -0,0 +1,68 @@
+ .. 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.
+
+
+Amazon QuickSight Operators
+========================================
+
+Prerequisite Tasks
+------------------
+
+.. include:: _partials/prerequisite_tasks.rst
+
+Overview
+--------
+
+Airflow to Amazon QuickSight integration allows to create and start the SPICE ingestion for dataset.
+
+  - :class:`~airflow.providers.amazon.aws.operators.quicksight.QuickSightCreateIngestionOperator`
+  - :class:`~airflow.providers.amazon.aws.sensor.quicksight.QuickSightSensor`
+
+Purpose
+"""""""
+
+This example DAG ``example_quicksight.py`` uses ``QuickSightCreateIngestionOperator`` for
+creating and starting the SPICE ingestion for the dataset configured to use SPICE. In the example,
+we created two ingestion. One of the ingestion waits for the SPICE ingestion to complete while
+other ingestion does not wait for completion and uses ``QuickSightSensor`` to check for ingestion
+status until it completes
+
+Defining tasks
+""""""""""""""
+
+In the following code we create and start a QuickSight SPICE ingestion for the dataset and wait
+for its completion.
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_quicksight.py
+    :language: python
+    :start-after: [START howto_operator_quicksight]
+    :end-before: [END howto_operator_quicksight]
+
+In the below, we create and start the SPICE ingestion but does not wait for completion. We use

Review comment:
       ```suggestion
   In the below example, we create and start the SPICE ingestion but do not wait for completion. We use
   ```

##########
File path: docs/apache-airflow-providers-amazon/operators/quicksight.rst
##########
@@ -0,0 +1,68 @@
+ .. 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.
+
+
+Amazon QuickSight Operators
+========================================
+
+Prerequisite Tasks
+------------------
+
+.. include:: _partials/prerequisite_tasks.rst
+
+Overview
+--------
+
+Airflow to Amazon QuickSight integration allows to create and start the SPICE ingestion for dataset.
+
+  - :class:`~airflow.providers.amazon.aws.operators.quicksight.QuickSightCreateIngestionOperator`
+  - :class:`~airflow.providers.amazon.aws.sensor.quicksight.QuickSightSensor`
+
+Purpose
+"""""""
+
+This example DAG ``example_quicksight.py`` uses ``QuickSightCreateIngestionOperator`` for
+creating and starting the SPICE ingestion for the dataset configured to use SPICE. In the example,
+we created two ingestion. One of the ingestion waits for the SPICE ingestion to complete while
+other ingestion does not wait for completion and uses ``QuickSightSensor`` to check for ingestion
+status until it completes
+
+Defining tasks
+""""""""""""""
+
+In the following code we create and start a QuickSight SPICE ingestion for the dataset and wait
+for its completion.
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_quicksight.py
+    :language: python
+    :start-after: [START howto_operator_quicksight]
+    :end-before: [END howto_operator_quicksight]
+
+In the below, we create and start the SPICE ingestion but does not wait for completion. We use
+sensor to poll for Ingestion status until it Completes.
+
+.. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_quicksight.py
+    :language: python

Review comment:
       ```suggestion
       :language: python
       :dedent: 4
   ```

##########
File path: docs/apache-airflow-providers-amazon/operators/quicksight.rst
##########
@@ -0,0 +1,68 @@
+ .. 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.
+
+
+Amazon QuickSight Operators
+========================================
+
+Prerequisite Tasks
+------------------
+
+.. include:: _partials/prerequisite_tasks.rst
+
+Overview
+--------
+
+Airflow to Amazon QuickSight integration allows to create and start the SPICE ingestion for dataset.
+
+  - :class:`~airflow.providers.amazon.aws.operators.quicksight.QuickSightCreateIngestionOperator`
+  - :class:`~airflow.providers.amazon.aws.sensor.quicksight.QuickSightSensor`
+
+Purpose
+"""""""
+
+This example DAG ``example_quicksight.py`` uses ``QuickSightCreateIngestionOperator`` for
+creating and starting the SPICE ingestion for the dataset configured to use SPICE. In the example,
+we created two ingestion. One of the ingestion waits for the SPICE ingestion to complete while

Review comment:
       ```suggestion
   we created two ingestions. One of the ingestions waits for the SPICE ingestion to complete while
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org