Merge pull request #50 from datafuselabs/feat/keep-query-id-header

feat: query with  X-Databend-Query-Id header

Author: hantmac

databend_py/connection.py

@@ -15,6 +15,8 @@
 
 headers = {'Content-Type': 'application/json', 'User-Agent': sdk_info(), 'Accept': 'application/json',
            'X-DATABEND-ROUTE': 'warehouse'}
+XDatabendQueryIDHeader = "X-Databend-Query-Id"
+QueryID = "id"
 
 
 class ServerInfo(object):
@@ -146,10 +148,13 @@ def query(self, statement):
         else:
             self.client_session = self.default_session()
             query_sql['session'] = self.client_session
+        if XDatabendQueryIDHeader in self.additional_headers:
+            del self.additional_headers[XDatabendQueryIDHeader]
         log.logger.debug(f"http headers {self.make_headers()}")
         try:
             resp_dict = self.do_query(url, query_sql)
             self.client_session = resp_dict.get("session", self.default_session())
+            self.additional_headers = {XDatabendQueryIDHeader: resp_dict.get(QueryID)}
             return resp_dict
         except Exception as err:
             log.logger.error(

tests/test_client.py

@@ -176,7 +176,7 @@ def test_cookies(self):
         client = Client.from_url(url_with_persist_cookies)
         client.execute("select 1")
         # self.assertIsNotNone(client.connection.cookies)
-    
+
     def test_null_to_none(self):
         client = Client.from_url(self.databend_url)
         _, data = client.execute("select NULL as test")
@@ -190,6 +190,15 @@ def test_null_to_none(self):
         _, data = client.execute("select NULL as test")
         self.assertIsNone(data[0][0])
 
+    def test_set_query_id_header(self):
+        client = Client.from_url(self.databend_url)
+        client.execute("select 1")
+        print(client.connection.additional_headers)
+        execute_query_id1 = client.connection.additional_headers["X-Databend-Query-Id"]
+        self.assertEqual("X-Databend-Query-Id" in client.connection.additional_headers, True)
+        client.execute("select 2")
+        self.assertNotEqual(execute_query_id1, client.connection.additional_headers["X-Databend-Query-Id"])
+
 
 if __name__ == '__main__':
     print("start test......")