From 193f8425b18ab6ff1526b2bff52b50343aacd98c Mon Sep 17 00:00:00 2001
From: Roman Chukov <r.chukov@corp.mail.ru>
Date: Thu, 8 Jun 2023 16:26:20 +0300
Subject: [PATCH] Adding methods for getting redash jobid for the query and for
 tracking redash job status by its id

---
 redashAPI/client.py | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/redashAPI/client.py b/redashAPI/client.py
index 477bfac..a4a5587 100644
--- a/redashAPI/client.py
+++ b/redashAPI/client.py
@@ -99,6 +99,44 @@ def generate_query_results(self, ds_id: int, qry: str, qry_id: int=None, max_age
 
         return res
 
+    def query_and_return_jobid(self, ds_id: int, query: str):
+        payload = {
+            'data_source_id': ds_id,
+            'query': query,
+            'max_age': 0
+        }
+
+        res = self.post('query_results', payload)
+        job_id = res.json().get('job', {}).get('id')
+
+        return job_id
+
+    def track_job_status_by_id(self, job_id):
+        print(F"Tracking job {job_id}")
+        stickArray = [ "|", "/", "-", "\\", ]
+        i = 0
+
+        while True:
+            job = self.get(f'jobs/{job_id}')
+            job = job.json().get('job', {})
+            if job.get('status') == 3:
+                query_result_id = job.get('query_result_id')
+                break
+            elif job.get('status') == 4:
+                print(F"Failed to execute job {job_id}")
+                return "FAILURE"
+            elif job.get('status') == 5:
+                print(F"Canceled job {job_id}")
+                return "CANCELLED"
+            else:
+                sys.stdout.write('\b') # move back the cursor
+                sys.stdout.write(stickArray[i % 4])
+                sys.stdout.flush()
+                i = i + 1
+                time.sleep(0.2)
+
+        return self.get(f'query_results/{query_result_id}')
+
     def query_and_wait_result(self, ds_id: int, query: str, timeout: int=60):
         payload = {
             'data_source_id': ds_id,