diff --git a/clickhouse_driver/connection.py b/clickhouse_driver/connection.py index 007a127..9255774 100644 --- a/clickhouse_driver/connection.py +++ b/clickhouse_driver/connection.py @@ -141,6 +141,10 @@ class Connection(object): Defaults to ``False``. :param client_revision: can be used for client version downgrading. Defaults to ``None``. + :param disable_reconnect: disable automatic reconnect in case of + failed ``ping``, helpful when every reconnect + need to be caught in calling code. + Defaults to ``False``. """ def __init__( @@ -161,7 +165,8 @@ def __init__( alt_hosts=None, settings_is_important=False, tcp_keepalive=False, - client_revision=None + client_revision=None, + disable_reconnect=False, ): if secure: default_port = defines.DEFAULT_SECURE_PORT @@ -187,6 +192,7 @@ def __init__( self.client_revision = min( client_revision or defines.CLIENT_REVISION, defines.CLIENT_REVISION ) + self.disable_reconnect = disable_reconnect self.secure_socket = secure self.verify_cert = verify @@ -258,6 +264,11 @@ def force_connect(self): self.connect() elif not self.ping(): + if self.disable_reconnect: + raise errors.NetworkError( + "Connection was closed, reconnect is disabled." + ) + logger.warning('Connection was closed, reconnecting.') self.connect() diff --git a/tests/test_connect.py b/tests/test_connect.py index 34cdc53..f7bea34 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -92,6 +92,28 @@ def test_transport_not_connection_on_disconnect(self): # Close newly created socket. connection.socket.close() + def test_transport_disable_reconnect(self): + with self.created_client(disable_reconnect=True) as client: + # Create connection. + client.execute('SELECT 1') + + connection = client.connection + + with patch.object(connection, 'ping') as mocked_ping: + mocked_ping.return_value = False + + self.assertTrue(connection.connected) + error = errors.NetworkError + with self.assertRaises(error) as e: + client.execute('SELECT 1') + + self.assertFalse(connection.connected) + + self.assertEqual( + str(e.exception), + 'Code: 210. Connection was closed, reconnect is disabled.' + ) + def test_socket_error_on_ping(self): self.client.execute('SELECT 1')