Add settings final clause (#292)

* Add settings final clause --------- Co-authored-by: Konstantin Lebedev <[email protected]>
xzkostyan · Mar 14, 2024 · 3576913 · 3576913
1 parent 7c05244
commit 3576913
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 1 deletion.
diff --git a/clickhouse_sqlalchemy/drivers/http/connector.py b/clickhouse_sqlalchemy/drivers/http/connector.py
@@ -101,6 +101,21 @@ def description(self):
     def close(self):
         pass
 
+    def _prepare(self, raw_sql, context=None):
+        if context:
+            execution_options = context.execution_options
+        else:
+            execution_options = {}
+
+        settings = execution_options.get('settings')
+        if settings:
+            raw_settings = ", ".join(
+                f"{key}={value}" for key, value in settings.items()
+            )
+            raw_sql = raw_sql + " SETTINGS " + raw_settings
+
+        return raw_sql
+
     def execute(self, operation, parameters=None, context=None):
         raw_sql = operation
 
@@ -112,6 +127,7 @@ def execute(self, operation, parameters=None, context=None):
 
         transport = self._connection.transport
         params = {'query_id': self._query_id}
+        raw_sql = self._prepare(raw_sql, context)
         response_gen = transport.execute(raw_sql, params=params)
 
         self._process_response(response_gen)

diff --git a/docs/features.rst b/docs/features.rst
@@ -906,7 +906,7 @@ becomes (respectively)
 FINAL
 +++++
 
-.. note:: Currently ``FINAL`` clause is supported only for table specified in ``FROM`` clause.
+.. note:: Currently ``FINAL`` clause is supported only for table specified in ``FROM`` clause. To apply ``FINAL`` modifier to all tables in a query, settings with "final=1" can be passed using execution options.
 
     .. code-block:: python
 

diff --git a/tests/drivers/http/test_cursor.py b/tests/drivers/http/test_cursor.py
@@ -1,6 +1,7 @@
 from sqlalchemy import text
 
 from tests.testcase import HttpSessionTestCase, HttpEngineTestCase
+from tests.util import require_server_version
 
 
 class CursorTestCase(HttpSessionTestCase, HttpEngineTestCase):
@@ -16,3 +17,14 @@ def test_check_iter_cursor_by_engine(self):
                 text('SELECT number FROM system.numbers LIMIT 5')
             )
             self.assertListEqual(list(rv), [(x,) for x in range(5)])
+
+    @require_server_version(23, 2, 1)
+    def test_with_settings_in_execution_options(self):
+        rv = self.session.execute(
+            text("SELECT number FROM system.numbers LIMIT 5"),
+            execution_options={"settings": {"final": 1}}
+        )
+        self.assertEqual(
+            dict(rv.context.execution_options), {"settings": {"final": 1}}
+        )
+        self.assertListEqual(list(rv), [(x,) for x in range(5)])
diff --git a/tests/drivers/native/test_cursor.py b/tests/drivers/native/test_cursor.py
@@ -1,6 +1,7 @@
 from sqlalchemy import text
 
 from tests.testcase import NativeSessionTestCase
+from tests.util import require_server_version
 
 
 class CursorTestCase(NativeSessionTestCase):
@@ -35,3 +36,14 @@ def test_with_stream_results(self):
                                   execution_options={"stream_results": True})
 
         self.assertEqual(len(rv.fetchall()), 1)
+
+    @require_server_version(23, 2, 1)
+    def test_with_settings_in_execution_options(self):
+        rv = self.session.execute(
+            text("SELECT * FROM system.numbers LIMIT 1"),
+            execution_options={"settings": {"final": 1}}
+        )
+        self.assertEqual(
+            dict(rv.context.execution_options), {"settings": {"final": 1}}
+        )
+        self.assertEqual(len(rv.fetchall()), 1)