Repository: apache/airflow

Issue:
Refactor poll_query_status in EMRContainerHook
### Body

The goal is to refactor the code so that we can remove this TODO
https://github.com/apache/airflow/blob/7640ba4e8ee239d6e2bbf950d53d624b9df93059/airflow/providers/amazon/aws/hooks/emr_containers.py#L174-L176

More information about the concerns can be found on https://github.com/apache/airflow/pull/16766#discussion_r668089559

### Committer

- [X] I acknowledge that I am a maintainer/committer of the Apache Airflow project.


Discussion:
Hi. I would like to work on this issue. Can you assign it to me?
Go ahead!
Thanks @patrikhv - let me know if I can help in any way! (original PR contributor and EMR team member)
Hi, I am working on that, but I do not have any idea how to refactor it. I already just did a more explicit call, if the query is in a completed stage. Can you give me any tips?
Had a chat with dacort.  They mentioned that the change they had in mind was that we don't actually confirm that the state is at a terminal state, only that it is not at an intermediate state.  I can't think of any other changes I'd make to it.
@patrikhv - It looks to me like adding in the explicit "is it a terminal state?" check is all that is needed here.  You want to put in the PR?
@ferruzzi since patrikhv is not working on it. Do you want to put a PR to solve this?
@eladkal I have a few other things I'm working on.  I can have a look when I get time, but it'll likely be a couple weeks unless somebody gets to it before then.
Hello! I am new to the community and would like to take this on as my first issue. May I submit a PR for this? @eladkal 
Feel fre to attempt it :). Maybe @ferruzzi will be able to help with reviewing it :)

Resolving patch:
diff --git a/airflow/providers/amazon/aws/hooks/emr.py b/airflow/providers/amazon/aws/hooks/emr.py
--- a/airflow/providers/amazon/aws/hooks/emr.py
+++ b/airflow/providers/amazon/aws/hooks/emr.py
@@ -114,6 +114,12 @@ class EmrContainerHook(AwsBaseHook):
         "CANCEL_PENDING",
     )
     SUCCESS_STATES = ("COMPLETED",)
+    TERMINAL_STATES = (
+        "COMPLETED",
+        "FAILED",
+        "CANCELLED",
+        "CANCEL_PENDING",
+    )
 
     def __init__(self, *args: Any, virtual_cluster_id: Optional[str] = None, **kwargs: Any) -> None:
         super().__init__(client_type="emr-containers", *args, **kwargs)  # type: ignore
@@ -228,19 +234,16 @@ def poll_query_status(
         try_number = 1
         final_query_state = None  # Query state when query reaches final state or max_tries reached
 
-        # TODO: Make this logic a little bit more robust.
-        # Currently this polls until the state is *not* one of the INTERMEDIATE_STATES
-        # While that should work in most cases...it might not. :)
         while True:
             query_state = self.check_query_status(job_id)
             if query_state is None:
                 self.log.info("Try %s: Invalid query state. Retrying again", try_number)
-            elif query_state in self.INTERMEDIATE_STATES:
-                self.log.info("Try %s: Query is still in an intermediate state - %s", try_number, query_state)
-            else:
+            elif query_state in self.TERMINAL_STATES:
                 self.log.info("Try %s: Query execution completed. Final state is %s", try_number, query_state)
                 final_query_state = query_state
                 break
+            else:
+                self.log.info("Try %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


Test patch:
diff --git a/tests/providers/amazon/aws/hooks/test_emr_containers.py b/tests/providers/amazon/aws/hooks/test_emr_containers.py
--- a/tests/providers/amazon/aws/hooks/test_emr_containers.py
+++ b/tests/providers/amazon/aws/hooks/test_emr_containers.py
@@ -28,6 +28,22 @@
     'virtualClusterId': 'vc1234',
 }
 
+JOB1_RUN_DESCRIPTION = {
+    'jobRun': {
+        'id': 'job123456',
+        'virtualClusterId': 'vc1234',
+        'state': 'COMPLETED',
+    }
+}
+
+JOB2_RUN_DESCRIPTION = {
+    'jobRun': {
+        'id': 'job123456',
+        'virtualClusterId': 'vc1234',
+        'state': 'RUNNING',
+    }
+}
+
 
 class TestEmrContainerHook(unittest.TestCase):
     def setUp(self):
@@ -55,3 +71,29 @@ def test_submit_job(self, mock_session):
             client_request_token="uuidtoken",
         )
         assert emr_containers_job == 'job123456'
+
+    @mock.patch("boto3.session.Session")
+    def test_query_status_polling_when_terminal(self, mock_session):
+        emr_client_mock = mock.MagicMock()
+        emr_session_mock = mock.MagicMock()
+        emr_session_mock.client.return_value = emr_client_mock
+        mock_session.return_value = emr_session_mock
+        emr_client_mock.describe_job_run.return_value = JOB1_RUN_DESCRIPTION
+
+        query_status = self.emr_containers.poll_query_status(job_id='job123456')
+        # should only poll once since query is already in terminal state
+        emr_client_mock.describe_job_run.assert_called_once()
+        assert query_status == 'COMPLETED'
+
+    @mock.patch("boto3.session.Session")
+    def test_query_status_polling_with_timeout(self, mock_session):
+        emr_client_mock = mock.MagicMock()
+        emr_session_mock = mock.MagicMock()
+        emr_session_mock.client.return_value = emr_client_mock
+        mock_session.return_value = emr_session_mock
+        emr_client_mock.describe_job_run.return_value = JOB2_RUN_DESCRIPTION
+
+        query_status = self.emr_containers.poll_query_status(job_id='job123456', max_tries=2)
+        # should poll until max_tries is reached since query is in non-terminal state
+        assert emr_client_mock.describe_job_run.call_count == 2
+        assert query_status == 'RUNNING'