fetch more pages for long-running queries

databend_py/connection.py

@@ -114,14 +114,24 @@ def query(self, statement, session):
             session = {"database": self.database}
             query_sql['session'] = session
         log.logger.debug(f"http headers {self.make_headers()}")
-        response = requests.post(url,
-                                 data=json.dumps(query_sql),
-                                 headers=self.make_headers(),
-                                 auth=HTTPBasicAuth(self.user, self.password),
-                                 verify=True)
 
         try:
-            return json.loads(response.content)
+            response = requests.post(url,
+                                    data=json.dumps(query_sql),
+                                    headers=self.make_headers(),
+                                    auth=HTTPBasicAuth(self.user, self.password),
+                                    verify=True)
+            response.raise_for_status()
+
+            j = response.json()
+            while j.get('state') == 'Running':
+                time.sleep(0.5)
+                response = self.next_page(j.get('next_uri'))
+                response.raise_for_status()
+                j = response.json()
+
+            return j
+
         except Exception as err:
             log.logger.error(
                 f"http error on {url}, SQL: {statement} content: {response.content} error msg:{str(err)}"