Add support for the export/import of Node extras (#2416)

The extras of `Nodes` can now be exported together with other data.
When importing an archive that contains extras for nodes, the behavior
depends on whether that node was already present in the database into
which the data is imported, namely:

 * The `Node` did not yet exist: In this case, the extras will be imported
   wholesale or not at all. The behavior can be toggled from `verdi import`
   with the flag `--extras-mode-new`.

 * The `Node` already exists: In this case, the various possibilities are a
   lot more complex, because the existing and archived node can have
   overlapping and non-overlapping extras. What to do in case of conflicts
   and in general can be controlled by the flag `--extras-mode-existing`.
   The currently implemented options are as follows:

    - ask: import all extras and prompt what to do for existing extras.
    - keep_existing: import all extras and keep original value of existing.
    - update_existing: import all extras and overwrite value of existing.
    - mirror: remove any pre-existing extras and import all archive extras.
    - none: do not import any extras.

aiida/backends/tests/export_and_import.py

@@ -2458,3 +2458,154 @@ def test_multiple_user_comments_for_single_node(self):
 
         finally:
             shutil.rmtree(tmp_folder, ignore_errors=True)
+
+class TestExtras(AiidaTestCase):
+    """Test extras import"""
+
+    @classmethod
+    def setUpClass(cls, *args, **kwargs):
+        """Only run to prepare an export file"""
+        import os 
+        import tempfile
+        from aiida.orm import Data
+        from aiida.orm.importexport import export
+        super(TestExtras, cls).setUpClass()
+        d = Data()
+        d.label = 'my_test_data_node'
+        d.store()
+        d.set_extras({'b': 2, 'c': 3})
+        tmp_folder = tempfile.mkdtemp()
+        cls.export_file = os.path.join(tmp_folder, 'export.aiida')
+        export([d], outfile=cls.export_file, silent=True)
+
+    def setUp(self):
+        """This function runs before every test execution"""
+        from aiida.orm.importexport import import_data
+        self.clean_db()
+        self.insert_data()
+
+    def import_extras(self, mode_new='import'):
+        """Import an aiida database"""
+        from aiida.orm import Data
+        from aiida.orm.querybuilder import QueryBuilder
+        import_data(self.export_file, silent=True, extras_mode_new=mode_new)
+        q = QueryBuilder().append(Data, filters={'label': 'my_test_data_node'})
+        self.assertEquals(q.count(), 1)
+        self.imported_node = q.all()[0][0]
+
+    def modify_extras(self, mode_existing):
+        """Import the same aiida database again"""
+        from aiida.orm import Data
+        from aiida.orm.querybuilder import QueryBuilder
+        self.imported_node.set_extra('a', 1)
+        self.imported_node.set_extra('b', 1000)
+        self.imported_node.del_extra('c')
+
+        import_data(self.export_file, silent=True, extras_mode_existing=mode_existing)
+
+        # Query again the database
+        q = QueryBuilder().append(Data, filters={'label': 'my_test_data_node'})
+        self.assertEquals(q.count(), 1)
+        return q.all()[0][0]
+
+    def tearDown(self):
+        pass
+
+
+    def test_import_of_extras(self):
+        """Check if extras are properly imported"""
+        self.import_extras()
+        self.assertEquals(self.imported_node.get_extra('b'), 2)
+        self.assertEquals(self.imported_node.get_extra('c'), 3)
+
+    def test_absence_of_extras(self):
+        """Check whether extras are not imported if the mode is set to 'none'"""
+        self.import_extras(mode_new='none')
+        with self.assertRaises(AttributeError):
+            # the extra 'b' should not exist
+            self.imported_node.get_extra('b')
+        with self.assertRaises(AttributeError):
+            # the extra 'c' should not exist
+            self.imported_node.get_extra('c')
+
+    def test_extras_import_mode_keep_existing(self):
+        """Check if old extras are not modified in case of name collision"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='kcl')
+
+        # Check that extras are imported correctly
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        self.assertEquals(imported_node.get_extra('b'), 1000)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+
+    def test_extras_import_mode_update_existing(self):
+        """Check if old extras are modified in case of name collision"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='kcu')
+
+        # Check that extras are imported correctly
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        self.assertEquals(imported_node.get_extra('b'), 2)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+
+    def test_extras_import_mode_mirror(self):
+        """Check if old extras are fully overwritten by the imported ones"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='ncu')
+
+        # Check that extras are imported correctly
+        with self.assertRaises(AttributeError): # the extra 
+            # 'a' should not exist, as the extras were fully mirrored with respect to
+            # the imported node
+            imported_node.get_extra('a')
+        self.assertEquals(imported_node.get_extra('b'), 2)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+
+    def test_extras_import_mode_none(self):
+        """Check if old extras are fully overwritten by the imported ones"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='knl')
+
+        # Check if extras are imported correctly
+        self.assertEquals(imported_node.get_extra('b'), 1000)
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        with self.assertRaises(AttributeError): # the extra
+            # 'c' should not exist, as the extras were keept untached
+            imported_node.get_extra('c')
+
+    def test_extras_import_mode_strange(self):
+        """Check a mode that is probably does not make much sense but is still available"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='kcd')
+
+        # Check if extras are imported correctly
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+        with self.assertRaises(AttributeError): # the extra
+            # 'b' should not exist, as the collided extras are deleted
+            imported_node.get_extra('b')
+
+    def test_extras_import_mode_correct(self):
+        """Test all possible import modes except 'ask' """
+        self.import_extras()
+        for l1 in ['k', 'n']: # keep or not keep old extras
+            for l2 in ['n', 'c']: # create or not create new extras
+                for l3 in ['l', 'u', 'd']: # leave old, update or delete collided extras
+                    mode = l1 + l2 + l3
+                    import_data(self.export_file, silent=True, extras_mode_existing=mode)
+
+    def test_extras_import_mode_wrong(self):
+        """Check a mode that is wrong"""
+        self.import_extras()
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='xnd') # first letter is wrong
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='nxd') # second letter is wrong
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='nnx') # third letter is wrong
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='n') # too short
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='nndnn') # too long
+        with self.assertRaises(TypeError):
+            import_data(self.export_file, silent=True, extras_mode_existing=5) # wrong type

aiida/cmdline/commands/cmd_import.py

@@ -11,6 +11,7 @@
 from __future__ import division
 from __future__ import print_function
 from __future__ import absolute_import
+from enum import Enum
 import click
 
 from aiida.cmdline.commands.cmd_verdi import verdi
@@ -19,6 +20,19 @@
 from aiida.cmdline.utils import decorators, echo
 from aiida.common import exceptions
 
+EXTRAS_MODE_EXISTING = ['keep_existing', 'update_existing', 'mirror', 'none', 'ask']
+EXTRAS_MODE_NEW = ['import', 'none']
+
+
+# pylint: disable=too-few-public-methods
+class ExtrasImportCode(Enum):
+    """Exit codes for the verdi command line."""
+    keep_existing = 'kcl'
+    update_existing = 'kcu'
+    mirror = 'ncu'
+    none = 'knl'
+    ask = 'kca'
+
 
 @verdi.command('import')
 @click.argument('archives', nargs=-1, type=click.Path(exists=True, readable=True))
@@ -35,8 +49,28 @@
     type=GroupParamType(create_if_not_exist=True),
     help='Specify group to which all the import nodes will be added. If such a group does not exist, it will be'
     ' created automatically.')
+@click.option(
+    '-e',
+    '--extras-mode-existing',
+    type=click.Choice(EXTRAS_MODE_EXISTING),
+    default='keep_existing',
+    help="Specify which extras from the export archive should be imported for nodes that are already contained in the "
+    "database: "
+    "ask: import all extras and prompt what to do for existing extras. "
+    "keep_existing: import all extras and keep original value of existing extras. "
+    "update_existing: import all extras and overwrite value of existing extras. "
+    "mirror: import all extras and remove any existing extras that are not present in the archive. "
+    "none: do not import any extras.")
+@click.option(
+    '-n',
+    '--extras-mode-new',
+    type=click.Choice(EXTRAS_MODE_NEW),
+    default='import',
+    help="Specify whether to import extras of new nodes: "
+    "import: import extras. "
+    "none: do not import extras.")
 @decorators.with_dbenv()
-def cmd_import(archives, webpages, group):
+def cmd_import(archives, webpages, group, extras_mode_existing, extras_mode_new):
     """Import one or multiple exported AiiDA archives
 
     The ARCHIVES can be specified by their relative or absolute file path, or their HTTP URL.
@@ -78,7 +112,11 @@ def cmd_import(archives, webpages, group):
         echo.echo_info('importing archive {}'.format(archive))
 
         try:
-            import_data(archive, group)
+            import_data(
+                archive,
+                group,
+                extras_mode_existing=ExtrasImportCode[extras_mode_existing].value,
+                extras_mode_new=extras_mode_new)
         except exceptions.IncompatibleArchiveVersionError as exception:
             echo.echo_warning('{} cannot be imported: {}'.format(archive, exception))
             echo.echo_warning('run `verdi export migrate {}` to update it'.format(archive))
@@ -105,7 +143,11 @@ def cmd_import(archives, webpages, group):
             echo.echo_success('archive downloaded, proceeding with import')
 
             try:
-                import_data(temp_folder.get_abs_path(temp_file), group)
+                import_data(
+                    temp_folder.get_abs_path(temp_file),
+                    group,
+                    extras_mode_existing=ExtrasImportCode[extras_mode_existing].value,
+                    extras_mode_new=extras_mode_new)
             except exceptions.IncompatibleArchiveVersionError as exception:
                 echo.echo_warning('{} cannot be imported: {}'.format(archive, exception))
                 echo.echo_warning('download the archive file and run `verdi export migrate` to update it')