diff --git a/CHANGELOG.md b/CHANGELOG.md index 7cbafb53..ec31de4e 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/clickhouse_driver/client.py b/clickhouse_driver/client.py index 5e9b50a6..a87c979a 100644 --- a/clickhouse_driver/client.py +++ b/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. diff --git a/docs/features.rst b/docs/features.rst index 0dbef6b0..e0dbcb4b 100644 --- a/docs/features.rst +++ b/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. diff --git a/tests/test_client.py b/tests/test_client.py index 8249a336..c6f79ea2 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,5 +1,4 @@ import ssl -from unittest import TestCase from clickhouse_driver import Client from clickhouse_driver.compression.lz4 import Compressor as LZ4Compressor @@ -7,9 +6,10 @@ 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)