From 63773a306180251647d599d7881a6171dfb1eb17 Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Thu, 16 Nov 2023 03:06:13 -0500
Subject: [PATCH] Allow watching existing task via cluster task

---
 client-cli/pvc/cli/cli.py     | 32 ++++++++++++++++++++++----
 client-cli/pvc/cli/waiters.py | 43 +++++++++++++++++++++--------------
 2 files changed, 53 insertions(+), 22 deletions(-)

diff --git a/client-cli/pvc/cli/cli.py b/client-cli/pvc/cli/cli.py
index 7fa7a576..cb8b9a38 100644
--- a/client-cli/pvc/cli/cli.py
+++ b/client-cli/pvc/cli/cli.py
@@ -539,7 +539,15 @@ def cli_cluster_maintenance_off():
 ###############################################################################
 @click.command(name="task", short_help="Show status of worker task.")
 @connection_req
-@click.argument("job_id", required=False, default=None)
+@click.argument("task_id", required=False, default=None)
+@click.option(
+    "--wait/--no-wait",
+    "wait_flag",
+    is_flag=True,
+    default=True,
+    show_default=True,
+    help="""Wait or don't wait for task to complete, showing progress. Requires TASK_ID; overrides "-f"/"--format".""",
+)
 @format_opt(
     {
         "pretty": cli_provisioner_status_format_pretty,
@@ -550,12 +558,26 @@ def cli_cluster_maintenance_off():
         "json-pretty": lambda d: jdumps(d, indent=2),
     }
 )
-def cli_cluster_task(job_id, format_function):
+def cli_cluster_task(task_id, wait_flag, format_function):
     """
-    Show the current status of worker task JOB_ID or a list of all active and pending tasks.
+    Show the current status of worker task TASK_ID or a list of all active and pending tasks.
     """
-    retcode, retdata = pvc.lib.common.task_status(CLI_CONFIG, job_id)
-    finish(retcode, retdata, format_function)
+    if task_id is None:
+        wait_flag = False
+
+    if wait_flag:
+        # First validate that this is actually a valid task that is running
+        retcode, retdata = pvc.lib.common.task_status(CLI_CONFIG, None)
+        if task_id in [i["id"] for i in retdata]:
+            retmsg = wait_for_celery_task(
+                CLI_CONFIG, {"task_id": task_id}, start_late=True
+            )
+        else:
+            retmsg = f"No task with ID {task_id} found."
+        finish(retcode, retmsg)
+    else:
+        retcode, retdata = pvc.lib.common.task_status(CLI_CONFIG, task_id)
+        finish(retcode, retdata, format_function)
 
 
 ###############################################################################
diff --git a/client-cli/pvc/cli/waiters.py b/client-cli/pvc/cli/waiters.py
index 70749b41..3f840ba7 100644
--- a/client-cli/pvc/cli/waiters.py
+++ b/client-cli/pvc/cli/waiters.py
@@ -65,36 +65,45 @@ def cli_node_waiter(config, node, state_field, state_value):
     echo(config, f" done. [{int(t_end - t_start)}s]")
 
 
-def wait_for_celery_task(CLI_CONFIG, task_detail):
+def wait_for_celery_task(CLI_CONFIG, task_detail, start_late=False):
     """
     Wait for a Celery task to complete
     """
 
     task_id = task_detail["task_id"]
-    run_on = task_detail["run_on"]
 
-    echo(CLI_CONFIG, f"Task ID: {task_id} assigned to node {run_on}")
-    echo(CLI_CONFIG, "")
+    if not start_late:
+        run_on = task_detail["run_on"]
 
-    # Wait for the task to start
-    echo(CLI_CONFIG, "Waiting for task to start...", newline=False)
-    while True:
-        sleep(0.5)
+        echo(CLI_CONFIG, f"Task ID: {task_id} assigned to node {run_on}")
+        echo(CLI_CONFIG, "")
+
+        # Wait for the task to start
+        echo(CLI_CONFIG, "Waiting for task to start...", newline=False)
+        while True:
+            sleep(0.5)
+            task_status = pvc.lib.common.task_status(
+                CLI_CONFIG, task_id=task_id, is_watching=True
+            )
+            if task_status.get("state") != "PENDING":
+                break
+            echo(CLI_CONFIG, ".", newline=False)
+        echo(CLI_CONFIG, " done.")
+        echo(CLI_CONFIG, "")
+
+        echo(
+            CLI_CONFIG,
+            task_status.get("status") + ":",
+        )
+    else:
         task_status = pvc.lib.common.task_status(
             CLI_CONFIG, task_id=task_id, is_watching=True
         )
-        if task_status.get("state") != "PENDING":
-            break
-        echo(CLI_CONFIG, ".", newline=False)
-    echo(CLI_CONFIG, " done.")
-    echo(CLI_CONFIG, "")
+
+        echo(CLI_CONFIG, f"Watching existing task {task_id}:")
 
     # Start following the task state, updating progress as we go
     total_task = task_status.get("total")
-    echo(
-        CLI_CONFIG,
-        task_status.get("status") + ":",
-    )
     with progressbar(length=total_task, show_eta=False) as bar:
         last_task = 0
         maxlen = 21