From b1b34f69da1bd43e1e6d97cee707b43c1e75c15c Mon Sep 17 00:00:00 2001 From: yakutovicha Date: Wed, 13 Feb 2019 20:14:02 +0100 Subject: [PATCH] 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 | 151 ++++++++++++ aiida/cmdline/commands/cmd_import.py | 48 +++- aiida/orm/importexport.py | 288 +++++++++++++++++++++- 3 files changed, 479 insertions(+), 8 deletions(-) diff --git a/aiida/backends/tests/export_and_import.py b/aiida/backends/tests/export_and_import.py index 5f04237857..92f67d85ba 100644 --- a/aiida/backends/tests/export_and_import.py +++ b/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 diff --git a/aiida/cmdline/commands/cmd_import.py b/aiida/cmdline/commands/cmd_import.py index 2742c2ccd3..ce31857f2e 100644 --- a/aiida/cmdline/commands/cmd_import.py +++ b/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') diff --git a/aiida/orm/importexport.py b/aiida/orm/importexport.py index 7c809fc52b..ecb2542bd9 100644 --- a/aiida/orm/importexport.py +++ b/aiida/orm/importexport.py @@ -13,6 +13,7 @@ import io +import click import six from six.moves import zip from six.moves.html_parser import HTMLParser @@ -382,17 +383,111 @@ def deserialize_field(k, v, fields_info, import_unique_ids_mappings, else: return ("{}_id".format(k), None) +def merge_extras(old_extras, new_extras, mode): + """ + :param old_extras: a dictionary containing the old extras of an already existing node + :param new_extras: a dictionary containing the new extras of an imported node + :param extras_mode_existing: 3 letter code that will identify what to do with the extras import. The first letter acts on + extras that are present in the original node and not present in the imported node. Can be + either k (keep it) or n (do not keep it). The second letter acts on the imported extras that + are not present in the original node. Can be either c (create it) or n (do not create it). The + third letter says what to do in case of a name collision. Can be l (leave the old value), u + (update with a new value), d (delete the extra), a (ask what to do if the content is + different). + """ + from six import string_types + if not isinstance(mode, string_types): + raise TypeError("Parameter 'mode' should be of string type, you provided '{}' type".format(type(mode))) + elif not len(mode) == 3: + raise ValueError("Parameter 'mode' should be a 3-letter string, you provided: '{}'".format(mode)) + + old_keys = set(old_extras.keys()) + new_keys = set(new_extras.keys()) + + collided_keys = old_keys.intersection(new_keys) + old_keys_only = old_keys.difference(collided_keys) + new_keys_only = new_keys.difference(collided_keys) + + final_extras = {} + + # Fast implementations for the common operations: + if mode == 'ncu': # 'mirror' operation: remove old extras, put only the new ones + return new_extras + + elif mode == 'knl': # 'none': keep old extras, do not add imported ones + return old_extras + + elif mode == 'kcu': # 'update_existing' operation: if an extra already exists, + # overwrite its new value with a new one + final_extras = new_extras + for key in old_keys_only: + final_extras[key] = old_extras[key] + + elif mode == 'kcl': # 'keep_existing': if an extra already exists, keep its original value + final_extras = old_extras + for key in new_keys_only: + final_extras[key] = new_extras[key] + + elif mode == 'kca': # 'ask': if an extra already exists ask a user whether + # to update its value + final_extras = old_extras + for key in new_keys_only: + final_extras[key] = new_extras[key] + for key in collided_keys: + if old_extras[key] != new_extras[key]: + if click.confirm('The extra {} collided, would you' + ' like to overwrite its value?\nOld value: {}\nNew value: {}\n'.format(key, + old_extras[key], new_extras[key])): + final_extras[key] = new_extras[key] + + # Slow, but more general implementation + else: + final_keys = set() + if mode[0] == 'k': + for key in old_keys_only: + final_extras[key] = old_extras[key] + elif mode[0] != 'n': + raise ValueError("Unknown first letter of the update extras mode: '{}'. Should be either 'k' or 'n'".format(mode)) + + if mode[1] == 'c': + for key in new_keys_only: + final_extras[key] = new_extras[key] + elif mode[1] != 'n': + raise ValueError("Unknown second letter of the update extras mode: '{}'. Should be either 'c' or 'n'".format(mode)) + + if mode[2] == 'u': + for key in collided_keys: + final_extras[key] = new_extras[key] + elif mode[2] == 'l': + for key in collided_keys: + final_extras[key] = old_extras[key] + elif mode[2] == 'a': + for key in collided_keys: + if old_extras[key] != new_extras[key]: + if click.confirm('The extra {} collided, would you' + ' like to overwrite its value?\nOld value: {}\nNew value: {}\n'.format(key, + old_extras[key], new_extras[key])): + final_extras[key] = new_extras[key] + else: + final_extras[key] = old_extras[key] + elif mode[2] != 'd': + raise ValueError("Unknown third letter of the update extras mode: '{}'. Should be one of 'u'/'l'/'a'/'d'".format(mode)) + + return final_extras + def import_data(in_path, group=None, ignore_unknown_nodes=False, - silent=False): + extras_mode_existing='kcl', extras_mode_new='import', silent=False): from aiida.backends.settings import BACKEND from aiida.backends.profile import BACKEND_DJANGO, BACKEND_SQLA if BACKEND == BACKEND_SQLA: return import_data_sqla(in_path, user_group=group, ignore_unknown_nodes=ignore_unknown_nodes, + extras_mode_existing=extras_mode_existing, extras_mode_new=extras_mode_new, silent=silent) elif BACKEND == BACKEND_DJANGO: return import_data_dj(in_path, user_group=group, ignore_unknown_nodes=ignore_unknown_nodes, + extras_mode_existing=extras_mode_existing, extras_mode_new=extras_mode_new, silent=silent) else: raise Exception("Unknown settings.BACKEND: {}".format( @@ -400,7 +495,7 @@ def import_data(in_path, group=None, ignore_unknown_nodes=False, def import_data_dj(in_path, user_group=None, ignore_unknown_nodes=False, - silent=False): + extras_mode_existing='kcl', extras_mode_new='import', silent=False): """ Import exported AiiDA environment to the AiiDA database. If the 'in_path' is a folder, calls extract_tree; otherwise, tries to @@ -408,6 +503,14 @@ def import_data_dj(in_path, user_group=None, ignore_unknown_nodes=False, correct function. :param in_path: the path to a file or folder that can be imported in AiiDA + :param extras_mode_existing: 3 letter code that will identify what to do with the extras import. The first letter acts on + extras that are present in the original node and not present in the imported node. Can be + either k (keep it) or n (do not keep it). The second letter acts on the imported extras that + are not present in the original node. Can be either c (create it) or n (do not create it). The + third letter says what to do in case of a name collision. Can be l (leave the old value), u + (update with a new value), d (delete the extra), a (ask what to do if the content is + different). + :param extras_mode_new: 'import' to import extras of new nodes, 'none' to ignore them """ import os import tarfile @@ -745,7 +848,7 @@ def import_data_dj(in_path, user_group=None, ignore_unknown_nodes=False, import_entry_id, new_pk)) - # For DbNodes, we also have to store Attributes! + # For DbNodes, we also have to store its attributes if model_name == NODE_ENTITY_NAME: if not silent: print("STORING NEW NODE ATTRIBUTES...") @@ -771,6 +874,81 @@ def import_data_dj(in_path, user_group=None, ignore_unknown_nodes=False, attributes=deserialized_attributes, with_transaction=False) + # For DbNodes, we also have to store its extras + if model_name == NODE_ENTITY_NAME: + if extras_mode_new == 'import': + if not silent: + print("STORING NEW NODE EXTRAS...") + for unique_id, new_pk in just_saved.items(): + import_entry_id = import_entry_ids[unique_id] + # Get extras from import file + try: + extras = data['node_extras'][ + str(import_entry_id)] + extras_conversion = data[ + 'node_extras_conversion'][ + str(import_entry_id)] + except KeyError: + raise ValueError("Unable to find extras info " + "for DbNode with UUID = {}".format( + unique_id)) + deserialized_extras = deserialize_attributes(extras, extras_conversion) + # TODO: remove when aiida extras will be moved somewhere else + # from here + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key.startswith('_aiida_')} + if models.DbNode.objects.filter(uuid=unique_id)[0].type.endswith('code.Code.'): + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key == 'hidden'} + # till here + models.DbExtra.reset_values_for_node( + dbnode=new_pk, + attributes=deserialized_extras, + with_transaction=False) + elif extras_mode_new == 'none': + if not silent: + print("SKIPPING NEW NODE EXTRAS...") + else: + raise ValueError("Unknown extras_mode_new value: {}, should be either 'import' or " + "'none'".format(extras_mode_new)) + + # For the existing DbNodes we may want to choose the import mode + if not silent: + print("UPDATING EXISTING NODE EXTRAS (mode: {})".format(extras_mode_existing)) + + for import_entry_id, entry_data in existing_entries[model_name].items(): + unique_id = entry_data[unique_identifier] + existing_entry_id = foreign_ids_reverse_mappings[model_name][unique_id] + # Get extras from import file + try: + extras = data['node_extras'][ + str(import_entry_id)] + extras_conversion = data[ + 'node_extras_conversion'][ + str(import_entry_id)] + except KeyError: + raise ValueError("Unable to find extras info " + "for DbNode with UUID = {}".format( + unique_id)) + + # Here I have to deserialize the extras + old_extras = models.DbExtra.get_all_values_for_nodepk(existing_entry_id) + deserialized_extras = deserialize_attributes(extras, extras_conversion) + # TODO: remove when aiida extras will be moved somewhere else + # from here + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key.startswith('_aiida_')} + if models.DbNode.objects.filter(uuid=unique_id)[0].type.endswith('code.Code.'): + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key == 'hidden'} + # till here + merged_extras = merge_extras(old_extras, deserialized_extras, extras_mode_existing) + + models.DbExtra.reset_values_for_node( + dbnode=existing_entry_id, + attributes=merged_extras, + with_transaction=False) + if not silent: print("STORING NODE LINKS...") ## TODO: check that we are not creating input links of an already @@ -935,7 +1113,8 @@ def validate_uuid(given_uuid): return str(parsed_uuid) == given_uuid -def import_data_sqla(in_path, user_group=None, ignore_unknown_nodes=False, silent=False): +def import_data_sqla(in_path, user_group=None, ignore_unknown_nodes=False, + extras_mode_existing='kcl', extras_mode_new='ignore', silent=False): """ Import exported AiiDA environment to the AiiDA database. If the 'in_path' is a folder, calls extract_tree; otherwise, tries to @@ -943,6 +1122,14 @@ def import_data_sqla(in_path, user_group=None, ignore_unknown_nodes=False, silen correct function. :param in_path: the path to a file or folder that can be imported in AiiDA + :param extras_mode_existing: 3 letter code that will identify what to do with the extras import. The first letter acts on + extras that are present in the original node and not present in the imported node. Can be + either k (keep it) or n (do not keep it). The second letter acts on the imported extras that + are not present in the original node. Can be either c (create it) or n (do not create it). The + third letter says what to do in case of a name collision. Can be l (leave the old value), u + (update with a new value), d (delete the extra), a (ask what to do if the content is + different). + :param extras_mode_new: 'import' to import extras of new nodes, 'none' to ignore them """ import os import tarfile @@ -952,6 +1139,8 @@ def import_data_sqla(in_path, user_group=None, ignore_unknown_nodes=False, silen from aiida.common import timezone from aiida.orm import Node, Group + from aiida.backends.sqlalchemy.models.node import DbNode + from aiida.backends.sqlalchemy.utils import flag_modified from aiida.common.archive import extract_tree, extract_tar, extract_zip, extract_cif from aiida.common.folders import SandboxFolder, RepositoryFolder from aiida.common.utils import get_object_from_string @@ -1346,11 +1535,80 @@ def import_data_sqla(in_path, user_group=None, ignore_unknown_nodes=False, silen attributes, attributes_conversion) if deserialized_attributes: - from sqlalchemy.dialects.postgresql import JSONB o.attributes = dict() for k, v in deserialized_attributes.items(): o.attributes[k] = v + # For DbNodes, we also have to store extras + # Get extras from import file + if extras_mode_new == 'import': + if not silent: + print("STORING NEW NODE EXTRAS...") + try: + extras = data['node_extras'][ + str(import_entry_id)] + + extras_conversion = data[ + 'node_extras_conversion'][ + str(import_entry_id)] + except KeyError: + raise ValueError( + "Unable to find extras info " + "for DbNode with UUID = {}".format( + unique_id)) + # Here I have to deserialize the extras + deserialized_extras = deserialize_attributes(extras, extras_conversion) + # TODO: remove when aiida extras will be moved somewhere else + # from here + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key.startswith('_aiida_')} + if o.type.endswith('code.Code.'): + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key == 'hidden'} + # till here + o.extras = dict() + for k, v in deserialized_extras.items(): + o.extras[k] = v + elif extras_mode_new == 'none': + if not silent: + print("SKIPPING NEW NODE EXTRAS...") + else: + raise ValueError("Unknown extras_mode_new value: {}, should be either 'import' or " + "'none'".format(extras_mode_new)) + + if not silent: + print("UPDATING EXISTING NODE EXTRAS (mode: {})".format(extras_mode_existing)) + + uuid_import_pk_match = {entry_data[unique_identifier]:import_entry_id for import_entry_id, entry_data in + existing_entries[entity_name].items()} + for db_node in session.query(DbNode).filter(DbNode.uuid.in_(uuid_import_pk_match)).distinct().all(): + import_entry_id = uuid_import_pk_match[str(db_node.uuid)] + # Get extras from import file + try: + extras = data['node_extras'][ + str(import_entry_id)] + extras_conversion = data[ + 'node_extras_conversion'][ + str(import_entry_id)] + except KeyError: + raise ValueError("Unable to find extras info " + "for DbNode with UUID = {}".format( + unique_id)) + + # Here I have to deserialize the extras + old_extras = db_node.extras + deserialized_extras = deserialize_attributes(extras, extras_conversion) + # TODO: remove when aiida extras will be moved somewhere else + # from here + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key.startswith('_aiida_')} + if db_node.type.endswith('code.Code.'): + deserialized_extras = {key:value for key,value in deserialized_extras.items() if not + key == 'hidden'} + # till here + db_node.extras = merge_extras(old_extras, deserialized_extras, extras_mode_existing) + flag_modified(db_node, "extras") + # Store them all in once; However, the PK # are not set in this way... if objects_to_create: @@ -2155,6 +2413,24 @@ def export_tree(what, folder, allowed_licenses=None, forbidden_licenses=None, node_attributes_conversion[str(n.pk)]) = serialize_dict( n.get_attrs(), track_conversion=True) + ## EXTRAS + if not silent: + print("STORING NODE EXTRAS...") + node_extras = {} + node_extras_conversion = {} + + # A second QueryBuilder query to get the extras. See if this can be + # optimized + if len(all_nodes_pk) > 0: + all_nodes_query = QueryBuilder() + all_nodes_query.append(Node, filters={"id": {"in": all_nodes_pk}}, + project=["*"]) + for res in all_nodes_query.iterall(): + n = res[0] + (node_extras[str(n.pk)], + node_extras_conversion[str(n.pk)]) = serialize_dict( + n.get_extras(), track_conversion=True) + if not silent: print("STORING NODE LINKS...") @@ -2349,6 +2625,8 @@ def export_tree(what, folder, allowed_licenses=None, forbidden_licenses=None, data = { 'node_attributes': node_attributes, 'node_attributes_conversion': node_attributes_conversion, + 'node_extras': node_extras, + 'node_extras_conversion': node_extras_conversion, 'export_data': export_data, 'links_uuid': links_uuid, 'groups_uuid': groups_uuid,