This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 76ffb4693fe Add deferrable mode to QuickSight operator and sensor
(#70218)
76ffb4693fe is described below
commit 76ffb4693feb30c5e1e9f7e0a5e630d0c59a5d63
Author: Emin Özata <[email protected]>
AuthorDate: Fri Sep 11 02:21:14 2026 +0300
Add deferrable mode to QuickSight operator and sensor (#70218)
The QuickSight create-ingestion wait path used a hand-written polling
loop that held a worker slot with time.sleep and had no upper bound, so
a stuck ingestion could block a worker indefinitely. Moving the wait to
a custom boto waiter lets both the operator and the sensor defer to the
triggerer with deferrable=True, matching the deferrable pattern used
across the Amazon provider, and bounds the synchronous wait as well.
AWS waits up to 36 hours for a SPICE ingestion in its own CloudFormation
wait policy, where 36 hours is both the default and the maximum.
Refreshes that run for several hours are common, so a shorter cap fails
Dags that previously waited without a limit.
---
generated/known_airflow_exceptions.txt | 1 -
providers/amazon/docs/operators/quicksight.rst | 6 ++
providers/amazon/provider.yaml | 3 +
.../providers/amazon/aws/hooks/quicksight.py | 85 +++++++++++++-----
.../providers/amazon/aws/operators/quicksight.py | 77 ++++++++++++++--
.../providers/amazon/aws/sensors/quicksight.py | 49 +++++++++-
.../providers/amazon/aws/triggers/quicksight.py | 87 ++++++++++++++++++
.../providers/amazon/aws/waiters/quicksight.json | 48 ++++++++++
.../airflow/providers/amazon/get_provider_info.py | 4 +
.../tests/unit/amazon/aws/hooks/test_quicksight.py | 77 +++++++++++++---
.../unit/amazon/aws/operators/test_quicksight.py | 100 +++++++++++++++++++--
.../unit/amazon/aws/sensors/test_quicksight.py | 46 +++++++++-
.../unit/amazon/aws/triggers/test_quicksight.py | 64 +++++++++++++
.../unit/amazon/aws/waiters/test_quicksight.py | 76 ++++++++++++++++
14 files changed, 668 insertions(+), 55 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index 2c619d8eaff..4269e6eb801 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -105,7 +105,6 @@
providers/amazon/src/airflow/providers/amazon/aws/sensors/kinesis_analytics.py::
providers/amazon/src/airflow/providers/amazon/aws/sensors/lambda_function.py::1
providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa.py::4
providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py::1
-providers/amazon/src/airflow/providers/amazon/aws/sensors/quicksight.py::1
providers/amazon/src/airflow/providers/amazon/aws/sensors/rds.py::1
providers/amazon/src/airflow/providers/amazon/aws/sensors/redshift_cluster.py::1
providers/amazon/src/airflow/providers/amazon/aws/sensors/s3.py::3
diff --git a/providers/amazon/docs/operators/quicksight.rst
b/providers/amazon/docs/operators/quicksight.rst
index 07cd48f7256..e5453cf2a8b 100644
--- a/providers/amazon/docs/operators/quicksight.rst
+++ b/providers/amazon/docs/operators/quicksight.rst
@@ -46,6 +46,9 @@ Amazon QuickSight create ingestion
The ``QuickSightCreateIngestionOperator`` creates and starts a new SPICE
ingestion for a dataset.
The operator also refreshes existing SPICE datasets.
+This operator can be run in deferrable mode by passing ``deferrable=True`` as
a parameter. This requires
+the aiobotocore module to be installed.
+
.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_quicksight.py
:language: python
:dedent: 4
@@ -62,6 +65,9 @@ Amazon QuickSight ingestion sensor
The ``QuickSightSensor`` waits for an Amazon QuickSight create ingestion until
it reaches a terminal state.
+This sensor can be run in deferrable mode by passing ``deferrable=True`` as a
parameter. This requires
+the aiobotocore module to be installed.
+
.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_quicksight.py
:language: python
:dedent: 4
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 00a365e552d..b7731887e17 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -898,6 +898,9 @@ triggers:
- integration-name: AWS Database Migration Service
python-modules:
- airflow.providers.amazon.aws.triggers.dms
+ - integration-name: Amazon QuickSight
+ python-modules:
+ - airflow.providers.amazon.aws.triggers.quicksight
transfers:
- source-integration-name: Amazon DynamoDB
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/hooks/quicksight.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/quicksight.py
index 0585213fff3..3441291aa56 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/quicksight.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/quicksight.py
@@ -18,10 +18,13 @@
from __future__ import annotations
import time
+import warnings
from botocore.exceptions import ClientError
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
from airflow.providers.common.compat.sdk import AirflowException
@@ -52,6 +55,7 @@ class QuickSightHook(AwsBaseHook):
wait_for_completion: bool = True,
check_interval: int = 30,
aws_account_id: str | None = None,
+ waiter_max_attempts: int = 4320,
) -> dict:
"""
Create and start a new SPICE ingestion for a dataset; refresh the
SPICE datasets.
@@ -66,31 +70,63 @@ class QuickSightHook(AwsBaseHook):
:param check_interval: the time interval in seconds which the operator
will check the status of QuickSight Ingestion
:param aws_account_id: An AWS Account ID, if set to ``None`` then use
associated AWS Account ID.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:return: Returns descriptive information about the created data
ingestion
having Ingestion ARN, HTTP status, ingestion ID and ingestion
status.
"""
aws_account_id = aws_account_id or self.account_id
self.log.info("Creating QuickSight Ingestion for data set id %s.",
data_set_id)
- try:
- create_ingestion_response = self.conn.create_ingestion(
- DataSetId=data_set_id,
- IngestionId=ingestion_id,
- IngestionType=ingestion_type,
- AwsAccountId=aws_account_id,
+ create_ingestion_response = self.conn.create_ingestion(
+ DataSetId=data_set_id,
+ IngestionId=ingestion_id,
+ IngestionType=ingestion_type,
+ AwsAccountId=aws_account_id,
+ )
+ if wait_for_completion:
+ self.wait_for_ingestion(
+ data_set_id=data_set_id,
+ ingestion_id=ingestion_id,
+ aws_account_id=aws_account_id,
+ waiter_delay=check_interval,
+ waiter_max_attempts=waiter_max_attempts,
)
+ return create_ingestion_response
+
+ def wait_for_ingestion(
+ self,
+ *,
+ data_set_id: str,
+ ingestion_id: str,
+ aws_account_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 4320,
+ ) -> None:
+ """
+ Poll a SPICE ingestion until it completes.
- if wait_for_completion:
- self.wait_for_state(
- aws_account_id=aws_account_id,
- data_set_id=data_set_id,
- ingestion_id=ingestion_id,
- target_state={"COMPLETED"},
- check_interval=check_interval,
- )
- return create_ingestion_response
- except Exception as general_error:
- self.log.error("Failed to run Amazon QuickSight create_ingestion
API, error: %s", general_error)
- raise
+ :param data_set_id: QuickSight Data Set ID
+ :param ingestion_id: QuickSight Ingestion ID
+ :param aws_account_id: An AWS Account ID, if set to ``None`` then use
associated AWS Account ID.
+ :param waiter_delay: The amount of time in seconds to wait between
attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :raises RuntimeError: If the ingestion fails, is cancelled or times
out.
+ """
+ try:
+ wait(
+ waiter=self.get_waiter("ingestion_complete"),
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ args={
+ "AwsAccountId": aws_account_id or self.account_id,
+ "DataSetId": data_set_id,
+ "IngestionId": ingestion_id,
+ },
+ failure_message="Amazon QuickSight SPICE ingestion failed.",
+ status_message="Status of Amazon QuickSight SPICE ingestion
is",
+ status_args=["Ingestion.IngestionStatus",
"Ingestion.ErrorInfo"],
+ )
+ except AirflowException as e:
+ raise RuntimeError(str(e)) from e
def get_status(self, aws_account_id: str | None, data_set_id: str,
ingestion_id: str) -> str:
"""
@@ -144,14 +180,23 @@ class QuickSightHook(AwsBaseHook):
"""
Check status of a QuickSight Create Ingestion API.
+ .. deprecated::
+ Use :meth:`wait_for_ingestion` instead.
+
:param aws_account_id: An AWS Account ID, if set to ``None`` then use
associated AWS Account ID.
:param data_set_id: QuickSight Data Set ID
:param ingestion_id: QuickSight Ingestion ID
:param target_state: Describes the QuickSight Job's Target State
:param check_interval: the time interval in seconds which the operator
will check the status of QuickSight Ingestion
- :return: response of describe_ingestion call after Ingestion is done
+ :return: the final status of the ingestion
"""
+ warnings.warn(
+ "`QuickSightHook.wait_for_state` is deprecated and will be removed
in a future release. "
+ "Use `QuickSightHook.wait_for_ingestion` instead.",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
aws_account_id = aws_account_id or self.account_id
while True:
@@ -162,7 +207,7 @@ class QuickSightHook(AwsBaseHook):
raise AirflowException(f"The Amazon QuickSight Ingestion
failed. Error info: {info}")
if status == "CANCELLED":
raise AirflowException("The Amazon QuickSight SPICE ingestion
cancelled!")
- if status not in self.NON_TERMINAL_STATES or status ==
target_state:
+ if status not in self.NON_TERMINAL_STATES:
break
time.sleep(check_interval)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/quicksight.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/quicksight.py
index 51b0b3b372e..64086104b79 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/quicksight.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/quicksight.py
@@ -16,12 +16,17 @@
# under the License.
from __future__ import annotations
+import warnings
from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.quicksight import
QuickSightIngestionCompletedTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
if TYPE_CHECKING:
from airflow.sdk import Context
@@ -39,10 +44,13 @@ class
QuickSightCreateIngestionOperator(AwsBaseOperator[QuickSightHook]):
:param ingestion_id: ID for the ingestion.
: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 Amazon 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 Amazon
QuickSight Ingestion
+ :param wait_for_completion: If True, wait for the ingestion to reach a
terminal state. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks.
(default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for
completion. (default: 4320)
+ :param check_interval: Deprecated, use ``waiter_delay`` instead.
+ :param deferrable: If True, the operator will wait asynchronously for the
ingestion to complete.
+ This implies waiting for completion. This mode requires aiobotocore
module to be installed.
+ (default: False, but can be overridden in config file by setting
default_deferrable to True)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used.
If
running Airflow in a distributed manner and aws_conn_id is None or
@@ -61,6 +69,8 @@ class
QuickSightCreateIngestionOperator(AwsBaseOperator[QuickSightHook]):
"ingestion_id",
"ingestion_type",
"wait_for_completion",
+ "waiter_delay",
+ "waiter_max_attempts",
"check_interval",
)
ui_color = "#ffd700"
@@ -71,7 +81,10 @@ class
QuickSightCreateIngestionOperator(AwsBaseOperator[QuickSightHook]):
ingestion_id: str,
ingestion_type: str = "FULL_REFRESH",
wait_for_completion: bool = True,
- check_interval: int = 30,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 4320,
+ check_interval: int | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
):
super().__init__(**kwargs)
@@ -79,14 +92,60 @@ class
QuickSightCreateIngestionOperator(AwsBaseOperator[QuickSightHook]):
self.ingestion_id = ingestion_id
self.ingestion_type = ingestion_type
self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
self.check_interval = check_interval
+ self.deferrable = deferrable
def execute(self, context: Context):
+ # check_interval may be templated, so it is only resolvable once
rendering has happened
+ if self.check_interval is not None:
+ warnings.warn(
+ "The `check_interval` parameter is deprecated and will be
removed in a future release. "
+ "Use `waiter_delay` instead. While `check_interval` is set, it
takes precedence over "
+ "`waiter_delay`.",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
self.log.info("Running the Amazon QuickSight SPICE Ingestion on
Dataset ID: %s", self.data_set_id)
- return self.hook.create_ingestion(
+ ingestion = self.hook.create_ingestion(
data_set_id=self.data_set_id,
ingestion_id=self.ingestion_id,
ingestion_type=self.ingestion_type,
- wait_for_completion=self.wait_for_completion,
- check_interval=self.check_interval,
+ wait_for_completion=False,
)
+ waiter_delay = int(self.waiter_delay if self.check_interval is None
else self.check_interval)
+ waiter_max_attempts = int(self.waiter_max_attempts)
+ if self.deferrable:
+ self.defer(
+ trigger=QuickSightIngestionCompletedTrigger(
+ data_set_id=self.data_set_id,
+ ingestion_id=self.ingestion_id,
+ aws_account_id=self.hook.account_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ botocore_config=self.botocore_config,
+ ),
+ method_name="execute_complete",
+ kwargs={"ingestion": ingestion},
+ )
+ elif self.wait_for_completion:
+ self.hook.wait_for_ingestion(
+ data_set_id=self.data_set_id,
+ ingestion_id=self.ingestion_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ )
+ return ingestion
+
+ def execute_complete(
+ self, context: Context, event: dict[str, Any] | None = None,
ingestion: dict[str, Any] | None = None
+ ) -> dict[str, Any] | None:
+ validated_event = validate_execute_complete_event(event)
+ if validated_event["status"] != "success":
+ raise RuntimeError(f"Error while running Amazon QuickSight SPICE
ingestion: {validated_event}")
+ self.log.info("Amazon QuickSight SPICE ingestion `%s` completed.",
validated_event["ingestion_id"])
+ return ingestion
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/sensors/quicksight.py
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/quicksight.py
index 4c1f9085033..6852d3a361a 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/quicksight.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/quicksight.py
@@ -18,11 +18,13 @@
from __future__ import annotations
from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.amazon.aws.triggers.quicksight import
QuickSightIngestionCompletedTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.common.compat.sdk import conf
if TYPE_CHECKING:
from airflow.sdk import Context
@@ -38,6 +40,10 @@ class QuickSightSensor(AwsBaseSensor[QuickSightHook]):
:param data_set_id: ID of the dataset used in the ingestion.
:param ingestion_id: ID for the ingestion.
+ :param max_retries: Number of times before returning the current state.
(default: 2160)
+ :param deferrable: If True, the sensor will operate in deferrable mode.
This mode requires
+ aiobotocore module to be installed.
+ (default: False, but can be overridden in config file by setting
default_deferrable to True)
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used.
If
running Airflow in a distributed manner and aws_conn_id is None or
@@ -53,13 +59,48 @@ class QuickSightSensor(AwsBaseSensor[QuickSightHook]):
aws_hook_class = QuickSightHook
template_fields: Sequence[str] = ("data_set_id", "ingestion_id",
"aws_conn_id")
- def __init__(self, *, data_set_id: str, ingestion_id: str, **kwargs):
+ def __init__(
+ self,
+ *,
+ data_set_id: str,
+ ingestion_id: str,
+ max_retries: int = 2160,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ **kwargs,
+ ):
super().__init__(**kwargs)
self.data_set_id = data_set_id
self.ingestion_id = ingestion_id
+ self.max_retries = max_retries
+ self.deferrable = deferrable
self.success_status = "COMPLETED"
self.errored_statuses = ("FAILED", "CANCELLED")
+ def execute(self, context: Context) -> Any:
+ if self.deferrable:
+ self.defer(
+ trigger=QuickSightIngestionCompletedTrigger(
+ data_set_id=self.data_set_id,
+ ingestion_id=self.ingestion_id,
+ aws_account_id=self.hook.account_id,
+ waiter_delay=int(self.poke_interval),
+ waiter_max_attempts=self.max_retries,
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ botocore_config=self.botocore_config,
+ ),
+ method_name="execute_complete",
+ )
+ else:
+ return super().execute(context=context)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ validated_event = validate_execute_complete_event(event)
+ if validated_event["status"] != "success":
+ raise RuntimeError(f"Error while waiting for the Amazon QuickSight
ingestion: {validated_event}")
+ self.log.info("Amazon QuickSight SPICE ingestion `%s` completed.",
validated_event["ingestion_id"])
+
def poke(self, context: Context) -> bool:
"""
Pokes until the QuickSight Ingestion has successfully finished.
@@ -72,5 +113,5 @@ class QuickSightSensor(AwsBaseSensor[QuickSightHook]):
self.log.info("QuickSight Status: %s", quicksight_ingestion_state)
if quicksight_ingestion_state in self.errored_statuses:
error = self.hook.get_error_info(None, self.data_set_id,
self.ingestion_id)
- raise AirflowException(f"The QuickSight Ingestion failed. Error
info: {error}")
+ raise RuntimeError(f"The QuickSight Ingestion failed. Error info:
{error}")
return quicksight_ingestion_state == self.success_status
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/quicksight.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/quicksight.py
new file mode 100644
index 00000000000..752af4b929e
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/quicksight.py
@@ -0,0 +1,87 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+
+class QuickSightIngestionCompletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when an Amazon QuickSight SPICE ingestion is complete.
+
+ :param data_set_id: ID of the dataset used in the ingestion.
+ :param ingestion_id: ID for the ingestion.
+ :param aws_account_id: The ID of the AWS account that owns the dataset.
+ :param waiter_delay: The amount of time in seconds to wait between
attempts. (default: 30)
+ :param waiter_max_attempts: The maximum number of attempts to be made.
(default: 4320)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: AWS region_name. If not specified then the default
boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates.
+ :param botocore_config: Configuration dictionary (key-values) for botocore
client.
+ """
+
+ def __init__(
+ self,
+ *,
+ data_set_id: str,
+ ingestion_id: str,
+ aws_account_id: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 4320,
+ aws_conn_id: str | None = "aws_default",
+ region_name: str | None = None,
+ verify: bool | str | None = None,
+ botocore_config: dict | None = None,
+ ) -> None:
+ super().__init__(
+ serialized_fields={
+ "data_set_id": data_set_id,
+ "ingestion_id": ingestion_id,
+ "aws_account_id": aws_account_id,
+ },
+ waiter_name="ingestion_complete",
+ waiter_args={
+ "AwsAccountId": aws_account_id,
+ "DataSetId": data_set_id,
+ "IngestionId": ingestion_id,
+ },
+ failure_message="Amazon QuickSight SPICE ingestion failed.",
+ status_message="Status of Amazon QuickSight SPICE ingestion is",
+ status_queries=["Ingestion.IngestionStatus",
"Ingestion.ErrorInfo"],
+ return_key="ingestion_id",
+ return_value=ingestion_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ verify=verify,
+ botocore_config=botocore_config,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return QuickSightHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/waiters/quicksight.json
b/providers/amazon/src/airflow/providers/amazon/aws/waiters/quicksight.json
new file mode 100644
index 00000000000..652c3b3d376
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/waiters/quicksight.json
@@ -0,0 +1,48 @@
+{
+ "version": 2,
+ "waiters": {
+ "ingestion_complete": {
+ "delay": 30,
+ "maxAttempts": 60,
+ "operation": "DescribeIngestion",
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "Ingestion.IngestionStatus",
+ "expected": "INITIALIZED",
+ "state": "retry"
+ },
+ {
+ "matcher": "path",
+ "argument": "Ingestion.IngestionStatus",
+ "expected": "QUEUED",
+ "state": "retry"
+ },
+ {
+ "matcher": "path",
+ "argument": "Ingestion.IngestionStatus",
+ "expected": "RUNNING",
+ "state": "retry"
+ },
+ {
+ "matcher": "path",
+ "argument": "Ingestion.IngestionStatus",
+ "expected": "COMPLETED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "Ingestion.IngestionStatus",
+ "expected": "FAILED",
+ "state": "failure"
+ },
+ {
+ "matcher": "path",
+ "argument": "Ingestion.IngestionStatus",
+ "expected": "CANCELLED",
+ "state": "failure"
+ }
+ ]
+ }
+ }
+}
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index ea3e1f49437..1038b5d698f 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -1022,6 +1022,10 @@ def get_provider_info():
"integration-name": "AWS Database Migration Service",
"python-modules":
["airflow.providers.amazon.aws.triggers.dms"],
},
+ {
+ "integration-name": "Amazon QuickSight",
+ "python-modules":
["airflow.providers.amazon.aws.triggers.quicksight"],
+ },
],
"transfers": [
{
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_quicksight.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_quicksight.py
index 8969eb72fb2..a6d93d81d06 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_quicksight.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_quicksight.py
@@ -22,6 +22,7 @@ from unittest import mock
import pytest
from botocore.exceptions import ClientError
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
from airflow.providers.common.compat.sdk import AirflowException
@@ -183,7 +184,10 @@ class TestQuicksight:
):
mocked_get_error_info.return_value = "Something Bad Happen"
hook = QuickSightHook(aws_conn_id=None, region_name="us-east-1")
- with pytest.raises(AirflowException, match="Error info: Something Bad
Happen"):
+ with (
+ pytest.warns(AirflowProviderDeprecationWarning,
match="wait_for_ingestion"),
+ pytest.raises(AirflowException, match="Error info: Something Bad
Happen"),
+ ):
hook.wait_for_state(
aws_account_id, "data_set_id", "ingestion_id",
target_state={"COMPLETED"}, check_interval=0
)
@@ -193,7 +197,10 @@ class TestQuicksight:
@mock.patch.object(QuickSightHook, "get_status", return_value="CANCELLED")
def test_wait_for_state_canceled(self, _):
hook = QuickSightHook(aws_conn_id=None, region_name="us-east-1")
- with pytest.raises(AirflowException, match="The Amazon QuickSight
SPICE ingestion cancelled"):
+ with (
+ pytest.warns(AirflowProviderDeprecationWarning,
match="wait_for_ingestion"),
+ pytest.raises(AirflowException, match="The Amazon QuickSight SPICE
ingestion cancelled"),
+ ):
hook.wait_for_state(
"aws_account_id", "data_set_id", "ingestion_id",
target_state={"COMPLETED"}, check_interval=0
)
@@ -202,12 +209,11 @@ class TestQuicksight:
def test_wait_for_state_completed(self, mocked_get_status):
mocked_get_status.side_effect = ["INITIALIZED", "QUEUED", "RUNNING",
"COMPLETED"]
hook = QuickSightHook(aws_conn_id=None, region_name="us-east-1")
- assert (
- hook.wait_for_state(
+ with pytest.warns(AirflowProviderDeprecationWarning,
match="wait_for_ingestion"):
+ status = hook.wait_for_state(
"aws_account_id", "data_set_id", "ingestion_id",
target_state={"COMPLETED"}, check_interval=0
)
- == "COMPLETED"
- )
+ assert status == "COMPLETED"
assert mocked_get_status.call_count == 4
@pytest.mark.parametrize(
@@ -220,7 +226,7 @@ class TestQuicksight:
mocked_client.create_ingestion.return_value =
MOCK_CREATE_INGESTION_RESPONSE
hook = QuickSightHook(aws_conn_id=None, region_name="us-east-1")
- with mock.patch.object(QuickSightHook, "wait_for_state") as
mocked_wait_for_state:
+ with mock.patch.object(QuickSightHook, "wait_for_ingestion") as
mocked_wait_for_ingestion:
assert (
hook.create_ingestion(
data_set_id="DemoDataSet",
@@ -229,23 +235,24 @@ class TestQuicksight:
aws_account_id=aws_account_id,
wait_for_completion=wait_for_completion,
check_interval=0,
+ waiter_max_attempts=2,
)
== MOCK_CREATE_INGESTION_RESPONSE
)
if wait_for_completion:
- mocked_wait_for_state.assert_called_once_with(
- aws_account_id=expected_account_id,
+ mocked_wait_for_ingestion.assert_called_once_with(
data_set_id="DemoDataSet",
ingestion_id="DemoDataSet_Ingestion",
- target_state={"COMPLETED"},
- check_interval=0,
+ aws_account_id=expected_account_id,
+ waiter_delay=0,
+ waiter_max_attempts=2,
)
else:
- mocked_wait_for_state.assert_not_called()
+ mocked_wait_for_ingestion.assert_not_called()
mocked_client.create_ingestion.assert_called_with(AwsAccountId=expected_account_id,
**MOCK_DATA)
- def test_create_ingestion_exception(self, mocked_account_id,
mocked_client, caplog):
+ def test_create_ingestion_exception(self, mocked_account_id,
mocked_client):
mocked_client.create_ingestion.side_effect = ValueError("Fake Error")
hook = QuickSightHook(aws_conn_id=None)
with pytest.raises(ValueError, match="Fake Error"):
@@ -254,4 +261,46 @@ class TestQuicksight:
ingestion_id="DemoDataSet_Ingestion",
ingestion_type="INCREMENTAL_REFRESH",
)
- assert "create_ingestion API, error: Fake Error" in caplog.text
+
+ @pytest.mark.parametrize(("aws_account_id", "expected_account_id"),
ACCOUNT_TEST_CASES)
+ @mock.patch("airflow.providers.amazon.aws.hooks.quicksight.wait")
+ @mock.patch.object(QuickSightHook, "get_waiter")
+ def test_wait_for_ingestion(
+ self, mocked_get_waiter, mocked_wait, aws_account_id,
expected_account_id, mocked_account_id
+ ):
+ hook = QuickSightHook(aws_conn_id=None, region_name="us-east-1")
+ hook.wait_for_ingestion(
+ data_set_id="DemoDataSet",
+ ingestion_id="DemoDataSet_Ingestion",
+ aws_account_id=aws_account_id,
+ waiter_delay=1,
+ waiter_max_attempts=2,
+ )
+ mocked_get_waiter.assert_called_once_with("ingestion_complete")
+ mocked_wait.assert_called_once_with(
+ waiter=mocked_get_waiter.return_value,
+ waiter_delay=1,
+ waiter_max_attempts=2,
+ args={
+ "AwsAccountId": expected_account_id,
+ "DataSetId": "DemoDataSet",
+ "IngestionId": "DemoDataSet_Ingestion",
+ },
+ failure_message="Amazon QuickSight SPICE ingestion failed.",
+ status_message="Status of Amazon QuickSight SPICE ingestion is",
+ status_args=["Ingestion.IngestionStatus", "Ingestion.ErrorInfo"],
+ )
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.hooks.quicksight.wait",
+ side_effect=AirflowException("Waiter error: max attempts reached"),
+ )
+ @mock.patch.object(QuickSightHook, "get_waiter")
+ def test_wait_for_ingestion_failure(self, mocked_get_waiter, mocked_wait):
+ hook = QuickSightHook(aws_conn_id=None, region_name="us-east-1")
+ with pytest.raises(RuntimeError, match="max attempts reached"):
+ hook.wait_for_ingestion(
+ data_set_id="DemoDataSet",
+ ingestion_id="DemoDataSet_Ingestion",
+ aws_account_id="123456789012",
+ )
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_quicksight.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_quicksight.py
index 5b0edd1ff5b..f3043c05073 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_quicksight.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_quicksight.py
@@ -19,8 +19,13 @@ from __future__ import annotations
from unittest import mock
+import pytest
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
from airflow.providers.amazon.aws.operators.quicksight import
QuickSightCreateIngestionOperator
+from airflow.providers.amazon.aws.triggers.quicksight import
QuickSightIngestionCompletedTrigger
+from airflow.providers.common.compat.sdk import TaskDeferred
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
@@ -38,6 +43,13 @@ MOCK_RESPONSE = {
}
[email protected]
+def mocked_account_id():
+ with mock.patch.object(QuickSightHook, "account_id",
new_callable=mock.PropertyMock) as m:
+ m.return_value = AWS_ACCOUNT_ID
+ yield m
+
+
class TestQuickSightCreateIngestionOperator:
def setup_method(self):
self.default_op_kwargs = {
@@ -71,16 +83,94 @@ class TestQuickSightCreateIngestionOperator:
assert op.hook._region_name is None
assert op.hook._verify is None
assert op.hook._config is None
+ assert op.wait_for_completion is True
+ assert op.waiter_delay == 30
+ assert op.waiter_max_attempts == 4320
+ assert op.deferrable is False
+ @pytest.mark.parametrize(
+ ("wait_for_completion", "deferrable"),
+ [
+ pytest.param(False, False, id="no_wait"),
+ pytest.param(True, False, id="wait"),
+ pytest.param(False, True, id="defer"),
+ ],
+ )
+ @mock.patch.object(QuickSightCreateIngestionOperator, "defer")
+ @mock.patch.object(QuickSightHook, "wait_for_ingestion")
@mock.patch.object(QuickSightHook, "create_ingestion",
return_value=MOCK_RESPONSE)
- def test_execute(self, mock_create_ingestion):
- QuickSightCreateIngestionOperator(**self.default_op_kwargs).execute({})
+ def test_execute_wait_combinations(
+ self,
+ mock_create_ingestion,
+ mock_wait_for_ingestion,
+ mock_defer,
+ wait_for_completion,
+ deferrable,
+ mocked_account_id,
+ ):
+ op = QuickSightCreateIngestionOperator(
+ **self.default_op_kwargs, wait_for_completion=wait_for_completion,
deferrable=deferrable
+ )
+
+ response = op.execute({})
+
+ assert response == MOCK_RESPONSE
mock_create_ingestion.assert_called_once_with(
data_set_id=DATA_SET_ID,
ingestion_id=INGESTION_ID,
- ingestion_type="FULL_REFRESH",
- wait_for_completion=True,
- check_interval=30,
+ ingestion_type=INGESTION_TYPE,
+ wait_for_completion=False,
+ )
+ assert mock_wait_for_ingestion.call_count == wait_for_completion
+ assert mock_defer.call_count == deferrable
+
+ @mock.patch.object(QuickSightHook, "create_ingestion",
return_value=MOCK_RESPONSE)
+ def test_execute_deferrable(self, mock_create_ingestion,
mocked_account_id):
+ op = QuickSightCreateIngestionOperator(
+ **self.default_op_kwargs, deferrable=True, waiter_delay=15,
waiter_max_attempts=3
+ )
+
+ with pytest.raises(TaskDeferred) as defer_exc:
+ op.execute({})
+
+ trigger = defer_exc.value.trigger
+ assert isinstance(trigger, QuickSightIngestionCompletedTrigger)
+ assert defer_exc.value.method_name == "execute_complete"
+ assert defer_exc.value.kwargs == {"ingestion": MOCK_RESPONSE}
+ _, kwargs = trigger.serialize()
+ assert kwargs["data_set_id"] == DATA_SET_ID
+ assert kwargs["ingestion_id"] == INGESTION_ID
+ assert kwargs["aws_account_id"] == AWS_ACCOUNT_ID
+ assert kwargs["waiter_delay"] == 15
+ assert kwargs["waiter_max_attempts"] == 3
+
+ def test_execute_complete(self):
+ op = QuickSightCreateIngestionOperator(**self.default_op_kwargs)
+ event = {"status": "success", "ingestion_id": INGESTION_ID}
+
+ assert op.execute_complete({}, event, ingestion=MOCK_RESPONSE) ==
MOCK_RESPONSE
+
+ def test_execute_complete_failure(self):
+ op = QuickSightCreateIngestionOperator(**self.default_op_kwargs)
+ event = {"status": "error", "message": "test failure", "ingestion_id":
INGESTION_ID}
+
+ with pytest.raises(RuntimeError, match="Error while running"):
+ op.execute_complete({}, event)
+
+ @pytest.mark.parametrize("check_interval", [10, "10"], ids=["int",
"templated_str"])
+ @mock.patch.object(QuickSightHook, "wait_for_ingestion")
+ @mock.patch.object(QuickSightHook, "create_ingestion",
return_value=MOCK_RESPONSE)
+ def test_check_interval_deprecation(self, mock_create_ingestion,
mock_wait_for_ingestion, check_interval):
+ op = QuickSightCreateIngestionOperator(**self.default_op_kwargs,
check_interval=check_interval)
+
+ with pytest.warns(AirflowProviderDeprecationWarning,
match="check_interval"):
+ op.execute({})
+
+ mock_wait_for_ingestion.assert_called_once_with(
+ data_set_id=DATA_SET_ID,
+ ingestion_id=INGESTION_ID,
+ waiter_delay=10,
+ waiter_max_attempts=4320,
)
def test_template_fields(self):
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_quicksight.py
b/providers/amazon/tests/unit/amazon/aws/sensors/test_quicksight.py
index 277749edd52..b65b8e6234a 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_quicksight.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_quicksight.py
@@ -23,10 +23,12 @@ import pytest
from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
from airflow.providers.amazon.aws.sensors.quicksight import QuickSightSensor
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.amazon.aws.triggers.quicksight import
QuickSightIngestionCompletedTrigger
+from airflow.providers.common.compat.sdk import TaskDeferred
DATA_SET_ID = "DemoDataSet"
INGESTION_ID = "DemoDataSet_Ingestion"
+AWS_ACCOUNT_ID = "123456789012"
@pytest.fixture
@@ -41,6 +43,13 @@ def mocked_get_error_info():
yield m
[email protected]
+def mocked_account_id():
+ with mock.patch.object(QuickSightHook, "account_id",
new_callable=mock.PropertyMock) as m:
+ m.return_value = AWS_ACCOUNT_ID
+ yield m
+
+
class TestQuickSightSensor:
def setup_method(self):
self.default_op_kwargs = {
@@ -74,6 +83,8 @@ class TestQuickSightSensor:
assert sensor.hook._region_name is None
assert sensor.hook._verify is None
assert sensor.hook._config is None
+ assert sensor.max_retries == 2160
+ assert sensor.deferrable is False
@pytest.mark.parametrize("status", ["COMPLETED"])
def test_poke_completed(self, status, mocked_get_status):
@@ -91,7 +102,38 @@ class TestQuickSightSensor:
def test_poke_terminated_status(self, status, mocked_get_status,
mocked_get_error_info):
mocked_get_status.return_value = status
mocked_get_error_info.return_value = "something bad happen"
- with pytest.raises(AirflowException, match="Error info: something bad
happen"):
+ with pytest.raises(RuntimeError, match="Error info: something bad
happen"):
QuickSightSensor(**self.default_op_kwargs).poke({})
mocked_get_status.assert_called_once_with(None, DATA_SET_ID,
INGESTION_ID)
mocked_get_error_info.assert_called_once_with(None, DATA_SET_ID,
INGESTION_ID)
+
+ def test_execute_deferrable(self, mocked_account_id):
+ sensor = QuickSightSensor(
+ **self.default_op_kwargs, deferrable=True, poke_interval=20.5,
max_retries=3
+ )
+
+ with pytest.raises(TaskDeferred) as defer_exc:
+ sensor.execute({})
+
+ trigger = defer_exc.value.trigger
+ assert isinstance(trigger, QuickSightIngestionCompletedTrigger)
+ assert defer_exc.value.method_name == "execute_complete"
+ _, kwargs = trigger.serialize()
+ assert kwargs["data_set_id"] == DATA_SET_ID
+ assert kwargs["ingestion_id"] == INGESTION_ID
+ assert kwargs["aws_account_id"] == AWS_ACCOUNT_ID
+ assert kwargs["waiter_delay"] == 20
+ assert kwargs["waiter_max_attempts"] == 3
+
+ def test_execute_complete(self):
+ sensor = QuickSightSensor(**self.default_op_kwargs)
+ event = {"status": "success", "ingestion_id": INGESTION_ID}
+
+ assert sensor.execute_complete({}, event) is None
+
+ def test_execute_complete_failure(self):
+ sensor = QuickSightSensor(**self.default_op_kwargs)
+ event = {"status": "error", "message": "Waiter error: max attempts
reached"}
+
+ with pytest.raises(RuntimeError, match="Error while waiting"):
+ sensor.execute_complete({}, event)
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_quicksight.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_quicksight.py
new file mode 100644
index 00000000000..bb4b61e50db
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_quicksight.py
@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+from airflow.providers.amazon.aws.triggers.quicksight import
QuickSightIngestionCompletedTrigger
+from airflow.triggers.base import TriggerEvent
+
+from unit.amazon.aws.utils.test_waiter import assert_expected_waiter_type
+
+BASE_TRIGGER_CLASSPATH = "airflow.providers.amazon.aws.triggers.quicksight."
+EXPECTED_WAITER_NAME = "ingestion_complete"
+DATA_SET_ID = "DemoDataSet"
+INGESTION_ID = "DemoDataSet_Ingestion"
+AWS_ACCOUNT_ID = "123456789012"
+
+
+class TestQuickSightIngestionCompletedTrigger:
+ def test_serialization(self):
+ """Assert that arguments and classpath are correctly serialized."""
+ trigger = QuickSightIngestionCompletedTrigger(
+ data_set_id=DATA_SET_ID, ingestion_id=INGESTION_ID,
aws_account_id=AWS_ACCOUNT_ID
+ )
+ classpath, kwargs = trigger.serialize()
+ assert classpath == BASE_TRIGGER_CLASSPATH +
"QuickSightIngestionCompletedTrigger"
+ assert kwargs.get("data_set_id") == DATA_SET_ID
+ assert kwargs.get("ingestion_id") == INGESTION_ID
+ assert kwargs.get("aws_account_id") == AWS_ACCOUNT_ID
+
+ @pytest.mark.asyncio
+ @mock.patch.object(QuickSightHook, "get_waiter")
+ @mock.patch.object(QuickSightHook, "get_async_conn")
+ async def test_run_success(self, mock_async_conn, mock_get_waiter):
+ mock_async_conn.__aenter__.return_value = mock.MagicMock()
+ mock_get_waiter().wait = AsyncMock()
+ trigger = QuickSightIngestionCompletedTrigger(
+ data_set_id=DATA_SET_ID, ingestion_id=INGESTION_ID,
aws_account_id=AWS_ACCOUNT_ID
+ )
+
+ generator = trigger.run()
+ response = await generator.asend(None)
+
+ assert response == TriggerEvent({"status": "success", "ingestion_id":
INGESTION_ID})
+ assert_expected_waiter_type(mock_get_waiter, EXPECTED_WAITER_NAME)
+ mock_get_waiter().wait.assert_called_once()
diff --git a/providers/amazon/tests/unit/amazon/aws/waiters/test_quicksight.py
b/providers/amazon/tests/unit/amazon/aws/waiters/test_quicksight.py
new file mode 100644
index 00000000000..9ccbe8e0edc
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/waiters/test_quicksight.py
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+
+import boto3
+import botocore
+import pytest
+
+from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
+
+INGESTION_WAITER_ARGS = {
+ "AwsAccountId": "123456789012",
+ "DataSetId": "DemoDataSet",
+ "IngestionId": "DemoDataSet_Ingestion",
+}
+
+
+class TestQuickSightCustomWaiters:
+ def test_service_waiters(self):
+ assert "ingestion_complete" in QuickSightHook().list_waiters()
+
+
+class TestQuickSightCustomWaitersBase:
+ @pytest.fixture(autouse=True)
+ def mock_conn(self, monkeypatch):
+ self.client = boto3.client("quicksight")
+ monkeypatch.setattr(QuickSightHook, "conn", self.client)
+
+
+class TestQuickSightIngestionCompleteWaiter(TestQuickSightCustomWaitersBase):
+ WAITER_NAME = "ingestion_complete"
+
+ @pytest.fixture
+ def mock_describe_ingestion(self):
+ with mock.patch.object(self.client, "describe_ingestion") as
mock_getter:
+ yield mock_getter
+
+ def test_ingestion_complete(self, mock_describe_ingestion):
+ mock_describe_ingestion.return_value = {"Ingestion":
{"IngestionStatus": "COMPLETED"}}
+
+
QuickSightHook().get_waiter(self.WAITER_NAME).wait(**INGESTION_WAITER_ARGS)
+
+ @pytest.mark.parametrize("state", ["FAILED", "CANCELLED"])
+ def test_ingestion_failed(self, state, mock_describe_ingestion):
+ mock_describe_ingestion.return_value = {"Ingestion":
{"IngestionStatus": state}}
+
+ with pytest.raises(botocore.exceptions.WaiterError):
+
QuickSightHook().get_waiter(self.WAITER_NAME).wait(**INGESTION_WAITER_ARGS)
+
+ def test_ingestion_wait(self, mock_describe_ingestion):
+ mock_describe_ingestion.side_effect = [
+ {"Ingestion": {"IngestionStatus": "INITIALIZED"}},
+ {"Ingestion": {"IngestionStatus": "QUEUED"}},
+ {"Ingestion": {"IngestionStatus": "RUNNING"}},
+ {"Ingestion": {"IngestionStatus": "COMPLETED"}},
+ ]
+
+ QuickSightHook().get_waiter(self.WAITER_NAME).wait(
+ **INGESTION_WAITER_ARGS, WaiterConfig={"Delay": 0.01,
"MaxAttempts": 4}
+ )
