From 9ef75ee114d550126b780428b22626121616c471 Mon Sep 17 00:00:00 2001
From: Yedrimas <fl.sebastien@hotmail.fr>
Date: Fri, 5 Jul 2024 12:14:04 +0200
Subject: [PATCH] feat: add query_id from execution_options for native driver
 (#322)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 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 <sebastienf@api.video>
---
 .../drivers/native/connector.py                  |  3 ++-
 tests/drivers/native/test_cursor.py              | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/clickhouse_sqlalchemy/drivers/native/connector.py b/clickhouse_sqlalchemy/drivers/native/connector.py
index 01481efa..94f4624a 100644
--- 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
index e38274df..5e5b4523 100644
--- 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)