Merge pull request #64 from datafuselabs/fix/additional_headers

fix: additional_headers

Author: hantmac

databend_py/connection.py

@@ -164,7 +164,11 @@ def query(self, statement):
         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)}
+            if self.additional_headers:
+                self.additional_headers.update(
+                    {XDatabendQueryIDHeader: resp_dict.get(QueryID)})
+            else:
+                self.additional_headers = {XDatabendQueryIDHeader: resp_dict.get(QueryID)}
             return self.wait_until_has_schema(resp_dict)
         except Exception as err:
             log.logger.error(

tests/test_client.py

@@ -199,9 +199,11 @@ def test_null_to_none(self):
         self.assertIsNone(data[0][0])
 
     def test_set_query_id_header(self):
+        os.environ["ADDITIONAL_HEADERS"] = "X-DATABENDCLOUD-TENANT=TENANT,X-DATABENDCLOUD-WAREHOUSE=WAREHOUSE"
         client = Client.from_url(self.databend_url)
+        self.assertEqual("X-DATABENDCLOUD-TENANT" in client.connection.additional_headers, True)
+        self.assertEqual(client.connection.additional_headers["X-DATABENDCLOUD-TENANT"], "TENANT")
         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")