Merge pull request #238 from wlhjason/client_context_manager

Implement Client as context manager

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+### Added
+- Support for using `Client` as context manager closing connection on exit. Solves issue [#237](https://github.com/mymarilyn/clickhouse-driver/issues/237).
+
 ## [0.2.1] - 2021-06-02
 ### Added
 - Linux wheels for AArch64. Pull request [#197](https://github.com/mymarilyn/clickhouse-driver/pull/197) by [odidev](https://github.com/odidev).

clickhouse_driver/client.py

@@ -90,6 +90,12 @@ def __init__(self, *args, **kwargs):
         self.reset_last_query()
         super(Client, self).__init__()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.disconnect()
+
     def disconnect(self):
         """
         Disconnects from the server.

docs/features.rst

@@ -423,3 +423,19 @@ Writing pandas DataFrame is also supported with `insert_dataframe`:
         ... )
         >>> client.insert_dataframe('INSERT INTO test VALUES', df)
         >>> 10000
+
+
+Automatic disposal
+------------------
+
+*New in version 0.2.2.*
+
+Each Client instance can be used as a context manager:
+
+    .. code-block:: python
+
+        >>> with Client('localhost') as client:
+        >>>     client.execute('SELECT 1')
+
+
+Upon exit, any established connection to the ClickHouse server will be closed automatically.

tests/test_client.py

@@ -1,15 +1,15 @@
 import ssl
-from unittest import TestCase
 
 from clickhouse_driver import Client
 from clickhouse_driver.compression.lz4 import Compressor as LZ4Compressor
 from clickhouse_driver.compression.lz4hc import Compressor as LZHC4Compressor
 from clickhouse_driver.compression.zstd import Compressor as ZSTDCompressor
 from clickhouse_driver.protocol import Compression
 from tests.numpy.util import check_numpy
+from tests.testcase import BaseTestCase
 
 
-class ClientFromUrlTestCase(TestCase):
+class ClientFromUrlTestCase(BaseTestCase):
     def assertHostsEqual(self, client, another, msg=None):
         self.assertEqual(list(client.connection.hosts), another, msg=msg)
 
@@ -227,3 +227,9 @@ def test_settings_is_important(self):
     def test_use_numpy(self):
         c = Client.from_url('clickhouse://host?use_numpy=true')
         self.assertTrue(c.connection.context.client_settings['use_numpy'])
+
+    def test_context_manager(self):
+        with self.client as c:
+            c.execute('SELECT 1')
+            self.assertTrue(c.connection.connected)
+        self.assertFalse(c.connection.connected)