fix for #286

xzkostyan · Jan 22, 2024 · b8b0619 · b8b0619
1 parent c34be61
commit b8b0619
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git 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
@@ -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)