def query_by_attrs(cls, query=None, **kwargs):
family_name = kwargs.pop('family_name', None)
if family_name:
group_filters = {'label': {'==': family_name}, 'type_string': {'==': cls.potcar_family_type_string}}
query = QueryBuilder()
query.append(Group, tag='family', filters=group_filters)
query.append(cls, tag=cls._query_label, with_group='family')
return super(PotcarData, cls).query_by_attrs(query=query, **kwargs)
def get_all_structure_from_structuregroup(structure_group):
if structure_group:
structure_group_name = structure_group.label
else:
return []
sqb = QueryBuilder()
sqb.append(Group, filters={'label': structure_group_name}, tag='g')
sqb.append(StructureData, member_of='g')
return [x[0] for x in sqb.all()]
def get_all_parents(self, node_pks, return_values=('id', )):
"""
Get all the parents of given nodes
:param node_pks: one node pk or an iterable of node pks
:return: a list of aiida objects with all the parents of the nodes
"""
from aiida.orm import Node, QueryBuilder
qb = QueryBuilder()
qb.append(Node, tag='low_node', filters={'id': {'in': node_pks}})
qb.append(Node, with_descendants='low_node', project=return_values)
return qb.all()
def process_pause(processes, all_entries, timeout, wait):
"""Pause running processes."""
from aiida.orm import ProcessNode, QueryBuilder
controller = get_manager().get_process_controller()
if processes and all_entries:
raise click.BadOptionUsage(
'all',
'cannot specify individual processes and the `--all` flag at the same time.'
)
if not processes and all_entries:
active_states = options.active_process_states()
builder = QueryBuilder().append(
ProcessNode,
filters={'attributes.process_state': {
'in': active_states
}})
processes = [entry[0] for entry in builder.all()]
futures = {}
for process in processes:
if process.is_terminated:
echo.echo_error('Process<{}> is already terminated'.format(
process.pk))
continue
try:
future = controller.pause_process(
process.pk, msg='Paused through `verdi process pause`')
except communications.UnroutableError:
echo.echo_error('Process<{}> is unreachable'.format(process.pk))
else:
futures[future] = process
process_actions(futures, 'pause', 'pausing', 'paused', wait, timeout)
def _get_unique_node_property(self, project: str) -> Union[Node, Any]:
query = QueryBuilder(limit=1)
query.append(self.AIIDA_ENTITY,
filters={"id": self._pk},
project=project)
if query.count() != 1:
raise AiidaEntityNotFound(
f"Could not find {self.AIIDA_ENTITY} with PK {self._pk}.")
return query.first()[0]
def group_remove_nodes(group, nodes, clear, force):
"""Remove nodes from a group."""
from aiida.orm import QueryBuilder, Group, Node
label = group.label
klass = group.__class__.__name__
if nodes and clear:
echo.echo_critical(
'Specify either the `--clear` flag to remove all nodes or the identifiers of the nodes you want to remove.'
)
if not force:
if nodes:
node_pks = [node.pk for node in nodes]
query = QueryBuilder()
query.append(Group, filters={'id': group.pk}, tag='group')
query.append(Node,
with_group='group',
filters={'id': {
'in': node_pks
}},
project='id')
group_node_pks = query.all(flat=True)
if not group_node_pks:
echo.echo_critical(
f'None of the specified nodes are in {klass}<{label}>.')
if len(node_pks) > len(group_node_pks):
node_pks = set(node_pks).difference(set(group_node_pks))
echo.echo_warning(
f'{len(node_pks)} nodes with PK {node_pks} are not in {klass}<{label}>.'
)
message = f'Are you sure you want to remove {len(group_node_pks)} nodes from {klass}<{label}>?'
elif clear:
message = f'Are you sure you want to remove ALL the nodes from {klass}<{label}>?'
else:
echo.echo_critical(
f'No nodes were provided for removal from {klass}<{label}>.')
click.confirm(message, abort=True)
if clear:
group.clear()
else:
group.remove_nodes(nodes)
def get(self):
from aiida.orm import QueryBuilder, Dict
qb = QueryBuilder()
qb.append(Dict, project=['id', 'ctime', 'attributes'], tag='pdata')
qb.order_by({'pdata': {'ctime': 'desc'}})
result = qb.first()
# Results are returned as a dictionary, datetime objects is
# serialized as ISO 8601
return dict(id=result[0],
ctime=result[1].isoformat(),
attributes=result[2])
def get_data_aiida(projections, sliders_dict, quantities, plot_info):
"""Query the AiiDA database"""
from figure.aiida import load_profile
from aiida.orm import QueryBuilder, Dict
load_profile()
filters = {}
def add_range_filter(bounds, label):
# a bit of cheating until this is resolved
# https://github.com/aiidateam/aiida_core/issues/1389
#filters['attributes.'+label] = {'>=':bounds[0]}
filters['attributes.' + label] = {
'and': [{
'>=': bounds[0]
}, {
'<': bounds[1]
}]
}
for k, v in sliders_dict.items():
# Note: filtering is costly, avoid if possible
if not v.value == quantities[k]['range']:
add_range_filter(v.value, k)
qb = QueryBuilder()
qb.append(
Dict,
filters=filters,
project=['attributes.' + p
for p in projections] + ['uuid', 'extras.cif_uuid'],
)
nresults = qb.count()
if nresults == 0:
plot_info.text = 'No matching COFs found.'
return data_empty
plot_info.text = '{} COFs found. Plotting...'.format(nresults)
# x,y position
x, y, clrs, uuids, names, cif_uuids = list(zip(*qb.all()))
plot_info.text = '{} COFs queried'.format(nresults)
x = list(map(float, x))
y = list(map(float, y))
cif_uuids = list(map(str, cif_uuids))
uuids = list(map(str, uuids))
if projections[2] == 'bond_type':
#clrs = map(lambda clr: bondtypes.index(clr), clrs)
clrs = list(map(str, clrs))
else:
clrs = list(map(float, clrs))
return dict(x=x, y=y, uuid=cif_uuids, color=clrs, name=names)
def retrieve_basis_sets(files, stop_if_existing):
"""Retrieve existing basis sets or create if them, if they do not exist.
:param files: list of basis set file paths
:param stop_if_existing: if True, check for the md5 of the files and,
if the file already exists in the DB, raises a MultipleObjectsError.
If False, simply adds the existing BasisSetData node to the group.
:return:
"""
basis_and_created = []
for basis_file in files:
_, content = parse_basis(basis_file)
md5sum = md5_from_string(content)
qb = QueryBuilder()
qb.append(BasisSetData, filters={"attributes.md5": {"==": md5sum}})
existing_basis = qb.first()
if existing_basis is None:
# return the basis set data instances, not stored
basisset, created = BasisSetData.get_or_create(
basis_file, use_first=True, store_basis=False
)
# to check whether only one basis set per element exists
# NOTE: actually, created has the meaning of "to_be_created"
basis_and_created.append((basisset, created))
else:
if stop_if_existing:
raise ValueError(
"A Basis Set with identical MD5 to "
" {} cannot be added with stop_if_existing"
"".format(basis_file)
)
existing_basis = existing_basis[0]
basis_and_created.append((existing_basis, False))
return basis_and_created
def _get_codes(self):
"""Query the list of available codes."""
user = User.objects.get_default()
return {
self._full_code_label(c[0]): c[0]
for c in QueryBuilder().append(
Code, filters={
"attributes.input_plugin": self.input_plugin
}).all() if c[0].computer.is_user_configured(user) and (
self.allow_hidden_codes or not c[0].hidden) and (
self.allow_disabled_computers
or c[0].computer.is_user_enabled(user))
}
def default_sssp_family(ctx, param, identifier): # pylint: disable=unused-argument
"""Determine the default if no value is specified."""
from aiida.common import exceptions
from aiida.orm import QueryBuilder
from aiida_sssp.groups import SsspFamily
if identifier is not None:
return identifier
try:
return QueryBuilder().append(SsspFamily).first()[0]
except exceptions.NotExistent:
raise click.BadParameter(
'failed to automatically detect an SSSP family: install it with `aiida-sssp install`.'
)
def _observe_node(self, _=None):
if self.node is None or self.eln is None:
return
if "eln" in self.node.extras:
info = self.node.extras["eln"]
else:
try:
q = QueryBuilder().append(
Node,
filters={"extras": {
"has_key": "eln"
}},
tag="source_node",
project="extras.eln",
)
q.append(Node,
filters={"uuid": self.node.uuid},
with_ancestors="source_node")
info = q.all(flat=True)[0]
except IndexError:
info = {}
self.eln.set_sample_config(**info)
def _set_shear_aiida_relaxes(self):
"""
Set list of AiidaRelaxWorkChain objects.
"""
rlx_wf = WorkflowFactory('vasp.relax')
qb = QueryBuilder()
qb.append(Node, filters={'id': {'==': self._pk}}, tag='wf')
qb.append(rlx_wf, with_incoming='wf', project=['id', 'label'])
qb_all = qb.all()
qb_all.sort(key=lambda qb_all: qb_all[1])
rlx_pks = [q[0] for q in qb_all if not load_node(q[0]).is_failed]
self._shear_aiida_relaxes = [
AiidaRelaxWorkChain(load_node(pk)) for pk in rlx_pks
]
def get(cls, **kwargs):
"""
Custom get for group which can be used to get a group with the given attributes
:param kwargs: the attributes to match the group to
:return: the group
:type nodes: :class:`aiida.orm.Node` or list
"""
from aiida.orm import QueryBuilder
if 'type_string' in kwargs:
message = '`type_string` is deprecated because it is determined automatically'
warnings.warn(message) # pylint: disable=no-member
type_check(kwargs['type_string'], str)
return QueryBuilder().append(cls, filters=kwargs).one()[0]
def get_allstructurenodes_fromgroup(structure_group):
from aiida.orm import Group
from aiida.orm import StructureData
from aiida.orm import QueryBuilder
if structure_group:
structure_group_label = structure_group.label
else:
return []
sqb = QueryBuilder()
sqb.append(Group, filters={'label': structure_group_label}, tag='g')
sqb.append(StructureData, with_group='g')
return [x[0] for x in sqb.all()]
def convert(self, value, param, ctx):
is_path = False
# Alternative one could check if int or uuid
# aiida allows also for shorten uuids
from aiida.orm import StructureData, QueryBuilder
try:
structure = types.DataParamType(
sub_classes=('aiida.data:structure', )).convert(
value, param, ctx)
except (NotExistent, click.exceptions.BadParameter) as er:
echo.echo(f'Tried to load node, could not fine one for {value}. '
'I will further check if it is a filepath.')
is_path = True
if is_path:
# If it is a path to a file try to convert the structure
pathtype = click.Path(exists=True,
dir_okay=False,
resolve_path=True)
filename = pathtype.convert(value, param, ctx)
try:
import ase.io
except ImportError:
echo.echo_critical(
'You have not installed the package ase. \nYou can install it with: pip install ase'
)
try:
asecell = ase.io.read(filename)
structure = StructureData(ase=asecell)
except ValueError as err:
echo.echo_critical(str(err))
# do not store structure, since this option is for calculation and workflow
# input, which will store the structure anyway.
# do not store again if structure is already there.
duplicate = QueryBuilder().append(StructureData,
filters={
'extras._aiida_hash':
structure._get_hash()
}).first() # pylint: disable=protected-access
if duplicate:
return duplicate[0]
return structure
def get_sssp_families_builder(version=None, functional=None, protocol=None):
"""Return a query builder that will query for SSSP families of the given configuration.
:param version: optional version filter
:param functional: optional functional filter
:param protocol: optional protocol filter
:return: `QueryBuilder` instance
"""
from aiida.orm import QueryBuilder
from aiida_sssp.groups import SsspFamily
label = 'SSSP/{version}/{functional}/{protocol}'
filters = {
'label': {
'like': label.format(version=version or '%', functional=functional or '%', protocol=protocol or '%')
}
}
builder = QueryBuilder().append(SsspFamily, filters=filters)
return builder
def generate_wf_option_node(computer=None, check_existence=True, **kwargs):
"""Create a option node for a certain workflow or calculation entry point.
:param wf_entry_point: string the entry point to create the node for, default='fleur.scf'
:param computer: dict {computername, queue} to provide computer dependend defaults
:param kwargs: dict, further key word argument by which the node content will be updated
:returns: AiiDA Dict node
"""
option_node_dict = generate_wf_option_dict(computer=computer, **kwargs)
option_node = orm.Dict(dict=option_node_dict)
if check_existence:
duplicate = QueryBuilder().append(orm.Dict,
filters={
'extras._aiida_hash':
option_node._get_hash()
}).first() # pylint: disable=protected-access
if duplicate:
option_node = duplicate[0]
return option_node
def cmd_install_openmx(version, protocol, hardness, traceback):
"""Install an OpenMX configuration.
The OpenMX configuration will be automatically downloaded from t-ozaki.issp.u-tokyo.ac.jp to create a new
`OpenmxBasisSet`.
"""
# pylint: disable=too-many-locals
configuration = OpenmxConfiguration(version, protocol, hardness)
if configuration not in OpenmxBasisSet.valid_configurations:
echo.echo_critical(f'{version} {hardness} {protocol} is not a valid OpenMX basis set configuration')
label = OpenmxBasisSet.format_configuration_label(configuration)
description = f'OpenMX 20{version} {hardness} {protocol} installed with aiida-basis v{__version__}'
metadata = OpenmxBasisSet.get_configuration_metadata(configuration)
if QueryBuilder().append(OpenmxBasisSet, filters={'label': label}).first():
echo.echo_critical(f'{OpenmxBasisSet.__name__}<{label}> is already installed')
with tempfile.TemporaryDirectory() as dirpath:
dirpath = pathlib.Path(dirpath)
download_openmx(configuration, dirpath)
with attempt('parsing PAOs... ', include_traceback=traceback):
basis_set = create_basis_set_from_directory(OpenmxBasisSet, label, dirpath)
orbital_configurations = {}
for element, values in metadata.items():
if basis_set.get_basis(element).md5 != values['md5']:
Group.objects.delete(basis_set.pk)
msg = f'md5 of PAO for element {element} does not match that of the metadata {values["md5"]}'
echo.echo_critical(msg)
orbital_configurations[element] = metadata[element]['orbital_configuration']
basis_set.description = description
basis_set.set_orbital_configurations(orbital_configurations)
echo.echo_success(f'installed `{label}` containing `{basis_set.count()}` orbital bases')
def get_members(
self,
elements: Optional[Sequence[str]] = None,
structure: Optional[StructureData] = None) -> Dict[str, List[_T]]:
"""
Return a dict of kind names/elements to a list of respective data nodes
for the given list of elements or structure.
:param elements: list of element symbols.
:param structure: the ``StructureData`` node.
"""
assert (elements is None) ^ (
structure is None
), "Exactly one of the parameters elements and structure must be specified"
assert isinstance(elements, Sequence) or isinstance(
structure, StructureData)
if structure:
elements = list(structure.get_symbols_set())
query = (QueryBuilder().append(
self.__class__, filters={
"id": self.pk
}, tag="group").append(
self.member_type,
with_group="group",
filters={"attributes.element": {
"in": elements
}}))
pseudos: Dict[str, List[_T]] = {}
for (pseudo, ) in query.iterall():
pseudos.setdefault(pseudo.element, []).append(pseudo)
return pseudos
def get_potcars_dict(cls, elements, family_name, mapping=None):
"""
Get a dictionary {element: ``PotcarData.full_name``} for all given symbols.
:param elements: The list of symbols to find POTCARs for
:param family_name: The POTCAR family to be used
:param mapping: A mapping[element] -> ``full_name``, for example: mapping={'In': 'In', 'As': 'As_d'}
Exceptions:
*If the mapping does not contain an item for a given element name, raise a ``ValueError``.
*If no POTCAR is found for a given element, a ``NotExistent`` error is raised.
If there are multiple POTCAR with the same ``full_name``, the first one
returned by ``PotcarData.find()`` will be used.
"""
if not mapping:
mapping = {element: element for element in elements}
group_filters = {'label': {'==': family_name}, 'type_string': {'==': cls.potcar_family_type_string}}
element_filters = {'attributes.full_name': {'in': [mapping[element] for element in elements]}}
query = QueryBuilder()
query.append(Group, tag='family', filters=group_filters)
query.append(cls, tag='potcar', with_group='family', filters=element_filters)
result_potcars = {}
for element in elements:
if element not in mapping:
raise ValueError('Potcar mapping must contain an item for each element in the structure, '
'with the full name of the POTCAR file (i.e. "In_d", "As_h").')
full_name = mapping[element]
potcars_of_kind = [potcar[0] for potcar in query.all() if potcar[0].full_name == full_name]
if not potcars_of_kind:
raise NotExistent('No POTCAR found for full name {} in family {}'.format(full_name, family_name))
if len(potcars_of_kind) > 1:
result_potcars[element] = cls.find(family=family_name, full_name=full_name)[0]
else:
result_potcars[element] = potcars_of_kind[0]
return result_potcars
def generate_wf_para_node(wf_entry_point='fleur.scf',
check_existence=True,
**kwargs):
"""Create a wf parameter node for a certain workflow or calculation entry point.
:param wf_entry_point: string the entry point to create the node for, default='fleur.scf'
:param kwargs: dict, further key word argument by which the node content will be updated
:returns: AiiDA Dict node
"""
wf_para_node_dict = generate_wf_para_dict(wf_entry_point=wf_entry_point,
check_existence=check_existence,
**kwargs)
wf_para_node = orm.Dict(dict=wf_para_node_dict)
if check_existence:
duplicate = QueryBuilder().append(orm.Dict,
filters={
'extras._aiida_hash':
wf_para_node._get_hash()
}).first() # pylint: disable=protected-access
if duplicate:
wf_para_node = duplicate[0]
return wf_para_node
def _set_structure_pks(self):
"""
Set structure pks.
"""
qb = QueryBuilder()
qb.append(Node, filters={'id': {'==': self._pk}}, tag='wf')
qb.append(
Node,
filters={'label': {
'==': 'get_twinboundary_shear_structure'
}},
project=['id'],
with_incoming='wf')
cf_pks = [q[0] for q in qb.all()]
cf_pks.sort()
shear_ratios = self._shear_strain_ratios
orders = list(np.argsort(shear_ratios))
orig_pks = []
input_pks = []
# for ix in orders:
for ix in range(len(self._shear_aiida_relaxes)):
cf = load_node(cf_pks[ix])
orig_pks.append(cf.outputs.twinboundary_shear_structure_orig.pk)
input_pks.append(cf.outputs.twinboundary_shear_structure.pk)
rlx_pks = []
for aiida_rlx, i_struct_pk in zip(self._shear_aiida_relaxes,
input_pks):
pks = aiida_rlx.get_pks()
assert pks['initial_structure_pk'] == i_struct_pk, \
"Input structure does not match."
rlx_pks.append(pks['final_structure_pk'])
self._structure_pks = {
'original_structures': orig_pks,
'input_structures': input_pks,
'relax_structures': rlx_pks,
}
def do_group_belong(self, args):
"""Command for list groups"""
from aiida.orm import QueryBuilder, Node, Group
from tabulate import tabulate
q = QueryBuilder()
q.append(Node, filters={'id': self._current_node.pk})
q.append(Group, with_node=Node, project=['*'])
projection_lambdas = {
'pk': lambda group: str(group.pk),
'label': lambda group: group.label,
'type_string': lambda group: group.type_string,
'count': lambda group: group.count(),
'user': lambda group: group.user.email.strip(),
'description': lambda group: group.description
}
table = []
projection_header = ['PK', 'Label', 'Type string', 'User']
projection_fields = ['pk', 'label', 'type_string', 'user']
if args.with_count:
projection_header.append('Node Count')
projection_fields.append('count')
if args.with_description:
projection_header.append('Description')
projection_fields.append('description')
for (group, ) in q.all():
table.append([
projection_lambdas[field](group) for field in projection_fields
])
self.poutput(tabulate(table, headers=projection_header))
def test_get_or_create_invalid_prefix(self):
"""Test the ``get_or_create_group`` method of ``Autogroup`` when there is already a group
with the same prefix, but followed by other non-underscore characters."""
label_prefix = 'new_test_prefix_TestAutogroup'
# I create a group with the same prefix, but followed by non-underscore
# characters. These should be ignored in the logic.
AutoGroup(label='{}xx'.format(label_prefix)).store()
# Check that there are no groups to begin with
queryb = QueryBuilder().append(AutoGroup,
filters={'label': label_prefix})
assert not list(queryb.all())
queryb = QueryBuilder().append(
AutoGroup,
filters={'label': {
'like': r'{}\_%'.format(label_prefix)
}})
assert not list(queryb.all())
# First group (no existing one)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))
# Second group (only one with no suffix existing)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix + '_1'
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))
def _add_nodes_to_groups(*, group_count: int,
group_uuids: Iterable[Tuple[str, Set[str]]],
foreign_ids_reverse_mappings: Dict[str, Dict[str,
int]]):
"""Add nodes to imported groups."""
if not group_count:
return
pbar_base_str = 'Groups - '
with get_progress_reporter()(total=group_count,
desc=pbar_base_str) as progress:
for groupuuid, groupnodes in group_uuids:
if not groupnodes:
progress.update()
continue
# TODO: cache these to avoid too many queries
qb_group = QueryBuilder().append(
Group, filters={'uuid': {
'==': groupuuid
}})
group_ = qb_group.first()[0]
progress.set_description_str(
f'{pbar_base_str}label={group_.label}', refresh=False)
progress.update()
nodes_to_store = [
foreign_ids_reverse_mappings[NODE_ENTITY_NAME][node_uuid]
for node_uuid in groupnodes
]
qb_nodes = QueryBuilder().append(
Node, filters={'id': {
'in': nodes_to_store
}})
# Adding nodes to group avoiding the SQLA ORM to increase speed
nodes_to_add = [n[0].backend_entity for n in qb_nodes.all()]
group_.backend_entity.add_nodes(nodes_to_add, skip_orm=True)
def get_or_create_group(self):
"""Return the current `AutoGroup`, or create one if None has been set yet.
This function implements a somewhat complex logic that is however needed
to make sure that, even if `verdi run` is called at the same time multiple
times, e.g. in a for loop in bash, there is never the risk that two ``verdi run``
Unix processes try to create the same group, with the same label, ending
up in a crash of the code (see PR #3650).
Here, instead, we make sure that if this concurrency issue happens,
one of the two will get a IntegrityError from the DB, and then recover
trying to create a group with a different label (with a numeric suffix appended),
until it manages to create it.
"""
from aiida.orm import QueryBuilder
# When this function is called, if it is the first time, just generate
# a new group name (later on, after this ``if`` block`).
# In that case, we will later cache in ``self._group_label`` the group label,
# So the group with the same name can be returned quickly in future
# calls of this method.
if self._group_label is not None:
builder = QueryBuilder().append(
AutoGroup, filters={'label': self._group_label})
results = [res[0] for res in builder.iterall()]
if results:
# If it is not empty, it should have only one result due to the uniqueness constraints
assert len(
results
) == 1, 'I got more than one autogroup with the same label!'
return results[0]
# There are no results: probably the group has been deleted.
# I continue as if it was not cached
self._group_label = None
label_prefix = self.get_group_label_prefix()
# Try to do a preliminary QB query to avoid to do too many try/except
# if many of the prefix_NUMBER groups already exist
queryb = QueryBuilder().append(
AutoGroup,
filters={
'or': [{
'label': {
'==': label_prefix
}
}, {
'label': {
'like': escape_for_sql_like(label_prefix + '_') + '%'
}
}]
},
project='label')
existing_group_labels = [
res[0][len(label_prefix):] for res in queryb.all()
]
existing_group_ints = []
for label in existing_group_labels:
if label == '':
# This is just the prefix without name - corresponds to counter = 0
existing_group_ints.append(0)
elif label.startswith('_'):
try:
existing_group_ints.append(int(label[1:]))
except ValueError:
# It's not an integer, so it will never collide - just ignore it
pass
if not existing_group_ints:
counter = 0
else:
counter = max(existing_group_ints) + 1
while True:
try:
label = label_prefix if counter == 0 else '{}_{}'.format(
label_prefix, counter)
group = AutoGroup(label=label).store()
self._group_label = group.label
except exceptions.IntegrityError:
counter += 1
else:
break
return group
def _retrieve_linked_nodes_query(current_node, input_type, output_type,
direction, link_type_value):
"""Helper function for :py:func:`~aiida.tools.importexport.dbexport.utils.retrieve_linked_nodes`
A general :py:class:`~aiida.orm.querybuilder.QueryBuilder` query, retrieving linked Nodes and returning link
information and the found Nodes.
:param current_node: The current Node's PK.
:type current_node: int
:param input_type: Source Node class for Link
:type input_type: :py:class:`~aiida.orm.nodes.data.data.Data`,
:py:class:`~aiida.orm.nodes.process.process.ProcessNode`.
:param output_type: Target Node class for Link
:type output_type: :py:class:`~aiida.orm.nodes.data.data.Data`,
:py:class:`~aiida.orm.nodes.process.process.ProcessNode`.
:param direction: Link direction, must be either ``'forward'`` or ``'backward'``.
:type direction: str
:param link_type_value: A :py:class:`~aiida.common.links.LinkType` value, e.g. ``LinkType.RETURN.value``.
:type link_type_value: str
:return: Dictionary of link information to be used for the export archive and set of found Nodes.
"""
found_nodes = set()
links_uuid_dict = {}
filters_input = {}
filters_output = {}
if direction == 'forward':
filters_input['id'] = current_node
elif direction == 'backward':
filters_output['id'] = current_node
else:
raise exceptions.ExportValidationError(
'direction must be either "forward" or "backward"')
builder = QueryBuilder()
builder.append(input_type,
project=['uuid', 'id'],
tag='input',
filters=filters_input)
builder.append(output_type,
project=['uuid', 'id'],
with_incoming='input',
filters=filters_output,
edge_filters={'type': link_type_value},
edge_project=['label', 'type'])
for input_uuid, input_pk, output_uuid, output_pk, link_label, link_type in builder.iterall(
):
links_uuid_entry = {
'input': str(input_uuid),
'output': str(output_uuid),
'label': str(link_label),
'type': str(link_type)
}
links_uuid_dict[frozenset(links_uuid_entry.items())] = links_uuid_entry
node_pk = output_pk if direction == 'forward' else input_pk
found_nodes.add(node_pk)
return links_uuid_dict, found_nodes
def delete_nodes(pks, verbosity=0, dry_run=False, force=False, **kwargs):
"""Delete nodes by a list of pks.
This command will delete not only the specified nodes, but also the ones that are
linked to these and should be also deleted in order to keep a consistent provenance
according to the rules explained in the concepts section of the documentation.
In summary:
1. If a DATA node is deleted, any process nodes linked to it will also be deleted.
2. If a CALC node is deleted, any incoming WORK node (callers) will be deleted as
well whereas any incoming DATA node (inputs) will be kept. Outgoing DATA nodes
(outputs) will be deleted by default but this can be disabled.
3. If a WORK node is deleted, any incoming WORK node (callers) will be deleted as
well, but all DATA nodes will be kept. Outgoing WORK or CALC nodes will be kept by
default, but deletion of either of both kind of connected nodes can be enabled.
These rules are 'recursive', so if a CALC node is deleted, then its output DATA
nodes will be deleted as well, and then any CALC node that may have those as
inputs, and so on.
:param pks: a list of the PKs of the nodes to delete
:param bool force: do not ask for confirmation to delete nodes.
:param int verbosity: 0 prints nothing,
1 prints just sums and total,
2 prints individual nodes.
:param kwargs: graph traversal rules. See :const:`aiida.common.links.GraphTraversalRules` what rule names
are toggleable and what the defaults are.
:param bool dry_run:
Just perform a dry run and do not delete anything. Print statistics according
to the verbosity level set.
:param bool force:
Do not ask for confirmation to delete nodes.
"""
# pylint: disable=too-many-arguments,too-many-branches,too-many-locals,too-many-statements
from aiida.backends.utils import delete_nodes_and_connections
from aiida.common import exceptions
from aiida.orm import Node, QueryBuilder, load_node
from aiida.tools.graph.graph_traversers import get_nodes_delete
starting_pks = []
for pk in pks:
try:
load_node(pk)
except exceptions.NotExistent:
echo.echo_warning(
f'warning: node with pk<{pk}> does not exist, skipping')
else:
starting_pks.append(pk)
# An empty set might be problematic for the queries done below.
if not starting_pks:
if verbosity:
echo.echo('Nothing to delete')
return
pks_set_to_delete = get_nodes_delete(starting_pks, **kwargs)['nodes']
if verbosity > 0:
echo.echo('I {} delete {} node{}'.format(
'would' if dry_run else 'will', len(pks_set_to_delete),
's' if len(pks_set_to_delete) > 1 else ''))
if verbosity > 1:
builder = QueryBuilder().append(
Node,
filters={'id': {
'in': pks_set_to_delete
}},
project=('uuid', 'id', 'node_type', 'label'))
echo.echo(f"The nodes I {'would' if dry_run else 'will'} delete:")
for uuid, pk, type_string, label in builder.iterall():
try:
short_type_string = type_string.split('.')[-2]
except IndexError:
short_type_string = type_string
echo.echo(f' {uuid} {pk} {short_type_string} {label}')
if dry_run:
if verbosity > 0:
echo.echo(
'\nThis was a dry run, exiting without deleting anything')
return
# Asking for user confirmation here
if force:
pass
else:
echo.echo_warning(
f'YOU ARE ABOUT TO DELETE {len(pks_set_to_delete)} NODES! THIS CANNOT BE UNDONE!'
)
if not click.confirm('Shall I continue?'):
echo.echo('Exiting without deleting')
return
# Recover the list of folders to delete before actually deleting the nodes. I will delete the folders only later,
# so that if there is a problem during the deletion of the nodes in the DB, I don't delete the folders
repositories = [load_node(pk)._repository for pk in pks_set_to_delete] # pylint: disable=protected-access
if verbosity > 0:
echo.echo('Starting node deletion...')
delete_nodes_and_connections(pks_set_to_delete)
if verbosity > 0:
echo.echo(
'Nodes deleted from database, deleting files from the repository now...'
)
# If we are here, we managed to delete the entries from the DB.
# I can now delete the folders
for repository in repositories:
repository.erase(force=True)
if verbosity > 0:
echo.echo('Deletion completed.')
def search(self, _=None):
"""Launch the search of structures in AiiDA database."""
self.preprocess()
qbuild = QueryBuilder()
# If the date range is valid, use it for the search
try:
start_date = datetime.datetime.strptime(
self.start_date_widget.value, '%Y-%m-%d')
end_date = datetime.datetime.strptime(
self.end_date_widget.value,
'%Y-%m-%d') + datetime.timedelta(hours=24)
# Otherwise revert to the standard (i.e. last 7 days)
except ValueError:
start_date = datetime.datetime.now() - datetime.timedelta(days=7)
end_date = datetime.datetime.now() + datetime.timedelta(hours=24)
self.start_date_widget.value = start_date.strftime('%Y-%m-%d')
self.end_date_widget.value = end_date.strftime('%Y-%m-%d')
filters = {}
filters['ctime'] = {'and': [{'>': start_date}, {'<=': end_date}]}
if self.mode.value == "uploaded":
qbuild2 = QueryBuilder().append(self.query_structure_type,
project=["id"],
tag='structures').append(
Node,
with_outgoing='structures')
processed_nodes = [n[0] for n in qbuild2.all()]
if processed_nodes:
filters['id'] = {"!in": processed_nodes}
qbuild.append(self.query_structure_type, filters=filters)
elif self.mode.value == "calculated":
if self.drop_label.value == 'All':
qbuild.append((CalcJobNode, WorkChainNode),
tag='calcjobworkchain')
else:
qbuild.append((CalcJobNode, WorkChainNode),
filters={'label': self.drop_label.value},
tag='calcjobworkchain')
qbuild.append(self.query_structure_type,
with_incoming='calcjobworkchain',
filters=filters)
elif self.mode.value == "edited":
qbuild.append(CalcFunctionNode)
qbuild.append(self.query_structure_type,
with_incoming=CalcFunctionNode,
filters=filters)
elif self.mode.value == "all":
qbuild.append(self.query_structure_type, filters=filters)
qbuild.order_by({self.query_structure_type: {'ctime': 'desc'}})
matches = {n[0] for n in qbuild.iterall()}
matches = sorted(matches, reverse=True, key=lambda n: n.ctime)
options = OrderedDict()
options["Select a Structure ({} found)".format(len(matches))] = False
for mch in matches:
label = "PK: {}".format(mch.id)
label += " | " + mch.ctime.strftime("%Y-%m-%d %H:%M")
label += " | " + mch.get_extra("formula")
label += " | " + mch.node_type.split('.')[-2]
label += " | " + mch.label
label += " | " + mch.description
options[label] = mch
self.results.options = options
