def find(self, query, fmt=DEFAULT_API_FORMAT):
"""
Query the database with a given dictionary of query parameters
:param query: a dictionary with the query parameters
"""
# pylint: disable=too-many-branches
if not isinstance(query, dict):
raise TypeError('The query argument should be a dictionary')
pagesize = self.pagesize
response = self.get(q=json.dumps(query),
fmt=ApiFormat.JSON,
pagesize=pagesize)
content = self.get_response_content(response, fmt=ApiFormat.JSON)
count = content['count']
npages = content['npages']
for page in range(npages):
response = self.get(q=json.dumps(query),
fmt=fmt,
pagesize=pagesize,
page=page)
content = self.get_response_content(response, fmt=fmt)
if fmt == ApiFormat.JSON:
if (page + 1) * pagesize > count:
last = count - (page * pagesize)
else:
last = pagesize
for i in range(last):
result = content['out'][i]
result['license'] = content['disclaimer']
yield result
elif fmt == ApiFormat.CIF:
lines = content.splitlines()
cif = []
for line in lines:
if cif:
if line.startswith('data_'):
text = '\n'.join(cif)
cif = [line]
yield text
else:
cif.append(line)
else:
if line.startswith('data_'):
cif.append(line)
if cif:
yield '\n'.join(cif)
def test_json(self):
"""Test loading and dumping from json."""
dictionary_01 = extendeddicts.AttributeDict({'x': 1, 'y': 2})
dictionary_02 = json.loads(json.dumps(dictionary_01))
# Note that here I am comparing a dictionary (dictionary_02) with a
# extendeddicts.AttributeDict (dictionary_02) and they still compare to equal
self.assertEqual(dictionary_01, dictionary_02)
def _repr_json_(self):
"""
Output in JSON format.
"""
obj = {'current_state': self.current_state}
if version_info[0] >= 3:
return obj
return json.dumps(obj)
def serialize(self):
"""
Serialize the current data (as obtained by ``self.get_dict()``) into a JSON string.
:return: A string with serialised representation of the current data.
"""
from aiida.common import json
return json.dumps(self.get_dict())
def serialize(self):
"""
Serialise the current data
:return: A serialised representation of the current data
"""
from aiida.common import json
ser_data = {
k: self.serialize_field(v, self._special_serializers.get(k, None))
for k, v in self.items()
}
return json.dumps(ser_data)
def delete_db(profile, non_interactive=True, verbose=False):
"""
Delete an AiiDA database associated with an AiiDA profile.
:param profile: AiiDA Profile
:type profile: :class:`aiida.manage.configuration.profile.Profile`
:param non_interactive: do not prompt for configuration values, fail if not all values are given as kwargs.
:type non_interactive: bool
:param verbose: if True, print parameters of DB connection
:type verbose: bool
"""
from aiida.manage.configuration import get_config
from aiida.manage.external.postgres import Postgres
from aiida.common import json
postgres = Postgres.from_profile(profile,
interactive=not non_interactive,
quiet=False)
if verbose:
echo.echo_info('Parameters used to connect to postgres:')
echo.echo(json.dumps(postgres.dbinfo, indent=4))
database_name = profile.database_name
if not postgres.db_exists(database_name):
echo.echo_info(
"Associated database '{}' does not exist.".format(database_name))
elif non_interactive or click.confirm(
"Delete associated database '{}'?\n"
'WARNING: All data will be lost.'.format(database_name)):
echo.echo_info("Deleting database '{}'.".format(database_name))
postgres.drop_db(database_name)
user = profile.database_username
config = get_config()
users = [
available_profile.database_username
for available_profile in config.profiles
]
if not postgres.dbuser_exists(user):
echo.echo_info(
"Associated database user '{}' does not exist.".format(user))
elif users.count(user) > 1:
echo.echo_info(
"Associated database user '{}' is used by other profiles "
'and will not be deleted.'.format(user))
elif non_interactive or click.confirm(
"Delete database user '{}'?".format(user)):
echo.echo_info("Deleting user '{}'.".format(user))
postgres.drop_dbuser(user)
def _format_dictionary_json_date(dictionary, sort_keys=True):
"""Return a dictionary formatted as a string using the json format and converting dates to strings."""
from aiida.common import json
def default_jsondump(data):
"""Function needed to decode datetimes, that would otherwise not be JSON-decodable."""
import datetime
from aiida.common import timezone
if isinstance(data, datetime.datetime):
return timezone.localtime(data).strftime('%Y-%m-%dT%H:%M:%S.%f%z')
raise TypeError(f'{repr(data)} is not JSON serializable')
return json.dumps(dictionary, indent=4, sort_keys=sort_keys, default=default_jsondump)
def read_text(self, path) -> str:
"""write text from the cache or base folder.
:param path: path relative to base folder
"""
if path not in self._cache:
return (self._path / path).read_text(self._encoding)
ctype, content = self._cache[path]
if ctype == 'text':
return content
if ctype == 'json':
return json.dumps(content)
raise TypeError(
f"content of type '{ctype}' could not be converted to text")
def compare_config_in_memory_and_on_disk(self, config, filepath):
"""Verify that the contents of `config` are identical to the contents of the file with path `filepath`.
:param config: instance of `Config`
:param filepath: absolute filepath to a configuration file
:raises AssertionError: if content of `config` is not equal to that of file on disk
"""
from aiida.manage.configuration.settings import DEFAULT_CONFIG_INDENT_SIZE
in_memory = json.dumps(config.dictionary,
indent=DEFAULT_CONFIG_INDENT_SIZE)
# Read the content stored on disk
with open(filepath, 'r') as handle:
on_disk = handle.read()
# Compare content of in memory config and the one on disk
self.assertEqual(in_memory, on_disk)
def presubmit(self, folder):
"""Prepares the calculation folder with all inputs, ready to be copied to the cluster.
:param folder: a SandboxFolder that can be used to write calculation input files and the scheduling script.
:type folder: :class:`aiida.common.folders.Folder`
:return calcinfo: the CalcInfo object containing the information needed by the daemon to handle operations.
:rtype calcinfo: :class:`aiida.common.CalcInfo`
"""
# pylint: disable=too-many-locals,too-many-statements,too-many-branches
import os
from aiida.common.exceptions import PluginInternalError, ValidationError, InvalidOperation, InputValidationError
from aiida.common import json
from aiida.common.utils import validate_list_of_string_tuples
from aiida.common.datastructures import CodeInfo, CodeRunMode
from aiida.orm import load_node, Code, Computer
from aiida.plugins import DataFactory
from aiida.schedulers.datastructures import JobTemplate
computer = self.node.computer
inputs = self.node.get_incoming(link_type=LinkType.INPUT_CALC)
if not self.inputs.metadata.dry_run and self.node.has_cached_links():
raise InvalidOperation('calculation node has unstored links in cache')
codes = [_ for _ in inputs.all_nodes() if isinstance(_, Code)]
for code in codes:
if not code.can_run_on(computer):
raise InputValidationError('The selected code {} for calculation {} cannot run on computer {}'.format(
code.pk, self.node.pk, computer.name))
if code.is_local() and code.get_local_executable() in folder.get_content_list():
raise PluginInternalError('The plugin created a file {} that is also the executable name!'.format(
code.get_local_executable()))
calc_info = self.prepare_for_submission(folder)
calc_info.uuid = str(self.node.uuid)
scheduler = computer.get_scheduler()
# I create the job template to pass to the scheduler
job_tmpl = JobTemplate()
job_tmpl.shebang = computer.get_shebang()
job_tmpl.submit_as_hold = False
job_tmpl.rerunnable = False
job_tmpl.job_environment = {}
# 'email', 'email_on_started', 'email_on_terminated',
job_tmpl.job_name = 'aiida-{}'.format(self.node.pk)
job_tmpl.sched_output_path = self.options.scheduler_stdout
if self.options.scheduler_stderr == self.options.scheduler_stdout:
job_tmpl.sched_join_files = True
else:
job_tmpl.sched_error_path = self.options.scheduler_stderr
job_tmpl.sched_join_files = False
# Set retrieve path, add also scheduler STDOUT and STDERR
retrieve_list = (calc_info.retrieve_list if calc_info.retrieve_list is not None else [])
if (job_tmpl.sched_output_path is not None and job_tmpl.sched_output_path not in retrieve_list):
retrieve_list.append(job_tmpl.sched_output_path)
if not job_tmpl.sched_join_files:
if (job_tmpl.sched_error_path is not None and job_tmpl.sched_error_path not in retrieve_list):
retrieve_list.append(job_tmpl.sched_error_path)
self.node.set_retrieve_list(retrieve_list)
retrieve_singlefile_list = (calc_info.retrieve_singlefile_list
if calc_info.retrieve_singlefile_list is not None else [])
# a validation on the subclasses of retrieve_singlefile_list
for _, subclassname, _ in retrieve_singlefile_list:
file_sub_class = DataFactory(subclassname)
if not issubclass(file_sub_class, orm.SinglefileData):
raise PluginInternalError(
'[presubmission of calc {}] retrieve_singlefile_list subclass problem: {} is '
'not subclass of SinglefileData'.format(self.node.pk, file_sub_class.__name__))
if retrieve_singlefile_list:
self.node.set_retrieve_singlefile_list(retrieve_singlefile_list)
# Handle the retrieve_temporary_list
retrieve_temporary_list = (calc_info.retrieve_temporary_list
if calc_info.retrieve_temporary_list is not None else [])
self.node.set_retrieve_temporary_list(retrieve_temporary_list)
# the if is done so that if the method returns None, this is
# not added. This has two advantages:
# - it does not add too many \n\n if most of the prepend_text are empty
# - most importantly, skips the cases in which one of the methods
# would return None, in which case the join method would raise
# an exception
prepend_texts = [computer.get_prepend_text()] + \
[code.get_prepend_text() for code in codes] + \
[calc_info.prepend_text, self.node.get_option('prepend_text')]
job_tmpl.prepend_text = '\n\n'.join(prepend_text for prepend_text in prepend_texts if prepend_text)
append_texts = [self.node.get_option('append_text'), calc_info.append_text] + \
[code.get_append_text() for code in codes] + \
[computer.get_append_text()]
job_tmpl.append_text = '\n\n'.join(append_text for append_text in append_texts if append_text)
# Set resources, also with get_default_mpiprocs_per_machine
resources = self.node.get_option('resources')
scheduler.preprocess_resources(resources, computer.get_default_mpiprocs_per_machine())
job_tmpl.job_resource = scheduler.create_job_resource(**resources)
subst_dict = {'tot_num_mpiprocs': job_tmpl.job_resource.get_tot_num_mpiprocs()}
for key, value in job_tmpl.job_resource.items():
subst_dict[key] = value
mpi_args = [arg.format(**subst_dict) for arg in computer.get_mpirun_command()]
extra_mpirun_params = self.node.get_option('mpirun_extra_params') # same for all codes in the same calc
# set the codes_info
if not isinstance(calc_info.codes_info, (list, tuple)):
raise PluginInternalError('codes_info passed to CalcInfo must be a list of CalcInfo objects')
codes_info = []
for code_info in calc_info.codes_info:
if not isinstance(code_info, CodeInfo):
raise PluginInternalError('Invalid codes_info, must be a list of CodeInfo objects')
if code_info.code_uuid is None:
raise PluginInternalError('CalcInfo should have '
'the information of the code '
'to be launched')
this_code = load_node(code_info.code_uuid, sub_classes=(Code,))
this_withmpi = code_info.withmpi # to decide better how to set the default
if this_withmpi is None:
if len(calc_info.codes_info) > 1:
raise PluginInternalError('For more than one code, it is '
'necessary to set withmpi in '
'codes_info')
else:
this_withmpi = self.node.get_option('withmpi')
if this_withmpi:
this_argv = (mpi_args + extra_mpirun_params + [this_code.get_execname()] +
(code_info.cmdline_params if code_info.cmdline_params is not None else []))
else:
this_argv = [this_code.get_execname()] + (code_info.cmdline_params
if code_info.cmdline_params is not None else [])
# overwrite the old cmdline_params and add codename and mpirun stuff
code_info.cmdline_params = this_argv
codes_info.append(code_info)
job_tmpl.codes_info = codes_info
# set the codes execution mode
if len(codes) > 1:
try:
job_tmpl.codes_run_mode = calc_info.codes_run_mode
except KeyError:
raise PluginInternalError('Need to set the order of the code execution (parallel or serial?)')
else:
job_tmpl.codes_run_mode = CodeRunMode.SERIAL
########################################################################
custom_sched_commands = self.node.get_option('custom_scheduler_commands')
if custom_sched_commands:
job_tmpl.custom_scheduler_commands = custom_sched_commands
job_tmpl.import_sys_environment = self.node.get_option('import_sys_environment')
job_tmpl.job_environment = self.node.get_option('environment_variables')
queue_name = self.node.get_option('queue_name')
account = self.node.get_option('account')
qos = self.node.get_option('qos')
if queue_name is not None:
job_tmpl.queue_name = queue_name
if account is not None:
job_tmpl.account = account
if qos is not None:
job_tmpl.qos = qos
priority = self.node.get_option('priority')
if priority is not None:
job_tmpl.priority = priority
max_memory_kb = self.node.get_option('max_memory_kb')
if max_memory_kb is not None:
job_tmpl.max_memory_kb = max_memory_kb
max_wallclock_seconds = self.node.get_option('max_wallclock_seconds')
if max_wallclock_seconds is not None:
job_tmpl.max_wallclock_seconds = max_wallclock_seconds
max_memory_kb = self.node.get_option('max_memory_kb')
if max_memory_kb is not None:
job_tmpl.max_memory_kb = max_memory_kb
submit_script_filename = self.node.get_option('submit_script_filename')
script_content = scheduler.get_submit_script(job_tmpl)
folder.create_file_from_filelike(io.StringIO(script_content), submit_script_filename, 'w', encoding='utf8')
subfolder = folder.get_subfolder('.aiida', create=True)
subfolder.create_file_from_filelike(io.StringIO(json.dumps(job_tmpl)), 'job_tmpl.json', 'w', encoding='utf8')
subfolder.create_file_from_filelike(io.StringIO(json.dumps(calc_info)), 'calcinfo.json', 'w', encoding='utf8')
if calc_info.local_copy_list is None:
calc_info.local_copy_list = []
if calc_info.remote_copy_list is None:
calc_info.remote_copy_list = []
# Some validation
this_pk = self.node.pk if self.node.pk is not None else '[UNSTORED]'
local_copy_list = calc_info.local_copy_list
try:
validate_list_of_string_tuples(local_copy_list, tuple_length=3)
except ValidationError as exc:
raise PluginInternalError('[presubmission of calc {}] '
'local_copy_list format problem: {}'.format(this_pk, exc))
remote_copy_list = calc_info.remote_copy_list
try:
validate_list_of_string_tuples(remote_copy_list, tuple_length=3)
except ValidationError as exc:
raise PluginInternalError('[presubmission of calc {}] '
'remote_copy_list format problem: {}'.format(this_pk, exc))
for (remote_computer_uuid, _, dest_rel_path) in remote_copy_list:
try:
Computer.objects.get(uuid=remote_computer_uuid) # pylint: disable=unused-variable
except exceptions.NotExistent:
raise PluginInternalError('[presubmission of calc {}] '
'The remote copy requires a computer with UUID={}'
'but no such computer was found in the '
'database'.format(this_pk, remote_computer_uuid))
if os.path.isabs(dest_rel_path):
raise PluginInternalError('[presubmission of calc {}] '
'The destination path of the remote copy '
'is absolute! ({})'.format(this_pk, dest_rel_path))
return calc_info
def create_value(cls,
key,
value,
subspecifier_value=None,
other_attribs={}):
"""
Create a new list of attributes, without storing them, associated
with the current key/value pair (and to the given subspecifier,
e.g. the DbNode for DbAttributes and DbExtras).
:note: No hits are done on the DB, in particular no check is done
on the existence of the given nodes.
:param key: a string with the key to create (can contain the
separator cls._sep if this is a sub-attribute: indeed, this
function calls itself recursively)
:param value: the value to store (a basic data type or a list or a dict)
:param subspecifier_value: must be None if this class has no
subspecifier set (e.g., the DbSetting class).
Must be the value of the subspecifier (e.g., the dbnode) for classes
that define it (e.g. DbAttribute and DbExtra)
:param other_attribs: a dictionary of other parameters, to store
only on the level-zero attribute (e.g. for description in DbSetting).
:return: always a list of class instances; it is the user
responsibility to store such entries (typically with a Django
bulk_create() call).
"""
import datetime
from aiida.common import json
from aiida.common.timezone import is_naive, make_aware, get_current_timezone
if cls._subspecifier_field_name is None:
if subspecifier_value is not None:
raise ValueError('You cannot specify a subspecifier value for '
'class {} because it has no subspecifiers'
''.format(cls.__name__))
if issubclass(cls, DbAttributeFunctionality):
new_entry = db_attribute_base_model(key=key, **other_attribs)
else:
new_entry = db_extra_base_model(key=key, **other_attribs)
else:
if subspecifier_value is None:
raise ValueError(
'You also have to specify a subspecifier value '
'for class {} (the {})'.format(
cls.__name__, cls._subspecifier_field_name))
further_params = other_attribs.copy()
further_params.update(
{cls._subspecifier_field_name: subspecifier_value})
# new_entry = cls(key=key, **further_params)
if issubclass(cls, DbAttributeFunctionality):
new_entry = db_attribute_base_model(key=key, **further_params)
else:
new_entry = db_extra_base_model(key=key, **further_params)
list_to_return = [new_entry]
if value is None:
new_entry.datatype = 'none'
new_entry.bval = None
new_entry.tval = ''
new_entry.ival = None
new_entry.fval = None
new_entry.dval = None
elif isinstance(value, bool):
new_entry.datatype = 'bool'
new_entry.bval = value
new_entry.tval = ''
new_entry.ival = None
new_entry.fval = None
new_entry.dval = None
elif isinstance(value, six.integer_types):
new_entry.datatype = 'int'
new_entry.ival = value
new_entry.tval = ''
new_entry.bval = None
new_entry.fval = None
new_entry.dval = None
elif isinstance(value, float):
new_entry.datatype = 'float'
new_entry.fval = value
new_entry.tval = ''
new_entry.ival = None
new_entry.bval = None
new_entry.dval = None
elif isinstance(value, six.string_types):
new_entry.datatype = 'txt'
new_entry.tval = value
new_entry.bval = None
new_entry.ival = None
new_entry.fval = None
new_entry.dval = None
elif isinstance(value, datetime.datetime):
# current timezone is taken from the settings file of django
if is_naive(value):
value_to_set = make_aware(value, get_current_timezone())
else:
value_to_set = value
new_entry.datatype = 'date'
# TODO: time-aware and time-naive datetime objects, see
# https://docs.djangoproject.com/en/dev/topics/i18n/timezones/#naive-and-aware-datetime-objects
new_entry.dval = value_to_set
new_entry.tval = ''
new_entry.bval = None
new_entry.ival = None
new_entry.fval = None
elif isinstance(value, (list, tuple)):
new_entry.datatype = 'list'
new_entry.dval = None
new_entry.tval = ''
new_entry.bval = None
new_entry.ival = len(value)
new_entry.fval = None
for i, subv in enumerate(value):
# I do not need get_or_create here, because
# above I deleted all children (and I
# expect no concurrency)
# NOTE: I do not pass other_attribs
list_to_return.extend(
cls.create_value(key=('{}{}{:d}'.format(key, cls._sep, i)),
value=subv,
subspecifier_value=subspecifier_value))
elif isinstance(value, dict):
new_entry.datatype = 'dict'
new_entry.dval = None
new_entry.tval = ''
new_entry.bval = None
new_entry.ival = len(value)
new_entry.fval = None
for subk, subv in value.items():
cls.validate_key(subk)
# I do not need get_or_create here, because
# above I deleted all children (and I
# expect no concurrency)
# NOTE: I do not pass other_attribs
list_to_return.extend(
cls.create_value(key='{}{}{}'.format(key, cls._sep, subk),
value=subv,
subspecifier_value=subspecifier_value))
else:
try:
jsondata = json.dumps(value)
except TypeError:
raise ValueError(
'Unable to store the value: it must be either a basic datatype, or json-serializable: {}'
.format(value))
new_entry.datatype = 'json'
new_entry.tval = jsondata
new_entry.bval = None
new_entry.ival = None
new_entry.fval = None
return list_to_return
def export_tree(what,
folder,
allowed_licenses=None,
forbidden_licenses=None,
silent=False,
include_comments=True,
include_logs=True,
**kwargs):
"""Export the entries passed in the 'what' list to a file tree.
:param what: a list of entity instances; they can belong to different models/entities.
:type what: list
:param folder: a temporary folder to build the archive before compression.
:type folder: :py:class:`~aiida.common.folders.Folder`
:param allowed_licenses: List or function. If a list, then checks whether all licenses of Data nodes are in the
list. If a function, then calls function for licenses of Data nodes expecting True if license is allowed, False
otherwise.
:type allowed_licenses: list
:param forbidden_licenses: List or function. If a list, then checks whether all licenses of Data nodes are in the
list. If a function, then calls function for licenses of Data nodes expecting True if license is allowed, False
otherwise.
:type forbidden_licenses: list
:param silent: suppress prints.
:type silent: bool
:param include_comments: In-/exclude export of comments for given node(s) in ``what``.
Default: True, *include* comments in export (as well as relevant users).
:type include_comments: bool
:param include_logs: In-/exclude export of logs for given node(s) in ``what``.
Default: True, *include* logs in export.
:type include_logs: bool
:param kwargs: graph traversal rules. See :const:`aiida.common.links.GraphTraversalRules` what rule names
are toggleable and what the defaults are.
:raises `~aiida.tools.importexport.common.exceptions.ArchiveExportError`: if there are any internal errors when
exporting.
:raises `~aiida.common.exceptions.LicensingException`: if any node is licensed under forbidden license.
"""
from collections import defaultdict
if not silent:
print('STARTING EXPORT...')
all_fields_info, unique_identifiers = get_all_fields_info()
entities_starting_set = defaultdict(set)
# The set that contains the nodes ids of the nodes that should be exported
given_data_entry_ids = set()
given_calculation_entry_ids = set()
given_group_entry_ids = set()
given_computer_entry_ids = set()
given_groups = set()
given_log_entry_ids = set()
given_comment_entry_ids = set()
# I store a list of the actual dbnodes
for entry in what:
# This returns the class name (as in imports). E.g. for a model node:
# aiida.backends.djsite.db.models.DbNode
# entry_class_string = get_class_string(entry)
# Now a load the backend-independent name into entry_entity_name, e.g. Node!
# entry_entity_name = schema_to_entity_names(entry_class_string)
if issubclass(entry.__class__, orm.Group):
entities_starting_set[GROUP_ENTITY_NAME].add(entry.uuid)
given_group_entry_ids.add(entry.id)
given_groups.add(entry)
elif issubclass(entry.__class__, orm.Node):
entities_starting_set[NODE_ENTITY_NAME].add(entry.uuid)
if issubclass(entry.__class__, orm.Data):
given_data_entry_ids.add(entry.pk)
elif issubclass(entry.__class__, orm.ProcessNode):
given_calculation_entry_ids.add(entry.pk)
elif issubclass(entry.__class__, orm.Computer):
entities_starting_set[COMPUTER_ENTITY_NAME].add(entry.uuid)
given_computer_entry_ids.add(entry.pk)
else:
raise exceptions.ArchiveExportError(
'I was given {} ({}), which is not a Node, Computer, or Group instance'
.format(entry, type(entry)))
# Add all the nodes contained within the specified groups
for group in given_groups:
for entry in group.nodes:
entities_starting_set[NODE_ENTITY_NAME].add(entry.uuid)
if issubclass(entry.__class__, orm.Data):
given_data_entry_ids.add(entry.pk)
elif issubclass(entry.__class__, orm.ProcessNode):
given_calculation_entry_ids.add(entry.pk)
for entity, entity_set in entities_starting_set.items():
entities_starting_set[entity] = list(entity_set)
# We will iteratively explore the AiiDA graph to find further nodes that
# should also be exported.
# At the same time, we will create the links_uuid list of dicts to be exported
if not silent:
print('RETRIEVING LINKED NODES AND STORING LINKS...')
to_be_exported, links_uuid, graph_traversal_rules = retrieve_linked_nodes(
given_calculation_entry_ids, given_data_entry_ids, **kwargs)
## Universal "entities" attributed to all types of nodes
# Logs
if include_logs and to_be_exported:
# Get related log(s) - universal for all nodes
builder = orm.QueryBuilder()
builder.append(orm.Log,
filters={'dbnode_id': {
'in': to_be_exported
}},
project='id')
res = {_[0] for _ in builder.all()}
given_log_entry_ids.update(res)
# Comments
if include_comments and to_be_exported:
# Get related log(s) - universal for all nodes
builder = orm.QueryBuilder()
builder.append(orm.Comment,
filters={'dbnode_id': {
'in': to_be_exported
}},
project='id')
res = {_[0] for _ in builder.all()}
given_comment_entry_ids.update(res)
# Here we get all the columns that we plan to project per entity that we
# would like to extract
given_entities = list()
if given_group_entry_ids:
given_entities.append(GROUP_ENTITY_NAME)
if to_be_exported:
given_entities.append(NODE_ENTITY_NAME)
if given_computer_entry_ids:
given_entities.append(COMPUTER_ENTITY_NAME)
if given_log_entry_ids:
given_entities.append(LOG_ENTITY_NAME)
if given_comment_entry_ids:
given_entities.append(COMMENT_ENTITY_NAME)
entries_to_add = dict()
for given_entity in given_entities:
project_cols = ['id']
# The following gets a list of fields that we need,
# e.g. user, mtime, uuid, computer
entity_prop = all_fields_info[given_entity].keys()
# Here we do the necessary renaming of properties
for prop in entity_prop:
# nprop contains the list of projections
nprop = (file_fields_to_model_fields[given_entity][prop] if prop
in file_fields_to_model_fields[given_entity] else prop)
project_cols.append(nprop)
# Getting the ids that correspond to the right entity
if given_entity == GROUP_ENTITY_NAME:
entry_ids_to_add = given_group_entry_ids
elif given_entity == NODE_ENTITY_NAME:
entry_ids_to_add = to_be_exported
elif given_entity == COMPUTER_ENTITY_NAME:
entry_ids_to_add = given_computer_entry_ids
elif given_entity == LOG_ENTITY_NAME:
entry_ids_to_add = given_log_entry_ids
elif given_entity == COMMENT_ENTITY_NAME:
entry_ids_to_add = given_comment_entry_ids
builder = orm.QueryBuilder()
builder.append(entity_names_to_entities[given_entity],
filters={'id': {
'in': entry_ids_to_add
}},
project=project_cols,
tag=given_entity,
outerjoin=True)
entries_to_add[given_entity] = builder
# TODO (Spyros) To see better! Especially for functional licenses
# Check the licenses of exported data.
if allowed_licenses is not None or forbidden_licenses is not None:
builder = orm.QueryBuilder()
builder.append(orm.Node,
project=['id', 'attributes.source.license'],
filters={'id': {
'in': to_be_exported
}})
# Skip those nodes where the license is not set (this is the standard behavior with Django)
node_licenses = list(
(a, b) for [a, b] in builder.all() if b is not None)
check_licenses(node_licenses, allowed_licenses, forbidden_licenses)
############################################################
##### Start automatic recursive export data generation #####
############################################################
if not silent:
print('STORING DATABASE ENTRIES...')
export_data = dict()
entity_separator = '_'
for entity_name, partial_query in entries_to_add.items():
foreign_fields = {
k: v
for k, v in all_fields_info[entity_name].items()
# all_fields_info[model_name].items()
if 'requires' in v
}
for value in foreign_fields.values():
ref_model_name = value['requires']
fill_in_query(partial_query, entity_name, ref_model_name,
[entity_name], entity_separator)
for temp_d in partial_query.iterdict():
for k in temp_d.keys():
# Get current entity
current_entity = k.split(entity_separator)[-1]
# This is a empty result of an outer join.
# It should not be taken into account.
if temp_d[k]['id'] is None:
continue
temp_d2 = {
temp_d[k]['id']:
serialize_dict(temp_d[k],
remove_fields=['id'],
rename_fields=model_fields_to_file_fields[
current_entity])
}
try:
export_data[current_entity].update(temp_d2)
except KeyError:
export_data[current_entity] = temp_d2
#######################################
# Manually manage attributes and extras
#######################################
# I use .get because there may be no nodes to export
all_nodes_pk = list()
if NODE_ENTITY_NAME in export_data:
all_nodes_pk.extend(export_data.get(NODE_ENTITY_NAME).keys())
if sum(len(model_data) for model_data in export_data.values()) == 0:
if not silent:
print('No nodes to store, exiting...')
return
if not silent:
print('Exporting a total of {} db entries, of which {} nodes.'.format(
sum(len(model_data) for model_data in export_data.values()),
len(all_nodes_pk)))
# ATTRIBUTES and EXTRAS
if not silent:
print('STORING NODE ATTRIBUTES AND EXTRAS...')
node_attributes = {}
node_extras = {}
# A second QueryBuilder query to get the attributes and extras. See if this can be optimized
if all_nodes_pk:
all_nodes_query = orm.QueryBuilder()
all_nodes_query.append(orm.Node,
filters={'id': {
'in': all_nodes_pk
}},
project=['id', 'attributes', 'extras'])
for res_pk, res_attributes, res_extras in all_nodes_query.iterall():
node_attributes[str(res_pk)] = res_attributes
node_extras[str(res_pk)] = res_extras
if not silent:
print('STORING GROUP ELEMENTS...')
groups_uuid = dict()
# If a group is in the exported date, we export the group/node correlation
if GROUP_ENTITY_NAME in export_data:
for curr_group in export_data[GROUP_ENTITY_NAME]:
group_uuid_qb = orm.QueryBuilder()
group_uuid_qb.append(entity_names_to_entities[GROUP_ENTITY_NAME],
filters={'id': {
'==': curr_group
}},
project=['uuid'],
tag='group')
group_uuid_qb.append(entity_names_to_entities[NODE_ENTITY_NAME],
project=['uuid'],
with_group='group')
for res in group_uuid_qb.iterall():
if str(res[0]) in groups_uuid:
groups_uuid[str(res[0])].append(str(res[1]))
else:
groups_uuid[str(res[0])] = [str(res[1])]
#######################################
# Final check for unsealed ProcessNodes
#######################################
process_nodes = set()
for node_pk, content in export_data.get(NODE_ENTITY_NAME, {}).items():
if content['node_type'].startswith('process.'):
process_nodes.add(node_pk)
check_process_nodes_sealed(process_nodes)
######################################
# Now I store
######################################
# subfolder inside the export package
nodesubfolder = folder.get_subfolder(NODES_EXPORT_SUBFOLDER,
create=True,
reset_limit=True)
if not silent:
print('STORING DATA...')
data = {
'node_attributes': node_attributes,
'node_extras': node_extras,
'export_data': export_data,
'links_uuid': links_uuid,
'groups_uuid': groups_uuid
}
# N.B. We're really calling zipfolder.open (if exporting a zipfile)
with folder.open('data.json', mode='w') as fhandle:
# fhandle.write(json.dumps(data, cls=UUIDEncoder))
fhandle.write(json.dumps(data))
# Add proper signature to unique identifiers & all_fields_info
# Ignore if a key doesn't exist in any of the two dictionaries
metadata = {
'aiida_version': get_version(),
'export_version': EXPORT_VERSION,
'all_fields_info': all_fields_info,
'unique_identifiers': unique_identifiers,
'export_parameters': {
'graph_traversal_rules': graph_traversal_rules,
'entities_starting_set': entities_starting_set,
'include_comments': include_comments,
'include_logs': include_logs
}
}
with folder.open('metadata.json', 'w') as fhandle:
fhandle.write(json.dumps(metadata))
if silent is not True:
print('STORING REPOSITORY FILES...')
# If there are no nodes, there are no repository files to store
if all_nodes_pk:
# Large speed increase by not getting the node itself and looping in memory in python, but just getting the uuid
uuid_query = orm.QueryBuilder()
uuid_query.append(orm.Node,
filters={'id': {
'in': all_nodes_pk
}},
project=['uuid'])
for res in uuid_query.all():
uuid = str(res[0])
sharded_uuid = export_shard_uuid(uuid)
# Important to set create=False, otherwise creates twice a subfolder. Maybe this is a bug of insert_path?
thisnodefolder = nodesubfolder.get_subfolder(sharded_uuid,
create=False,
reset_limit=True)
# Make sure the node's repository folder was not deleted
src = RepositoryFolder(section=Repository._section_name, uuid=uuid) # pylint: disable=protected-access
if not src.exists():
raise exceptions.ArchiveExportError(
'Unable to find the repository folder for Node with UUID={} in the local repository'
.format(uuid))
# In this way, I copy the content of the folder, and not the folder itself
thisnodefolder.insert_path(src=src.abspath, dest_name='.')
def test_raise_wrong_metadata_type_error(self):
"""
Test a TypeError exception is thrown with string metadata.
Also test that metadata is correctly created.
"""
from aiida.common import json
# Create CalculationNode
calc = orm.CalculationNode().store()
# dict metadata
correct_metadata_format = {
'msg': 'Life is like riding a bicycle.',
'args': '()',
'name': 'aiida.orm.node.process.calculation.CalculationNode'
}
# str of dict metadata
wrong_metadata_format = str(correct_metadata_format)
# JSON-serialized-deserialized dict metadata
json_metadata_format = json.loads(json.dumps(correct_metadata_format))
# Check an error is raised when creating a Log with wrong metadata
with self.assertRaises(TypeError):
Log(now(),
'loggername',
logging.getLevelName(LOG_LEVEL_REPORT),
calc.id,
'To keep your balance, you must keep moving',
metadata=wrong_metadata_format)
# Check no error is raised when creating a Log with dict metadata
correct_metadata_log = Log(
now(),
'loggername',
logging.getLevelName(LOG_LEVEL_REPORT),
calc.id,
'To keep your balance, you must keep moving',
metadata=correct_metadata_format)
# Check metadata is correctly created
self.assertEqual(correct_metadata_log.metadata,
correct_metadata_format)
# Create Log with json metadata, making sure TypeError is NOT raised
json_metadata_log = Log(now(),
'loggername',
logging.getLevelName(LOG_LEVEL_REPORT),
calc.id,
'To keep your balance, you must keep moving',
metadata=json_metadata_format)
# Check metadata is correctly created
self.assertEqual(json_metadata_log.metadata, json_metadata_format)
# Check no error is raised if no metadata is given
no_metadata_log = Log(now(),
'loggername',
logging.getLevelName(LOG_LEVEL_REPORT),
calc.id,
'To keep your balance, you must keep moving',
metadata=None)
# Check metadata is an empty dict for no_metadata_log
self.assertEqual(no_metadata_log.metadata, {})
def write_to_archive(
folder: Union[Folder, ZipFolder],
metadata: dict,
all_node_uuids: Set[str],
export_data: Dict[str, Dict[int, dict]],
node_attributes: Dict[str, dict],
node_extras: Dict[str, dict],
groups_uuid: Dict[str, List[str]],
links_uuid: List[dict],
silent: bool,
) -> None:
"""Store data to the archive."""
######################################
# Now collecting and storing
######################################
# subfolder inside the export package
nodesubfolder = folder.get_subfolder(NODES_EXPORT_SUBFOLDER,
create=True,
reset_limit=True)
EXPORT_LOGGER.debug("ADDING DATA TO EXPORT ARCHIVE...")
data = {
"node_attributes": node_attributes,
"node_extras": node_extras,
"export_data": export_data,
"links_uuid": links_uuid,
"groups_uuid": groups_uuid,
}
# N.B. We're really calling zipfolder.open (if exporting a zipfile)
with folder.open("data.json", mode="w") as fhandle:
# fhandle.write(json.dumps(data, cls=UUIDEncoder))
fhandle.write(json.dumps(data))
with folder.open("metadata.json", "w") as fhandle:
fhandle.write(json.dumps(metadata))
EXPORT_LOGGER.debug("ADDING REPOSITORY FILES TO EXPORT ARCHIVE...")
# If there are no nodes, there are no repository files to store
if all_node_uuids:
progress_bar = get_progress_bar(total=len(all_node_uuids),
disable=silent)
pbar_base_str = "Exporting repository - "
for uuid in all_node_uuids:
sharded_uuid = export_shard_uuid(uuid)
progress_bar.set_description_str(
f"{pbar_base_str}UUID={uuid.split('-')[0]}", refresh=False)
progress_bar.update()
# Important to set create=False, otherwise creates twice a subfolder.
# Maybe this is a bug of insert_path?
thisnodefolder = nodesubfolder.get_subfolder(sharded_uuid,
create=False,
reset_limit=True)
# Make sure the node's repository folder was not deleted
src = RepositoryFolder(section=Repository._section_name, uuid=uuid) # pylint: disable=protected-access
if not src.exists():
raise exceptions.ArchiveExportError(
f"Unable to find the repository folder for Node with UUID={uuid} "
"in the local repository")
# In this way, I copy the content of the folder, and not the folder itself
thisnodefolder.insert_path(src=src.abspath, dest_name=".")
def export_tree(entities=None,
folder=None,
allowed_licenses=None,
forbidden_licenses=None,
silent=False,
include_comments=True,
include_logs=True,
**kwargs):
"""Export the entries passed in the 'entities' list to a file tree.
.. deprecated:: 1.2.1
Support for the parameter `what` will be removed in `v2.0.0`. Please use `entities` instead.
:param entities: a list of entity instances; they can belong to different models/entities.
:type entities: list
:param folder: a temporary folder to build the archive before compression.
:type folder: :py:class:`~aiida.common.folders.Folder`
:param allowed_licenses: List or function. If a list, then checks whether all licenses of Data nodes are in the
list. If a function, then calls function for licenses of Data nodes expecting True if license is allowed, False
otherwise.
:type allowed_licenses: list
:param forbidden_licenses: List or function. If a list, then checks whether all licenses of Data nodes are in the
list. If a function, then calls function for licenses of Data nodes expecting True if license is allowed, False
otherwise.
:type forbidden_licenses: list
:param silent: suppress console prints and progress bar.
:type silent: bool
:param include_comments: In-/exclude export of comments for given node(s) in ``entities``.
Default: True, *include* comments in export (as well as relevant users).
:type include_comments: bool
:param include_logs: In-/exclude export of logs for given node(s) in ``entities``.
Default: True, *include* logs in export.
:type include_logs: bool
:param kwargs: graph traversal rules. See :const:`aiida.common.links.GraphTraversalRules` what rule names
are toggleable and what the defaults are.
:raises `~aiida.tools.importexport.common.exceptions.ArchiveExportError`: if there are any internal errors when
exporting.
:raises `~aiida.common.exceptions.LicensingException`: if any node is licensed under forbidden license.
"""
from collections import defaultdict
from aiida.tools.graph.graph_traversers import get_nodes_export
if silent:
logging.disable(level=logging.CRITICAL)
EXPORT_LOGGER.debug('STARTING EXPORT...')
# Backwards-compatibility
entities = deprecated_parameters(
old={
'name': 'what',
'value': kwargs.pop('what', None)
},
new={
'name': 'entities',
'value': entities
},
)
type_check(
entities, (list, tuple, set),
msg='`entities` must be specified and given as a list of AiiDA entities'
)
entities = list(entities)
type_check(
folder, (Folder, ZipFolder),
msg='`folder` must be specified and given as an AiiDA Folder entity')
all_fields_info, unique_identifiers = get_all_fields_info()
entities_starting_set = defaultdict(set)
# The set that contains the nodes ids of the nodes that should be exported
given_node_entry_ids = set()
given_log_entry_ids = set()
given_comment_entry_ids = set()
# Instantiate progress bar - go through list of `entities`
pbar_total = len(entities) + 1 if entities else 1
progress_bar = get_progress_bar(total=pbar_total,
leave=False,
disable=silent)
progress_bar.set_description_str('Collecting chosen entities',
refresh=False)
# I store a list of the actual dbnodes
for entry in entities:
progress_bar.update()
# This returns the class name (as in imports). E.g. for a model node:
# aiida.backends.djsite.db.models.DbNode
# entry_class_string = get_class_string(entry)
# Now a load the backend-independent name into entry_entity_name, e.g. Node!
# entry_entity_name = schema_to_entity_names(entry_class_string)
if issubclass(entry.__class__, orm.Group):
entities_starting_set[GROUP_ENTITY_NAME].add(entry.uuid)
elif issubclass(entry.__class__, orm.Node):
entities_starting_set[NODE_ENTITY_NAME].add(entry.uuid)
given_node_entry_ids.add(entry.pk)
elif issubclass(entry.__class__, orm.Computer):
entities_starting_set[COMPUTER_ENTITY_NAME].add(entry.uuid)
else:
raise exceptions.ArchiveExportError(
'I was given {} ({}), which is not a Node, Computer, or Group instance'
.format(entry, type(entry)))
# Add all the nodes contained within the specified groups
if GROUP_ENTITY_NAME in entities_starting_set:
progress_bar.set_description_str('Retrieving Nodes from Groups ...',
refresh=True)
# Use single query instead of given_group.nodes iterator for performance.
qh_groups = orm.QueryBuilder().append(
orm.Group,
filters={
'uuid': {
'in': entities_starting_set[GROUP_ENTITY_NAME]
}
},
tag='groups').queryhelp
# Delete this import once the dbexport.zip module has been renamed
from builtins import zip # pylint: disable=redefined-builtin
node_results = orm.QueryBuilder(**qh_groups).append(
orm.Node, project=['id', 'uuid'], with_group='groups').all()
if node_results:
pks, uuids = map(list, zip(*node_results))
entities_starting_set[NODE_ENTITY_NAME].update(uuids)
given_node_entry_ids.update(pks)
del node_results, pks, uuids
progress_bar.update()
# We will iteratively explore the AiiDA graph to find further nodes that should also be exported.
# At the same time, we will create the links_uuid list of dicts to be exported
progress_bar = get_progress_bar(total=1, disable=silent)
progress_bar.set_description_str(
'Getting provenance and storing links ...', refresh=True)
traverse_output = get_nodes_export(starting_pks=given_node_entry_ids,
get_links=True,
**kwargs)
node_ids_to_be_exported = traverse_output['nodes']
graph_traversal_rules = traverse_output['rules']
# A utility dictionary for mapping PK to UUID.
if node_ids_to_be_exported:
qbuilder = orm.QueryBuilder().append(
orm.Node,
project=('id', 'uuid'),
filters={'id': {
'in': node_ids_to_be_exported
}},
)
node_pk_2_uuid_mapping = dict(qbuilder.all())
else:
node_pk_2_uuid_mapping = {}
# The set of tuples now has to be transformed to a list of dicts
links_uuid = [{
'input': node_pk_2_uuid_mapping[link.source_id],
'output': node_pk_2_uuid_mapping[link.target_id],
'label': link.link_label,
'type': link.link_type
} for link in traverse_output['links']]
progress_bar.update()
# Progress bar initialization - Entities
progress_bar = get_progress_bar(total=1, disable=silent)
progress_bar.set_description_str('Initializing export of all entities',
refresh=True)
## Universal "entities" attributed to all types of nodes
# Logs
if include_logs and node_ids_to_be_exported:
# Get related log(s) - universal for all nodes
builder = orm.QueryBuilder()
builder.append(orm.Log,
filters={'dbnode_id': {
'in': node_ids_to_be_exported
}},
project='uuid')
res = set(builder.all(flat=True))
given_log_entry_ids.update(res)
# Comments
if include_comments and node_ids_to_be_exported:
# Get related log(s) - universal for all nodes
builder = orm.QueryBuilder()
builder.append(orm.Comment,
filters={'dbnode_id': {
'in': node_ids_to_be_exported
}},
project='uuid')
res = set(builder.all(flat=True))
given_comment_entry_ids.update(res)
# Here we get all the columns that we plan to project per entity that we would like to extract
given_entities = set(entities_starting_set.keys())
if node_ids_to_be_exported:
given_entities.add(NODE_ENTITY_NAME)
if given_log_entry_ids:
given_entities.add(LOG_ENTITY_NAME)
if given_comment_entry_ids:
given_entities.add(COMMENT_ENTITY_NAME)
progress_bar.update()
if given_entities:
progress_bar = get_progress_bar(total=len(given_entities),
disable=silent)
pbar_base_str = 'Preparing entities'
entries_to_add = dict()
for given_entity in given_entities:
progress_bar.set_description_str(pbar_base_str +
' - {}s'.format(given_entity),
refresh=False)
progress_bar.update()
project_cols = ['id']
# The following gets a list of fields that we need,
# e.g. user, mtime, uuid, computer
entity_prop = all_fields_info[given_entity].keys()
# Here we do the necessary renaming of properties
for prop in entity_prop:
# nprop contains the list of projections
nprop = (file_fields_to_model_fields[given_entity][prop] if prop
in file_fields_to_model_fields[given_entity] else prop)
project_cols.append(nprop)
# Getting the ids that correspond to the right entity
entry_uuids_to_add = entities_starting_set.get(given_entity, set())
if not entry_uuids_to_add:
if given_entity == LOG_ENTITY_NAME:
entry_uuids_to_add = given_log_entry_ids
elif given_entity == COMMENT_ENTITY_NAME:
entry_uuids_to_add = given_comment_entry_ids
elif given_entity == NODE_ENTITY_NAME:
entry_uuids_to_add.update(
{node_pk_2_uuid_mapping[_]
for _ in node_ids_to_be_exported})
builder = orm.QueryBuilder()
builder.append(entity_names_to_entities[given_entity],
filters={'uuid': {
'in': entry_uuids_to_add
}},
project=project_cols,
tag=given_entity,
outerjoin=True)
entries_to_add[given_entity] = builder
# TODO (Spyros) To see better! Especially for functional licenses
# Check the licenses of exported data.
if allowed_licenses is not None or forbidden_licenses is not None:
builder = orm.QueryBuilder()
builder.append(orm.Node,
project=['id', 'attributes.source.license'],
filters={'id': {
'in': node_ids_to_be_exported
}})
# Skip those nodes where the license is not set (this is the standard behavior with Django)
node_licenses = list(
(a, b) for [a, b] in builder.all() if b is not None)
check_licenses(node_licenses, allowed_licenses, forbidden_licenses)
############################################################
##### Start automatic recursive export data generation #####
############################################################
EXPORT_LOGGER.debug('GATHERING DATABASE ENTRIES...')
if entries_to_add:
progress_bar = get_progress_bar(total=len(entries_to_add),
disable=silent)
export_data = defaultdict(dict)
entity_separator = '_'
for entity_name, partial_query in entries_to_add.items():
progress_bar.set_description_str('Exporting {}s'.format(entity_name),
refresh=False)
progress_bar.update()
foreign_fields = {
k: v
for k, v in all_fields_info[entity_name].items() if 'requires' in v
}
for value in foreign_fields.values():
ref_model_name = value['requires']
fill_in_query(partial_query, entity_name, ref_model_name,
[entity_name], entity_separator)
for temp_d in partial_query.iterdict():
for key in temp_d:
# Get current entity
current_entity = key.split(entity_separator)[-1]
# This is a empty result of an outer join.
# It should not be taken into account.
if temp_d[key]['id'] is None:
continue
export_data[current_entity].update({
temp_d[key]['id']:
serialize_dict(temp_d[key],
remove_fields=['id'],
rename_fields=model_fields_to_file_fields[
current_entity])
})
# Close progress up until this point in order to print properly
close_progress_bar(leave=False)
#######################################
# Manually manage attributes and extras
#######################################
# Pointer. Renaming, since Nodes have now technically been retrieved and "stored"
all_node_pks = node_ids_to_be_exported
model_data = sum(len(model_data) for model_data in export_data.values())
if not model_data:
EXPORT_LOGGER.log(msg='Nothing to store, exiting...',
level=LOG_LEVEL_REPORT)
return
EXPORT_LOGGER.log(
msg='Exporting a total of {} database entries, of which {} are Nodes.'.
format(model_data, len(all_node_pks)),
level=LOG_LEVEL_REPORT)
# Instantiate new progress bar
progress_bar = get_progress_bar(total=1, leave=False, disable=silent)
# ATTRIBUTES and EXTRAS
EXPORT_LOGGER.debug('GATHERING NODE ATTRIBUTES AND EXTRAS...')
node_attributes = {}
node_extras = {}
# Another QueryBuilder query to get the attributes and extras. TODO: See if this can be optimized
if all_node_pks:
all_nodes_query = orm.QueryBuilder().append(
orm.Node,
filters={'id': {
'in': all_node_pks
}},
project=['id', 'attributes', 'extras'])
progress_bar = get_progress_bar(total=all_nodes_query.count(),
disable=silent)
progress_bar.set_description_str('Exporting Attributes and Extras',
refresh=False)
for node_pk, attributes, extras in all_nodes_query.iterall():
progress_bar.update()
node_attributes[str(node_pk)] = attributes
node_extras[str(node_pk)] = extras
EXPORT_LOGGER.debug('GATHERING GROUP ELEMENTS...')
groups_uuid = defaultdict(list)
# If a group is in the exported data, we export the group/node correlation
if GROUP_ENTITY_NAME in export_data:
group_uuids_with_node_uuids = orm.QueryBuilder().append(
orm.Group,
filters={
'id': {
'in': export_data[GROUP_ENTITY_NAME]
}
},
project='uuid',
tag='groups').append(orm.Node, project='uuid', with_group='groups')
# This part is _only_ for the progress bar
total_node_uuids_for_groups = group_uuids_with_node_uuids.count()
if total_node_uuids_for_groups:
progress_bar = get_progress_bar(total=total_node_uuids_for_groups,
disable=silent)
progress_bar.set_description_str('Exporting Groups ...',
refresh=False)
for group_uuid, node_uuid in group_uuids_with_node_uuids.iterall():
progress_bar.update()
groups_uuid[group_uuid].append(node_uuid)
#######################################
# Final check for unsealed ProcessNodes
#######################################
process_nodes = set()
for node_pk, content in export_data.get(NODE_ENTITY_NAME, {}).items():
if content['node_type'].startswith('process.'):
process_nodes.add(node_pk)
check_process_nodes_sealed(process_nodes)
######################################
# Now collecting and storing
######################################
# subfolder inside the export package
nodesubfolder = folder.get_subfolder(NODES_EXPORT_SUBFOLDER,
create=True,
reset_limit=True)
EXPORT_LOGGER.debug('ADDING DATA TO EXPORT ARCHIVE...')
data = {
'node_attributes': node_attributes,
'node_extras': node_extras,
'export_data': export_data,
'links_uuid': links_uuid,
'groups_uuid': groups_uuid
}
# N.B. We're really calling zipfolder.open (if exporting a zipfile)
with folder.open('data.json', mode='w') as fhandle:
# fhandle.write(json.dumps(data, cls=UUIDEncoder))
fhandle.write(json.dumps(data))
# Turn sets into lists to be able to export them as JSON metadata.
for entity, entity_set in entities_starting_set.items():
entities_starting_set[entity] = list(entity_set)
metadata = {
'aiida_version': get_version(),
'export_version': EXPORT_VERSION,
'all_fields_info': all_fields_info,
'unique_identifiers': unique_identifiers,
'export_parameters': {
'graph_traversal_rules': graph_traversal_rules,
'entities_starting_set': entities_starting_set,
'include_comments': include_comments,
'include_logs': include_logs
}
}
with folder.open('metadata.json', 'w') as fhandle:
fhandle.write(json.dumps(metadata))
EXPORT_LOGGER.debug('ADDING REPOSITORY FILES TO EXPORT ARCHIVE...')
# If there are no nodes, there are no repository files to store
if all_node_pks:
all_node_uuids = {node_pk_2_uuid_mapping[_] for _ in all_node_pks}
progress_bar = get_progress_bar(total=len(all_node_uuids),
disable=silent)
pbar_base_str = 'Exporting repository - '
for uuid in all_node_uuids:
sharded_uuid = export_shard_uuid(uuid)
progress_bar.set_description_str(
pbar_base_str + 'UUID={}'.format(uuid.split('-')[0]),
refresh=False)
progress_bar.update()
# Important to set create=False, otherwise creates twice a subfolder. Maybe this is a bug of insert_path?
thisnodefolder = nodesubfolder.get_subfolder(sharded_uuid,
create=False,
reset_limit=True)
# Make sure the node's repository folder was not deleted
src = RepositoryFolder(section=Repository._section_name, uuid=uuid) # pylint: disable=protected-access
if not src.exists():
raise exceptions.ArchiveExportError(
'Unable to find the repository folder for Node with UUID={} in the local repository'
.format(uuid))
# In this way, I copy the content of the folder, and not the folder itself
thisnodefolder.insert_path(src=src.abspath, dest_name='.')
close_progress_bar(leave=False)
# Reset logging level
if silent:
logging.disable(level=logging.NOTSET)
def dumps_json(dictionary):
"""Transforms all datetime object into isoformat and then returns the JSON."""
return json.dumps(recursive_datetime_to_isoformat(dictionary))
