Add wait_for_job to FlowRun

Author: Jordan Woods

tableauserverclient/server/endpoint/exceptions.py

@@ -64,3 +64,16 @@ def __str__(self):
         from pprint import pformat
 
         return pformat(self.error)
+
+
+class FlowRunFailedException(Exception):
+    def __init__(self, flow_run):
+        self.background_job_id = flow_run.background_job_id
+        self.flow_run = flow_run
+
+    def __str__(self):
+        return f"FlowRun {self.flow_run.id} failed with job id {self.background_job_id}"
+
+
+class FlowRunCanceledException(FlowRunFailedException):
+    pass 

tableauserverclient/server/endpoint/flow_runs_endpoint.py

@@ -1,10 +1,16 @@
 from .endpoint import Endpoint, QuerysetEndpoint, api
 from .. import FlowRunItem, PaginationItem
+from .exceptions import FlowRunFailedException, FlowRunCanceledException
+import time
 
 import logging
 
 logger = logging.getLogger("tableau.endpoint.flowruns")
 
+# Polling for job completion uses exponential backoff for the sleep intervals between polls
+ASYNC_JOB_POLL_MIN_INTERVAL=0.5
+ASYNC_JOB_POLL_MAX_INTERVAL=30
+ASYNC_JOB_POLL_BACKOFF_FACTOR=1.4
 
 class FlowRuns(QuerysetEndpoint):
     def __init__(self, parent_srv):
@@ -46,3 +52,37 @@ def cancel(self, flow_run_id):
         url = "{0}/{1}".format(self.baseurl, id_)
         self.put_request(url)
         logger.info("Deleted single flow (ID: {0})".format(id_))
+
+
+    @api(version="3.10")
+    def wait_for_job(self, flow_run_id, *, timeout=None):
+        id_ = getattr(flow_run_id, "id", flow_run_id)
+        wait_start_time = time.time()
+        logger.debug(f"Waiting for job {id_}")
+
+        current_sleep_interval = ASYNC_JOB_POLL_MIN_INTERVAL
+        flow_run = self.get_by_id(id_)
+        while flow_run.completed_at is None:
+            max_sleep_time = ASYNC_JOB_POLL_MAX_INTERVAL
+
+            if timeout is not None:
+                elapsed = (time.time() - wait_start_time)
+                if elapsed >= timeout:
+                    raise TimeoutError(f"Timeout after {elapsed} seconds waiting for asynchronous flow run: {id_}")
+                max_sleep_time = max(ASYNC_JOB_POLL_MIN_INTERVAL, timeout - elapsed)
+
+            time.sleep(min(current_sleep_interval, max_sleep_time))
+            job = self.get_by_id(id_)
+            current_sleep_interval *= ASYNC_JOB_POLL_BACKOFF_FACTOR
+            logger.debug(f"\tFlowRun {id_} progress={flow_run.progress}")
+
+        logger.info("FlowRun {} Completed: Status: {}".format(id_, flow_run.status))
+
+        if flow_run.status == "Success":
+            return flow_run
+        elif flow_run.status == "Failed":
+            raise FlowRunFailedException(flow_run)
+        elif flow_run.finish_code in ["Canceled", "Cancelled"]:
+            raise FlowRunCanceledException(flow_run)
+        else:
+            raise AssertionError("Unexpected status in flow_run", flow_run)