diff --git a/databend_py/connection.py b/databend_py/connection.py index c4d1820..93c0b55 100644 --- a/databend_py/connection.py +++ b/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( diff --git a/tests/test_client.py b/tests/test_client.py index 833f288..3716c56 100644 --- a/tests/test_client.py +++ b/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......")