From 8b28130d2765dc9945ac41fc8f0e0538c239ca40 Mon Sep 17 00:00:00 2001
From: Vladimir Rudev <rudev@azurgames.com>
Date: Fri, 6 Sep 2024 09:43:15 +0300
Subject: [PATCH] disable reconnect toggle

---
 clickhouse_driver/connection.py | 13 ++++++++++++-
 tests/test_connect.py           | 22 ++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

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')