Merge pull request influxdata#125 from influxdb/catch_connection_error

aviau · aviau · commit a85226d248ba · 2015-03-18T15:32:03.000-04:00
Now catch requests connection errors
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -5,6 +5,7 @@
 import json
 import socket
 import requests
+import requests.exceptions
 import warnings
 
 from influxdb import chunked_json
@@ -158,15 +159,25 @@ def request(self, url, method='GET', params=None, data=None,
         if data is not None and not isinstance(data, str):
             data = json.dumps(data)
 
-        response = session.request(
-            method=method,
-            url=url,
-            params=params,
-            data=data,
-            headers=self._headers,
-            verify=self._verify_ssl,
-            timeout=self._timeout
-        )
+        # Try to send the request a maximum of three times. (see #103)
+        # TODO (aviau): Make this configurable.
+        for i in range(0, 3):
+            try:
+                response = session.request(
+                    method=method,
+                    url=url,
+                    params=params,
+                    data=data,
+                    headers=self._headers,
+                    verify=self._verify_ssl,
+                    timeout=self._timeout
+                )
+                break
+            except requests.exceptions.ConnectionError as e:
+                if i < 2:
+                    continue
+                else:
+                    raise e
 
         if response.status_code == expected_response_code:
             return response
diff --git a/tests/influxdb/client_test.py b/tests/influxdb/client_test.py
@@ -4,12 +4,14 @@
 """
 import json
 import requests
+import requests.exceptions
 import socket
 import unittest
 import requests_mock
 from nose.tools import raises
 from mock import patch
 import warnings
+import mock
 
 from influxdb import InfluxDBClient
 from influxdb.client import session
@@ -640,3 +642,51 @@ def test_delete_database_user(self):
     def test_update_permission(self):
         cli = InfluxDBClient('host', 8086, 'username', 'password', 'db')
         cli.update_permission('admin', [])
+
+    @mock.patch('requests.Session.request')
+    def test_request_retry(self, mock_request):
+        """Tests that two connection errors will be handled"""
+
+        class CustomMock(object):
+            i = 0
+
+            def connection_error(self, *args, **kwargs):
+                self.i += 1
+
+                if self.i < 3:
+                    raise requests.exceptions.ConnectionError
+                else:
+                    r = requests.Response()
+                    r.status_code = 200
+                    return r
+
+        mock_request.side_effect = CustomMock().connection_error
+
+        cli = InfluxDBClient(database='db')
+        cli.write_points(
+            self.dummy_points
+        )
+
+    @mock.patch('requests.Session.request')
+    def test_request_retry_raises(self, mock_request):
+        """Tests that three connection errors will not be handled"""
+
+        class CustomMock(object):
+            i = 0
+
+            def connection_error(self, *args, **kwargs):
+                self.i += 1
+
+                if self.i < 4:
+                    raise requests.exceptions.ConnectionError
+                else:
+                    r = requests.Response()
+                    r.status_code = 200
+                    return r
+
+        mock_request.side_effect = CustomMock().connection_error
+
+        cli = InfluxDBClient(database='db')
+
+        with self.assertRaises(requests.exceptions.ConnectionError):
+            cli.write_points(self.dummy_points)
