Add TCP keepalive

clickhouse_driver/__init__.py

@@ -3,7 +3,7 @@
 from .dbapi import connect
 
 
-VERSION = (0, 2, 5)
+VERSION = (0, 2, 6)
 __version__ = '.'.join(str(x) for x in VERSION)
 
 __all__ = ['Client', 'connect']

clickhouse_driver/client.py

@@ -778,6 +778,15 @@ def from_url(cls, url):
             elif name == 'settings_is_important':
                 kwargs[name] = asbool(value)
 
+            elif name == 'tcp_keepalive':
+                try:
+                    kwargs[name] = asbool(value)
+                except ValueError:
+                    parts = value.split(',')
+                    kwargs[name] = (
+                        float(parts[0]), float(parts[1]), int(parts[2])
+                    )
+
             # ssl
             elif name == 'verify':
                 kwargs[name] = asbool(value)

clickhouse_driver/connection.py

@@ -3,6 +3,7 @@
 import ssl
 from collections import deque
 from contextlib import contextmanager
+from sys import platform
 from time import time
 from urllib.parse import urlparse
 
@@ -124,6 +125,13 @@ class Connection(object):
                                   ignored, ``True`` means that the query will
                                   fail with UNKNOWN_SETTING error.
                                   Defaults to ``False``.
+    :param tcp_keepalive: enables `TCP keepalive <https://tldp.org/HOWTO/
+                          TCP-Keepalive-HOWTO/overview.html>`_ on established
+                          connection. If is set to ``True``` system keepalive
+                          settings are used. You can also specify custom
+                          keepalive setting with tuple:
+                          ``(idle_time_sec, interval_sec, probes)``.
+                          Defaults to ``False``.
     """
 
     def __init__(
@@ -143,6 +151,7 @@ def __init__(
             server_hostname=None,
             alt_hosts=None,
             settings_is_important=False,
+            tcp_keepalive=False
     ):
         if secure:
             default_port = defines.DEFAULT_SECURE_PORT
@@ -164,6 +173,7 @@ def __init__(
         self.send_receive_timeout = send_receive_timeout
         self.sync_request_timeout = sync_request_timeout
         self.settings_is_important = settings_is_important
+        self.tcp_keepalive = tcp_keepalive
 
         self.secure_socket = secure
         self.verify_cert = verify
@@ -310,6 +320,8 @@ def _init_connection(self, host, port):
 
         # performance tweak
         self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        if self.tcp_keepalive:
+            self._set_keepalive()
 
         self.fin = BufferedSocketReader(self.socket, defines.BUFFER_SIZE)
         self.fout = BufferedSocketWriter(self.socket, defines.BUFFER_SIZE)
@@ -321,6 +333,34 @@ def _init_connection(self, host, port):
         self.block_in_raw = BlockInputStream(self.fin, self.context)
         self.block_out = self.get_block_out_stream()
 
+    def _set_keepalive(self):
+        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+
+        if not isinstance(self.tcp_keepalive, tuple):
+            return
+
+        idle_time_sec, interval_sec, probes = self.tcp_keepalive
+
+        if platform == 'linux' or platform == 'win32':
+            # This should also work for Windows
+            # starting with Windows 10, version 1709.
+            self.socket.setsockopt(
+                socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, idle_time_sec
+            )
+            self.socket.setsockopt(
+                socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, interval_sec
+            )
+            self.socket.setsockopt(
+                socket.IPPROTO_TCP, socket.TCP_KEEPCNT, probes
+            )
+
+        elif platform == 'darwin':
+            TCP_KEEPALIVE = 0x10
+            # Only interval is available in mac os.
+            self.socket.setsockopt(
+                socket.IPPROTO_TCP, TCP_KEEPALIVE, interval_sec
+            )
+
     def _format_connection_error(self, e, host, port):
         err = (e.strerror + ' ') if e.strerror else ''
         return err + '({}:{})'.format(host, port)

docs/features.rst

@@ -546,3 +546,43 @@ Each Client instance can be used as a context manager:
 
 Upon exit, any established connection to the ClickHouse server will be closed
 automatically.
+
+
+TCP keepalive
+-------------
+
+*New in version 0.2.6.*
+
+You can enable `TCP keepalive
+<https://tldp.org/HOWTO/TCP-Keepalive-HOWTO/overview.html>`_ on connection with
+ClickHouse server. This setting is disabled by default. When parameter
+``tcp_keepalive`` is set to ``True`` system TCP keepalive settings are used.
+
+    .. code-block:: python
+
+        >>> client = Client('localhost', tcp_keepalive=True)
+
+
+For Linux default TCP keepalive settings can be found in:
+
+    .. code-block:: bash
+
+        idle_time_sec - /proc/sys/net/ipv4/tcp_keepalive_time
+        interval_sec - /proc/sys/net/ipv4/tcp_keepalive_intvl
+        probes - /proc/sys/net/ipv4/tcp_keepalive_probes
+
+You can also specify custom keepalive settings with tuple
+``(idle_time_sec, interval_sec, probes)``:
+
+    .. code-block:: python
+
+        >>> client = Client('localhost', tcp_keepalive=(60.5, 5.1, 2))
+
+    .. note::
+
+        For Linux and Windows all parameters: idle time, interval and probes
+        can be changed for socket.
+
+        For Mac OS only the second parameter ``interval_sec`` can be changed
+        for socket. ``idle_time_sec``, ``probes`` are not used, but should be
+        specified for uniformity.

tests/test_client.py

@@ -275,3 +275,12 @@ def test_round_robin(self):
             'clickhouse://host?round_robin=true&alt_hosts=host2'
         )
         self.assertEqual(len(c.connections), 1)
+
+    def test_tcp_keepalive(self):
+        c = Client.from_url('clickhouse://host?tcp_keepalive=true')
+        self.assertTrue(c.connection.tcp_keepalive)
+
+        c = Client.from_url('clickhouse://host?tcp_keepalive=10.5,2.5,3')
+        self.assertEqual(
+            c.connection.tcp_keepalive, (10.5, 2.5, 3)
+        )

tests/test_connect.py

@@ -264,6 +264,19 @@ def test_round_robin(self):
             self.assertFalse(client.connection.connected)
             self.assertFalse(list(client.connections)[0].connected)
 
+    def test_tcp_keepalive(self):
+        self.assertFalse(self.client.connection.tcp_keepalive)
+
+        with self.created_client(tcp_keepalive=True) as client:
+            self.assertTrue(client.connection.tcp_keepalive)
+
+            client.execute('SELECT 1')
+
+        with self.created_client(tcp_keepalive=(100, 20, 2)) as client:
+            self.assertEqual(client.connection.tcp_keepalive, (100, 20, 2))
+
+            client.execute('SELECT 1')
+
 
 class FakeBufferedReader(BufferedReader):
     def __init__(self, inputs, bufsize=128):