DAS-2427: Adds retry logic to harmony requests Session (#119)

Author: flamingbear

harmony/client.py

@@ -28,6 +28,8 @@
 import platform
 from uuid import UUID
 from requests import Response
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 from requests.exceptions import JSONDecodeError
 import requests.models
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -164,6 +166,20 @@ def _session(self):
                 self.session = create_session(self.config, token=self.token)
             else:
                 self.session = create_session(self.config, auth=self.auth)
+            # Add retry logic
+            retry_strategy = Retry(
+                total=3,
+                backoff_factor=1,  # Wait 1, 2, 4 seconds between retries
+                backoff_jitter=0.5,
+                status_forcelist=[429, 500, 502, 503, 504],
+                allowed_methods=["GET"],
+                raise_on_status=False,
+
+            )
+            adapter = HTTPAdapter(max_retries=retry_strategy)
+            self.session.mount("https://", adapter)
+            self.session.mount("http://", adapter)
+
         return self.session
 
     def _http_method(self, request: BaseRequest) -> str:

tests/test_client.py

@@ -10,7 +10,9 @@
 import dateutil.parser
 import pytest
 import responses
+from responses import registries
 
+from urllib3.util.retry import Retry
 from harmony.request import BBox, Collection, LinkType, Request, Dimension, CapabilitiesRequest, \
     AddLabelsRequest, DeleteLabelsRequest, JobsRequest
 from harmony.client import Client, ProcessingFailedException, DEFAULT_JOB_LABEL
@@ -1482,12 +1484,6 @@ def test_handle_error_response_with_description_key():
         status=500,
         json=error
     )
-    responses.add(
-        responses.GET,
-        expected_status_url(job_id),
-        status=500,
-        json=error
-    )
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).submit(request)
     assert str(e.value) == f"('Internal Server Error', '{error['description']}')"
@@ -1499,6 +1495,7 @@ def test_handle_error_response_with_description_key():
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).progress(job_id)
     assert str(e.value) == f"('Internal Server Error', '{error['description']}')"
+    assert len(responses.calls) == 9
 
 @responses.activate
 def test_handle_error_response_no_description_key():
@@ -1521,12 +1518,6 @@ def test_handle_error_response_no_description_key():
         status=500,
         json=error
     )
-    responses.add(
-        responses.GET,
-        expected_status_url(job_id),
-        status=500,
-        json=error
-    )
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).submit(request)
     assert "500 Server Error: Internal Server Error for url" in str(e.value)
@@ -1538,6 +1529,8 @@ def test_handle_error_response_no_description_key():
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).progress(job_id)
     assert "500 Server Error: Internal Server Error for url" in str(e.value)
+    # Check retries
+    assert len(responses.calls) == 9 # (1 + 4 + 4) # Post isn't retried.
 
 @responses.activate
 def test_handle_error_response_no_json():
@@ -1559,12 +1552,6 @@ def test_handle_error_response_no_json():
         status=500,
         body='error'
     )
-    responses.add(
-        responses.GET,
-        expected_status_url(job_id),
-        status=500,
-        body='error'
-    )
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).submit(request)
     assert "500 Server Error: Internal Server Error for url" in str(e.value)
@@ -1576,6 +1563,9 @@ def test_handle_error_response_no_json():
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).progress(job_id)
     assert "500 Server Error: Internal Server Error for url" in str(e.value)
+    # Check retries
+    assert len(responses.calls) == 9 # (1 + 4 + 4)
+
 
 @responses.activate
 def test_handle_error_response_invalid_json():
@@ -1597,12 +1587,6 @@ def test_handle_error_response_invalid_json():
         status=500,
         json='error'
     )
-    responses.add(
-        responses.GET,
-        expected_status_url(job_id),
-        status=500,
-        json='error'
-    )
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).submit(request)
     assert "500 Server Error: Internal Server Error for url" in str(e.value)
@@ -1614,6 +1598,87 @@ def test_handle_error_response_invalid_json():
     with pytest.raises(Exception) as e:
         Client(should_validate_auth=False).progress(job_id)
     assert "500 Server Error: Internal Server Error for url" in str(e.value)
+    # Check retries
+    assert len(responses.calls) == 9 # (1 + 4 + 4)
+
+@responses.activate
+def test_handle_non_transient_error_no_retry():
+    job_id = '89733-badc-1324'
+    collection = Collection(id='F229040468263-STRP')
+    request = Request(
+        collection=collection,
+        spatial=BBox(-107, 40, -105, 42)
+    )
+    responses.add(
+        responses.POST,
+        expected_submit_url(collection.id),
+        status=200,
+        json=expected_job(collection.id, 'abcd-1234')
+    )
+    responses.add(
+        responses.GET,
+        expected_status_url(job_id),
+        status=431,
+        json='no-retry error'
+    )
+
+    Client(should_validate_auth=False).submit(request)
+
+    with pytest.raises(Exception) as e:
+        Client(should_validate_auth=False).status(job_id)
+    assert "431 Client Error: Request Header Fields Too Large for url" in str(e.value)
+
+    # Check no retries
+    assert len(responses.calls) == 2
+
+@responses.activate(registry=registries.OrderedRegistry)
+def test_handle_transient_error_responses():
+    job_id = '3141592653-acbd-1234'
+    collection = Collection(id='C1342468263-ANYTHING')
+    request = Request(
+        collection=collection,
+        spatial=BBox(-107, 40, -105, 42)
+    )
+    responses.add(
+        responses.POST,
+        expected_submit_url(collection.id),
+        status=200,
+        json=expected_job(collection.id, 'abcd-1234'),
+    )
+    first_success = responses.add(
+        responses.GET,
+        expected_status_url(job_id),
+        status=200,
+        json=expected_job(collection.id, job_id)
+    )
+    first_retry = responses.add(
+        responses.GET,
+        expected_status_url(job_id),
+        status=502,
+        json='error'
+    )
+    second_retry = responses.add(
+        responses.GET,
+        expected_status_url(job_id),
+        status=502,
+        json='error'
+    )
+    last_success = responses.add(
+        responses.GET,
+        expected_status_url(job_id),
+        status=200,
+        json=expected_job(collection.id, job_id)
+    )
+
+    Client(should_validate_auth=False).submit(request)
+    Client(should_validate_auth=False).status(job_id)
+    Client(should_validate_auth=False).status(job_id)
+
+    assert first_success.call_count == 1
+    assert first_retry.call_count == 1
+    assert second_retry.call_count == 1
+    assert last_success.call_count == 1
+
 
 def test_request_as_curl_get():
     collection = Collection(id='C1940468263-POCLOUD')
@@ -1845,4 +1910,4 @@ def test_client_custom_environment_not_affected_by_env_var():
     assert client.config.environment == Environment.UAT
     assert client.config.harmony_hostname == 'harmony.uat.earthdata.nasa.gov'
 
-    del os.environ['ENVIRONMENT']
+    del os.environ['ENVIRONMENT']