X-Git-Url: https://code.grnet.gr/git/snf-image-creator/blobdiff_plain/74149d07caa5363494b071483e55af861b7d953f..aa486e935452c7120582057caf9db0c77b88d68a:/image_creator/rsync.py

diff --git a/image_creator/rsync.py b/image_creator/rsync.py
index 22d3ad7..08bc27a 100644
--- a/image_creator/rsync.py
+++ b/image_creator/rsync.py
@@ -32,40 +32,110 @@
 # or implied, of GRNET S.A.
 
 import subprocess
+import time
+import signal
+
+from image_creator.util import FatalError
 
 
 class Rsync:
-    def __init__(self, src, dest, exclude=[]):
-        self.src = src
-        self.dest = dest
-        self.exclude = exclude
-        self.options = ['-v']
+    """Wrapper class for the rsync command"""
+
+    def __init__(self, output):
+        """Create an instance """
+        self._out = output
+        self._exclude = []
+        self._options = ['-v']
 
     def archive(self):
-        self.options.append('-a')
+        """Enable the archive option"""
+        self._options.append('-a')
         return self
 
-    def run(self, out, msg):
+    def xattrs(self):
+        """Preserve extended attributes"""
+        self._options.append('-X')
+        return self
+
+    def hard_links(self):
+        """Preserve hard links"""
+        self._options.append('-H')
+        return self
+
+    def acls(self):
+        """Preserve ACLs"""
+        self._options.append('-A')
+        return self
+
+    def sparse(self):
+        """Handle sparse files efficiently"""
+        self._options.append('-S')
+        return self
+
+    def exclude(self, pattern):
+        """Add an exclude pattern"""
+        self._exclude.append(pattern)
+        return self
+
+    def reset(self):
+        """Reset all rsync options"""
+        self._exclude = []
+        self._options = ['-v']
+
+    def run(self, src, dest, slabel='source', dlabel='destination'):
+        """Run the actual command"""
         cmd = []
         cmd.append('rsync')
-        cmd.extend(self.options)
-        for i in self.exclude:
+        cmd.extend(self._options)
+        for i in self._exclude:
             cmd.extend(['--exclude', i])
 
-        dry_run = subprocess.Popen(cmd + ['-n', self.src, self.dest],
-                                   shell=False, stdout=subprocess.PIPE,
-                                   bufsize=0)
-        print "%r" % (cmd + ['-n', self.src, self.dest])
-        total = 0
-        for line in iter(dry_run.stdout.readline, b''):
-            total += 1
+        self._out.output("Calculating total number of %s files ..." % slabel,
+                         False)
 
-        progress = out.Progress(total, msg)
-        run = subprocess.Popen(cmd + [self.src, self.dest], shell=False,
+        # If you don't specify a destination, rsync will list the source files.
+        dry_run = subprocess.Popen(cmd + [src], shell=False,
+                                   stdout=subprocess.PIPE, bufsize=0)
+        try:
+            total = 0
+            for line in iter(dry_run.stdout.readline, b''):
+                total += 1
+        finally:
+            dry_run.communicate()
+            if dry_run.returncode != 0:
+                raise FatalError("rsync failed")
+
+        self._out.success("%d" % total)
+
+        progress = self._out.Progress(total, "Copying files to %s" % dlabel)
+        run = subprocess.Popen(cmd + [src, dest], shell=False,
                                stdout=subprocess.PIPE, bufsize=0)
-        for line in iter(run.stdout.readline, b''):
-            progress.next()
+        try:
+            t = time.time()
+            i = 0
+            for line in iter(run.stdout.readline, b''):
+                i += 1
+                current = time.time()
+                if current - t > 0.1:
+                    t = current
+                    progress.goto(i)
+
+            progress.success('done')
+
+        finally:
+            def handler(signum, frame):
+                run.terminate()
+                time.sleep(1)
+                run.poll()
+                if run.returncode is None:
+                    run.kill()
+                run.wait()
 
-        progress.success('done')
+            signal.signal(signal.SIGALRM, handler)
+            signal.alarm(2)
+            run.communicate()
+            signal.alarm(0)
+            if run.returncode != 0:
+                raise FatalError("rsync failed")
 
 # vim: set sta sts=4 shiftwidth=4 sw=4 et ai :