Skip to content

Commit 6ba88c0

Browse files
rhajekrhajek
committed
feat: Add custom socket_options
1 parent dec2f05 commit 6ba88c0

File tree

2 files changed

+28
-15
lines changed

2 files changed

+28
-15
lines changed

‎influxdb/client.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ class InfluxDBClient(object):
8888
:param headers: headers to add to Requests, will add 'Content-Type'
8989
and 'Accept' unless these are already present, defaults to {}
9090
:type headers: dict
91-
:param socket_options: use custom tcp socket options, If not specified, then defaults are loaded from
91+
:param socket_options: use custom tcp socket options,
92+
If not specified, then defaults are loaded from
9293
``HTTPConnection.default_socket_options``
9394
:type socket_options: list
9495
@@ -134,7 +135,7 @@ def __init__(self,
134135
session = requests.Session()
135136

136137
self._session = session
137-
adapter = SocketOptionsAdapter(
138+
adapter = _SocketOptionsAdapter(
138139
pool_connections=int(pool_size),
139140
pool_maxsize=int(pool_size),
140141
socket_options=socket_options
@@ -1258,12 +1259,14 @@ def _msgpack_parse_hook(code, data):
12581259
return msgpack.ExtType(code, data)
12591260

12601261

1261-
class SocketOptionsAdapter(HTTPAdapter):
1262+
class _SocketOptionsAdapter(HTTPAdapter):
1263+
"""_SocketOptionsAdapter injects socket_options into HTTP Adapter."""
1264+
12621265
def __init__(self, *args, **kwargs):
12631266
self.socket_options = kwargs.pop("socket_options", None)
1264-
super(SocketOptionsAdapter, self).__init__(*args, **kwargs)
1267+
super(_SocketOptionsAdapter, self).__init__(*args, **kwargs)
12651268

12661269
def init_poolmanager(self, *args, **kwargs):
12671270
if self.socket_options is not None:
12681271
kwargs["socket_options"] = self.socket_options
1269-
super(SocketOptionsAdapter, self).init_poolmanager(*args, **kwargs)
1272+
super(_SocketOptionsAdapter, self).init_poolmanager(*args, **kwargs)

‎influxdb/tests/client_test.py

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1500,28 +1500,38 @@ def test_auth_token(self):
15001500
"my-token")
15011501

15021502
def test_custom_socket_options(self):
1503-
test_socket_options = HTTPConnection.default_socket_options + [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
1504-
(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
1505-
(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 15)]
1503+
"""Test custom socket options."""
1504+
test_socket_options = HTTPConnection.default_socket_options + \
1505+
[(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
1506+
(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60),
1507+
(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 15)]
15061508

1507-
cli = InfluxDBClient(username=None, password=None, socket_options=test_socket_options)
1509+
cli = InfluxDBClient(username=None, password=None,
1510+
socket_options=test_socket_options)
15081511

1509-
self.assertEquals(cli._session.adapters.get("http://").socket_options, test_socket_options)
1510-
self.assertEquals(cli._session.adapters.get("http://").poolmanager.connection_pool_kw.get("socket_options"),
1512+
self.assertEquals(cli._session.adapters.get("http://").socket_options,
1513+
test_socket_options)
1514+
self.assertEquals(cli._session.adapters.get("http://").poolmanager.
1515+
connection_pool_kw.get("socket_options"),
15111516
test_socket_options)
15121517

1513-
connection_pool = cli._session.adapters.get("http://").poolmanager.connection_from_url(
1518+
connection_pool = cli._session.adapters.get("http://").poolmanager \
1519+
.connection_from_url(
15141520
url="http://localhost:8086")
15151521
new_connection = connection_pool._new_conn()
15161522
self.assertEquals(new_connection.socket_options, test_socket_options)
15171523

15181524
def test_none_socket_options(self):
1525+
"""Test default socket options."""
15191526
cli = InfluxDBClient(username=None, password=None)
1520-
self.assertEquals(cli._session.adapters.get("http://").socket_options, None)
1521-
connection_pool = cli._session.adapters.get("http://").poolmanager.connection_from_url(
1527+
self.assertEquals(cli._session.adapters.get("http://").socket_options,
1528+
None)
1529+
connection_pool = cli._session.adapters.get("http://").poolmanager \
1530+
.connection_from_url(
15221531
url="http://localhost:8086")
15231532
new_connection = connection_pool._new_conn()
1524-
self.assertEquals(new_connection.socket_options, HTTPConnection.default_socket_options)
1533+
self.assertEquals(new_connection.socket_options,
1534+
HTTPConnection.default_socket_options)
15251535

15261536

15271537
class FakeClient(InfluxDBClient):

0 commit comments

Comments
 (0)