Refactor: Add WriteExtension.create_local_system method

This allows code sharing amongst all the places that create a
system in a raw disk image. This also adds the creation of
a factory-run subvolume, and fixes error messages for errors
that happen during a disk image creation.

Suggested-By: Richard Maw
Suggested-By: Sam Thursfield

4 files changed, 52 insertions, 69 deletions

diff --git a/morphlib/exts/kvm.write b/morphlib/exts/kvm.write
index e52e5553..09a7d224 100755
--- a/morphlib/exts/kvm.write
+++ b/morphlib/exts/kvm.write
@@ -18,8 +18,11 @@
 '''A Morph deployment write extension for deploying to KVM+libvirt.'''
 
 
+import cliapp
 import os
 import re
+import sys
+import tempfile
 import urlparse
 
 import morphlib.writeexts
@@ -54,35 +57,15 @@ class KvmPlusSshWriteExtension(morphlib.writeexts.WriteExtension):
         temp_root, location = args
         ssh_host, vm_name, vm_path = self.parse_location(location)
         
-        size = self.get_disk_size()
-        self.status(msg='Disk size is %(size)d bytes', size=size)
+        fd, raw_disk = tempfile.mkstemp()
+        os.close(fd)
+        self.create_local_system(temp_root, raw_disk)
         
-        raw_disk = tempfile.mkstemp()
-        self.create_raw_disk_image(raw_disk, size)
         try:
-            self.mkfs_btrfs(raw_disk)
-            mp = self.mount(raw_disk)
-        except BaseException:
-            self.status(msg='EEEK')
-            os.remove(raw_disk)
-            raise
-        try:
-            self.create_factory(mp, temp_root)
-            self.create_fstab(mp)
-            self.install_extlinux(mp)
-        except BaseException, e:
-            self.status(msg='EEK')
-            self.unmount(mp)
-            os.remove(raw_disk)
-            raise
-        else:
-            self.unmount(mp)
-
-        try:
-            self.transfer(raw_disk, size, ssh_host, vm_path)
+            self.transfer(raw_disk, ssh_host, vm_path)
             self.create_libvirt_guest(ssh_host, vm_name, vm_path)
         except BaseException:
-            self.status(msg='EEEK')
+            sys.stderr.write('Error deploying to libvirt')
             os.remove(raw_disk)
             raise
         else:
@@ -104,7 +87,7 @@ class KvmPlusSshWriteExtension(morphlib.writeexts.WriteExtension):
             raise cliapp.AppException('Cannot parse location %s' % location)
         return x.netloc, m.group('guest'), m.group('path')
 
-    def transfer(self, raw_disk, size, ssh_host, vm_path):
+    def transfer(self, raw_disk, ssh_host, vm_path):
         '''Transfer raw disk image to libvirt host.'''
 
         self.status(msg='Transfer disk image')
diff --git a/morphlib/exts/rawdisk.write b/morphlib/exts/rawdisk.write
index 89e9e82b..a55473f2 100755
--- a/morphlib/exts/rawdisk.write
+++ b/morphlib/exts/rawdisk.write
@@ -19,6 +19,7 @@
 
 
 import os
+import sys
 import time
 import tempfile
 
@@ -40,24 +41,8 @@ class RawDiskWriteExtension(morphlib.writeexts.WriteExtension):
         
         temp_root, location = args
         
-        size = self.get_disk_size()
-        self.status(msg='Disk size is %(size)d bytes', size=size)
-        
-        self.create_raw_disk_image(location, size)
-        self.mkfs_btrfs(location)
-        mp = self.mount(location)
-        try:
-            self.create_factory(mp, temp_root)
-            self.create_fstab(mp)
-            self.install_extlinux(mp)
-        except BaseException, e:
-            self.status(msg='EEK')
-            self.unmount(mp)
-            raise
-        else:
-            self.unmount(mp)
-            
-        self.status(msg='Disk image has been created')
+        self.create_local_system(temp_root, location)
+        self.status(msg='Disk image has been created at %s' % location)
 
 
 RawDiskWriteExtension().run()
diff --git a/morphlib/exts/virtualbox-ssh.write b/morphlib/exts/virtualbox-ssh.write
index dbfb90a7..5d118ec4 100755
--- a/morphlib/exts/virtualbox-ssh.write
+++ b/morphlib/exts/virtualbox-ssh.write
@@ -23,6 +23,7 @@ accessed over ssh. The machine gets created, but not started.
 '''
 
 
+import cliapp
 import os
 import re
 import time
@@ -61,36 +62,16 @@ class VirtualBoxPlusSshWriteExtension(morphlib.writeexts.WriteExtension):
         temp_root, location = args
         ssh_host, vm_name, vdi_path = self.parse_location(location)
         
-        size = self.get_disk_size()
-        self.status(msg='Disk size is %(size)d bytes', size=size)
-        
-        raw_disk = tempfile.mkstemp()
-        self.create_raw_disk_image(raw_disk, size)
-        try:
-            self.mkfs_btrfs(raw_disk)
-            mp = self.mount(raw_disk)
-        except BaseException:
-            self.status('EEEK')
-            os.remove(raw_disk)
-            raise
-        try:
-            self.create_factory(mp, temp_root)
-            self.create_fstab(mp)
-            self.install_extlinux(mp)
-        except BaseException, e:
-            self.status(msg='EEK')
-            self.unmount(mp)
-            os.remove(raw_disk)
-            raise
-        else:
-            self.unmount(mp)
+        fd, raw_disk = tempfile.mkstemp()
+        os.close(fd)
+        self.create_local_system(temp_root)
 
         try:
             self.transfer_and_convert_to_vdi(
                 raw_disk, size, ssh_host, vdi_path)
             self.create_virtualbox_guest(ssh_host, vm_name, vdi_path)
         except BaseException:
-            self.status('EEEK')
+            sys.stderr.write('Error deploying to VirtualBox')
             os.remove(raw_disk)
             raise
         else:
diff --git a/morphlib/writeexts.py b/morphlib/writeexts.py
index fae0d5d9..60848345 100755
--- a/morphlib/writeexts.py
+++ b/morphlib/writeexts.py
@@ -47,6 +47,31 @@ class WriteExtension(cliapp.Application):
         '''
         
         self.output.write('%s\n' % (kwargs['msg'] % kwargs))
+    
+    def create_local_system(self, temp_root, raw_disk):
+        '''Create a raw system image locally.'''
+        
+        size = self.get_disk_size()
+        self.create_raw_disk_image(raw_disk, size)
+        try:
+            self.mkfs_btrfs(raw_disk)
+            mp = self.mount(raw_disk)
+        except BaseException:
+            sys.stderr.write('Error creating disk image')
+            os.remove(raw_disk)
+            raise
+        try:
+            self.create_factory(mp, temp_root)
+            self.create_fstab(mp)
+            self.create_factory_run(mp)
+            self.install_extlinux(mp)
+        except BaseException, e:
+            sys.stderr.write('Error creating disk image')
+            self.unmount(mp)
+            os.remove(raw_disk)
+            raise
+        else:
+            self.unmount(mp)
 
     def get_disk_size(self):
         '''Parse disk size from environment.'''
@@ -120,6 +145,15 @@ class WriteExtension(cliapp.Application):
         factory_boot = os.path.join(factory, 'boot')
         root_boot = os.path.join(real_root, 'boot')
         cliapp.runcmd(['cp', '-a', factory_boot, root_boot])
+        
+    def create_factory_run(self, real_root):
+        '''Create the 'factory-run' snapshot.'''
+
+        self.status(msg='Creating factory-run subvolume')
+        factory = os.path.join(real_root, 'factory')
+        factory_run = factory + '-run'
+        cliapp.runcmd(
+            ['btrfs', 'subvolume', 'snapshot', factory, factory_run])
 
     def create_fstab(self, real_root):
         '''Create an fstab.'''
@@ -139,7 +173,7 @@ class WriteExtension(cliapp.Application):
             f.write('timeout 1\n')
             f.write('label linux\n')
             f.write('kernel /boot/vmlinuz\n')
-            f.write('append root=/dev/sda rootflags=subvol=factory '
+            f.write('append root=/dev/sda rootflags=subvol=factory-run '
                     'init=/sbin/init rw\n')
 
         self.status(msg='Installing extlinux')