Add option for okdata status to wait

Add a `--watch` option to the `okdata status` to watch the status of a
dataset upload.

Author: simenheg

CHANGELOG.md

@@ -1,3 +1,8 @@
+## ?.?.?
+
+* The `okdata status` command can now be set to watch the status of a dataset
+  upload with the `--watch` option.
+
 ## 3.1.0
 
 * Added support for Python 3.12.

okdata/cli/commands/datasets/datasets.py

@@ -354,8 +354,8 @@ def upload_file(self, source, target):
         summary = [f"Uploaded file to dataset: {dataset_id}", str(out)]
         if res["trace_id"]:
             summary += [
-                "\nYou can check the data processing status by running:\n",
-                f"  okdata status {res['trace_id']}",
+                "\nYou can watch the data processing status by running:\n",
+                f"  okdata status {res['trace_id']} --watch",
             ]
         self.print("\n".join(summary))
 

okdata/cli/commands/status.py

@@ -1,3 +1,5 @@
+from time import sleep
+
 from okdata.sdk.status import Status
 
 from okdata.cli.command import BaseCommand, BASE_COMMAND_OPTIONS
@@ -8,15 +10,17 @@ class StatusCommand(BaseCommand):
     __doc__ = f"""Oslo :: Status
 
 Usage:
-  okdata status <trace_id> [options --history]
+  okdata status <trace_id> [options]
 
 Examples:
   okdata status trace-id-from-system
-  okdata status trace-id-from-system --format=json | jq ".done"
+  okdata status trace-id-from-system --watch
   okdata status trace-id-from-system --history
+  okdata status trace-id-from-system --format=json | jq ".done"
 
 Options:{BASE_COMMAND_OPTIONS}
   --history
+  --watch
     """
 
     def __init__(self):
@@ -28,7 +32,7 @@ def login(self):
     def handler(self):
         self.log.info("StatusCommand.handler()")
         if self.arg("trace_id"):
-            self.status_for_id()
+            self.status_for_id(self.arg("trace_id"))
         else:
             self.print("Invalid command")
 
@@ -84,11 +88,35 @@ def full_history_for_status(self, trace_id, trace_events):
                 {"error": 1, "message": "No trace events found"},
             )
 
-    def status_for_id(self):
-        trace_id = self.arg("trace_id")
+    def get_trace_events(self, trace_id):
         self.log.info(f"Looking up status: {trace_id}")
         trace_events = self.sdk.get_status(trace_id)
-        trace_events = StatusCommand.get_error_messages(trace_events)
+        return StatusCommand.get_error_messages(trace_events)
+
+    def wait_until_done(self, trace_id):
+        """Wait until status for `trace_id` is ready, then return the trace."""
+        trace_events = self.get_trace_events(trace_id)
+        trace_status = StatusCommand.find_latest_event(trace_events)["trace_status"]
+
+        if trace_status != "FINISHED":
+            print("Waiting for upload to finish", end="", flush=True)
+
+        while trace_status != "FINISHED":
+            print(".", end="", flush=True)
+            sleep(2)
+            trace_events = self.get_trace_events(trace_id)
+            trace_status = StatusCommand.find_latest_event(trace_events)["trace_status"]
+        print()
+
+        return trace_events
+
+    def status_for_id(self, trace_id):
+        trace_events = (
+            self.wait_until_done(trace_id)
+            if self.opt("watch")
+            else self.get_trace_events(trace_id)
+        )
+
         if self.opt("history"):
             self.full_history_for_status(trace_id, trace_events)
         else: