Merge branch 'baserock/richardmaw/S8847/branch-from-image-v3'

Reviewed-by: Lars Wirzenius <lars.wirzenius@codethink.co.uk>
Reviewed-by: Jonathan Maw <jonathan.maw@codethink.co.uk>

6 files changed, 313 insertions, 47 deletions

diff --git a/morphlib/__init__.py b/morphlib/__init__.py
index b1e3c7c3..7eb3f975 100644
--- a/morphlib/__init__.py
+++ b/morphlib/__init__.py
@@ -78,6 +78,7 @@ import sourcepool
 import stagingarea
 import stopwatch
 import sysbranchdir
+import systemmetadatadir
 import tempdir
 import util
 import workspace
diff --git a/morphlib/plugins/branch_and_merge_new_plugin.py b/morphlib/plugins/branch_and_merge_new_plugin.py
index edda7d9c..3d0a71a5 100644
--- a/morphlib/plugins/branch_and_merge_new_plugin.py
+++ b/morphlib/plugins/branch_and_merge_new_plugin.py
@@ -15,6 +15,7 @@
 
 
 import cliapp
+import contextlib
 import glob
 import logging
 import os
@@ -55,6 +56,15 @@ class SimpleBranchAndMergePlugin(cliapp.Plugin):
                                 arg_synopsis='-- COMMAND [ARGS...]')
         self.app.add_subcommand('status', self.status,
                                 arg_synopsis='')
+        self.app.add_subcommand('branch-from-image', self.branch_from_image,
+                                arg_synopsis='BRANCH')
+        group_branch = 'Branching Options'
+        self.app.settings.string(['metadata-dir'],
+                                  'Set metadata location for branch-from-image'
+                                  ' (default: /baserock)',
+                                  metavar='DIR',
+                                  default='/baserock',
+                                  group=group_branch)
 
     def disable(self):
         pass
@@ -97,6 +107,40 @@ class SimpleBranchAndMergePlugin(cliapp.Plugin):
         ws = morphlib.workspace.open('.')
         self.app.output.write('%s\n' % ws.root)
 
+    # TODO: Move this somewhere nicer
+    @contextlib.contextmanager
+    def _initializing_system_branch(self, ws, root_url, system_branch,
+                                    cached_repo, base_ref):
+        '''A context manager for system branches under construction.
+
+        The purpose of this context manager is to factor out the branch
+        cleanup code for if an exception occurs while a branch is being
+        constructed.
+
+        This could be handled by a higher order function which takes
+        a function to initialize the branch as a parameter, but with
+        statements look nicer and are more obviously about resource
+        cleanup.
+
+        '''
+        root_dir = ws.get_default_system_branch_directory_name(system_branch)
+        try:
+            sb = morphlib.sysbranchdir.create(
+                root_dir, root_url, system_branch)
+            gd = sb.clone_cached_repo(cached_repo, base_ref)
+
+            yield (sb, gd)
+
+            gd.update_submodules(self.app)
+            gd.update_remotes()
+
+        except BaseException as e:
+            # Oops. Clean up.
+            logging.error('Caught exception: %s' % str(e))
+            logging.info('Removing half-finished branch %s' % system_branch)
+            self._remove_branch_dir_safe(ws.root, root_dir)
+            raise
+
     def checkout(self, args):
         '''Check out an existing system branch.
 
@@ -124,6 +168,7 @@ class SimpleBranchAndMergePlugin(cliapp.Plugin):
 
         root_url = args[0]
         system_branch = args[1]
+        base_ref = system_branch
 
         self._require_git_user_config()
 
@@ -139,27 +184,12 @@ class SimpleBranchAndMergePlugin(cliapp.Plugin):
         # Check the git branch exists.
         cached_repo.resolve_ref(system_branch)
 
-        root_dir = ws.get_default_system_branch_directory_name(system_branch)
-
-        try:
-            # Create the system branch directory. This doesn't yet clone
-            # the root repository there.
-            sb = morphlib.sysbranchdir.create(
-                root_dir, root_url, system_branch)
-
-            gd = sb.clone_cached_repo(cached_repo, system_branch)
+        with self._initializing_system_branch(
+            ws, root_url, system_branch, cached_repo, base_ref) as (sb, gd):
 
             if not self._checkout_has_systems(gd):
-                raise BranchRootHasNoSystemsError(root_url, system_branch)
+                raise BranchRootHasNoSystemsError(base_ref)
 
-            gd.update_submodules(self.app)
-            gd.update_remotes()
-        except BaseException as e:
-            # Oops. Clean up.
-            logging.error('Caught exception: %s' % str(e))
-            logging.info('Removing half-finished branch %s' % system_branch)
-            self._remove_branch_dir_safe(ws.root, root_dir)
-            raise
 
     def branch(self, args):
         '''Create a new system branch.
@@ -212,29 +242,14 @@ class SimpleBranchAndMergePlugin(cliapp.Plugin):
         # Make sure the base_ref exists.
         cached_repo.resolve_ref(base_ref)
 
-        root_dir = ws.get_default_system_branch_directory_name(system_branch)
-
-        try:
-            # Create the system branch directory. This doesn't yet clone
-            # the root repository there.
-            sb = morphlib.sysbranchdir.create(
-                root_dir, root_url, system_branch)
+        with self._initializing_system_branch(
+            ws, root_url, system_branch, cached_repo, base_ref) as (sb, gd):
 
-            gd = sb.clone_cached_repo(cached_repo, base_ref)
             gd.branch(system_branch, base_ref)
             gd.checkout(system_branch)
 
             if not self._checkout_has_systems(gd):
-                raise BranchRootHasNoSystemsError(root_url, base_ref)
-
-            gd.update_submodules(self.app)
-            gd.update_remotes()
-        except BaseException as e:
-            # Oops. Clean up.
-            logging.error('Caught exception: %s' % str(e))
-            logging.info('Removing half-finished branch %s' % system_branch)
-            self._remove_branch_dir_safe(ws.root, root_dir)
-            raise
+                raise BranchRootHasNoSystemsError(base_ref)
 
     def _save_dirty_morphologies(self, loader, sb, morphs):
         logging.debug('Saving dirty morphologies: start')
@@ -789,3 +804,98 @@ class SimpleBranchAndMergePlugin(cliapp.Plugin):
 
         if not has_uncommitted_changes:
             self.app.output.write("\nNo repos have outstanding changes.\n")
+
+    def branch_from_image(self, args):
+        '''Produce a branch of an existing system image.
+
+        Given the metadata specified by --metadata-dir, create a new
+        branch then petrify it to the state of the commits at the time
+        the system was built.
+
+        If --metadata-dir is not specified, it defaults to your currently
+        running system.
+
+        '''
+        if len(args) != 1:
+            raise cliapp.AppException(
+                "branch-from-image needs exactly 1 argument "
+                "of the new system branch's name")
+        system_branch = args[0]
+        metadata_path = self.app.settings['metadata-dir']
+        alias_resolver = morphlib.repoaliasresolver.RepoAliasResolver(
+            self.app.settings['repo-alias'])
+
+        self._require_git_user_config()
+
+        ws = morphlib.workspace.open('.')
+
+        system, metadata = self._load_system_metadata(metadata_path)
+        resolved_refs = dict(self._resolve_refs_from_metadata(alias_resolver,
+                                                              metadata))
+        logging.debug('Resolved refs: %r' % resolved_refs)
+        base_ref = system['sha1']
+        # The previous version would fall back to deducing this from the repo
+        # url and the repo alias resolver, but this does not always work, and
+        # new systems always have repo-alias in the metadata
+        root_url = system['repo-alias']
+
+        lrc, rrc = morphlib.util.new_repo_caches(self.app)
+        cached_repo = lrc.get_updated_repo(root_url)
+
+
+        with self._initializing_system_branch(
+            ws, root_url, system_branch, cached_repo, base_ref) as (sb, gd):
+
+            # TODO: It's nasty to clone to a sha1 then create a branch
+            #       of that sha1 then check it out, a nicer API may be the
+            #       initial clone not checking out a branch at all, then
+            #       the user creates and checks out their own branches
+            gd.branch(system_branch, base_ref)
+            gd.checkout(system_branch)
+
+            loader = morphlib.morphloader.MorphologyLoader()
+            morphs = self._load_all_sysbranch_morphologies(sb, loader)
+
+            morphs.repoint_refs(sb.root_repository_url,
+                                sb.system_branch_name)
+
+            morphs.petrify_chunks(resolved_refs)
+
+            self._save_dirty_morphologies(loader, sb, morphs.morphologies)
+
+    @staticmethod
+    def _load_system_metadata(path):
+        '''Load all metadata in `path` corresponding to a single System.
+        '''
+
+        smd = morphlib.systemmetadatadir.SystemMetadataDir(path)
+        metadata = smd.values()
+        systems = [md for md in metadata
+                   if 'kind' in md and md['kind'] == 'system']
+
+        if not systems:
+            raise cliapp.AppException(
+                'Metadata directory does not contain any systems.')
+        if len(systems) > 1:
+            raise cliapp.AppException(
+                'Metadata directory contains multiple systems.')
+        system_metadatum = systems[0]
+    
+        metadata_cache_id_lookup = dict((md['cache-key'], md)
+                                        for md in metadata)
+
+        return system_metadatum, metadata_cache_id_lookup
+
+    @staticmethod
+    def _resolve_refs_from_metadata(alias_resolver, metadata):
+        '''Pre-resolve a set of refs from existing metadata.
+
+        Given the metadata, generate a mapping of all the (repo, ref)
+        pairs defined in the metadata and the commit id they resolved to.
+
+        '''
+        for md in metadata.itervalues():
+            repourls = set((md['repo-alias'], md['repo']))
+            repourls.update(alias_resolver.aliases_from_url(md['repo']))
+            for repourl in repourls:
+                yield ((repourl, md['original_ref']), md['sha1'])
diff --git a/morphlib/plugins/branch_and_merge_plugin.py b/morphlib/plugins/branch_and_merge_plugin.py
index 37d5e40c..2f8560d0 100644
--- a/morphlib/plugins/branch_and_merge_plugin.py
+++ b/morphlib/plugins/branch_and_merge_plugin.py
@@ -68,15 +68,9 @@ class BranchAndMergePlugin(cliapp.Plugin):
         self.app.add_subcommand('build', self.build,
                                 arg_synopsis='SYSTEM')
         self.app.add_subcommand('old-status', self.status)
-        self.app.add_subcommand('branch-from-image', self.branch_from_image,
-                                 arg_synopsis='REPO BRANCH')
-        group_branch = 'Branching Options'
-        self.app.settings.string(['metadata-dir'],
-                                  'Set metadata location for branch-from-image'
-                                  ' (default: /baserock)',
-                                  metavar='DIR',
-                                  default='/baserock',
-                                  group=group_branch)
+        self.app.add_subcommand('old-branch-from-image',
+                                self.branch_from_image,
+                                arg_synopsis='REPO BRANCH')
 
         # Advanced commands
         self.app.add_subcommand('old-foreach', self.foreach,
diff --git a/morphlib/systemmetadatadir.py b/morphlib/systemmetadatadir.py
new file mode 100644
index 00000000..eac5b446
--- /dev/null
+++ b/morphlib/systemmetadatadir.py
@@ -0,0 +1,87 @@
+# Copyright (C) 2013  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+#
+# =*= License: GPL-2 =*=
+
+
+import collections
+import glob
+import json
+import os
+
+
+class SystemMetadataDir(collections.MutableMapping):
+
+    '''An abstraction over the /baserock metadata directory.
+
+    This allows methods of iterating over it, and accessing it like
+    a dict.
+
+    The /baserock metadata directory contains information about all of
+    the chunks in a built system. It exists to provide traceability from
+    the input sources to the output.
+
+    If you create the object with smd = SystemMetadataDir('/baserock')
+    data = smd['key'] will read /baserock/key.meta and return its JSON
+    encoded contents as native python objects.
+
+    smd['key'] = data will write data to /baserock/key.meta as JSON
+
+    The key may not have '\0' characters in it since the underlying
+    system calls don't support embedded NUL bytes.
+
+    The key may not have '/' characters in it since we do not support
+    morphologies with slashes in their names.
+
+    '''
+
+    def __init__(self, metadata_path):
+        collections.MutableMapping.__init__(self)
+        self._metadata_path = metadata_path
+
+    def _join_path(self, *args):
+        return os.path.join(self._metadata_path, *args)
+
+    def _raw_path_iter(self):
+        return glob.iglob(self._join_path('*.meta'))
+
+    @staticmethod
+    def _check_key(key):
+        if any(c in key for c in "\0/"):
+            raise KeyError(key)
+
+    def __getitem__(self, key):
+        self._check_key(key)
+        try:
+            with open(self._join_path('%s.meta' % key), 'r') as f:
+                return json.load(f)
+        except IOError:
+            raise KeyError(key)
+
+    def __setitem__(self, key, value):
+        self._check_key(key)
+        with open(self._join_path('%s.meta' % key), 'w') as f:
+            json.dump(value, f, indent=4, sort_keys=True)
+
+    def __delitem__(self, key):
+        self._check_key(key)
+        os.unlink(self._join_path('%s.meta' % key))
+
+    def __iter__(self):
+        return (os.path.basename(fn)[:-len('.meta')]
+                for fn in self._raw_path_iter())
+
+    def __len__(self):
+        return len(list(self._raw_path_iter()))
diff --git a/morphlib/systemmetadatadir_tests.py b/morphlib/systemmetadatadir_tests.py
new file mode 100644
index 00000000..0126f862
--- /dev/null
+++ b/morphlib/systemmetadatadir_tests.py
@@ -0,0 +1,75 @@
+# Copyright (C) 2013  Codethink Limited
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+#
+# =*= License: GPL-2 =*=
+
+
+import operator
+import os
+import shutil
+import tempfile
+import unittest
+
+import morphlib
+
+
+class SystemMetadataDirTests(unittest.TestCase):
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.metadatadir = os.path.join(self.tempdir, 'baserock')
+        os.mkdir(self.metadatadir)
+        self.smd = morphlib.systemmetadatadir.SystemMetadataDir(
+            self.metadatadir)
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+    def test_add_new(self):
+        self.smd['key'] = {'foo': 'bar'}
+        self.assertEqual(self.smd['key']['foo'], 'bar')
+
+    def test_replace(self):
+        self.smd['key'] = {'foo': 'bar'}
+        self.smd['key'] = {'foo': 'baz'}
+        self.assertEqual(self.smd['key']['foo'], 'baz')
+
+    def test_remove(self):
+        self.smd['key'] = {'foo': 'bar'}
+        del self.smd['key']
+        self.assertTrue('key' not in self.smd)
+
+    def test_iterate(self):
+        self.smd['build-essential'] = "Some data"
+        self.smd['core'] = "More data"
+        self.smd['foundation'] = "Yet more data"
+        self.assertEqual(sorted(self.smd.keys()),
+                         ['build-essential', 'core', 'foundation'])
+        self.assertEqual(dict(self.smd.iteritems()),
+                         {
+                             'build-essential': "Some data",
+                             'core': "More data",
+                             'foundation': "Yet more data",
+                         })
+
+    def test_raises_KeyError(self):
+        self.assertRaises(KeyError, operator.getitem, self.smd, 'key')
+
+    def test_validates_keys(self):
+        for key in ('foo/bar', 'baz\0quux'):
+            self.assertRaises(KeyError, operator.getitem, self.smd, key)
+            self.assertRaises(KeyError, operator.setitem,
+                              self.smd, key, 'value')
+            self.assertRaises(KeyError, operator.delitem, self.smd, key)
diff --git a/tests.as-root/branch-from-image-works.script b/tests.as-root/branch-from-image-works.script
index 9f82f629..942301e8 100755
--- a/tests.as-root/branch-from-image-works.script
+++ b/tests.as-root/branch-from-image-works.script
@@ -48,8 +48,7 @@ git commit --quiet -m 'Make hello say goodbye'
 workspace="$DATADIR/workspace"
 "$SRCDIR/scripts/test-morph" init "$workspace"
 cd "$workspace"
-"$SRCDIR/scripts/test-morph" branch-from-image \
-                             test:morphs mybranch \
+"$SRCDIR/scripts/test-morph" branch-from-image  mybranch \
                              --metadata-dir="$extracted/baserock"
 cd mybranch/test:morphs
 grep -qFe "$hello_chunk_commit" hello-stratum.morph