From d1451919dd695cd6bcb16179ed11851c7fa40c36 Mon Sep 17 00:00:00 2001 From: hantmac Date: Fri, 22 Mar 2024 17:01:51 +0800 Subject: [PATCH] fix --- databend_py/connection.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/databend_py/connection.py b/databend_py/connection.py index 362b771..383507c 100644 --- a/databend_py/connection.py +++ b/databend_py/connection.py @@ -165,7 +165,7 @@ def query(self, statement): 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 + return self.wait_until_has_schema(resp_dict) except Exception as err: log.logger.error( f"http error on {url}, SQL: {statement} error msg:{str(err)}" @@ -182,6 +182,19 @@ def format_url(self): def reset_session(self): self.client_session = dict() + def wait_until_has_schema(self, raw_data_dict): + resp_schema = raw_data_dict.get("schema") + while resp_schema is not None and len(resp_schema) == 0: + if raw_data_dict['next_uri'] is None: + break + resp = self.next_page(raw_data_dict['next_uri']) + resp_dict = json.loads(resp.content) + raw_data_dict = resp_dict + resp_schema = raw_data_dict.get("schema") + if resp_schema is not None and (len(resp_schema) != 0 or len(raw_data_dict.get("data")) != 0): + break + return raw_data_dict + def next_page(self, next_uri): url = "{}://{}:{}{}".format(self.schema, self.host, self.port, next_uri) return self.requests_session.get(url=url, headers=self.make_headers(), cookies=self.cookies)