Improve stage handling

Run start() at the beginning, and leverage the new tweaks to the CLI to
update the total steps later. Allows errors to be handled gracefully

api-daemon/pvcapid/flaskapi.py

@@ -3194,7 +3194,23 @@ class API_VM_Snapshot(Resource):
               id: Message
         """
         snapshot_name = reqargs.get("snapshot_name", None)
-        return api_helper.create_vm_snapshot(vm, snapshot_name=snapshot_name)
+
+        task = run_celery_task(
+            "vm.create_snapshot",
+            domain=vm,
+            snapshot_name=snapshot_name,
+            run_on="primary",
+        )
+
+        return (
+            {
+                "task_id": task.id,
+                "task_name": "vm.create_snapshot",
+                "run_on": f"{get_primary_node()} (primary)",
+            },
+            202,
+            {"Location": Api.url_for(api, API_Tasks_Element, task_id=task.id)},
+        )
 
     @RequestParser(
         [
@@ -3236,7 +3252,23 @@ class API_VM_Snapshot(Resource):
               id: Message
         """
         snapshot_name = reqargs.get("snapshot_name", None)
-        return api_helper.remove_vm_snapshot(vm, snapshot_name)
+
+        task = run_celery_task(
+            "vm.remove_snapshot",
+            domain=vm,
+            snapshot_name=snapshot_name,
+            run_on="primary",
+        )
+
+        return (
+            {
+                "task_id": task.id,
+                "task_name": "vm.remove_snapshot",
+                "run_on": f"{get_primary_node()} (primary)",
+            },
+            202,
+            {"Location": Api.url_for(api, API_Tasks_Element, task_id=task.id)},
+        )
 
 
 api.add_resource(API_VM_Snapshot, "/vm/<vm>/snapshot")
@@ -3284,7 +3316,23 @@ class API_VM_Snapshot_Rollback(Resource):
               id: Message
         """
         snapshot_name = reqargs.get("snapshot_name", None)
-        return api_helper.rollback_vm_snapshot(vm, snapshot_name)
+
+        task = run_celery_task(
+            "vm.rollback_snapshot",
+            domain=vm,
+            snapshot_name=snapshot_name,
+            run_on="primary",
+        )
+
+        return (
+            {
+                "task_id": task.id,
+                "task_name": "vm.rollback_snapshot",
+                "run_on": f"{get_primary_node()} (primary)",
+            },
+            202,
+            {"Location": Api.url_for(api, API_Tasks_Element, task_id=task.id)},
+        )
 
 
 api.add_resource(API_VM_Snapshot_Rollback, "/vm/<vm>/snapshot/rollback")
@@ -3354,8 +3402,24 @@ class API_VM_Snapshot_Export(Resource):
         snapshot_name = reqargs.get("snapshot_name", None)
         export_path = reqargs.get("export_path", None)
         incremental_parent = reqargs.get("incremental_parent", None)
-        return api_helper.export_vm_snapshot(
-            vm, snapshot_name, export_path, incremental_parent
+
+        task = run_celery_task(
+            "vm.export_snapshot",
+            domain=vm,
+            snapshot_name=snapshot_name,
+            export_path=export_path,
+            incremental_parent=incremental_parent,
+            run_on="primary",
+        )
+
+        return (
+            {
+                "task_id": task.id,
+                "task_name": "vm.export_snapshot",
+                "run_on": f"{get_primary_node()} (primary)",
+            },
+            202,
+            {"Location": Api.url_for(api, API_Tasks_Element, task_id=task.id)},
         )
 
 
@@ -3427,8 +3491,24 @@ class API_VM_Snapshot_Import(Resource):
         snapshot_name = reqargs.get("snapshot_name", None)
         import_path = reqargs.get("import_path", None)
         retain_snapshot = bool(strtobool(reqargs.get("retain_snapshot", "True")))
-        return api_helper.import_vm_snapshot(
-            vm, snapshot_name, import_path, retain_snapshot
+
+        task = run_celery_task(
+            "vm.import_snapshot",
+            domain=vm,
+            snapshot_name=snapshot_name,
+            import_path=import_path,
+            retain_snapshot=retain_snapshot,
+            run_on="primary",
+        )
+
+        return (
+            {
+                "task_id": task.id,
+                "task_name": "vm.import_snapshot",
+                "run_on": f"{get_primary_node()} (primary)",
+            },
+            202,
+            {"Location": Api.url_for(api, API_Tasks_Element, task_id=task.id)},
         )
 
 

client-cli/pvc/cli/cli.py

@@ -1758,7 +1758,7 @@ def cli_vm_flush_locks(domain, wait_flag):
     NOTE: This is a task-based command. The "--wait" flag (default) will block and show progress. Specifying the "--no-wait" flag will return immediately with a job ID instead, which can be queried externally later.
     """
 
-    retcode, retmsg = pvc.lib.vm.vm_locks(CLI_CONFIG, domain, wait_flag)
+    retcode, retmsg = pvc.lib.vm.vm_locks(CLI_CONFIG, domain, wait_flag=wait_flag)
 
     if retcode and wait_flag:
         retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
@@ -1787,7 +1787,15 @@ def cli_vm_snapshot():
 @connection_req
 @click.argument("domain")
 @click.argument("snapshot_name", required=False, default=None)
-def cli_vm_snapshot_create(domain, snapshot_name):
+@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",
+)
+def cli_vm_snapshot_create(domain, snapshot_name, wait_flag):
     """
     Create a snapshot of the disks and XML configuration of virtual machine DOMAIN, with the
     optional name SNAPSHOT_NAME. DOMAIN may be a UUID or name.
@@ -1797,18 +1805,12 @@ def cli_vm_snapshot_create(domain, snapshot_name):
     VM at the moment of the snapshot.
     """
 
-    echo(
-        CLI_CONFIG,
-        f"Taking snapshot of VM '{domain}'... ",
-        newline=False,
-    )
     retcode, retmsg = pvc.lib.vm.vm_create_snapshot(
-        CLI_CONFIG, domain, snapshot_name=snapshot_name
+        CLI_CONFIG, domain, snapshot_name=snapshot_name, wait_flag=wait_flag
     )
-    if retcode:
-        echo(CLI_CONFIG, "done.")
-    else:
-        echo(CLI_CONFIG, "failed.")
+
+    if retcode and wait_flag:
+        retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
     finish(retcode, retmsg)
 
 
@@ -1819,23 +1821,27 @@ def cli_vm_snapshot_create(domain, snapshot_name):
 @connection_req
 @click.argument("domain")
 @click.argument("snapshot_name")
+@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",
+)
 @confirm_opt("Remove shapshot {snapshot_name} of VM {domain}")
-def cli_vm_snapshot_remove(domain, snapshot_name):
+def cli_vm_snapshot_remove(domain, snapshot_name, wait_flag):
     """
     Remove the snapshot SNAPSHOT_NAME of the disks and XML configuration of virtual machine DOMAIN,
     DOMAIN may be a UUID or name.
     """
 
-    echo(
-        CLI_CONFIG,
-        f"Removing snapshot '{snapshot_name}' of VM '{domain}'... ",
-        newline=False,
+    retcode, retmsg = pvc.lib.vm.vm_remove_snapshot(
+        CLI_CONFIG, domain, snapshot_name, wait_flag=wait_flag
     )
-    retcode, retmsg = pvc.lib.vm.vm_remove_snapshot(CLI_CONFIG, domain, snapshot_name)
-    if retcode:
-        echo(CLI_CONFIG, "done.")
-    else:
-        echo(CLI_CONFIG, "failed.")
+
+    if retcode and wait_flag:
+        retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
     finish(retcode, retmsg)
 
 
@@ -1848,25 +1854,29 @@ def cli_vm_snapshot_remove(domain, snapshot_name):
 @connection_req
 @click.argument("domain")
 @click.argument("snapshot_name")
+@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",
+)
 @confirm_opt(
     "Roll back to snapshot {snapshot_name} of {domain} and lose all data and changes since this snapshot"
 )
-def cli_vm_snapshot_rollback(domain, snapshot_name):
+def cli_vm_snapshot_rollback(domain, snapshot_name, wait_flag):
     """
     Roll back to the snapshot SNAPSHOT_NAME of the disks and XML configuration of virtual machine DOMAIN,
     DOMAIN may be a UUID or name.
     """
 
-    echo(
-        CLI_CONFIG,
-        f"Rolling back to snapshot '{snapshot_name}' of VM '{domain}'... ",
-        newline=False,
+    retcode, retmsg = pvc.lib.vm.vm_rollback_snapshot(
+        CLI_CONFIG, domain, snapshot_name, wait_flag=wait_flag
     )
-    retcode, retmsg = pvc.lib.vm.vm_rollback_snapshot(CLI_CONFIG, domain, snapshot_name)
-    if retcode:
-        echo(CLI_CONFIG, "done.")
-    else:
-        echo(CLI_CONFIG, "failed.")
+
+    if retcode and wait_flag:
+        retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
     finish(retcode, retmsg)
 
 
@@ -1887,7 +1897,17 @@ def cli_vm_snapshot_rollback(domain, snapshot_name):
     default=None,
     help="Perform an incremental volume export from this parent snapshot.",
 )
-def cli_vm_snapshot_export(domain, snapshot_name, export_path, incremental_parent):
+@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",
+)
+def cli_vm_snapshot_export(
+    domain, snapshot_name, export_path, incremental_parent, wait_flag
+):
     """
     Export the (existing) snapshot SNAPSHOT_NAME of virtual machine DOMAIN to the absolute path EXPORT_PATH on the current PVC primary coordinator.
     DOMAIN may be a UUID or name.
@@ -1901,19 +1921,17 @@ def cli_vm_snapshot_export(domain, snapshot_name, export_path, incremental_paren
     Full export volume images are sparse-allocated, however it is recommended for safety to consider their maximum allocated size when allocated space for the EXPORT_PATH. Incremental volume images are generally small but are dependent entirely on the rate of data change in each volume.
     """
 
-    _, primary_node = pvc.lib.cluster.get_primary_node(CLI_CONFIG)
-    echo(
-        CLI_CONFIG,
-        f'Exporting snapshot "{snapshot_name}" of VM "{domain}" to "{primary_node}:{export_path}"... ',
-        newline=False,
-    )
     retcode, retmsg = pvc.lib.vm.vm_export_snapshot(
-        CLI_CONFIG, domain, snapshot_name, export_path, incremental_parent
+        CLI_CONFIG,
+        domain,
+        snapshot_name,
+        export_path,
+        incremental_parent=incremental_parent,
+        wait_flag=wait_flag,
     )
-    if retcode:
-        echo(CLI_CONFIG, "done.")
-    else:
-        echo(CLI_CONFIG, "failed.")
+
+    if retcode and wait_flag:
+        retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
     finish(retcode, retmsg)
 
 
@@ -1933,7 +1951,17 @@ def cli_vm_snapshot_export(domain, snapshot_name, export_path, incremental_paren
     default=True,
     help="Retain or remove restored (parent, if incremental) snapshot in Ceph.",
 )
-def cli_vm_snapshot_import(domain, snapshot_name, import_path, retain_snapshot):
+@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",
+)
+def cli_vm_snapshot_import(
+    domain, snapshot_name, import_path, retain_snapshot, wait_flag
+):
     """
     Import the snapshot SNAPSHOT_NAME of virtual machine DOMAIN from the absolute path IMPORT_PATH on the current PVC primary coordinator.
     DOMAIN may be a UUID or name.
@@ -1942,25 +1970,24 @@ def cli_vm_snapshot_import(domain, snapshot_name, import_path, retain_snapshot):
 
     The import will include the VM configuration, metainfo, and the point-in-time snapshot of all attached RBD volumes. Incremental imports will be automatically handled.
 
-    A VM named DOMAIN or with the same UUID must not exist; if a VM with the same name or UUID already exists, it must be removed, or renamed and then undefined (to preserve volumes), before importing.
+    A VM named DOMAIN or with the same UUID must not exist; if a VM with the same name or UUID already exists, it must be removed (or renamed and then undefined, to preserve volumes while freeing the UUID) before importing.
 
     If the "-r"/"--retain-snapshot" option is specified (the default), for incremental imports, only the parent snapshot is kept; for full imports, the imported snapshot is kept. If the "-R"/"--remove-snapshot" option is specified, the imported snapshot is removed.
 
     WARNING: The "-R"/"--remove-snapshot" option will invalidate any existing incremental snapshots based on the same incremental parent for the imported VM.
     """
 
-    echo(
-        CLI_CONFIG,
-        f"Importing snapshot '{snapshot_name}' of VM '{domain}'... ",
-        newline=False,
-    )
     retcode, retmsg = pvc.lib.vm.vm_import_snapshot(
-        CLI_CONFIG, domain, snapshot_name, import_path, retain_snapshot
+        CLI_CONFIG,
+        domain,
+        snapshot_name,
+        import_path,
+        retain_snapshot=retain_snapshot,
+        wait_flag=wait_flag,
     )
-    if retcode:
-        echo(CLI_CONFIG, "done.")
-    else:
-        echo(CLI_CONFIG, "failed.")
+
+    if retcode and wait_flag:
+        retmsg = wait_for_celery_task(CLI_CONFIG, retmsg)
     finish(retcode, retmsg)
 
 
@@ -2055,7 +2082,7 @@ def cli_vm_backup_restore(domain, backup_datestring, backup_path, retain_snapsho
 
     The restore will import the VM configuration, metainfo, and the point-in-time snapshot of all attached RBD volumes. Incremental backups will be automatically handled.
 
-    A VM named DOMAIN or with the same UUID must not exist; if a VM with the same name or UUID already exists, it must be removed, or renamed and then undefined (to preserve volumes), before restoring.
+    A VM named DOMAIN or with the same UUID must not exist; if a VM with the same name or UUID already exists, it must be removed (or renamed and then undefined, to preserve volumes while freeing the UUID) before importing.
 
     If the "-r"/"--retain-snapshot" option is specified (the default), for incremental restores, only the parent snapshot is kept; for full restores, the restored snapshot is kept. If the "-R"/"--remove-snapshot" option is specified, the imported snapshot is removed.
 

client-cli/pvc/cli/waiters.py

@@ -121,6 +121,8 @@ def wait_for_celery_task(CLI_CONFIG, task_detail, start_late=False):
                 break
             if task_status.get("current") > last_task:
                 current_task = int(task_status.get("current"))
+                total_task = int(task_status.get("total"))
+                bar.length = total_task
                 bar.update(current_task - last_task)
                 last_task = current_task
                 # The extensive spaces at the end cause this to overwrite longer previous messages

client-cli/pvc/lib/vm.py

@@ -421,7 +421,7 @@ def vm_node(config, vm, target_node, action, force=False, wait=False, force_live
     return retstatus, response.json().get("message", "")
 
 
-def vm_locks(config, vm, wait_flag):
+def vm_locks(config, vm, wait_flag=True):
     """
     Flush RBD locks of (stopped) VM
 
@@ -498,7 +498,7 @@ def vm_restore(config, vm, backup_path, backup_datestring, retain_snapshot=False
         return True, response.json().get("message", "")
 
 
-def vm_create_snapshot(config, vm, snapshot_name=None):
+def vm_create_snapshot(config, vm, snapshot_name=None, wait_flag=True):
     """
     Take a snapshot of a VM's disks and configuration
 
@@ -513,13 +513,10 @@ def vm_create_snapshot(config, vm, snapshot_name=None):
         config, "post", "/vm/{vm}/snapshot".format(vm=vm), params=params
     )
 
-    if response.status_code != 200:
-        return False, response.json().get("message", "")
-    else:
-        return True, response.json().get("message", "")
+    return get_wait_retdata(response, wait_flag)
 
 
-def vm_remove_snapshot(config, vm, snapshot_name):
+def vm_remove_snapshot(config, vm, snapshot_name, wait_flag=True):
     """
     Remove a snapshot of a VM's disks and configuration
 
@@ -532,13 +529,10 @@ def vm_remove_snapshot(config, vm, snapshot_name):
         config, "delete", "/vm/{vm}/snapshot".format(vm=vm), params=params
     )
 
-    if response.status_code != 200:
-        return False, response.json().get("message", "")
-    else:
-        return True, response.json().get("message", "")
+    return get_wait_retdata(response, wait_flag)
 
 
-def vm_rollback_snapshot(config, vm, snapshot_name):
+def vm_rollback_snapshot(config, vm, snapshot_name, wait_flag=True):
     """
     Roll back to a snapshot of a VM's disks and configuration
 
@@ -551,13 +545,12 @@ def vm_rollback_snapshot(config, vm, snapshot_name):
         config, "post", "/vm/{vm}/snapshot/rollback".format(vm=vm), params=params
     )
 
-    if response.status_code != 200:
-        return False, response.json().get("message", "")
-    else:
-        return True, response.json().get("message", "")
+    return get_wait_retdata(response, wait_flag)
 
 
-def vm_export_snapshot(config, vm, snapshot_name, export_path, incremental_parent):
+def vm_export_snapshot(
+    config, vm, snapshot_name, export_path, incremental_parent=None, wait_flag=True
+):
     """
     Export an (existing) snapshot of a VM's disks and configuration to export_path, optionally
     incremental with incremental_parent
@@ -577,13 +570,12 @@ def vm_export_snapshot(config, vm, snapshot_name, export_path, incremental_paren
         config, "post", "/vm/{vm}/snapshot/export".format(vm=vm), params=params
     )
 
-    if response.status_code != 200:
-        return False, response.json().get("message", "")
-    else:
-        return True, response.json().get("message", "")
+    return get_wait_retdata(response, wait_flag)
 
 
-def vm_import_snapshot(config, vm, snapshot_name, import_path, retain_snapshot=False):
+def vm_import_snapshot(
+    config, vm, snapshot_name, import_path, retain_snapshot=False, wait_flag=True
+):
     """
     Import a snapshot of {vm} and its volumes from a local primary coordinator filesystem path
 
@@ -600,10 +592,7 @@ def vm_import_snapshot(config, vm, snapshot_name, import_path, retain_snapshot=F
         config, "post", "/vm/{vm}/snapshot/import".format(vm=vm), params=params
     )
 
-    if response.status_code != 200:
-        return False, response.json().get("message", "")
-    else:
-        return True, response.json().get("message", "")
+    return get_wait_retdata(response, wait_flag)
 
 
 def vm_vcpus_set(config, vm, vcpus, topology, restart):

worker-daemon/pvcworkerd/Daemon.py

@@ -28,6 +28,11 @@ from daemon_lib.vm import (
     vm_worker_flush_locks,
     vm_worker_attach_device,
     vm_worker_detach_device,
+    vm_worker_create_snapshot,
+    vm_worker_remove_snapshot,
+    vm_worker_rollback_snapshot,
+    vm_worker_export_snapshot,
+    vm_worker_import_snapshot,
 )
 from daemon_lib.ceph import (
     osd_worker_add_osd,
@@ -123,6 +128,87 @@ def vm_device_detach(self, domain=None, xml=None, run_on=None):
     return run_vm_device_detach(self, domain, xml)
 
 
+@celery.task(name="vm.create_snapshot", bind=True, routing_key="run_on")
+def vm_create_snapshot(self, domain=None, snapshot_name=None, run_on="primary"):
+    @ZKConnection(config)
+    def run_vm_create_snapshot(zkhandler, self, domain, snapshot_name):
+        return vm_worker_create_snapshot(zkhandler, self, domain, snapshot_name)
+
+    return run_vm_create_snapshot(self, domain, snapshot_name)
+
+
+@celery.task(name="vm.remove_snapshot", bind=True, routing_key="run_on")
+def vm_remove_snapshot(self, domain=None, snapshot_name=None, run_on="primary"):
+    @ZKConnection(config)
+    def run_vm_remove_snapshot(zkhandler, self, domain, snapshot_name):
+        return vm_worker_remove_snapshot(zkhandler, self, domain, snapshot_name)
+
+    return run_vm_remove_snapshot(self, domain, snapshot_name)
+
+
+@celery.task(name="vm.rollback_snapshot", bind=True, routing_key="run_on")
+def vm_rollback_snapshot(self, domain=None, snapshot_name=None, run_on="primary"):
+    @ZKConnection(config)
+    def run_vm_rollback_snapshot(zkhandler, self, domain, snapshot_name):
+        return vm_worker_rollback_snapshot(zkhandler, self, domain, snapshot_name)
+
+    return run_vm_rollback_snapshot(self, domain, snapshot_name)
+
+
+@celery.task(name="vm.export_snapshot", bind=True, routing_key="run_on")
+def vm_export_snapshot(
+    self,
+    domain=None,
+    snapshot_name=None,
+    export_path=None,
+    incremental_parent=None,
+    run_on="primary",
+):
+    @ZKConnection(config)
+    def run_vm_export_snapshot(
+        zkhandler, self, domain, snapshot_name, export_path, incremental_parent=None
+    ):
+        return vm_worker_export_snapshot(
+            zkhandler,
+            self,
+            domain,
+            snapshot_name,
+            export_path,
+            incremental_parent=incremental_parent,
+        )
+
+    return run_vm_export_snapshot(
+        self, domain, snapshot_name, export_path, incremental_parent=incremental_parent
+    )
+
+
+@celery.task(name="vm.import_snapshot", bind=True, routing_key="run_on")
+def vm_import_snapshot(
+    self,
+    domain=None,
+    snapshot_name=None,
+    import_path=None,
+    retain_snapshot=True,
+    run_on="primary",
+):
+    @ZKConnection(config)
+    def run_vm_import_snapshot(
+        zkhandler, self, domain, snapshot_name, import_path, retain_snapshot=True
+    ):
+        return vm_worker_import_snapshot(
+            zkhandler,
+            self,
+            domain,
+            snapshot_name,
+            import_path,
+            retain_snapshot=retain_snapshot,
+        )
+
+    return run_vm_import_snapshot(
+        self, domain, snapshot_name, import_path, retain_snapshot=retain_snapshot
+    )
+
+
 @celery.task(name="osd.add", bind=True, routing_key="run_on")
 def osd_add(
     self,