From 2a5c0a8231c7f2d4cefb8392a1c5f9539aafc986 Mon Sep 17 00:00:00 2001 From: victorphoenix3 Date: Thu, 10 Feb 2022 16:09:51 +0530 Subject: [PATCH] refactors polling logic for athena queries --- airflow/providers/amazon/aws/hooks/athena.py | 13 ++++++++----- tests/providers/amazon/aws/operators/test_athena.py | 6 +++--- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/athena.py b/airflow/providers/amazon/aws/hooks/athena.py index 1c5c42fa7a06d..bbd08f9034529 100644 --- a/airflow/providers/amazon/aws/hooks/athena.py +++ b/airflow/providers/amazon/aws/hooks/athena.py @@ -54,6 +54,11 @@ class AthenaHook(AwsBaseHook): 'CANCELLED', ) SUCCESS_STATES = ('SUCCEEDED',) + TERMINAL_STATES = ( + "SUCCEEDED", + "FAILED", + "CANCELLED", + ) def __init__(self, *args: Any, sleep_time: int = 30, **kwargs: Any) -> None: super().__init__(client_type='athena', *args, **kwargs) # type: ignore @@ -200,16 +205,14 @@ def poll_query_status(self, query_execution_id: str, max_tries: Optional[int] = query_state = self.check_query_status(query_execution_id) if query_state is None: self.log.info('Trial %s: Invalid query state. Retrying again', try_number) - elif query_state in self.INTERMEDIATE_STATES: - self.log.info( - 'Trial %s: Query is still in an intermediate state - %s', try_number, query_state - ) - else: + elif query_state in self.TERMINAL_STATES: self.log.info( 'Trial %s: Query execution completed. Final state is %s}', try_number, query_state ) final_query_state = query_state break + else: + self.log.info('Trial %s: Query is still in non-terminal state - %s', try_number, query_state) if max_tries and try_number >= max_tries: # Break loop if max_tries reached final_query_state = query_state break diff --git a/tests/providers/amazon/aws/operators/test_athena.py b/tests/providers/amazon/aws/operators/test_athena.py index 060bfb40c182a..a1cfc8478d9e8 100644 --- a/tests/providers/amazon/aws/operators/test_athena.py +++ b/tests/providers/amazon/aws/operators/test_athena.py @@ -72,7 +72,7 @@ def test_init(self): assert self.athena.hook.sleep_time == 0 - @mock.patch.object(AthenaHook, 'check_query_status', side_effect=("SUCCESS",)) + @mock.patch.object(AthenaHook, 'check_query_status', side_effect=("SUCCEEDED",)) @mock.patch.object(AthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AthenaHook, 'get_conn') def test_hook_run_small_success_query(self, mock_conn, mock_run_query, mock_check_query_status): @@ -92,7 +92,7 @@ def test_hook_run_small_success_query(self, mock_conn, mock_run_query, mock_chec side_effect=( "RUNNING", "RUNNING", - "SUCCESS", + "SUCCEEDED", ), ) @mock.patch.object(AthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @@ -202,7 +202,7 @@ def test_hook_run_failed_query_with_max_tries(self, mock_conn, mock_run_query, m ) assert mock_check_query_status.call_count == 3 - @mock.patch.object(AthenaHook, 'check_query_status', side_effect=("SUCCESS",)) + @mock.patch.object(AthenaHook, 'check_query_status', side_effect=("SUCCEEDED",)) @mock.patch.object(AthenaHook, 'run_query', return_value=ATHENA_QUERY_ID) @mock.patch.object(AthenaHook, 'get_conn') def test_return_value(self, mock_conn, mock_run_query, mock_check_query_status):