Merge branch 'baserock/pedroalvarez/rawdisk-to-device8'

Reviewed-by:
- Richard Maw
- Sam Thursfield

3 files changed, 50 insertions, 49 deletions

diff --git a/openstack.write b/openstack.write
index 516fe367..b1941d3c 100755
--- a/openstack.write
+++ b/openstack.write
@@ -79,8 +79,7 @@ class OpenStackWriteExtension(morphlib.writeexts.WriteExtension):
     def set_extlinux_root_to_virtio(self, raw_disk):
         '''Re-configures extlinux to use virtio disks'''
         self.status(msg='Updating extlinux.conf')
-        mp = self.mount(raw_disk)
-        try:
+        with self.mount(raw_disk) as mp:
             path = os.path.join(mp, 'extlinux.conf')
 
             with open(path) as f:
@@ -91,9 +90,6 @@ class OpenStackWriteExtension(morphlib.writeexts.WriteExtension):
             with open(path, "w") as f:
                 f.write(extlinux_conf)
 
-        finally:
-            self.unmount(mp)
-
     def get_openstack_parameters(self):
         '''Get the environment variables needed.
 
diff --git a/rawdisk.check b/rawdisk.check
index acdc4de1..094adb72 100755
--- a/rawdisk.check
+++ b/rawdisk.check
@@ -33,10 +33,11 @@ class RawdiskCheckExtension(morphlib.writeexts.WriteExtension):
         location = args[0]
         upgrade = self.get_environment_boolean('UPGRADE')
         if upgrade:
-            if not os.path.isfile(location):
-                raise cliapp.AppException(
-                    'Cannot upgrade %s: it is not an existing disk image' %
-                    location)
+            if not self.is_device(location):
+                if not os.path.isfile(location):
+                    raise cliapp.AppException(
+                        'Cannot upgrade %s: it is not an existing disk image' %
+                        location)
 
             version_label = os.environ.get('VERSION_LABEL')
             if version_label is None:
@@ -44,9 +45,10 @@ class RawdiskCheckExtension(morphlib.writeexts.WriteExtension):
                     'VERSION_LABEL was not given. It is required when '
                     'upgrading an existing system.')
         else:
-            if os.path.exists(location):
-                raise cliapp.AppException(
-                    'Target %s already exists. Use `morph upgrade` if you '
-                    'want to update an existing image.' % location)
+            if not self.is_device(location):
+                if os.path.exists(location):
+                    raise cliapp.AppException(
+                        'Target %s already exists. Use `morph upgrade` if you '
+                        'want to update an existing image.' % location)
 
 RawdiskCheckExtension().run()
diff --git a/rawdisk.write b/rawdisk.write
index 12db4398..b17f8aa7 100755
--- a/rawdisk.write
+++ b/rawdisk.write
@@ -43,67 +43,70 @@ class RawDiskWriteExtension(morphlib.writeexts.WriteExtension):
             raise cliapp.AppException('Wrong number of command line args')
         
         temp_root, location = args
-        if os.path.isfile(location):
+        upgrade = self.get_environment_boolean('UPGRADE')
+
+        if upgrade:
             self.upgrade_local_system(location, temp_root)
         else:
             try:
-                self.create_local_system(temp_root, location)
-                self.status(msg='Disk image has been created at %s' % location)
+                if not self.is_device(location):
+                    with self.created_disk_image(location):
+                        self.format_btrfs(location)
+                        self.create_system(temp_root, location)
+                    self.status(msg='Disk image has been created at %s' %
+                                     location)
+                else:
+                    self.format_btrfs(location)
+                    self.create_system(temp_root, location)
+                    self.status(msg='System deployed to %s' % location)
             except Exception:
-                self.status(msg='Failure to create disk image at %s' %
+                self.status(msg='Failure to deploy system to %s' %
                                  location)
-                if os.path.exists(location):
-                    os.remove(location)
                 raise
 
     def upgrade_local_system(self, raw_disk, temp_root):
         self.complete_fstab_for_btrfs_layout(temp_root)
 
-        mp = self.mount(raw_disk)
+        with self.mount(raw_disk) as mp:
+            version_label = self.get_version_label(mp)
+            self.status(msg='Updating image to a new version with label %s' %
+                        version_label)
 
-        version_label = self.get_version_label(mp)
-        self.status(msg='Updating image to a new version with label %s' %
-                    version_label)
+            version_root = os.path.join(mp, 'systems', version_label)
+            os.mkdir(version_root)
 
-        version_root = os.path.join(mp, 'systems', version_label)
-        os.mkdir(version_root)
+            old_orig = os.path.join(mp, 'systems', 'default', 'orig')
+            new_orig = os.path.join(version_root, 'orig')
+            cliapp.runcmd(
+                ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])
 
-        old_orig = os.path.join(mp, 'systems', 'default', 'orig')
-        new_orig = os.path.join(version_root, 'orig')
-        cliapp.runcmd(
-            ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])
+            cliapp.runcmd(
+                ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
+                 temp_root + os.path.sep, new_orig])
 
-        cliapp.runcmd(
-            ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
-             temp_root + os.path.sep, new_orig])
+            self.create_run(version_root)
 
-        self.create_run(version_root)
+            default_path = os.path.join(mp, 'systems', 'default')
+            if os.path.exists(default_path):
+                os.remove(default_path)
+            else:
+                # we are upgrading and old system that does
+                # not have an updated extlinux config file
+                if self.bootloader_config_is_wanted():
+                    self.generate_bootloader_config(mp)
+                self.install_bootloader(mp)
+            os.symlink(version_label, default_path)
 
-        default_path = os.path.join(mp, 'systems', 'default')
-        if os.path.exists(default_path):
-            os.remove(default_path)
-        else:
-            # we are upgrading and old system that does
-            # not have an updated extlinux config file
             if self.bootloader_config_is_wanted():
-                self.generate_bootloader_config(mp)
-            self.install_bootloader(mp)
-        os.symlink(version_label, default_path)
-
-        if self.bootloader_config_is_wanted():
-            self.install_kernel(version_root, temp_root)
-
-        self.unmount(mp)
+                self.install_kernel(version_root, temp_root)
 
     def get_version_label(self, mp):
         version_label = os.environ.get('VERSION_LABEL')
 
         if version_label is None:
-            self.unmount(mp)
             raise cliapp.AppException('VERSION_LABEL was not given')
 
         if os.path.exists(os.path.join(mp, 'systems', version_label)):
-            self.unmount(mp)
             raise cliapp.AppException('VERSION_LABEL %s already exists'
                                       % version_label)