feat: Add custom socket_options

Author: rhajek

influxdb/client.py

@@ -88,7 +88,8 @@ class InfluxDBClient(object):
     :param headers: headers to add to Requests, will add 'Content-Type'
         and 'Accept' unless these are already present, defaults to {}
     :type headers: dict
-    :param socket_options: use custom tcp socket options, If not specified, then defaults are loaded from
+    :param socket_options: use custom tcp socket options,
+        If not specified, then defaults are loaded from
         ``HTTPConnection.default_socket_options``
     :type socket_options: list
 
@@ -134,7 +135,7 @@ def __init__(self,
             session = requests.Session()
 
         self._session = session
-        adapter = SocketOptionsAdapter(
+        adapter = _SocketOptionsAdapter(
             pool_connections=int(pool_size),
             pool_maxsize=int(pool_size),
             socket_options=socket_options
@@ -1258,12 +1259,14 @@ def _msgpack_parse_hook(code, data):
     return msgpack.ExtType(code, data)
 
 
-class SocketOptionsAdapter(HTTPAdapter):
+class _SocketOptionsAdapter(HTTPAdapter):
+    """_SocketOptionsAdapter injects socket_options into HTTP Adapter."""
+
     def __init__(self, *args, **kwargs):
         self.socket_options = kwargs.pop("socket_options", None)
-        super(SocketOptionsAdapter, self).__init__(*args, **kwargs)
+        super(_SocketOptionsAdapter, self).__init__(*args, **kwargs)
 
     def init_poolmanager(self, *args, **kwargs):
         if self.socket_options is not None:
             kwargs["socket_options"] = self.socket_options
-        super(SocketOptionsAdapter, self).init_poolmanager(*args, **kwargs)
+        super(_SocketOptionsAdapter, self).init_poolmanager(*args, **kwargs)

influxdb/tests/client_test.py

@@ -1500,28 +1500,38 @@ def test_auth_token(self):
                              "my-token")
 
     def test_custom_socket_options(self):
-        test_socket_options = HTTPConnection.default_socket_options + [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
-                                                                       (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
-                                                                       (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 15)]
+        """Test custom socket options."""
+        test_socket_options = HTTPConnection.default_socket_options + \
+            [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+             (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
+             (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 15)]
 
-        cli = InfluxDBClient(username=None, password=None, socket_options=test_socket_options)
+        cli = InfluxDBClient(username=None, password=None,
+                             socket_options=test_socket_options)
 
-        self.assertEquals(cli._session.adapters.get("http://").socket_options, test_socket_options)
-        self.assertEquals(cli._session.adapters.get("http://").poolmanager.connection_pool_kw.get("socket_options"),
+        self.assertEquals(cli._session.adapters.get("http://").socket_options,
+                          test_socket_options)
+        self.assertEquals(cli._session.adapters.get("http://").poolmanager.
+                          connection_pool_kw.get("socket_options"),
                           test_socket_options)
 
-        connection_pool = cli._session.adapters.get("http://").poolmanager.connection_from_url(
+        connection_pool = cli._session.adapters.get("http://").poolmanager \
+            .connection_from_url(
             url="http://localhost:8086")
         new_connection = connection_pool._new_conn()
         self.assertEquals(new_connection.socket_options, test_socket_options)
 
     def test_none_socket_options(self):
+        """Test default socket options."""
         cli = InfluxDBClient(username=None, password=None)
-        self.assertEquals(cli._session.adapters.get("http://").socket_options, None)
-        connection_pool = cli._session.adapters.get("http://").poolmanager.connection_from_url(
+        self.assertEquals(cli._session.adapters.get("http://").socket_options,
+                          None)
+        connection_pool = cli._session.adapters.get("http://").poolmanager \
+            .connection_from_url(
             url="http://localhost:8086")
         new_connection = connection_pool._new_conn()
-        self.assertEquals(new_connection.socket_options, HTTPConnection.default_socket_options)
+        self.assertEquals(new_connection.socket_options,
+                          HTTPConnection.default_socket_options)
 
 
 class FakeClient(InfluxDBClient):