Add support for rate limiting (#21)

* Add granular control over the conditions under which we retry a request

* Add retry functionality and allow config

* Update Readme

* Bump version to 1.2.0

* fixup! Add retry functionality and allow config

* Review fixes

Author: spyrbri

README.md

@@ -65,6 +65,18 @@ This throws error or returns `<Ping{data='pong!'}>`
 You can also pass to the Config initializer:
 * `request_timeout=` sets timeout for requests (seconds), default: none (see [requests docs](http://docs.python-requests.org/en/master/user/quickstart/#timeouts) for details)
 
+### Rate Limits & Exponential Backoff
+The library will keep retrying if the request exceeds the rate limit or if there's any network related error.
+By default, the request will be retried for 20 times (approximately 15 minutes) before finally giving up.
+
+You can change the retry count from the Config initializer:
+
+ * `max_retries=` sets the maximum number of retries for failed requests, default: 20
+ * `backoff_factor=` sets the exponential backoff factor, default: 2
+
+Set max_retries 0 to disable it.
+Set backoff_factor 0 to disable it.
+
 ## Usage
 
 The library is based on [promises](https://pypi.python.org/pypi/promise) (mechanism similar to futures).

chartmogul/__init__.py

@@ -29,7 +29,7 @@
 """
 
 __title__ = 'chartmogul'
-__version__ = '1.1.8'
+__version__ = '1.2.0'
 __build__ = 0x000000
 __author__ = 'ChartMogul Ltd'
 __license__ = 'MIT'

chartmogul/api/config.py

@@ -5,6 +5,8 @@
 class Config:
     uri = API_BASE + "/" + VERSION
 
-    def __init__(self, account_token, secret_key, request_timeout=None):
+    def __init__(self, account_token, secret_key, request_timeout=None, max_retries=20, backoff_factor=2):
         self.auth = (account_token, secret_key)
         self.request_timeout = request_timeout
+        self.max_retries = max_retries
+        self.backoff_factor = backoff_factor

chartmogul/resource.py

@@ -2,6 +2,7 @@
 from json import dumps
 from promise import Promise
 from uritemplate import URITemplate
+from .retry_request import requests_retry_session
 from .errors import APIError, ConfigurationError, ArgumentMissingError, annotateHTTPError
 from .api.config import Config
 from datetime import datetime, date
@@ -119,14 +120,14 @@ def _request(cls, config, method, http_verb, path, data=None, **kwargs):
                 data = dumps(data, default=json_serial)
 
         return Promise(lambda resolve, _:
-                       resolve(getattr(requests, http_verb)(
-                           config.uri + path,
-                           data=data,
-                           headers={'content-type': 'application/json'},
-                           params=params,
-                           auth=config.auth,
-                           timeout=config.request_timeout)
-                       )).then(cls._load).catch(annotateHTTPError)
+            resolve(getattr(requests_retry_session(config.max_retries, config.backoff_factor), http_verb)(
+                config.uri + path,
+                data=data,
+                headers={'content-type': 'application/json'},
+                params=params,
+                auth=config.auth,
+                timeout=config.request_timeout)
+            )).then(cls._load).catch(annotateHTTPError)
 
     @classmethod
     def _expandPath(cls, path, kwargs):
@@ -160,7 +161,6 @@ def fc(cls, config, **kwargs):
             return cls._request(config, method, http_verb, pathTemp, **kwargs)
         return fc
 
-
 def _add_method(cls, method, http_verb, path=None):
     """
     Dynamically define all possible actions.

chartmogul/retry_request.py

@@ -0,0 +1,24 @@
+import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+
+METHOD_WHITELIST = ['HEAD', 'GET', 'POST', 'PUT', 'DELETE', 'OPTIONS', 'TRACE']
+STATUS_FORCELIST = (429, 500, 502, 503, 504, 520, 524)
+
+def requests_retry_session(retries=20, backoff_factor=2, session=None,):
+    session = session or requests.Session()
+    adapter = _retry_adapter(retries, backoff_factor)
+    session.mount('https://', adapter)
+    return session
+
+def _retry_adapter(retries, backoff_factor):
+    retry = Retry(
+        total=retries,
+        read=retries,
+        connect=retries,
+        status=retries,
+        method_whitelist=METHOD_WHITELIST,
+        status_forcelist=STATUS_FORCELIST,
+        backoff_factor=backoff_factor,
+    )
+    return HTTPAdapter(max_retries=retry)

setup.py

@@ -20,7 +20,7 @@
     'marshmallow>=2.12.1',
     'future>=0.16.0',
 ]
-test_requirements = ['mock>=1.0.1', 'requests-mock>=1.3.0', 'vcrpy>=1.11.1']
+test_requirements = ['mock>=1.0.1', 'requests-mock>=1.3.0', 'vcrpy>=1.11.1', 'httpretty>=0.9.5']
 
 with open('chartmogul/__init__.py', 'r') as fd:
     version = re.search(r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]',

test/api/test_retry_request.py

@@ -0,0 +1,49 @@
+import unittest
+
+import httpretty
+import chartmogul
+from chartmogul import Config, DataSource
+from datetime import date, datetime
+from requests.exceptions import RetryError
+from chartmogul.retry_request import requests_retry_session
+
+class RetryRequestTestCase(unittest.TestCase):
+
+    @httpretty.activate
+    def test_retry_request(self):
+        httpretty.register_uri(
+            httpretty.GET,
+            "https://example:444/testing",
+            responses=[
+                httpretty.Response(body='{}', status=500),
+                httpretty.Response(body='{}', status=200),
+            ]
+        )
+
+        with self.assertRaises(RetryError):
+            requests_retry_session(0).get('https://example:444/testing')
+
+        response = requests_retry_session(2, 0).get('https://example:444/testing')
+        self.assertEqual(response.text, '{}')
+
+    @httpretty.activate
+    def test_requests_retry_session_on_resource(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            "https://api.chartmogul.com/v1/data_sources",
+            responses=[
+                httpretty.Response(body='{}', status=500),
+                httpretty.Response(body='{}', status=500),
+                httpretty.Response(body='{}', status=500),
+                httpretty.Response(body='{}', status=500),
+                httpretty.Response(body='{}', status=200),
+            ]
+        )
+
+        # max_retries set as 4
+        # backoff_factor set as 0 to avoid waiting while testing
+        config = Config("token", "secret", None, 4, 0)
+        try:
+            DataSource.create(config, data={ "test_date": date(2015, 1, 1) }).get()
+        except RetryError:
+            self.fail("request raised retryError unexpectedly!")