From b8b06199858a59cbade834856a460731925b5031 Mon Sep 17 00:00:00 2001 From: Alik Kurdyukov Date: Mon, 22 Jan 2024 20:10:50 +0400 Subject: [PATCH 1/2] fix for #286 --- clickhouse_sqlalchemy/drivers/native/connector.py | 2 +- tests/drivers/native/test_cursor.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/clickhouse_sqlalchemy/drivers/native/connector.py b/clickhouse_sqlalchemy/drivers/native/connector.py index 06632239..cce86189 100644 --- a/clickhouse_sqlalchemy/drivers/native/connector.py +++ b/clickhouse_sqlalchemy/drivers/native/connector.py @@ -132,7 +132,7 @@ def _prepare(self, context=None): if self._stream_results and execute_iter: execute = execute_iter settings = settings or {} - settings['max_block_size'] = execution_options['max_row_buffer'] + settings['max_block_size'] = execution_options.get('max_row_buffer', 1000) execute_kwargs = { 'settings': settings, diff --git a/tests/drivers/native/test_cursor.py b/tests/drivers/native/test_cursor.py index e4239e8f..10b6acce 100644 --- a/tests/drivers/native/test_cursor.py +++ b/tests/drivers/native/test_cursor.py @@ -22,3 +22,14 @@ def test_check_iter_cursor(self): text('SELECT number FROM system.numbers LIMIT 5') ) self.assertListEqual(list(rv), [(x,) for x in range(5)]) + + def test_execute_with_stream(self): + rv = self.session.execute(text("SELECT * FROM system.numbers LIMIT 1")).yield_per(10) + + self.assertEqual(len(rv.fetchall()), 1) + + def test_with_stream_results(self): + rv = self.session.execute(text("SELECT * FROM system.numbers LIMIT 1"), + execution_options={"stream_results": True}) + + self.assertEqual(len(rv.fetchall()), 1) From 5d59385e5865007aa2ebbb3e094fa5b839304e7e Mon Sep 17 00:00:00 2001 From: Alik Kurdyukov Date: Mon, 22 Jan 2024 20:15:56 +0400 Subject: [PATCH 2/2] flake8 fix --- clickhouse_sqlalchemy/drivers/native/connector.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/clickhouse_sqlalchemy/drivers/native/connector.py b/clickhouse_sqlalchemy/drivers/native/connector.py index cce86189..01481efa 100644 --- a/clickhouse_sqlalchemy/drivers/native/connector.py +++ b/clickhouse_sqlalchemy/drivers/native/connector.py @@ -132,7 +132,8 @@ def _prepare(self, context=None): if self._stream_results and execute_iter: execute = execute_iter settings = settings or {} - settings['max_block_size'] = execution_options.get('max_row_buffer', 1000) + settings['max_block_size'] = ( + execution_options.get('max_row_buffer', 1000)) execute_kwargs = { 'settings': settings,