def test_v02_to_newest(self, temp_dir):
"""Test migration of exported files from v0.2 to newest export version"""
# Get export file with export version 0.2
input_file = get_archive_file('export_v0.2.aiida',
**self.external_archive)
output_file = os.path.join(temp_dir, 'output_file.aiida')
# Perform the migration
migrate_archive(input_file, output_file)
metadata, _ = get_json_files(output_file)
verify_metadata_version(metadata, version=newest_version)
# Load the migrated file
import_data(output_file, silent=True)
# Do the necessary checks
self.assertEqual(orm.QueryBuilder().append(orm.Node).count(),
self.node_count)
# Verify that CalculationNodes have non-empty attribute dictionaries
builder = orm.QueryBuilder().append(orm.CalculationNode)
for [calculation] in builder.iterall():
self.assertIsInstance(calculation.attributes, dict)
self.assertNotEqual(len(calculation.attributes), 0)
# Verify that the StructureData nodes maintained their (same) label, cell, and kinds
builder = orm.QueryBuilder().append(orm.StructureData)
self.assertEqual(
builder.count(),
self.struct_count,
msg='There should be {} StructureData, instead {} were/was found'.
format(self.struct_count, builder.count()))
for structures in builder.all():
structure = structures[0]
self.assertEqual(structure.label, self.known_struct_label)
self.assertEqual(structure.cell, self.known_cell)
builder = orm.QueryBuilder().append(orm.StructureData,
project=['attributes.kinds'])
for [kinds] in builder.iterall():
self.assertEqual(len(kinds), len(self.known_kinds))
for kind in kinds:
self.assertIn(kind,
self.known_kinds,
msg="Kind '{}' not found in: {}".format(
kind, self.known_kinds))
# Check that there is a StructureData that is an input of a CalculationNode
builder = orm.QueryBuilder()
builder.append(orm.StructureData, tag='structure')
builder.append(orm.CalculationNode, with_incoming='structure')
self.assertGreater(len(builder.all()), 0)
# Check that there is a RemoteData that is the output of a CalculationNode
builder = orm.QueryBuilder()
builder.append(orm.CalculationNode, tag='parent')
builder.append(orm.RemoteData, with_incoming='parent')
self.assertGreater(len(builder.all()), 0)
def test_migrate_recursively(self):
"""Test function 'migrate_recursively'"""
import io
import tarfile
import zipfile
from aiida.common.exceptions import NotExistent
from aiida.common.folders import SandboxFolder
from aiida.common.json import load as jsonload
from aiida.tools.importexport.common.archive import extract_tar, extract_zip
# Get metadata.json and data.json as dicts from v0.1 file archive
# Cannot use 'get_json_files' for 'export_v0.1_simple.aiida',
# because we need to pass the SandboxFolder to 'migrate_recursively'
dirpath_archive = get_archive_file('export_v0.1_simple.aiida',
**self.core_archive)
with SandboxFolder(sandbox_in_repo=False) as folder:
if zipfile.is_zipfile(dirpath_archive):
extract_zip(dirpath_archive, folder, silent=True)
elif tarfile.is_tarfile(dirpath_archive):
extract_tar(dirpath_archive, folder, silent=True)
else:
raise ValueError(
'invalid file format, expected either a zip archive or gzipped tarball'
)
try:
with io.open(folder.get_abs_path('data.json'),
'r',
encoding='utf8') as fhandle:
data = jsonload(fhandle)
with io.open(folder.get_abs_path('metadata.json'),
'r',
encoding='utf8') as fhandle:
metadata = jsonload(fhandle)
except IOError:
raise NotExistent(
'export archive does not contain the required file {}'.
format(fhandle.filename))
verify_metadata_version(metadata, version='0.1')
# Migrate to newest version
new_version = migrate_recursively(metadata, data, folder)
verify_metadata_version(metadata, version=newest_version)
self.assertEqual(new_version, newest_version)
def migrate(input_file, output_file, force, silent, archive_format):
# pylint: disable=too-many-locals,too-many-statements,too-many-branches
"""
Migrate an old export archive file to the most recent format.
"""
import tarfile
import zipfile
from aiida.common import json
from aiida.common.folders import SandboxFolder
from aiida.tools.importexport import migration, extract_zip, extract_tar
if os.path.exists(output_file) and not force:
echo.echo_critical('the output file already exists')
with SandboxFolder(sandbox_in_repo=False) as folder:
if zipfile.is_zipfile(input_file):
extract_zip(input_file, folder, silent=silent)
elif tarfile.is_tarfile(input_file):
extract_tar(input_file, folder, silent=silent)
else:
echo.echo_critical('invalid file format, expected either a zip archive or gzipped tarball')
try:
with io.open(folder.get_abs_path('data.json'), 'r', encoding='utf8') as fhandle:
data = json.load(fhandle)
with io.open(folder.get_abs_path('metadata.json'), 'r', encoding='utf8') as fhandle:
metadata = json.load(fhandle)
except IOError:
echo.echo_critical('export archive does not contain the required file {}'.format(fhandle.filename))
old_version = migration.verify_metadata_version(metadata)
new_version = migration.migrate_recursively(metadata, data, folder)
with io.open(folder.get_abs_path('data.json'), 'wb') as fhandle:
json.dump(data, fhandle, indent=4)
with io.open(folder.get_abs_path('metadata.json'), 'wb') as fhandle:
json.dump(metadata, fhandle)
if archive_format in ['zip', 'zip-uncompressed']:
compression = zipfile.ZIP_DEFLATED if archive_format == 'zip' else zipfile.ZIP_STORED
with zipfile.ZipFile(output_file, mode='w', compression=compression, allowZip64=True) as archive:
src = folder.abspath
for dirpath, dirnames, filenames in os.walk(src):
relpath = os.path.relpath(dirpath, src)
for filename in dirnames + filenames:
real_src = os.path.join(dirpath, filename)
real_dest = os.path.join(relpath, filename)
archive.write(real_src, real_dest)
elif archive_format == 'tar.gz':
with tarfile.open(output_file, 'w:gz', format=tarfile.PAX_FORMAT, dereference=True) as archive:
archive.add(folder.abspath, arcname='')
if not silent:
echo.echo_success('migrated the archive from version {} to {}'.format(old_version, new_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 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')
with SandboxFolder(sandbox_in_repo=False) as folder:
if zipfile.is_zipfile(input_file):
extract_zip(input_file, folder, silent=silent)
elif tarfile.is_tarfile(input_file):
extract_tar(input_file, folder, silent=silent)
else:
echo.echo_critical('invalid file format, expected either a zip archive or gzipped tarball')
try:
with open(folder.get_abs_path('data.json'), 'r', encoding='utf8') as fhandle:
data = json.load(fhandle)
with open(folder.get_abs_path('metadata.json'), 'r', encoding='utf8') as fhandle:
metadata = json.load(fhandle)
except IOError:
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:
echo.echo_critical(str(exception))
with open(folder.get_abs_path('data.json'), 'wb') as fhandle:
json.dump(data, fhandle, indent=4)
with open(folder.get_abs_path('metadata.json'), 'wb') as fhandle:
json.dump(metadata, fhandle)
if archive_format in ['zip', 'zip-uncompressed']:
compression = zipfile.ZIP_DEFLATED if archive_format == 'zip' else zipfile.ZIP_STORED
with zipfile.ZipFile(output_file, mode='w', compression=compression, allowZip64=True) as archive:
src = folder.abspath
for dirpath, dirnames, filenames in os.walk(src):
relpath = os.path.relpath(dirpath, src)
for filename in dirnames + filenames:
real_src = os.path.join(dirpath, filename)
real_dest = os.path.join(relpath, filename)
archive.write(real_src, real_dest)
elif archive_format == 'tar.gz':
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))
