Add progress bar in pithos upload & compute md5sum
authorNikos Skalkotos <skalkoto@grnet.gr>
Fri, 20 Apr 2012 14:38:40 +0000 (17:38 +0300)
committerNikos Skalkotos <skalkoto@grnet.gr>
Fri, 20 Apr 2012 14:38:40 +0000 (17:38 +0300)
image_creator/disk.py
image_creator/kamaki_wrapper.py
image_creator/main.py
image_creator/util.py

index ef53799..1c9d04e 100644 (file)
@@ -179,7 +179,8 @@ class DiskDevice(object):
 
     def enable(self):
         """Enable a newly created DiskDevice"""
-        self.progressbar = progress("Launching helper VM: ")
+        new_progress = progress("Launching helper VM: ")
+        self.progressbar = new_progress()
         self.progressbar.next()
         eh = self.g.set_event_callback(self.progress_callback,
                                                     guestfs.EVENT_PROGRESS)
@@ -317,8 +318,8 @@ class DiskDevice(object):
         blocksize = 2 ** 22  # 4MB
         size = self.size()
         progress_size = (size + 2 ** 20 - 1) // 2 ** 20  # in MB
-        progressbar = progress("Dumping image file: ", progress_size)
-
+        new_progress = progress("Dumping image file: ")
+        progressbar = new_progress(progress_size)
         source = open(self.device, "r")
         try:
             dest = open(outfile, "w")
index b7fce51..5be0c33 100644 (file)
@@ -38,7 +38,7 @@ from kamaki.clients import ClientError
 from kamaki.clients.image import ImageClient
 from kamaki.clients.pithos import PithosClient
 
-from image_creator.util import FatalError
+from image_creator.util import FatalError, progress
 
 CONTAINER = "images"
 
@@ -73,7 +73,10 @@ class Kamaki:
                     raise FatalError("Pithos client: %d %s" % \
                                                         (e.status, e.message))
             try:
-                self.pithos_client.create_object(remote_path, f, size)
+               hash_progress = progress("(1/2)  Calculating block hashes:")
+               upload_progress = progress("(2/2)  Uploading missing blocks:")
+                self.pithos_client.create_object(remote_path, f, size,
+                                               hash_progress, upload_progress)
                 self.uploaded_object = "pithos://%s/%s/%s" % \
                                 (self.account, self.container, remote_path)
             except ClientError as e:
index d073946..af236d6 100755 (executable)
@@ -37,7 +37,8 @@ from image_creator import get_os_class
 from image_creator import __version__ as version
 from image_creator import util
 from image_creator.disk import Disk
-from image_creator.util import get_command, error, success, output, FatalError
+from image_creator.util import get_command, error, success, output, \
+                                                    FatalError, progress, md5
 from image_creator.kamaki_wrapper import Kamaki
 import sys
 import os
@@ -153,7 +154,7 @@ def image_creator():
                         % os.path.basename(sys.argv[0]))
 
     if not options.force and options.outfile is not None:
-        for extension in ('', '.meta'):
+        for extension in ('', '.meta', '.md5sum'):
             filename = "%s%s" % (options.outfile, extension)
             if os.path.exists(filename):
                 raise FatalError("Output file %s exists "
@@ -193,6 +194,12 @@ def image_creator():
         size = options.shrink and dev.shrink() or dev.size()
         metadata['SIZE'] = str(size // 2 ** 20)
 
+        #Calculating MD5sum
+        output("Calculating md5sum...", False)
+        checksum = md5(snapshot, size)
+        success(checksum)
+        output()
+
         if options.outfile is not None:
             f = open('%s.%s' % (options.outfile, 'meta'), 'w')
             try:
@@ -207,10 +214,9 @@ def image_creator():
         disk.destroy_device(dev)
 
         if options.upload:
-            output("Uploading image to pithos...", False)
+            output("Uploading image to pithos:")
             kamaki = Kamaki(options.account, options.token)
             kamaki.upload(snapshot, size, options.upload)
-            output("done")
 
     finally:
         output('cleaning up...')
index 1b4994a..2805de7 100644 (file)
@@ -33,6 +33,7 @@
 
 import sys
 import pbs
+import hashlib
 from clint.textui import colored, progress as uiprogress
 
 
@@ -84,19 +85,34 @@ def output(msg="", new_line=True):
             sys.stdout.flush()
 
 
-def progress(label='', n=100):
+def progress(message=''):
 
     PROGRESS_LENGTH = 32
     MESSAGE_LENGTH = 32
 
-    def progress_generator(label, n):
+    def progress_generator(n=100):
         position = 0
-        for i in uiprogress.bar(range(n), label.ljust(MESSAGE_LENGTH), \
-                                                    PROGRESS_LENGTH, silent):
+        msg = message.ljust(MESSAGE_LENGTH)
+        for i in uiprogress.bar(range(n), msg, PROGRESS_LENGTH, silent):
             if i < position:
                 continue
             position = yield
         yield  # suppress the StopIteration exception
-    return progress_generator(label, n)
+    return progress_generator
+
+def md5(filename, size, progress = None):
+
+    BLOCKSIZE = 2^22  # 4MB
+
+    md5 = hashlib.md5()
+    with open(filename, "r") as src:
+        left = size
+        while left > 0:
+            length = min(left, BLOCKSIZE)
+            data = src.read(length)
+            md5.update(data)
+            left -= length
+
+    return md5.hexdigest()
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :