Merge pull request #287 from akurdyukov/stream_results

fix for #286
xzkostyan · Jan 23, 2024 · e129569 · e129569
2 parents c34be61 + 5d59385
commit e129569
Show file tree

Hide file tree

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