Add parameter chunk_size and default as 1 (#315)

* Add parameter chunk_size
mymarilyn · Jun 7, 2022 · 76b4ed6 · 76b4ed6
1 parent aeae5c4
commit 76b4ed6
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/clickhouse_driver/client.py b/clickhouse_driver/client.py
@@ -350,7 +350,7 @@ def execute_with_progress(
     def execute_iter(
             self, query, params=None, with_column_types=False,
             external_tables=None, query_id=None, settings=None,
-            types_check=False):
+            types_check=False, chunk_size=1):
         """
         *New in version 0.0.14.*
 
@@ -372,15 +372,16 @@ def execute_iter(
                          Defaults to ``None`` (no additional settings).
         :param types_check: enables type checking of data for INSERT queries.
                             Causes additional overhead. Defaults to ``False``.
+        :param chunk_size: chunk query results.
         :return: :ref:`iter-query-result` proxy.
         """
-
         with self.disconnect_on_error(query, settings):
-            return self.iter_process_ordinary_query(
+            rv = self.iter_process_ordinary_query(
                 query, params=params, with_column_types=with_column_types,
                 external_tables=external_tables,
                 query_id=query_id, types_check=types_check
             )
+            return chunks(rv, chunk_size) if chunk_size > 1 else rv
 
     def query_dataframe(
             self, query, params=None, external_tables=None, query_id=None,

diff --git a/tests/test_blocks.py b/tests/test_blocks.py
@@ -4,6 +4,7 @@
 from clickhouse_driver.errors import ServerException
 from tests.testcase import BaseTestCase, file_config
 from tests.util import capture_logging
+from clickhouse_driver.util.helpers import chunks
 
 
 class BlocksTestCase(BaseTestCase):
@@ -149,10 +150,27 @@ def test_select_with_iter(self):
             'SELECT number FROM system.numbers LIMIT 10'
         )
         self.assertIsInstance(result, types.GeneratorType)
+        self.assertEqual(list(result), list(zip(range(10))))
+        self.assertEqual(list(result), [])
 
+    def test_select_with_chunk_one_iter(self):
+        result = self.client.execute_iter(
+            'SELECT number FROM system.numbers LIMIT 10',
+            chunk_size=1
+        )
+        self.assertIsInstance(result, types.GeneratorType)
         self.assertEqual(list(result), list(zip(range(10))))
         self.assertEqual(list(result), [])
 
+    def test_select_with_chunk_some_iter(self):
+        result = self.client.execute_iter(
+            'SELECT number FROM system.numbers LIMIT 10',
+            chunk_size=3
+        )
+        self.assertIsInstance(result, types.GeneratorType)
+        self.assertEqual(list(result), list(chunks(zip(range(10)), 3)))
+        self.assertEqual(list(result), [])
+
     def test_select_with_iter_with_column_types(self):
         result = self.client.execute_iter(
             'SELECT CAST(number AS UInt32) as number '