verdi export migrate: add --in-place flag to migrate archive in p…

…lace (#4220)

When an export archive needs to be migrated, one often does not care
about the original archive and simply wants to overwrite the existing
file with the migrated archive. The new flag `--in-place` saves users
from having to specify a temporary filename and copying it over the
original file after it has been migrated.

This commit also changes the exit code for migrating an export file that
already is up to date from an error to success (0).

aiida/cmdline/commands/cmd_export.py

@@ -11,6 +11,7 @@
 """`verdi export` command."""
 
 import os
+import tempfile
 
 import click
 import tabulate
@@ -141,31 +142,45 @@ def create(
 
 @verdi_export.command('migrate')
 @arguments.INPUT_FILE()
-@arguments.OUTPUT_FILE()
+@arguments.OUTPUT_FILE(required=False)
 @options.ARCHIVE_FORMAT()
 @options.FORCE(help='overwrite output file if it already exists')
+@click.option('-i', '--in-place', is_flag=True, help='Migrate the archive in place, overwriting the original file.')
 @options.SILENT()
 @click.option(
     '-v',
     '--version',
     type=click.STRING,
     required=False,
     metavar='VERSION',
-    help='Specify an exact archive version to migrate to. By default the most recent version is taken.'
+    # Note: Adding aiida.tools.EXPORT_VERSION as a default value explicitly would result in a slow import of
+    # aiida.tools and, as a consequence, aiida.orm. As long as this is the case, better determine the latest export
+    # version inside the function when needed.
+    help='Archive format version to migrate to (defaults to latest version).',
 )
-def migrate(input_file, output_file, force, silent, archive_format, version):
+def migrate(input_file, output_file, force, silent, in_place, archive_format, version):
     # pylint: disable=too-many-locals,too-many-statements,too-many-branches
     """Migrate an export archive to a more recent format version."""
     import tarfile
     import zipfile
 
     from aiida.common import json
     from aiida.common.folders import SandboxFolder
-    from aiida.tools.importexport import EXPORT_VERSION, migration, extract_zip, extract_tar, ArchiveMigrationError
+    from aiida.tools.importexport import migration, extract_zip, extract_tar, ArchiveMigrationError, EXPORT_VERSION
 
     if version is None:
         version = EXPORT_VERSION
 
+    if in_place:
+        if output_file:
+            echo.echo_critical('output file specified together with --in-place flag')
+        tempdir = tempfile.TemporaryDirectory()
+        output_file = os.path.join(tempdir.name, 'archive.aiida')
+    elif not output_file:
+        echo.echo_critical(
+            'no output file specified. Please add --in-place flag if you would like to migrate in place.'
+        )
+
     if os.path.exists(output_file) and not force:
         echo.echo_critical('the output file already exists')
 
@@ -187,6 +202,10 @@ def migrate(input_file, output_file, force, silent, archive_format, version):
             echo.echo_critical('export archive does not contain the required file {}'.format(fhandle.filename))
 
         old_version = migration.verify_metadata_version(metadata)
+        if version <= old_version:
+            echo.echo_success('nothing to be done - archive already at version {} >= {}'.format(old_version, version))
+            return
+
         try:
             new_version = migration.migrate_recursively(metadata, data, folder, version)
         except ArchiveMigrationError as exception:
@@ -212,5 +231,9 @@ def migrate(input_file, output_file, force, silent, archive_format, version):
             with tarfile.open(output_file, 'w:gz', format=tarfile.PAX_FORMAT, dereference=True) as archive:
                 archive.add(folder.abspath, arcname='')
 
+        if in_place:
+            os.rename(output_file, input_file)
+            tempdir.cleanup()
+
         if not silent:
             echo.echo_success('migrated the archive from version {} to {}'.format(old_version, new_version))

aiida/tools/importexport/migration/__init__.py

@@ -52,8 +52,8 @@ def migrate_recursively(metadata, data, folder, version=EXPORT_VERSION):
 
     try:
         if old_version == version:
-            raise ArchiveMigrationError('Your export file is already at the version {}'.format(version))
-        elif old_version > version:
+            return old_version
+        if old_version > version:
             raise ArchiveMigrationError('Backward migrations are not supported')
         elif old_version in MIGRATE_FUNCTIONS:
             MIGRATE_FUNCTIONS[old_version](metadata, data, folder)

tests/cmdline/commands/test_export.py

@@ -181,18 +181,6 @@ def test_migrate_version_specific(self):
         finally:
             delete_temporary_file(filename_output)
 
-    def test_migrate_versions_recent(self):
-        """Migrating an archive with the current version should exit with non-zero status."""
-        filename_input = get_archive_file(self.newest_archive, filepath=self.fixture_archive)
-        filename_output = next(tempfile._get_candidate_names())  # pylint: disable=protected-access
-
-        try:
-            options = [filename_input, filename_output]
-            result = self.cli_runner.invoke(cmd_export.migrate, options)
-            self.assertIsNotNone(result.exception)
-        finally:
-            delete_temporary_file(filename_output)
-
     def test_migrate_force(self):
         """Test that passing the -f/--force option will overwrite the output file even if it exists."""
         filename_input = get_archive_file(self.penultimate_archive, filepath=self.fixture_archive)
@@ -213,6 +201,39 @@ def test_migrate_force(self):
                 self.assertTrue(os.path.isfile(filename_output))
                 self.assertEqual(zipfile.ZipFile(filename_output).testzip(), None)
 
+    def test_migrate_in_place(self):
+        """Test that passing the -i/--in-place option will overwrite the passed file."""
+        archive = 'export_v0.1_simple.aiida'
+        target_version = '0.2'
+        filename_input = get_archive_file(archive, filepath=self.fixture_archive)
+        filename_tmp = next(tempfile._get_candidate_names())  # pylint: disable=protected-access
+
+        try:
+            # copy file (don't want to overwrite test data)
+            shutil.copy(filename_input, filename_tmp)
+
+            # specifying both output and in-place should except
+            options = [filename_tmp, '--in-place', '--output-file', 'test.aiida']
+            result = self.cli_runner.invoke(cmd_export.migrate, options)
+            self.assertIsNotNone(result.exception, result.output)
+
+            # specifying neither output nor in-place should except
+            options = [filename_tmp]
+            result = self.cli_runner.invoke(cmd_export.migrate, options)
+            self.assertIsNotNone(result.exception, result.output)
+
+            # check that in-place migration produces a valid archive in place of the old file
+            options = [filename_tmp, '--in-place', '--version', target_version]
+            result = self.cli_runner.invoke(cmd_export.migrate, options)
+            self.assertIsNone(result.exception, result.output)
+            self.assertTrue(os.path.isfile(filename_tmp))
+            # check that files in zip file are ok
+            self.assertEqual(zipfile.ZipFile(filename_tmp).testzip(), None)
+            with Archive(filename_tmp) as archive_object:
+                self.assertEqual(archive_object.version_format, target_version)
+        finally:
+            os.remove(filename_tmp)
+
     def test_migrate_silent(self):
         """Test that the captured output is an empty string when the -s/--silent option is passed."""
         filename_input = get_archive_file(self.penultimate_archive, filepath=self.fixture_archive)

tests/tools/importexport/migration/test_migration.py

@@ -114,9 +114,7 @@ def test_migrate_recursively_specific_version(self):
             with self.assertRaises(ArchiveMigrationError):
                 migrate_recursively(archive.meta_data, archive.data, None, version='0.2')
 
-            # Same version will also raise
-            with self.assertRaises(ArchiveMigrationError):
-                migrate_recursively(archive.meta_data, archive.data, None, version='0.3')
+            migrate_recursively(archive.meta_data, archive.data, None, version='0.3')
 
             migrated_version = '0.5'
             version = migrate_recursively(archive.meta_data, archive.data, None, version=migrated_version)
@@ -186,17 +184,11 @@ def test_wrong_versions(self):
                 )
 
     def test_migrate_newest_version(self):
-        """Test that an exception is raised when an export file with the newest export version is migrated."""
+        """Test that  migrating the latest version runs without complaints."""
         metadata = {'export_version': newest_version}
 
-        with self.assertRaises(ArchiveMigrationError):
-            new_version = migrate_recursively(metadata, {}, None)
-
-            self.assertIsNone(
-                new_version,
-                msg='migrate_recursively should not return anything, '
-                "hence the 'return' should be None, but instead it is {}".format(new_version)
-            )
+        new_version = migrate_recursively(metadata, {}, None)
+        self.assertEqual(new_version, newest_version)
 
     @with_temp_dir
     def test_v02_to_newest(self, temp_dir):