This is an automated email from the ASF dual-hosted git repository.
onikolas 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 4e61b654c3 Fetch status in while loop so as to not exit too early
(#31804)
4e61b654c3 is described below
commit 4e61b654c39d47bfe72e3004be45ba1bd97be660
Author: Niko Oliveira <[email protected]>
AuthorDate: Thu Jun 8 14:56:51 2023 -0700
Fetch status in while loop so as to not exit too early (#31804)
This fixes a condition described in #30626 for the AWS Quicksight hook
---
airflow/providers/amazon/aws/hooks/quicksight.py | 10 ++++------
tests/providers/amazon/aws/hooks/test_quicksight.py | 12 +++++++++++-
2 files changed, 15 insertions(+), 7 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/quicksight.py
b/airflow/providers/amazon/aws/hooks/quicksight.py
index 9ba840ea02..2c61252b74 100644
--- a/airflow/providers/amazon/aws/hooks/quicksight.py
+++ b/airflow/providers/amazon/aws/hooks/quicksight.py
@@ -153,19 +153,17 @@ class QuickSightHook(AwsBaseHook):
will check the status of QuickSight Ingestion
:return: response of describe_ingestion call after Ingestion is is done
"""
- sec = 0
- status = self.get_status(aws_account_id, data_set_id, ingestion_id)
- while status in self.NON_TERMINAL_STATES and status != target_state:
+ while True:
+ status = self.get_status(aws_account_id, data_set_id, ingestion_id)
self.log.info("Current status is %s", status)
if status in self.FAILED_STATES:
info = self.get_error_info(aws_account_id, data_set_id,
ingestion_id)
raise AirflowException(f"The Amazon QuickSight Ingestion
failed. Error info: {info}")
if status == "CANCELLED":
raise AirflowException("The Amazon QuickSight SPICE ingestion
cancelled!")
- # wait and try again
+ if status not in self.NON_TERMINAL_STATES or status ==
target_state:
+ break
time.sleep(check_interval)
- sec += check_interval
- status = self.get_status(aws_account_id, data_set_id, ingestion_id)
self.log.info("QuickSight Ingestion completed")
return status
diff --git a/tests/providers/amazon/aws/hooks/test_quicksight.py
b/tests/providers/amazon/aws/hooks/test_quicksight.py
index 5918c20c54..9c8ef16ce8 100644
--- a/tests/providers/amazon/aws/hooks/test_quicksight.py
+++ b/tests/providers/amazon/aws/hooks/test_quicksight.py
@@ -22,12 +22,12 @@ from unittest import mock
import pytest
from botocore.exceptions import ClientError
+from airflow.exceptions import AirflowException
from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
from airflow.providers.amazon.aws.hooks.sts import StsHook
AWS_ACCOUNT_ID = "123456789012"
-
MOCK_DATA = {
"DataSetId": "DemoDataSet",
"IngestionId": "DemoDataSet_Ingestion",
@@ -99,6 +99,16 @@ class TestQuicksight:
mock_conn.return_value.create_ingestion.assert_called_with(**expected_call_params)
assert result == MOCK_CREATE_INGESTION_RESPONSE
+ @mock.patch.object(QuickSightHook, "get_conn")
+ @mock.patch.object(QuickSightHook, "get_status")
+ def test_fast_failing_ingestion(self, mock_get_status, mock_conn):
+ quicksight_hook = QuickSightHook(aws_conn_id="aws_default",
region_name="us-east-1")
+ mock_get_status.return_value = "FAILED"
+ with pytest.raises(AirflowException):
+ quicksight_hook.wait_for_state(
+ "account_id", "data_set_id", "ingestion_id",
target_state={"COMPLETED"}, check_interval=1
+ )
+
@mock.patch.object(StsHook, "get_conn")
@mock.patch.object(StsHook, "get_account_number")
def test_create_ingestion_exception(self, mock_get_account_number,
sts_conn):
