Add progress and handle abort for import disk

disk.py

@@ -266,11 +266,26 @@ 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])
+        percent = 0
+        try:
             with open(downloaded_file, 'wb') as f:
             with open(downloaded_file, 'wb') as f:
-            f.write(r.content)
+                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}
+                        )
 
 
             cmdline = ['qemu-img',
             cmdline = ['qemu-img',
                        'convert',
                        'convert',
@@ -278,7 +293,17 @@ class Disk(object):
                        downloaded_file,
                        downloaded_file,
                        self.get_path()]
                        self.get_path()]
             subprocess.check_output(cmdline)
             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)
             os.unlink(downloaded_file)
 
 
             if not self.check_valid_image():
             if not self.check_valid_image():

storagedriver.py

@@ -41,11 +41,19 @@ class download(AbortableTask):
                 'checksum': disk.checksum, }
                 'checksum': disk.checksum, }
 
 
 
 
-@celery.task()
-def import_disk(disk_desc, url):
+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 = Disk.deserialize(disk_desc)
-    disk.import_disk(url)
-    return disk.size
+        disk.import_disk(self, url, parent_id)
+        return {
+            "size": disk.size,
+            "checksum": disk.checksum
+        }
 
 
 
 
 @celery.task()
 @celery.task()