feat: add query_id from execution_options for native driver (#322)

* feat: add query_id from execution_options for native driver * fix: move import at top of the file * fix: flake8 --------- Co-authored-by: Sébastien Fleury <[email protected]>
xzkostyan · Jul 5, 2024 · 9ef75ee · 9ef75ee
1 parent b62512f
commit 9ef75ee
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/clickhouse_sqlalchemy/drivers/native/connector.py b/clickhouse_sqlalchemy/drivers/native/connector.py
@@ -138,7 +138,8 @@ def _prepare(self, context=None):
         execute_kwargs = {
             'settings': settings,
             'external_tables': external_tables,
-            'types_check': execution_options.get('types_check', False)
+            'types_check': execution_options.get('types_check', False),
+            'query_id': execution_options.get('query_id', None)
         }
 
         return execute, execute_kwargs

diff --git a/tests/drivers/native/test_cursor.py b/tests/drivers/native/test_cursor.py
@@ -1,3 +1,5 @@
+import uuid
+
 from sqlalchemy import text
 
 from tests.testcase import NativeSessionTestCase
@@ -47,3 +49,17 @@ def test_with_settings_in_execution_options(self):
             dict(rv.context.execution_options), {"settings": {"final": 1}}
         )
         self.assertEqual(len(rv.fetchall()), 1)
+
+    def test_set_query_id(self):
+        query_id = str(uuid.uuid4())
+        self.session.execute(
+            text("SELECT 1"),
+            execution_options={'query_id': query_id}
+        ).first()
+        self.session.execute(text("SYSTEM FLUSH LOGS"))
+        rv = self.session.execute(
+            text(f"SELECT COUNT(*) "
+                 f"FROM system.query_log "
+                 f"WHERE query_id = '{query_id}'")
+        ).first()
+        self.assertEqual(rv[0], 2)