make max number of retries a param

kbcstorage/base.py

@@ -9,7 +9,7 @@
 .. _Storage API documentation:
     http://docs.keboola.apiary.io/
 """
-from . import retry_requests
+from kbcstorage.retry_requests import MAX_RETRIES_DEFAULT, RetryRequests
 import requests
 
 
@@ -22,7 +22,7 @@ class Endpoint:
         base_url (str): The base URL for this endpoint.
         token (str): A key for the Storage API.
     """
-    def __init__(self, root_url, path_component, token):
+    def __init__(self, root_url, path_component, token, max_requests_retries=MAX_RETRIES_DEFAULT):
         """
         Create an endpoint.
 
@@ -45,6 +45,7 @@ def __init__(self, root_url, path_component, token):
         self._auth_header = {'X-StorageApi-Token': self.token,
                              'Accept-Encoding': 'gzip',
                              'User-Agent': 'Keboola Storage API Python Client'}
+        self.requests = RetryRequests(max_requests_retries)
 
     def _get_raw(self, url, params=None, **kwargs):
         """
@@ -67,7 +68,7 @@ def _get_raw(self, url, params=None, **kwargs):
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
 
-        r = retry_requests.get(url, params=params, headers=headers, **kwargs)
+        r = self.requests.get(url, params=params, headers=headers, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:
@@ -112,7 +113,7 @@ def _post(self, *args, **kwargs):
         """
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
-        r = retry_requests.post(headers=headers, *args, **kwargs)
+        r = self.requests.post(headers=headers, *args, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:
@@ -138,7 +139,7 @@ def _put(self, *args, **kwargs):
         """
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
-        r = retry_requests.put(headers=headers, *args, **kwargs)
+        r = self.requests.put(headers=headers, *args, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:
@@ -164,7 +165,7 @@ def _delete(self, *args, **kwargs):
         """
         headers = kwargs.pop('headers', {})
         headers.update(self._auth_header)
-        r = retry_requests.delete(headers=headers, *args, **kwargs)
+        r = self.requests.delete(headers=headers, *args, **kwargs)
         try:
             r.raise_for_status()
         except requests.HTTPError:

kbcstorage/retry_requests.py

@@ -1,35 +1,35 @@
 import time
 import requests
 
-MAX_RETRIES = 11
+MAX_RETRIES_DEFAULT = 11
 BACKOFF_FACTOR = 1.0
 
 
 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(MAX_RETRIES - 1):
-        if response.status_code == 501 or response.status_code < 500:
-            return response
-        time.sleep(_get_backoff_time(retry_count))
-        response = request_func(url, **kwargs)
-    return response
+class RetryRequests:
+    def __init__(self, max_requests_retries=MAX_RETRIES_DEFAULT) -> None:
+        self.max_retries = max_requests_retries
 
+    def _retry_request(self, request_func, url, *args, **kwargs):
+        response = request_func(url, *args, **kwargs)
+        for retry_count in range(self.max_retries - 1):
+            if response.status_code == 501 or response.status_code < 500:
+                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 get(self, url, *args, **kwargs):
+        return self._retry_request(requests.get, url, *args, **kwargs)
 
+    def post(self, url, *args, **kwargs):
+        return self._retry_request(requests.post, url, *args, **kwargs)
 
-def post(url, *args, **kwargs):
-    return _retry_request(requests.post, url, *args, **kwargs)
+    def put(self, url, *args, **kwargs):
+        return self._retry_request(requests.put, 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)
+    def delete(self, url, *args, **kwargs):
+        return self._retry_request(requests.delete, url, *args, **kwargs)