Add progress and handle abort for import disk

disk.py

@@ -266,26 +266,51 @@ class Disk(object):
                 raise Exception("Invalid file format. Only qcow and "
                 raise Exception("Invalid file format. Only qcow and "
                                 "iso files are allowed. Image from: %s" % url)
                                 "iso files are allowed. Image from: %s" % url)
 
 
-    def import_disk(self, url):
+    def import_disk(self, task, url, parent_id):
         r = requests.get(url, stream=True)
         r = requests.get(url, stream=True)
+        clen = int(r.headers.get('content-length'))
         downloaded_file = os.path.join(self.dir, url.split('/')[-1])
         downloaded_file = os.path.join(self.dir, url.split('/')[-1])
-        with open(downloaded_file, 'wb') as f:
-            f.write(r.content)
-
-        cmdline = ['qemu-img',
-                   'convert',
-                   '-O', 'qcow2',
-                   downloaded_file,
-                   self.get_path()]
-        subprocess.check_output(cmdline)
+        percent = 0
+        try:
+            with open(downloaded_file, 'wb') as f:
+                for chunk in r.iter_content(chunk_size=256 * 1024):
+                    f.write(chunk)
+                    current_size = f.tell()
+                    new_percent = current_size * 100 / clen
+                    if task.is_aborted():
+                        raise AbortException()
+                    if new_percent > percent:
+                        percent = new_percent
+                        task.update_state(
+                            task_id=parent_id,
+                            state=task.AsyncResult(parent_id).state,
+                            meta={'percent': percent}
+                        )
 
 
-        os.unlink(downloaded_file)
+            cmdline = ['qemu-img',
+                       'convert',
+                       '-O', 'qcow2',
+                       downloaded_file,
+                       self.get_path()]
+            subprocess.check_output(cmdline)
+        except AbortException:
+            os.unlink(downloaded_file)
+            if os.path.exists(self.get_path()):
+                os.unlink(self.get_path())
+            logger.info("Import of disk %s aborted" % self.name)
+        except:
+            os.unlink(downloaded_file)
+            if os.path.exists(self.get_path()):
+                os.unlink(self.get_path())
+            raise
+        else:
+            os.unlink(downloaded_file)
 
 
-        if not self.check_valid_image():
-            os.unlink(self.get_path())
-            raise Exception("Invalid file format.")
+            if not self.check_valid_image():
+                os.unlink(self.get_path())
+                raise Exception("Invalid file format.")
 
 
-        self.size = Disk.get(self.dir, self.name).size
+            self.size = Disk.get(self.dir, self.name).size
 
 
     def export(self, format, exported_name, upload_link):
     def export(self, format, exported_name, upload_link):
         exported_path = self.get_path() + '.' + format
         exported_path = self.get_path() + '.' + format

storagedriver.py

@@ -41,11 +41,19 @@ class download(AbortableTask):
                 'checksum': disk.checksum, }
                 'checksum': disk.checksum, }
 
 
 
 
-@celery.task()
-def import_disk(disk_desc, url):
-    disk = Disk.deserialize(disk_desc)
-    disk.import_disk(url)
-    return disk.size
+class import_disk(AbortableTask):
+    time_limit = 18000
+
+    def run(self, **kwargs):
+        disk_desc = kwargs["disk_desc"]
+        url = kwargs["url"]
+        parent_id = kwargs["task"]
+        disk = Disk.deserialize(disk_desc)
+        disk.import_disk(self, url, parent_id)
+        return {
+            "size": disk.size,
+            "checksum": disk.checksum
+        }
 
 
 
 
 @celery.task()
 @celery.task()