def test_simple_import(self): """ This is a very simple test which checks that an export file with nodes that are not associated to a computer is imported correctly. In Django when such nodes are exported, there is an empty set for computers in the export file. In SQLA there is such a set only when a computer is associated with the exported nodes. When an empty computer set is found at the export file (when imported to an SQLA profile), the SQLA import code used to crash. This test demonstrates this problem. :return: """ import inspect import os curr_path = inspect.getfile(inspect.currentframe()) folder_path = os.path.dirname(curr_path) relative_folder_path = ("export_import_test_files/" "SSSP_parameters_2.aiida") test_file_path = os.path.join(folder_path, relative_folder_path) # Clean the database self.clean_db() # Insert the default data to the database self.insert_data() # Import the needed data import_data(test_file_path, silent=True)
def parse_data(content, name, date): # pylint: disable=unused-argument if content is None: return '' content_type, content_string = content.split(',') # pylint: disable=unused-variable decoded = base64.b64decode(content_string) fd, path = tempfile.mkstemp() try: with os.fdopen(fd, 'wb') as tmp: # do stuff with temp file tmp.write(decoded) try: from aiida import load_dbenv, is_dbenv_loaded if not is_dbenv_loaded(): load_dbenv() from aiida.orm.importexport import import_data #from aiida.common import exceptions import_data(path) except Exception: msg = 'an exception occurred while importing the archive {}'.format( name) msg += traceback.format_exc() else: msg = 'Success: imported archive {}'.format(name) finally: os.remove(path) print(msg) return msg
def test_1(self): import os import shutil import tempfile from aiida.orm import DataFactory from aiida.orm import load_node from aiida.orm.calculation.job import JobCalculation from aiida.orm.importexport import export # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: StructureData = DataFactory('structure') sd = StructureData() sd.store() calc = JobCalculation() calc.set_computer(self.computer) calc.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc.store() calc.add_link_from(sd) pks = [sd.pk, calc.pk] attrs = {} for pk in pks: node = load_node(pk) attrs[node.uuid] = dict() for k in node.attrs(): attrs[node.uuid][k] = node.get_attr(k) filename = os.path.join(temp_folder, "export.tar.gz") export([calc.dbnode], outfile=filename, silent=True) self.clean_db() # NOTE: it is better to load new nodes by uuid, rather than assuming # that they will have the first 3 pks. In fact, a recommended policy in # databases is that pk always increment, even if you've deleted elements import_data(filename, silent=True) for uuid in attrs.keys(): node = load_node(uuid) # for k in node.attrs(): for k in attrs[uuid].keys(): self.assertEquals(attrs[uuid][k], node.get_attr(k)) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def test_3(self): """ Test importing of nodes, that have links to unknown nodes. """ import json import tarfile import os import shutil import tempfile from aiida.orm import DataFactory from aiida.orm.importexport import export from aiida.common.folders import SandboxFolder # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: StructureData = DataFactory('structure') sd = StructureData() sd.store() filename = os.path.join(temp_folder, "export.tar.gz") export([sd.dbnode], outfile=filename, silent=True) unpack = SandboxFolder() with tarfile.open(filename, "r:gz", format=tarfile.PAX_FORMAT) as tar: tar.extractall(unpack.abspath) with open(unpack.get_abs_path('data.json'), 'r') as f: metadata = json.load(f) metadata['links_uuid'].append({ 'output': sd.uuid, 'input': 'non-existing-uuid', 'label': 'parent' }) with open(unpack.get_abs_path('data.json'), 'w') as f: json.dump(metadata, f) with tarfile.open(filename, "w:gz", format=tarfile.PAX_FORMAT) as tar: tar.add(unpack.abspath, arcname="") self.clean_db() with self.assertRaises(ValueError): import_data(filename, silent=True) import_data(filename, ignore_unknown_nodes=True, silent=True) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def test_2(self): """ Test the check for the export format version. """ import json import tarfile import os import shutil import tempfile from aiida.orm import DataFactory from aiida.orm.importexport import export # Creating a folder for the import/export files export_file_tmp_folder = tempfile.mkdtemp() unpack_tmp_folder = tempfile.mkdtemp() try: StructureData = DataFactory('structure') sd = StructureData() sd.store() filename = os.path.join(export_file_tmp_folder, "export.tar.gz") export([sd.dbnode], outfile=filename, silent=True) with tarfile.open(filename, "r:gz", format=tarfile.PAX_FORMAT) as tar: tar.extractall(unpack_tmp_folder) with open(os.path.join(unpack_tmp_folder, 'metadata.json'), 'r') as f: metadata = json.load(f) metadata['export_version'] = 0.0 with open(os.path.join(unpack_tmp_folder, 'metadata.json'), 'w') as f: json.dump(metadata, f) with tarfile.open(filename, "w:gz", format=tarfile.PAX_FORMAT) as tar: tar.add(unpack_tmp_folder, arcname="") self.tearDownClass() self.setUpClass() with self.assertRaises(ValueError): import_data(filename, silent=True) finally: # Deleting the created temporary folders shutil.rmtree(export_file_tmp_folder, ignore_errors=True) shutil.rmtree(unpack_tmp_folder, ignore_errors=True)
def import_archive_fixture(filepath): """ Import an archive fixture, which is an AiiDA export archive :param filepath: the relative path of the archive file within the fixture directory """ from aiida.orm.importexport import import_data filepath_current = os.path.dirname(os.path.realpath(__file__)) filepath_fixtures = os.path.join(filepath_current, os.pardir, 'fixtures') filepath_archive = os.path.join(filepath_fixtures, filepath) if not os.path.isfile(filepath_archive): raise ValueError('archive {} does not exist in the fixture directory {}'.format(filepath, filepath_fixtures)) import_data(filepath_archive, silent=True)
def read_test(self, outfolder): import os import importlib import json from aiida.orm import JobCalculation from aiida.orm.utils import load_node from aiida.orm.importexport import import_data imported = import_data(outfolder, ignore_unknown_nodes=True, silent=True) calc = None for _, pk in imported['Node']['new']: c = load_node(pk) if issubclass(c.__class__, JobCalculation): calc = c break retrieved = calc.out.retrieved retrieve_temporary_list = calc.get_attr('retrieve_temporary_list', None) if retrieve_temporary_list: from aiida.orm.data.folder import FolderData retrieved_temporary_folder = FolderData() retrieved_temporary_folder.replace_with_folder( retrieved.get_abs_path('.'), overwrite=True) else: retrieved_temporary_folder = None try: with open(os.path.join(outfolder, '_aiida_checks.json')) as f: tests = json.load(f) except IOError: raise ValueError("This test does not provide a check file!") except ValueError: raise ValueError( "This test does provide a check file, but it cannot " "be JSON-decoded!") mod_path = 'aiida.backends.tests.parser_tests.{}'.format( os.path.split(outfolder)[1]) skip_test = False try: m = importlib.import_module(mod_path) skip_test = m.skip_condition() except Exception: pass if skip_test: raise SkipTestException return calc, tests, retrieved_temporary_folder
def read_test(outfolder): """ Read a test folder created by output_test. .. note:: This method should only be called in the testing environment, because it's importing data in the current database. """ import os import importlib import json from aiida.common.exceptions import NotExistent from aiida.orm import JobCalculation from aiida.orm.utils import load_node from aiida.orm.importexport import import_data imported = import_data(outfolder, ignore_unknown_nodes=True, silent=True) calc = None for _, pk in imported['aiida.backends.djsite.db.models.DbNode']['new']: c = load_node(pk) if issubclass(c.__class__, JobCalculation): calc = c break retrieved = calc.out.retrieved try: with open(os.path.join(outfolder, '_aiida_checks.json')) as f: tests = json.load(f) except IOError: raise ValueError("This test does not provide a check file!") except ValueError: raise ValueError("This test does provide a check file, but it cannot " "be JSON-decoded!") mod_path = 'aiida.backends.djsite.db.subtests.parser_tests.{}'.format( os.path.split(outfolder)[1]) skip_test = False try: m = importlib.import_module(mod_path) skip_test = m.skip_condition() except Exception: pass if skip_test: raise SkipTestException return calc, {'retrieved': retrieved}, tests
def read_test(self, outfolder): import os import importlib import json from aiida.orm import JobCalculation from aiida.orm.utils import load_node from aiida.orm.importexport import import_data imported = import_data(outfolder, ignore_unknown_nodes=True, silent=True) calc = None for _, pk in imported['aiida.backends.djsite.db.models.DbNode']['new']: c = load_node(pk) if issubclass(c.__class__, JobCalculation): calc = c break retrieved = calc.out.retrieved try: with open(os.path.join(outfolder, '_aiida_checks.json')) as f: tests = json.load(f) except IOError: raise ValueError("This test does not provide a check file!") except ValueError: raise ValueError( "This test does provide a check file, but it cannot " "be JSON-decoded!") mod_path = 'aiida.backends.tests.parser_tests.{}'.format( os.path.split(outfolder)[1]) skip_test = False try: m = importlib.import_module(mod_path) skip_test = m.skip_condition() except Exception: pass if skip_test: raise SkipTestException return calc, {'retrieved': retrieved}, tests
def test_different_computer_same_name_import(self): """ This test checks that if there is a name collision, the imported computers are renamed accordingly. """ import os import shutil import tempfile from aiida.orm.importexport import export from aiida.orm.querybuilder import QueryBuilder from aiida.orm.computer import Computer from aiida.orm.calculation.job import JobCalculation from aiida.orm.importexport import COMP_DUPL_SUFFIX # Creating a folder for the import/export files export_file_tmp_folder = tempfile.mkdtemp() unpack_tmp_folder = tempfile.mkdtemp() try: # Set the computer name comp1_name = "localhost_1" self.computer.set_name(comp1_name) # Store a calculation calc1_label = "calc1" calc1 = JobCalculation() calc1.set_computer(self.computer) calc1.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc1.label = calc1_label calc1.store() calc1._set_state(u'RETRIEVING') # Export the first job calculation filename1 = os.path.join(export_file_tmp_folder, "export1.tar.gz") export([calc1.dbnode], outfile=filename1, silent=True) # Reset the database self.clean_db() self.insert_data() # Set the computer name to the same name as before self.computer.set_name(comp1_name) # Store a second calculation calc2_label = "calc2" calc2 = JobCalculation() calc2.set_computer(self.computer) calc2.set_resources({ "num_machines": 2, "num_mpiprocs_per_machine": 2 }) calc2.label = calc2_label calc2.store() calc2._set_state(u'RETRIEVING') # Export the second job calculation filename2 = os.path.join(export_file_tmp_folder, "export2.tar.gz") export([calc2.dbnode], outfile=filename2, silent=True) # Reset the database self.clean_db() self.insert_data() # Set the computer name to the same name as before self.computer.set_name(comp1_name) # Store a third calculation calc3_label = "calc3" calc3 = JobCalculation() calc3.set_computer(self.computer) calc3.set_resources({ "num_machines": 2, "num_mpiprocs_per_machine": 2 }) calc3.label = calc3_label calc3.store() calc3._set_state(u'RETRIEVING') # Export the third job calculation filename3 = os.path.join(export_file_tmp_folder, "export3.tar.gz") export([calc3.dbnode], outfile=filename3, silent=True) # Clean the local database self.clean_db() # Check that there are no computers qb = QueryBuilder() qb.append(Computer, project=['*']) self.assertEqual( qb.count(), 0, "There should not be any computers" "in the database at this point.") # Check that there are no calculations qb = QueryBuilder() qb.append(JobCalculation, project=['*']) self.assertEqual( qb.count(), 0, "There should not be any " "calculations in the database at " "this point.") # Import all the calculations import_data(filename1, silent=True) import_data(filename2, silent=True) import_data(filename3, silent=True) # Retrieve the calculation-computer pairs qb = QueryBuilder() qb.append(JobCalculation, project=['label'], tag='jcalc') qb.append(Computer, project=['name'], computer_of='jcalc') self.assertEqual(qb.count(), 3, "Three combinations expected.") res = qb.all() self.assertIn([calc1_label, comp1_name], res, "Calc-Computer combination not found.") self.assertIn( [calc2_label, comp1_name + COMP_DUPL_SUFFIX.format(0)], res, "Calc-Computer combination not found.") self.assertIn( [calc3_label, comp1_name + COMP_DUPL_SUFFIX.format(1)], res, "Calc-Computer combination not found.") finally: # Deleting the created temporary folders shutil.rmtree(export_file_tmp_folder, ignore_errors=True) shutil.rmtree(unpack_tmp_folder, ignore_errors=True)
def test_import(self): from aiida.orm.querybuilder import QueryBuilder from aiida.orm.node import Node from aiida.orm.calculation import Calculation from aiida.orm.data.structure import StructureData import inspect import os curr_path = inspect.getfile(inspect.currentframe()) folder_path = os.path.dirname(curr_path) relative_folder_path = ("export_import_test_files/" "parents_of_6537645.aiida") test_file_path = os.path.join(folder_path, relative_folder_path) # Clean the database self.clean_db() # Insert the default data to the database self.insert_data() # Import the needed data import_data(test_file_path, silent=True) # Check that the number of nodes if correct qb = QueryBuilder() qb.append(Node, project=["id"]) self.assertEquals(qb.count(), 83, "The number of Nodes is not the " "expected one.") # Check the number of calculations and that the attributes were # imported correctly qb = QueryBuilder() qb.append(Calculation, project=["*"]) self.assertEquals( qb.count(), 19, "The number of Calculations is not " "the expected one.") for [calc] in qb.all(): attr = calc.get_attrs() self.assertIsInstance(attr, dict, "A dictionary should be " "returned") self.assertNotEquals(len(attr), 0, "The attributes should not be " "empty.") # Check the number of the structure data and that the label is the # expected one qb = QueryBuilder() qb.append(StructureData, project=["*"]) self.assertEquals( qb.count(), 7, "The number of StructureData is not " "the expected one.") for [struct] in qb.all(): self.assertEquals(struct.label, "3D_with_2D_substructure", "A label is not correct") # TO BE SEEN WITH MOUNET # print "<================= ParameterData attributes.energy ====================>" # # from aiida.orm.data.parameter import ParameterData # qb = QueryBuilder() # # qb.append(Calculation, filters={ # # 'id': {"==": 6525492}}, project=["id"], tag="res") # qb.append(ParameterData, project=["attributes"], tag="res") # print qb.all() # for [struct] in qb.all(): # print struct # # print struct.get_attrs() # # print struct.uuid # # print struct.label # print "==============" # TO BE SEEN WITH MOUNET # Check that the cell attributes of the structure data is not empty. qb = QueryBuilder() qb.append(StructureData, project=["attributes.cell"]) for [cell] in qb.all(): self.assertNotEquals(len(cell), 0, "There should be cells.") # Check that the cell of specific structure data is the expected one qb = QueryBuilder() qb.append( StructureData, project=["attributes.cell"], filters={'uuid': { "==": "45670237-dc1e-4300-8e0b-4d3639dc77cf" }}) for [cell] in qb.all(): #print cell self.assertEquals( cell, [[8.34, 0.0, 0.0], [0.298041701839357, 8.53479766274308, 0.0], [0.842650688117053, 0.47118495164127, 10.6965192730702]], "The cell is not the expected one.") # Check that the kind attributes are the correct ones. qb = QueryBuilder() qb.append(StructureData, project=["attributes.kinds"], tag="res") for [kinds] in qb.all(): self.assertEqual(len(kinds), 2, "Attributes kinds should be of " "length 2") self.assertIn( { u'symbols': [u'Fe'], u'weights': [1.0], u'mass': 55.847, u'name': u'Fe' }, kinds) self.assertIn( { u'symbols': [u'S'], u'weights': [1.0], u'mass': 32.066, u'name': u'S' }, kinds) # Check that there are StructureData that are outputs of Calculations qb = QueryBuilder() qb.append(Calculation, project=["uuid"], tag="res") qb.append(StructureData, output_of="res") self.assertGreater(len(qb.all()), 0, "There should be results for the" "query.") # Check that there are RemoteData that are children and # parents of Calculations from aiida.orm.data.remote import RemoteData qb = QueryBuilder() qb.append(Calculation, tag="c1") qb.append(RemoteData, project=["uuid"], output_of="c1", tag='r1') qb.append(Calculation, output_of="r1", tag="c2") self.assertGreater(len(qb.all()), 0, "There should be results for the" "query.") # TO BE SEEN WITH MOUNET # from aiida.orm.data.array.trajectory import TrajectoryData # qb = QueryBuilder() # qb.append(TrajectoryData, project=["*"], tag="res") # print qb.all() # for [struct] in qb.all(): # print struct # print struct.get_attrs() # # print struct.uuid # # print struct.label # print "==============" # TO BE SEEN WITH MOUNET # Check that a specific UUID exists qb = QueryBuilder() qb.append( Node, filters={'uuid': { "==": "45670237-dc1e-4300-8e0b-4d3639dc77cf" }}, project=["*"], tag="res") self.assertGreater(len(qb.all()), 0, "There should be results for the" "query.")
def test_same_computer_import(self): """ Test that you can import nodes in steps without any problems. In this test we will import a first calculation and then a second one. The import should work as expected and have in the end two job calculations. Each calculation is related to the same computer. In the end we should have only one computer """ import os import shutil import tempfile from aiida.orm.importexport import export from aiida.orm.querybuilder import QueryBuilder from aiida.orm.computer import Computer from aiida.orm.calculation.job import JobCalculation # Creating a folder for the import/export files export_file_tmp_folder = tempfile.mkdtemp() unpack_tmp_folder = tempfile.mkdtemp() try: # Store two job calculation related to the same computer calc1_label = "calc1" calc1 = JobCalculation() calc1.set_computer(self.computer) calc1.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc1.label = calc1_label calc1.store() calc1._set_state(u'RETRIEVING') calc2_label = "calc2" calc2 = JobCalculation() calc2.set_computer(self.computer) calc2.set_resources({ "num_machines": 2, "num_mpiprocs_per_machine": 2 }) calc2.label = calc2_label calc2.store() calc2._set_state(u'RETRIEVING') # Store locally the computer name comp_name = unicode(self.computer.name) comp_uuid = unicode(self.computer.uuid) # Export the first job calculation filename1 = os.path.join(export_file_tmp_folder, "export1.tar.gz") export([calc1.dbnode], outfile=filename1, silent=True) # Export the second job calculation filename2 = os.path.join(export_file_tmp_folder, "export2.tar.gz") export([calc2.dbnode], outfile=filename2, silent=True) # Clean the local database self.clean_db() # Check that there are no computers qb = QueryBuilder() qb.append(Computer, project=['*']) self.assertEqual( qb.count(), 0, "There should not be any computers" "in the database at this point.") # Check that there are no calculations qb = QueryBuilder() qb.append(JobCalculation, project=['*']) self.assertEqual( qb.count(), 0, "There should not be any " "calculations in the database at " "this point.") # Import the first calculation import_data(filename1, silent=True) # Check that the calculation computer is imported correctly. qb = QueryBuilder() qb.append(JobCalculation, project=['label']) self.assertEqual(qb.count(), 1, "Only one calculation should be " "found.") self.assertEqual(unicode(qb.first()[0]), calc1_label, "The calculation label is not correct.") # Check that the referenced computer is imported correctly. qb = QueryBuilder() qb.append(Computer, project=['name', 'uuid', 'id']) self.assertEqual(qb.count(), 1, "Only one computer should be " "found.") self.assertEqual(unicode(qb.first()[0]), comp_name, "The computer name is not correct.") self.assertEqual(unicode(qb.first()[1]), comp_uuid, "The computer uuid is not correct.") # Store the id of the computer comp_id = qb.first()[2] # Import the second calculation import_data(filename2, silent=True) # Check that the number of computers remains the same and its data # did not change. qb = QueryBuilder() qb.append(Computer, project=['name', 'uuid', 'id']) self.assertEqual(qb.count(), 1, "Only one computer should be " "found.") self.assertEqual(unicode(qb.first()[0]), comp_name, "The computer name is not correct.") self.assertEqual(unicode(qb.first()[1]), comp_uuid, "The computer uuid is not correct.") self.assertEqual(qb.first()[2], comp_id, "The computer id is not correct.") # Check that now you have two calculations attached to the same # computer. qb = QueryBuilder() qb.append(Computer, tag='comp') qb.append(JobCalculation, has_computer='comp', project=['label']) self.assertEqual(qb.count(), 2, "Two calculations should be " "found.") ret_labels = set(_ for [_] in qb.all()) self.assertEqual( ret_labels, set([calc1_label, calc2_label]), "The labels of the calculations are not correct.") finally: # Deleting the created temporary folders shutil.rmtree(export_file_tmp_folder, ignore_errors=True) shutil.rmtree(unpack_tmp_folder, ignore_errors=True)
def test_same_computer_different_name_import(self): """ This test checks that if the computer is re-imported with a different name to the same database, then the original computer will not be renamed. It also checks that the names were correctly imported (without any change since there is no computer name collision) """ import os import shutil import tempfile from aiida.orm.importexport import export from aiida.orm.querybuilder import QueryBuilder from aiida.orm.computer import Computer from aiida.orm.calculation.job import JobCalculation # Creating a folder for the import/export files export_file_tmp_folder = tempfile.mkdtemp() unpack_tmp_folder = tempfile.mkdtemp() try: # Store a calculation calc1_label = "calc1" calc1 = JobCalculation() calc1.set_computer(self.computer) calc1.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc1.label = calc1_label calc1.store() calc1._set_state(u'RETRIEVING') # Store locally the computer name comp1_name = unicode(self.computer.name) # Export the first job calculation filename1 = os.path.join(export_file_tmp_folder, "export1.tar.gz") export([calc1.dbnode], outfile=filename1, silent=True) # Rename the computer self.computer.set_name(comp1_name + "_updated") # Store a second calculation calc2_label = "calc2" calc2 = JobCalculation() calc2.set_computer(self.computer) calc2.set_resources({ "num_machines": 2, "num_mpiprocs_per_machine": 2 }) calc2.label = calc2_label calc2.store() calc2._set_state(u'RETRIEVING') # Export the second job calculation filename2 = os.path.join(export_file_tmp_folder, "export2.tar.gz") export([calc2.dbnode], outfile=filename2, silent=True) # Clean the local database self.clean_db() # Check that there are no computers qb = QueryBuilder() qb.append(Computer, project=['*']) self.assertEqual( qb.count(), 0, "There should not be any computers" "in the database at this point.") # Check that there are no calculations qb = QueryBuilder() qb.append(JobCalculation, project=['*']) self.assertEqual( qb.count(), 0, "There should not be any " "calculations in the database at " "this point.") # Import the first calculation import_data(filename1, silent=True) # Check that the calculation computer is imported correctly. qb = QueryBuilder() qb.append(JobCalculation, project=['label']) self.assertEqual(qb.count(), 1, "Only one calculation should be " "found.") self.assertEqual(unicode(qb.first()[0]), calc1_label, "The calculation label is not correct.") # Check that the referenced computer is imported correctly. qb = QueryBuilder() qb.append(Computer, project=['name', 'uuid', 'id']) self.assertEqual(qb.count(), 1, "Only one computer should be " "found.") self.assertEqual(unicode(qb.first()[0]), comp1_name, "The computer name is not correct.") # Import the second calculation import_data(filename2, silent=True) # Check that the number of computers remains the same and its data # did not change. qb = QueryBuilder() qb.append(Computer, project=['name']) self.assertEqual(qb.count(), 1, "Only one computer should be " "found.") self.assertEqual(unicode(qb.first()[0]), comp1_name, "The computer name is not correct.") finally: # Deleting the created temporary folders shutil.rmtree(export_file_tmp_folder, ignore_errors=True) shutil.rmtree(unpack_tmp_folder, ignore_errors=True)
def test_complex_graph_import_export(self): """ This test checks that a small and bit complex graph can be correctly exported and imported. It will create the graph, store it to the database, export it to a file and import it. In the end it will check if the initial nodes are present at the imported graph. """ import tempfile import shutil import os from aiida.orm.calculation.job import JobCalculation from aiida.orm.data.folder import FolderData from aiida.orm.data.parameter import ParameterData from aiida.orm.data.remote import RemoteData from aiida.common.links import LinkType from aiida.orm.importexport import export, import_data from aiida.orm.utils import load_node from aiida.common.exceptions import NotExistent temp_folder = tempfile.mkdtemp() try: calc1 = JobCalculation() calc1.set_computer(self.computer) calc1.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc1.label = "calc1" calc1.store() calc1._set_state(u'RETRIEVING') pd1 = ParameterData() pd1.label = "pd1" pd1.store() pd2 = ParameterData() pd2.label = "pd2" pd2.store() rd1 = RemoteData() rd1.label = "rd1" rd1.set_remote_path("/x/y.py") rd1.set_computer(self.computer) rd1.store() rd1.add_link_from(calc1, link_type=LinkType.CREATE) calc2 = JobCalculation() calc2.set_computer(self.computer) calc2.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc2.label = "calc2" calc2.store() calc2.add_link_from(pd1, link_type=LinkType.INPUT) calc2.add_link_from(pd2, link_type=LinkType.INPUT) calc2.add_link_from(rd1, link_type=LinkType.INPUT) calc2._set_state(u'SUBMITTING') fd1 = FolderData() fd1.label = "fd1" fd1.store() fd1.add_link_from(calc2, link_type=LinkType.CREATE) node_uuids_labels = { calc1.uuid: calc1.label, pd1.uuid: pd1.label, pd2.uuid: pd2.label, rd1.uuid: rd1.label, calc2.uuid: calc2.label, fd1.uuid: fd1.label } filename = os.path.join(temp_folder, "export.tar.gz") export([fd1.dbnode], outfile=filename, silent=True) self.clean_db() import_data(filename, silent=True, ignore_unknown_nodes=True) for uuid, label in node_uuids_labels.iteritems(): try: load_node(uuid) except NotExistent: self.fail("Node with UUID {} and label {} was not " "found.".format(uuid, label)) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def cmd_import(archives, webpages): """Import one or multiple exported AiiDA archives The ARCHIVES can be specified by their relative or absolute file path, or their HTTP URL. """ # pylint: disable=too-many-branches,broad-except import traceback import urllib2 from aiida.common.folders import SandboxFolder from aiida.orm.importexport import get_valid_import_links, import_data archives_url = [] archives_file = [] for archive in archives: if archive.startswith('http://') or archive.startswith('https://'): archives_url.append(archive) else: archives_file.append(archive) if webpages is not None: for webpage in webpages: try: echo.echo_info( 'retrieving archive URLS from {}'.format(webpage)) urls = get_valid_import_links(webpage) except Exception: echo.echo_error( 'an exception occurred while trying to discover archives at URL {}' .format(webpage)) echo.echo(traceback.format_exc()) click.confirm('do you want to continue?', abort=True) else: echo.echo_success( '{} archive URLs discovered and added'.format(len(urls))) archives_url += urls if not archives_url + archives_file: echo.echo_critical('no valid exported archives were found') for archive in archives_file: try: echo.echo_info('importing archive {}'.format(archive)) import_data(archive) except Exception: echo.echo_error( 'an exception occurred while importing the archive {}'.format( archive)) echo.echo(traceback.format_exc()) click.confirm('do you want to continue?', abort=True) else: echo.echo_success('imported archive {}'.format(archive)) for archive in archives_url: try: echo.echo_info('downloading archive {}'.format(archive)) response = urllib2.urlopen(archive) with SandboxFolder() as temp_folder: temp_file = 'importfile.tar.gz' temp_folder.create_file_from_filelike(response, temp_file) echo.echo_success('archive downloaded, proceeding with import') import_data(temp_folder.get_abs_path(temp_file)) except Exception: echo.echo_error( 'an exception occurred while importing the archive {}'.format( archive)) echo.echo(traceback.format_exc()) click.confirm('do you want to continue?', abort=True) else: echo.echo_success('imported archive {}'.format(archive))
def run(self, *args): load_dbenv() import argparse import traceback import urllib2 from aiida.common.folders import SandboxFolder from aiida.orm.importexport import get_valid_import_links, import_data from aiida.common.exceptions import ContentNotExistent parser = argparse.ArgumentParser( prog=self.get_full_command_name(), description='Import data in the DB.') parser.add_argument('-w', '--webpage', nargs='+', type=str, dest='webpages', metavar='URL', help="Download all URLs in the given HTTP web " "page with extension .aiida") parser.add_argument(nargs='*', type=str, dest='files', metavar='URL_OR_PATH', help="Import the given files or URLs") parsed_args = parser.parse_args(args) all_args = [] if parsed_args.files is None else parsed_args.files urls = [] files = [] for path in all_args: if path.startswith('http://') or path.startswith('https://'): urls.append(path) else: files.append(path) webpages = [] if parsed_args.webpages is None else parsed_args.webpages for webpage in webpages: try: print "**** Getting links from {}".format(webpage) found_urls = get_valid_import_links(webpage) print " `-> {} links found.".format(len(found_urls)) urls += found_urls except Exception: traceback.print_exc() print "" print "> There has been an exception during the import of webpage" print "> {}".format(webpage) answer = raw_input("> Do you want to continue (c) or stop " "(S, default)? ") if answer.lower() == 'c': continue else: return if not (urls + files): print >> sys.stderr, ("Pass at least one file or URL from which " "you want to import data.") sys.exit(1) for filename in files: try: print "**** Importing file {}".format(filename) try: import_data(filename) except ContentNotExistent as ce: print ce except Exception: traceback.print_exc() print "" print "> There has been an exception during the import of file" print "> {}".format(filename) answer = raw_input("> Do you want to continue (c) or stop " "(S, default)? ") if answer.lower() == 'c': continue else: return download_file_name = 'importfile.tar.gz' for url in urls: try: print "**** Downloading url {}".format(url) response = urllib2.urlopen(url) with SandboxFolder() as temp_download_folder: temp_download_folder.create_file_from_filelike( response, download_file_name) print " `-> File downloaded. Importing it..." import_data(temp_download_folder.get_abs_path( download_file_name)) except Exception: traceback.print_exc() print "" print "> There has been an exception during the import of url" print "> {}".format(url) answer = raw_input("> Do you want to continue (c) or stop " "(S, default)? ") if answer.lower() == 'c': continue else: return
def test_5(self): """ This test checks that nodes belonging to different users are correctly exported & imported. """ import os import shutil import tempfile from aiida.orm import load_node from aiida.orm.calculation.job import JobCalculation from aiida.orm.data.structure import StructureData from aiida.orm.importexport import export from aiida.common.datastructures import calc_states from aiida.common.links import LinkType from aiida.orm.user import User from aiida.common.utils import get_configured_user_email # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: # Create another user new_email = "[email protected]" user = User(email=new_email) user.force_save() # Create a structure data node that has a calculation as output sd1 = StructureData() sd1.dbnode.user = user._dbuser sd1.label = 'sd1' sd1.store() jc1 = JobCalculation() jc1.set_computer(self.computer) jc1.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc1.dbnode.user = user._dbuser jc1.label = 'jc1' jc1.store() jc1.add_link_from(sd1) jc1._set_state(calc_states.PARSING) # Create some nodes from a different user sd2 = StructureData() sd2.dbnode.user = user._dbuser sd2.label = 'sd2' sd2.store() sd2.add_link_from(jc1, label='l1', link_type=LinkType.RETURN) jc2 = JobCalculation() jc2.set_computer(self.computer) jc2.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc2.label = 'jc2' jc2.store() jc2.add_link_from(sd2, label='l2') jc2._set_state(calc_states.PARSING) sd3 = StructureData() sd3.label = 'sd3' sd3.store() sd3.add_link_from(jc2, label='l3', link_type=LinkType.RETURN) uuids_u1 = [sd1.uuid, jc1.uuid, sd2.uuid] uuids_u2 = [jc2.uuid, sd3.uuid] filename = os.path.join(temp_folder, "export.tar.gz") export([sd3.dbnode], outfile=filename, silent=True) self.clean_db() import_data(filename, silent=True) # Check that the imported nodes are correctly imported and that # the user assigned to the nodes is the right one for uuid in uuids_u1: self.assertEquals(load_node(uuid).get_user().email, new_email) for uuid in uuids_u2: self.assertEquals(load_node(uuid).get_user().email, get_configured_user_email()) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def test_6(self): """ This test checks that nodes belonging to user A (which is not the default user) can be correctly exported, imported, enriched with nodes from the default user, re-exported & re-imported and that in the end all the nodes that have been finally imported belonging to the right users. """ import os import shutil import tempfile from aiida.orm import load_node from aiida.orm.calculation.job import JobCalculation from aiida.orm.data.structure import StructureData from aiida.orm.importexport import export from aiida.common.datastructures import calc_states from aiida.common.links import LinkType from aiida.common.utils import get_configured_user_email from aiida.orm.user import User # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: # Create another user new_email = "[email protected]" user = User(email=new_email) user.force_save() # Create a structure data node that has a calculation as output sd1 = StructureData() sd1.dbnode.user = user._dbuser sd1.label = 'sd1' sd1.store() jc1 = JobCalculation() jc1.set_computer(self.computer) jc1.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc1.dbnode.user = user._dbuser jc1.label = 'jc1' jc1.store() jc1.add_link_from(sd1) jc1._set_state(calc_states.PARSING) # Create some nodes from a different user sd2 = StructureData() sd2.dbnode.user = user._dbuser sd2.label = 'sd2' sd2.store() sd2.add_link_from(jc1, label='l1', link_type=LinkType.RETURN) # Set the jc1 to FINISHED jc1._set_state(calc_states.FINISHED) # At this point we export the generated data filename1 = os.path.join(temp_folder, "export1.tar.gz") export([sd2.dbnode], outfile=filename1, silent=True) uuids1 = [sd1.uuid, jc1.uuid, sd2.uuid] self.clean_db() self.insert_data() import_data(filename1, silent=True) # Check that the imported nodes are correctly imported and that # the user assigned to the nodes is the right one for uuid in uuids1: self.assertEquals(load_node(uuid).get_user().email, new_email) # Now we continue to generate more data based on the imported # data sd2_imp = load_node(sd2.uuid) jc2 = JobCalculation() jc2.set_computer(self.computer) jc2.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc2.label = 'jc2' jc2.store() jc2.add_link_from(sd2_imp, label='l2') jc2._set_state(calc_states.PARSING) sd3 = StructureData() sd3.label = 'sd3' sd3.store() sd3.add_link_from(jc2, label='l3', link_type=LinkType.RETURN) # Set the jc2 to FINISHED jc2._set_state(calc_states.FINISHED) # Store the UUIDs of the nodes that should be checked # if they can be imported correctly. uuids2 = [jc2.uuid, sd3.uuid] filename2 = os.path.join(temp_folder, "export2.tar.gz") export([sd3.dbnode], outfile=filename2, silent=True) self.clean_db() self.insert_data() import_data(filename2, silent=True) # Check that the imported nodes are correctly imported and that # the user assigned to the nodes is the right one for uuid in uuids1: self.assertEquals(load_node(uuid).get_user().email, new_email) for uuid in uuids2: self.assertEquals(load_node(uuid).get_user().email, get_configured_user_email()) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)