vincbeck commented on code in PR #70218:
URL: https://github.com/apache/airflow/pull/70218#discussion_r3658259138
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/quicksight.py:
##########
@@ -66,31 +71,63 @@ def create_ingestion(
: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 = 60,
+ ) -> 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 QuickSightIngestionFailedError: 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 QuickSightIngestionFailedError(str(e)) from e
Review Comment:
Do we need a custom exception? `RuntimeError` suffice here no?
##########
providers/amazon/docs/changelog.rst:
##########
@@ -26,6 +26,13 @@
Changelog
---------
+.. warning::
Review Comment:
Do not update changelog. This will be done by the release manager
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/quicksight.py:
##########
@@ -71,22 +82,73 @@ def __init__(
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 = 60,
+ check_interval: int | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
):
super().__init__(**kwargs)
+ if check_interval is not None:
Review Comment:
This is kind of weird, you set a default value and then you check whether a
value is assigned to it. In other words, anyone using this operator without
setting `check_interval` will get a deprecation warning. I would leave the
default value to `None` and keep following the deprecation warning process
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
