Merge pull request influxdata#435 from xginn8/define_retries (Thanks @xginn8!)

aviau · web-flow · commit f93998011d6c · 2017-04-12T02:59:11.000-04:00
Allow user-configurable number of retries when submitting request to …
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -53,6 +53,9 @@ class InfluxDBClient(object):
     :param timeout: number of seconds Requests will wait for your client to
         establish a connection, defaults to None
     :type timeout: int
+    :param retries: number of retries your client will try before aborting,
+        defaults to 3. 0 indicates try until success
+    :type retries: int
     :param use_udp: use UDP to connect to InfluxDB, defaults to False
     :type use_udp: bool
     :param udp_port: UDP port to connect to InfluxDB, defaults to 4444
@@ -70,6 +73,7 @@ def __init__(self,
                  ssl=False,
                  verify_ssl=False,
                  timeout=None,
+                 retries=3,
                  use_udp=False,
                  udp_port=4444,
                  proxies=None,
@@ -81,6 +85,7 @@ def __init__(self,
         self._password = password
         self._database = database
         self._timeout = timeout
+        self._retries = retries
 
         self._verify_ssl = verify_ssl
 
@@ -225,9 +230,10 @@ def request(self, url, method='GET', params=None, data=None,
         if isinstance(data, (dict, list)):
             data = json.dumps(data)
 
-        # Try to send the request a maximum of three times. (see #103)
-        # TODO (aviau): Make this configurable.
-        for i in range(0, 3):
+        # Try to send the request more than once by default (see #103)
+        retry = True
+        _try = 0
+        while retry:
             try:
                 response = self._session.request(
                     method=method,
@@ -242,10 +248,12 @@ def request(self, url, method='GET', params=None, data=None,
                 )
                 break
             except requests.exceptions.ConnectionError as e:
-                if i < 2:
-                    continue
-                else:
-                    raise e
+                _try += 1
+                if self._retries != 0:
+                    retry = _try < self._retries
+
+        else:
+            raise requests.exceptions.ConnectionError
 
         if 500 <= response.status_code < 600:
             raise InfluxDBServerError(response.content)
diff --git a/influxdb/influxdb08/client.py b/influxdb/influxdb08/client.py
@@ -55,6 +55,9 @@ class InfluxDBClient(object):
     :param verify_ssl: verify SSL certificates for HTTPS requests, defaults is
         False
     :type verify_ssl: boolean
+    :param retries: number of retries your client will try before aborting,
+        defaults to 3. 0 indicates try until success
+    :type retries: int
     :param timeout: number of seconds Requests will wait for your client to
         establish a connection, defaults to None
     :type timeout: int
@@ -73,6 +76,7 @@ def __init__(self,
                  ssl=False,
                  verify_ssl=False,
                  timeout=None,
+                 retries=3,
                  use_udp=False,
                  udp_port=4444):
         """
@@ -84,6 +88,7 @@ def __init__(self,
         self._password = password
         self._database = database
         self._timeout = timeout
+        self._retries = retries
 
         self._verify_ssl = verify_ssl
 
@@ -228,9 +233,10 @@ def request(self, url, method='GET', params=None, data=None,
         if data is not None and not isinstance(data, str):
             data = json.dumps(data)
 
-        # Try to send the request a maximum of three times. (see #103)
-        # TODO (aviau): Make this configurable.
-        for i in range(0, 3):
+        retry = True
+        _try = 0
+        # Try to send the request more than once by default (see #103)
+        while retry:
             try:
                 response = session.request(
                     method=method,
@@ -244,10 +250,11 @@ def request(self, url, method='GET', params=None, data=None,
                 break
             except (requests.exceptions.ConnectionError,
                     requests.exceptions.Timeout) as e:
-                if i < 2:
-                    continue
-                else:
-                    raise e
+                _try += 1
+                if self._retries != 0:
+                    retry = _try < self._retries
+        else:
+            raise requests.exceptions.ConnectionError
 
         if response.status_code == expected_response_code:
             return response
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -645,6 +645,61 @@ def connection_error(self, *args, **kwargs):
         with self.assertRaises(requests.exceptions.ConnectionError):
             cli.write_points(self.dummy_points)
 
+    @mock.patch('requests.Session.request')
+    def test_random_request_retry(self, mock_request):
+        """Tests that a random number of connection errors will be handled"""
+
+        class CustomMock(object):
+            def __init__(self, retries):
+                self.i = 0
+                self.retries = retries
+
+            def connection_error(self, *args, **kwargs):
+                self.i += 1
+
+                if self.i < self.retries:
+                    raise requests.exceptions.ConnectionError
+                else:
+                    r = requests.Response()
+                    r.status_code = 204
+                    return r
+
+        retries = random.randint(1, 100)
+        mock_request.side_effect = CustomMock(retries).connection_error
+
+        cli = InfluxDBClient(database='db', retries=retries)
+        cli.write_points(
+            self.dummy_points
+        )
+
+    @mock.patch('requests.Session.request')
+    def test_random_request_retry_raises(self, mock_request):
+        """Tests that a random number of connection errors plus one \
+                will not be handled"""
+
+        class CustomMock(object):
+            def __init__(self, retries):
+                self.i = 0
+                self.retries = retries
+
+            def connection_error(self, *args, **kwargs):
+                self.i += 1
+
+                if self.i < self.retries + 1:
+                    raise requests.exceptions.ConnectionError
+                else:
+                    r = requests.Response()
+                    r.status_code = 200
+                    return r
+
+        retries = random.randint(1, 100)
+        mock_request.side_effect = CustomMock(retries).connection_error
+
+        cli = InfluxDBClient(database='db', retries=retries)
+
+        with self.assertRaises(requests.exceptions.ConnectionError):
+            cli.write_points(self.dummy_points)
+
     def test_get_list_users(self):
         example_response = (
             '{"results":[{"series":[{"columns":["user","admin"],'
