Retry requests

kbcstorage/base.py

@@ -9,6 +9,7 @@
 .. _Storage API documentation:
     http://docs.keboola.apiary.io/
 """
+from . import retry_requests
 import requests
 
 
@@ -66,7 +67,7 @@ def _get_raw(self, url, params=None, **kwargs):
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
 
-        r = requests.get(url, params, headers=headers, **kwargs)
+        r = retry_requests.get(url, params=params, headers=headers, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:
@@ -111,7 +112,7 @@ def _post(self, *args, **kwargs):
         """
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
-        r = requests.post(headers=headers, *args, **kwargs)
+        r = retry_requests.post(headers=headers, *args, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:
@@ -137,7 +138,7 @@ def _put(self, *args, **kwargs):
         """
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
-        r = requests.put(headers=headers, *args, **kwargs)
+        r = retry_requests.put(headers=headers, *args, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:
@@ -163,7 +164,7 @@ def _delete(self, *args, **kwargs):
         """
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
-        r = requests.delete(headers=headers, *args, **kwargs)
+        r = retry_requests.delete(headers=headers, *args, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:

kbcstorage/retry_requests.py

@@ -0,0 +1,31 @@
+import time
+import requests
+
+MAX_RETRIES = 5
+BACKOFF_FACTOR = 1.0
+RETRY_STATUS_CODES = {500, 502, 503, 504}
+
+
+def _get_backoff_time(retry_count):
+    return BACKOFF_FACTOR * (2 ** retry_count)
+
+def _retry_request(request_func, url, *args, **kwargs):
+    response = request_func(url, *args, **kwargs)
+    for retry_count in range(1, MAX_RETRIES):
+        if response.status_code not in RETRY_STATUS_CODES:
+            return response
+        time.sleep(_get_backoff_time(retry_count))
+        response = request_func(url, **kwargs)
+    return response
+
+def get(url, *args, **kwargs):
+    return _retry_request(requests.get, url, *args, **kwargs)
+
+def post(url, *args, **kwargs):
+    return _retry_request(requests.post, url, *args, **kwargs)
+
+def put(url, *args, **kwargs):
+    return _retry_request(requests.put, url, *args, **kwargs)
+
+def delete(url, *args, **kwargs):
+    return _retry_request(requests.delete, url, *args, **kwargs)

tests/mocks/test_retry.py

@@ -0,0 +1,88 @@
+"""
+Test basic functionality of the Tables endpoint
+"""
+import time
+import unittest
+from unittest.mock import patch
+from urllib.error import HTTPError
+
+import requests
+import responses
+
+from kbcstorage.tables import Tables
+
+from .table_responses import list_response
+
+
+class TestRequestRetry(unittest.TestCase):
+    """
+    Test that requests are retried.
+    """
+    def setUp(self):
+        token = 'dummy_token'
+        base_url = 'https://connection.keboola.com/'
+        self.tables = Tables(base_url, token)
+
+    @responses.activate
+    @patch('time.sleep', return_value=None)
+    def test_ok(self, sleep_mock):
+        """
+        Retry will try at least 5 times.
+        """
+        for _ in range(4):
+            responses.add(
+                responses.Response(
+                    method='GET',
+                    url='https://connection.keboola.com/v2/storage/tables',
+                    json=list_response,
+                    status=502
+                )
+            )
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tables',
+                json=list_response,
+            )
+        )
+        tables_list = self.tables.list()
+        assert isinstance(tables_list, list)
+
+    @responses.activate
+    @patch('time.sleep', return_value=None)
+    def test_raises_error_many_tries(self, sleep_mock):
+        """
+        Retry will fail if it gets enough of error responses.
+        """
+        for _ in range(20):
+            responses.add(
+                responses.Response(
+                    method='GET',
+                    url='https://connection.keboola.com/v2/storage/tables',
+                    json=list_response,
+                    status=502
+                )
+            )
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tables',
+                json=list_response,
+            )
+        )
+        with self.assertRaises(requests.exceptions.HTTPError):
+            self.tables.list()
+
+    @responses.activate
+    @patch('time.sleep', return_value=None)
+    def test_raises_error_on_4xx(self, sleep_mock):
+        responses.add(
+            responses.Response(
+                method='GET',
+                url='https://connection.keboola.com/v2/storage/tables',
+                json=list_response,
+                status=401
+            )
+        )
+        with self.assertRaises(requests.exceptions.HTTPError):
+            self.tables.list()