def get_upf_groups(cls, filter_elements=None, user=None):
"""Return all names of groups of type UpfFamily, possibly with some filters.
:param filter_elements: A string or a list of strings.
If present, returns only the groups that contains one UPF for every element present in the list. The default
is `None`, meaning that all families are returned.
:param user: if None (default), return the groups for all users.
If defined, it should be either a `User` instance or the user email.
:return: list of `Group` entities of type UPF.
"""
from aiida.orm import UpfFamily
from aiida.orm import QueryBuilder
from aiida.orm import User
builder = QueryBuilder()
builder.append(UpfFamily, tag='group', project='*')
if user:
builder.append(User,
filters={'email': {
'==': user
}},
with_group='group')
if isinstance(filter_elements, str):
filter_elements = [filter_elements]
if filter_elements is not None:
builder.append(
UpfData,
filters={'attributes.element': {
'in': filter_elements
}},
with_group='group')
builder.order_by({UpfFamily: {'id': 'asc'}})
return [group for group, in builder.all()]
def get_data_aiida(quantitites):
"""Query the AiiDA database for a list of quantities (other appl results).
Return list of entries like [mat_id, mat_name, mat_class, quantity-1, quantity-2, ..., quantity-n].
"""
qb = QueryBuilder()
qb.append(Group,
filters={'label': {
'like': GROUP_DIR + "%"
}},
tag='g',
project=[
'extras.mat_id', 'extras.name_conventional',
'extras.class_material'
])
for q in quantitites:
qb.append(Dict,
project=['attributes.{}'.format(q['key'])],
filters={'extras.{}'.format(TAG_KEY): q['dict']},
with_group='g')
return qb.all()
def node_show(nodes, print_groups):
"""Show generic information on one or more nodes."""
from aiida.cmdline.utils.common import get_node_info
for node in nodes:
# pylint: disable=fixme
# TODO: Add a check here on the node type, otherwise it might try to access
# attributes such as code which are not necessarily there
echo.echo(get_node_info(node))
if print_groups:
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.groups import Group
from aiida.orm import Node # pylint: disable=redefined-outer-name
# pylint: disable=invalid-name
qb = QueryBuilder()
qb.append(Node, tag='node', filters={'id': {'==': node.pk}})
qb.append(Group,
tag='groups',
with_node='node',
project=['id', 'label', 'type_string'])
echo.echo('#### GROUPS:')
if qb.count() == 0:
echo.echo(f'Node {node.pk} does not belong to any group')
else:
echo.echo(f'Node {node.pk} belongs to the following groups:')
res = qb.iterdict()
table = [(gr['groups']['id'], gr['groups']['label'],
gr['groups']['type_string']) for gr in res]
table.sort()
echo.echo(
tabulate.tabulate(table,
headers=['PK', 'Label', 'Group type']))
def querry_for_ref_structure(element_string):
"""
This methods finds StructureData nodes with the following extras:
extra.type = 'bulk', # Should be done by looking at pbc, but I could not
get querry to work.
extra.specific = 'reference',
'extra.elemental' = True,
extra.structure = element_string
param: element_string: string of an element
return: the latest StructureData node that was found
"""
#query db
q = QueryBuilder()
q.append(StructureData,
filters={
'extras.type': {
'==': 'bulk'
},
'extras.specification': {
'==': 'reference'
},
'extras.elemental': {
'==': True
},
'extras.element': {
'==': element_string
}
})
q.order_by({StructureData: 'ctime'}) #always use the most recent
structures = q.all()
if structures:
return structures[-1][0]
else:
return None
def _build_query(order_by=None, limit=None, past_days=None):
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.calculation import Calculation
import aiida.utils.timezone as timezone
import datetime
from aiida.orm.mixins import SealableMixin
_SEALED_ATTRIBUTE_KEY = 'attributes.{}'.format(SealableMixin.SEALED_KEY)
# The things that we want to get out
calculation_projections = \
['id', 'ctime', 'type', _SEALED_ATTRIBUTE_KEY]
now = timezone.now()
# The things to filter by
calculation_filters = {}
if past_days is not None:
n_days_ago = now - datetime.timedelta(days=past_days)
calculation_filters['ctime'] = {'>': n_days_ago}
qb = QueryBuilder()
# ORDER
if order_by is not None:
qb.order_by({'calculation': [order_by]})
# LIMIT
if limit is not None:
qb.limit(limit)
# Build the quiery
qb.append(cls=Calculation,
filters=calculation_filters,
project=calculation_projections,
tag='calculation')
return qb.iterdict()
def mock_vasp(aiida_env, localhost):
"""Points to a mock-up of a VASP executable."""
from aiida.orm import Code
from aiida.orm.querybuilder import QueryBuilder
query_builder = QueryBuilder()
query_builder.append(Code, tag='code')
query_builder.add_filter('code', {'label': {'==': 'mock-vasp'}})
query_results = query_builder.all()
if query_results:
code = query_results[0][0]
else:
os_env = os.environ.copy()
if not localhost.pk:
localhost.store()
mock_vasp_path = sp.check_output(['which', 'mock-vasp'],
env=os_env).strip()
code = Code()
code.label = 'mock-vasp'
code.description = 'Mock VASP for tests'
code.set_remote_computer_exec((localhost, mock_vasp_path))
code.set_input_plugin_name('vasp.vasp')
aiidapath = py_path.local(aiida_env.root_dir).join('.aiida')
code.set_prepend_text('export AIIDA_PATH={}'.format(aiidapath))
return code
def _retrieve_basis_sets(files, stop_if_existing):
""" get existing basis sets or create if not
: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:
"""
from aiida.orm.querybuilder import QueryBuilder
basis_and_created = []
for f in files:
_, content = parse_basis(f)
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(f,
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(f))
existing_basis = existing_basis[0]
basis_and_created.append((existing_basis, False))
return basis_and_created
def get_data_aiida(inp_list):
"""Query the AiiDA database: find info in the README.
Note: Removing the group filter speeds up queries on 70k COFs from ~10s to ~1s.
"""
qb = QueryBuilder()
qb.append(Group, filters={'label': {'like': 'group_%'}}, tag='group')
qb.append(CifData,
with_group='group',
filters={'extras.group_tag': 'orig_cif'},
project=['uuid', 'label'])
for inp in inp_list:
if 'henry_coefficient_average' in inp:
proj = 'henry_coefficient_average' # take out _co2, _n2, _ht
else:
proj = inp
qb.append(Dict,
with_group='group',
filters={'extras.group_tag': get_tag[inp]},
project=['attributes.{}'.format(proj)])
return qb.all()
def preprocess_spm_calcs(
workchain_list=['STMWorkChain', 'PdosWorkChain', 'AfmWorkChain']):
qb = QueryBuilder()
qb.append(WorkCalculation,
filters={
'attributes._process_label': {
'in': workchain_list
},
'or': [
{
'extras': {
'!has_key': 'preprocess_version'
}
},
{
'extras.preprocess_version': {
'<': PREPROCESS_VERSION
}
},
],
})
qb.order_by({WorkCalculation: {'ctime': 'asc'}})
for m in qb.all():
n = m[0]
## ---------------------------------------------------------------
## calculation not finished
if not n.is_sealed:
print("Skipping underway workchain PK %d" % n.pk)
continue
calc_states = [out.get_state() for out in n.get_outputs()]
if 'WITHSCHEDULER' in calc_states:
print("Skipping underway workchain PK %d" % n.pk)
continue
## ---------------------------------------------------------------
if 'obsolete' not in n.get_extras():
n.set_extra('obsolete', False)
if n.get_extra('obsolete'):
continue
wc_name = n.get_attrs()['_process_label']
try:
if not all(
[calc.get_state() == 'FINISHED' for calc in n.get_outputs()]):
raise (Exception("Not all calculations are 'FINISHED'"))
preprocess_one(n)
print("Preprocessed PK %d (%s)" % (n.pk, wc_name))
n.set_extra('preprocess_successful', True)
n.set_extra('preprocess_version', PREPROCESS_VERSION)
if 'preprocess_error' in n.get_extras():
n.del_extra('preprocess_error')
except Exception as e:
n.set_extra('preprocess_successful', False)
n.set_extra('preprocess_error', str(e))
n.set_extra('preprocess_version', PREPROCESS_VERSION)
print("Failed to preprocess PK %d (%s): %s" % (n.pk, wc_name, e))
parser = argparse.ArgumentParser()
parser.add_argument("code",
help="code and machine where you would like to run")
parser.add_argument("json_hpc", help="json file with HPC parameters")
parser.add_argument("json_pw", help="json file with PW parameters")
args = parser.parse_args()
StructureData = DataFactory('structure')
UpfData = DataFactory('upf')
ParameterData = DataFactory('parameter')
KpointsData = DataFactory('array.kpoints')
with open(args.json_hpc) as data_file:
json_hpc = json.load(data_file)
qb = QueryBuilder()
qb.append(JobCalculation, tag="mycalculation", project=["*"])
qb.append(Group,
filters={"name": json_hpc["query_group"]},
group_of="mycalculation")
calcs_list = qb.all()
pseudo_family = json_hpc['pseudo']
structures_wf = []
kpoints_wf = []
pw_parameters_wf = []
hpc_workflow_params = {}
keys = []
count = 0
#for i in calcs_list[0:1]:
def test_querybuilder_classifications(self):
"""
This tests the classifications of the QueryBuilder u. the django backend.
"""
from aiida.backends.querybuild.dummy_model import (
DbNode, DbUser, DbComputer,
DbGroup,
)
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.utils import (DataFactory, CalculationFactory)
from aiida.orm.data.structure import StructureData
from aiida.orm.implementation.django.node import Node
from aiida.orm import Group, User, Node, Computer, Data, Calculation
from aiida.common.exceptions import InputValidationError
qb = QueryBuilder()
with self.assertRaises(InputValidationError):
qb._get_ormclass(None, 'data')
with self.assertRaises(InputValidationError):
qb._get_ormclass(None, 'data.Data')
with self.assertRaises(InputValidationError):
qb._get_ormclass(None, '.')
for cls, clstype, query_type_string in (
qb._get_ormclass(StructureData, None),
qb._get_ormclass(None, 'data.structure.StructureData.'),
):
self.assertEqual(clstype, 'data.structure.StructureData.')
self.assertTrue(issubclass(cls, DbNode))
self.assertEqual(clstype, 'data.structure.StructureData.')
self.assertEqual(query_type_string,
StructureData._query_type_string)
for cls, clstype, query_type_string in (
qb._get_ormclass(Node, None),
qb._get_ormclass(DbNode, None),
qb._get_ormclass(None, '')
):
self.assertEqual(clstype, Node._plugin_type_string)
self.assertEqual(query_type_string, Node._query_type_string)
self.assertTrue(issubclass(cls, DbNode))
for cls, clstype, query_type_string in (
qb._get_ormclass(DbGroup, None),
qb._get_ormclass(Group, None),
qb._get_ormclass(None, 'group'),
qb._get_ormclass(None, 'Group'),
):
self.assertEqual(clstype, 'group')
self.assertEqual(query_type_string, None)
self.assertTrue(issubclass(cls, DbGroup))
for cls, clstype, query_type_string in (
qb._get_ormclass(DbUser, None),
qb._get_ormclass(DbUser, None),
qb._get_ormclass(None, "user"),
qb._get_ormclass(None, "User"),
):
self.assertEqual(clstype, 'user')
self.assertEqual(query_type_string, None)
self.assertTrue(issubclass(cls, DbUser))
for cls, clstype, query_type_string in (
qb._get_ormclass(DbComputer, None),
qb._get_ormclass(Computer, None),
qb._get_ormclass(None, 'computer'),
qb._get_ormclass(None, 'Computer'),
):
self.assertEqual(clstype, 'computer')
self.assertEqual(query_type_string, None)
self.assertTrue(issubclass(cls, DbComputer))
for cls, clstype, query_type_string in (
qb._get_ormclass(Data, None),
qb._get_ormclass(None, 'data.Data.'),
):
self.assertEqual(clstype, Data._plugin_type_string)
self.assertEqual(query_type_string, Data._query_type_string)
self.assertTrue(issubclass(cls, DbNode))
def test_simple_query_django_2(self):
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm import Node
from datetime import datetime
from aiida.backends.querybuild.dummy_model import (
DbNode, DbLink, DbAttribute, session
)
n0 = DbNode(
label='hello',
type='',
description='', user_id=1,
)
n1 = DbNode(
label='foo',
type='',
description='I am FoO', user_id=2,
)
n2 = DbNode(
label='bar',
type='',
description='I am BaR', user_id=3,
)
DbAttribute(
key='foo',
datatype='txt',
tval='bar',
dbnode=n0
)
l1 = DbLink(input=n0, output=n1, label='random_1', type='')
l2 = DbLink(input=n1, output=n2, label='random_2', type='')
session.add_all([n0, n1, n2, l1, l2])
qb1 = QueryBuilder()
qb1.append(
DbNode,
filters={
'label': 'hello',
}
)
self.assertEqual(len(list(qb1.all())), 1)
qh = {
'path': [
{
'cls': Node,
'tag': 'n1'
},
{
'cls': Node,
'tag': 'n2',
'output_of': 'n1'
}
],
'filters': {
'n1': {
'label': {'ilike': '%foO%'},
},
'n2': {
'label': {'ilike': 'bar%'},
}
},
'project': {
'n1': ['id', 'uuid', 'ctime', 'label'],
'n2': ['id', 'description', 'label'],
}
}
qb2 = QueryBuilder(**qh)
resdict = qb2.dict()
self.assertEqual(len(resdict), 1)
resdict = resdict[0]
self.assertTrue(isinstance(resdict['n1']['ctime'], datetime))
self.assertEqual(resdict['n2']['label'], 'bar')
qh = {
'path': [
{
'cls': Node,
'label': 'n1'
},
{
'cls': Node,
'label': 'n2',
'output_of': 'n1'
}
],
'filters': {
'n1--n2': {'label': {'like': '%_2'}}
}
}
qb = QueryBuilder(**qh)
self.assertEqual(qb.count(), 1)
# Test the hashing:
query1 = qb.get_query()
qb.add_filter('n2', {'label': 'nonexistentlabel'})
self.assertEqual(qb.count(), 0)
query2 = qb.get_query()
query3 = qb.get_query()
self.assertTrue(id(query1) != id(query2))
self.assertTrue(id(query2) == id(query3))
def create(outfile, computers, groups, nodes, group_names, input_forward,
create_reversed, return_reversed, call_reversed, overwrite,
archive_format):
"""
Export nodes and groups of nodes to an archive file for backup or sharing purposes
"""
import sys
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
# TODO: Replace with aiida.cmdline.utils.decorators.with_dbenv decocator
# TODO: when we merge to develop
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm import Group, Node, Computer
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.importexport import export, export_zip
node_id_set = set(nodes)
group_dict = dict()
if group_names:
qb = QueryBuilder()
qb.append(Group, tag='group', project=['*'], filters={'name': {'in': group_names}})
qb.append(Node, tag='node', member_of='group', project=['id'])
res = qb.dict()
group_dict.update(
{group['group']['*'].id: group['group']['*'] for group in res})
node_id_set.update([node['node']['id'] for node in res])
if groups:
qb = QueryBuilder()
qb.append(Group, tag='group', project=['*'], filters={'id': {'in': groups}})
qb.append(Node, tag='node', member_of='group', project=['id'])
res = qb.dict()
group_dict.update(
{group['group']['*'].id: group['group']['*'] for group in res})
node_id_set.update([node['node']['id'] for node in res])
groups_list = group_dict.values()
# Getting the nodes that correspond to the ids that were found above
if len(node_id_set) > 0:
qb = QueryBuilder()
qb.append(Node, tag='node', project=['*'], filters={'id': {'in': node_id_set}})
node_list = [node[0] for node in qb.all()]
else:
node_list = list()
# Check if any of the nodes wasn't found in the database.
missing_nodes = node_id_set.difference(node.id for node in node_list)
for node_id in missing_nodes:
print >> sys.stderr, ('WARNING! Node with pk={} not found, skipping'.format(node_id))
if computers:
qb = QueryBuilder()
qb.append(Computer, tag='comp', project=['*'], filters={'id': {'in': set(computers)}})
computer_list = [computer[0] for computer in qb.all()]
missing_computers = set(computers).difference(computer.id for computer in computer_list)
for computer_id in missing_computers:
print >> sys.stderr, ('WARNING! Computer with pk={} not found, skipping'.format(computer_id))
else:
computer_list = []
what_list = node_list + computer_list + groups_list
additional_kwargs = dict()
if archive_format == 'zip':
export_function = export_zip
additional_kwargs.update({'use_compression': True})
elif archive_format == 'zip-uncompressed':
export_function = export_zip
additional_kwargs.update({'use_compression': False})
elif archive_format == 'tar.gz':
export_function = export
else:
print >> sys.stderr, 'invalid --archive-format value {}'.format(
archive_format)
sys.exit(1)
try:
export_function(
what=what_list, input_forward=input_forward,
create_reversed=create_reversed,
return_reversed=return_reversed,
call_reversed=call_reversed, outfile=outfile,
overwrite=overwrite, **additional_kwargs
)
except IOError as e:
print >> sys.stderr, 'IOError: {}'.format(e.message)
sys.exit(1)
def get_io_tree(self, nodeId, maxDepth=None):
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.node import Node
def addNodes(nodeId, maxDepth, nodes, addedNodes, addedEdges,
edgeType):
qb = QueryBuilder()
qb.append(Node, tag="main", filters={"id": {"==": nodeId}})
if edgeType == "ancestors":
qb.append(Node,
tag=edgeType,
project=['id', 'type'],
edge_project=['path', 'depth'],
ancestor_of_beta='main',
edge_filters={'depth': {
'<=': maxDepth
}})
elif edgeType == "desc":
qb.append(Node,
tag=edgeType,
project=['id', 'type'],
edge_project=['path', 'depth'],
descendant_of_beta='main',
edge_filters={'depth': {
'<=': maxDepth
}})
if (qb.count() > 0):
qbResults = qb.get_results_dict()
for resultDict in qbResults:
if resultDict[edgeType]["id"] not in addedNodes:
nodes.append({
"id":
len(addedNodes),
"nodeid":
resultDict[edgeType]["id"],
"nodetype":
resultDict[edgeType]["type"],
"group":
edgeType + "-" +
str(resultDict["main--" + edgeType]["depth"])
})
addedNodes.append(resultDict[edgeType]["id"])
path = resultDict["main--" + edgeType]["path"]
if edgeType == "ancestors":
startEdge = path[0]
endEdge = path[1]
elif edgeType == "desc":
startEdge = path[-2]
endEdge = path[-1]
if startEdge not in addedEdges.keys():
addedEdges[startEdge] = [endEdge]
elif endEdge not in addedEdges[startEdge]:
addedEdges[startEdge].append(endEdge)
return nodes, addedNodes, addedEdges
def addEdges(edges, addedNodes, addedEdges):
for fromNodeId in addedEdges.keys():
fromNodeIdIndex = addedNodes.index(fromNodeId)
for toNodeId in addedEdges[fromNodeId]:
toNodeIdIndex = addedNodes.index(toNodeId)
edges.append({
"from": fromNodeIdIndex,
"to": toNodeIdIndex,
"arrows": "to",
"color": {
"inherit": 'from'
}
})
return edges
nodes = []
edges = []
addedNodes = []
addedEdges = {}
if maxDepth is None:
from aiida.restapi.common.config import MAX_TREE_DEPTH
maxDepth = MAX_TREE_DEPTH
qb = QueryBuilder()
qb.append(Node,
tag="main",
project=["id", "type"],
filters={"id": {
"==": nodeId
}})
if qb.count() > 0:
mainNode = qb.first()
nodes.append({
"id": 0,
"nodeid": mainNode[0],
"nodetype": mainNode[1],
"group": "mainNode"
})
addedNodes.append(mainNode[0])
# get all descendents
nodes, addedNodes, addedEdges = addNodes(nodeId, maxDepth, nodes,
addedNodes, addedEdges,
"ancestors")
nodes, addedNodes, addedEdges = addNodes(nodeId, maxDepth, nodes,
addedNodes, addedEdges,
"desc")
edges = addEdges(edges, addedNodes, addedEdges)
return {"nodes": nodes, "edges": edges}
def load_node(node_id=None,
pk=None,
uuid=None,
parent_class=None,
query_with_dashes=True):
"""
Return an AiiDA node given PK or UUID.
:param node_id: PK (integer) or UUID (string) or a node
:param pk: PK of a node
:param uuid: UUID of a node, or the beginning of the uuid
:param parent_class: if specified, checks whether the node loaded is a
subclass of parent_class
:param bool query_with_dashes: Specific if uuid is passed, allows to put the uuid in the correct form.
Default=True
:return: an AiiDA node
:raise InputValidationError: if none or more than one of parameters is supplied
:raise TypeError: I the wrong types are provided
:raise NotExistent: if no matching Node is found.
:raise MultipleObjectsError: If more than one Node was fouuund
"""
from aiida.common.exceptions import NotExistent, MultipleObjectsError, NotExistent, InputValidationError
# This must be done inside here, because at import time the profile
# must have been already loaded. If you put it at the module level,
# the implementation is frozen to the default one at import time.
from aiida.orm.implementation import Node
from aiida.orm.querybuilder import QueryBuilder
# First checking if the inputs are valid:
inputs_provided = [val is not None
for val in (node_id, pk, uuid)].count(True)
if inputs_provided == 0:
raise InputValidationError(
"one of the parameters 'node_id', 'pk' and 'uuid' "
"has to be supplied")
elif inputs_provided > 1:
raise InputValidationError(
"only one of parameters 'node_id', 'pk' and 'uuid' "
"has to be supplied")
class_ = parent_class or Node
if not issubclass(class_, Node):
raise TypeError("{} is not a subclass of {}".format(class_, Node))
# The logic is as follows: If pk is specified I will look for the pk
# if UUID is specified for the uuid.
# node_id can either be string -> uuid or an integer -> pk
# Checking first if node_id specified
if node_id is not None:
if isinstance(node_id, (str, unicode)):
uuid = node_id
elif isinstance(node_id, int):
pk = node_id
else:
raise TypeError("'node_id' has to be either string, unicode or "
"integer, {} given".format(type(node_id)))
#I am checking whether uuid, if supplied, is a string
if uuid is not None:
if not isinstance(uuid, (str, unicode)):
raise TypeError("'uuid' has to be string or unicode")
# or whether the pk, if provided, is an integer
elif pk is not None:
if not isinstance(pk, int):
raise TypeError("'pk' has to be an integer")
else:
# I really shouldn't get here
assert True, "Neither pk nor uuid was provided"
qb = QueryBuilder()
qb.append(class_, tag='node', project='*')
if pk:
qb.add_filter('node', {'id': pk})
elif uuid:
# Put back dashes in the right place
start_uuid = uuid.replace('-', '')
if query_with_dashes:
# Essential that this is ordered from largest to smallest!
for dash_pos in [20, 16, 12, 8]:
if len(start_uuid) > dash_pos:
start_uuid = "{}-{}".format(start_uuid[:dash_pos],
start_uuid[dash_pos:])
qb.add_filter('node', {'uuid': {'like': "{}%".format(start_uuid)}})
try:
return qb.one()[0]
except MultipleObjectsError:
raise MultipleObjectsError("More than one node found")
except NotExistent:
raise NotExistent("No node was found")
def process_dummy_data(cls):
"""
This functions prepare atomic chunks of typical responses from the
RESTapi and puts them into class attributes
"""
#TODO: Storing the different nodes as lists and accessing them
# by their list index is very fragile and a pain to debug.
# Please change this!
computer_projections = [
"id", "uuid", "name", "hostname", "transport_type",
"scheduler_type"
]
computers = QueryBuilder().append(
Computer, tag="comp", project=computer_projections).order_by({
'comp': [{
'name': {
'order': 'asc'
}
}]
}).dict()
# Cast UUID into a string (e.g. in sqlalchemy it comes as a UUID object)
computers = [_['comp'] for _ in computers]
for comp in computers:
if comp['uuid'] is not None:
comp['uuid'] = str(comp['uuid'])
cls._dummy_data["computers"] = computers
calculation_projections = ["id", "uuid", "user_id", "type"]
calculations = QueryBuilder().append(
Calculation, tag="calc",
project=calculation_projections).order_by({
'calc': [{
'id': {
'order': 'desc'
}
}]
}).dict()
calculations = [_['calc'] for _ in calculations]
for calc in calculations:
if calc['uuid'] is not None:
calc['uuid'] = str(calc['uuid'])
cls._dummy_data["calculations"] = calculations
data_projections = ["id", "uuid", "user_id", "type"]
data_types = {
'cifdata': CifData,
'parameterdata': ParameterData,
'structuredata': StructureData,
'data': Data,
}
for label, dataclass in data_types.iteritems():
data = QueryBuilder().append(dataclass,
tag="data",
project=data_projections).order_by({
'data': [{
'id': {
'order': 'desc'
}
}]
}).dict()
data = [_['data'] for _ in data]
for datum in data:
if datum['uuid'] is not None:
datum['uuid'] = str(datum['uuid'])
cls._dummy_data[label] = data
def upload_psf_family(folder,
group_name,
group_description,
stop_if_existing=True):
"""
Upload a set of PSF files in a given group.
:param folder: a path containing all PSF files to be added.
Only files ending in .PSF (case-insensitive) are considered.
:param group_name: the name of the group to create. If it exists and is
non-empty, a UniquenessError is raised.
:param group_description: a string to be set as the group description.
Overwrites previous descriptions, if the group was existing.
: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 PsfData node to the group.
"""
import os
import aiida.common
from aiida.common import aiidalogger
from aiida.orm import Group
from aiida.common.exceptions import UniquenessError, NotExistent
from aiida.backends.utils import get_automatic_user
from aiida.orm.querybuilder import QueryBuilder
if not os.path.isdir(folder):
raise ValueError("folder must be a directory")
# only files, and only those ending with .psf or .PSF;
# go to the real file if it is a symlink
files = [
os.path.realpath(os.path.join(folder, i)) for i in os.listdir(folder)
if os.path.isfile(os.path.join(folder, i))
and i.lower().endswith('.psf')
]
nfiles = len(files)
try:
group = Group.get(name=group_name, type_string=PSFGROUP_TYPE)
group_created = False
except NotExistent:
group = Group(name=group_name,
type_string=PSFGROUP_TYPE,
user=get_automatic_user())
group_created = True
if group.user != get_automatic_user():
raise UniquenessError("There is already a PsfFamily group with name {}"
", but it belongs to user {}, therefore you "
"cannot modify it".format(
group_name, group.user.email))
# Always update description, even if the group already existed
group.description = group_description
# NOTE: GROUP SAVED ONLY AFTER CHECKS OF UNICITY
pseudo_and_created = []
for f in files:
md5sum = aiida.common.utils.md5_file(f)
qb = QueryBuilder()
qb.append(PsfData, filters={'attributes.md5': {'==': md5sum}})
existing_psf = qb.first()
#existing_psf = PsfData.query(dbattributes__key="md5",
# dbattributes__tval = md5sum)
if existing_psf is None:
# return the psfdata instances, not stored
pseudo, created = PsfData.get_or_create(f,
use_first=True,
store_psf=False)
# to check whether only one psf per element exists
# NOTE: actually, created has the meaning of "to_be_created"
pseudo_and_created.append((pseudo, created))
else:
if stop_if_existing:
raise ValueError("A PSF with identical MD5 to "
" {} cannot be added with stop_if_existing"
"".format(f))
existing_psf = existing_psf[0]
pseudo_and_created.append((existing_psf, False))
# check whether pseudo are unique per element
elements = [(i[0].element, i[0].md5sum) for i in pseudo_and_created]
# If group already exists, check also that I am not inserting more than
# once the same element
if not group_created:
for aiida_n in group.nodes:
# Skip non-pseudos
if not isinstance(aiida_n, PsfData):
continue
elements.append((aiida_n.element, aiida_n.md5sum))
elements = set(elements) # Discard elements with the same MD5, that would
# not be stored twice
elements_names = [e[0] for e in elements]
if not len(elements_names) == len(set(elements_names)):
duplicates = set(
[x for x in elements_names if elements_names.count(x) > 1])
duplicates_string = ", ".join(i for i in duplicates)
raise UniquenessError("More than one PSF found for the elements: " +
duplicates_string + ".")
# At this point, save the group, if still unstored
if group_created:
group.store()
# save the psf in the database, and add them to group
for pseudo, created in pseudo_and_created:
if created:
pseudo.store()
aiidalogger.debug("New node {} created for file {}".format(
pseudo.uuid, pseudo.filename))
else:
aiidalogger.debug("Reusing node {} for file {}".format(
pseudo.uuid, pseudo.filename))
# Add elements to the group all togetehr
group.add_nodes(pseudo for pseudo, created in pseudo_and_created)
nuploaded = len([_ for _, created in pseudo_and_created if created])
return nfiles, nuploaded
def test_different_computer_same_name_import(self):
"""
This test checks that if there is a name collision, the imported
computers are renamed accordingly.
"""
import os
import shutil
import tempfile
from aiida.orm.importexport import export
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.computer import Computer
from aiida.orm.calculation.job import JobCalculation
from aiida.orm.importexport import COMP_DUPL_SUFFIX
# Creating a folder for the import/export files
export_file_tmp_folder = tempfile.mkdtemp()
unpack_tmp_folder = tempfile.mkdtemp()
try:
# Set the computer name
comp1_name = "localhost_1"
self.computer.set_name(comp1_name)
# Store a calculation
calc1_label = "calc1"
calc1 = JobCalculation()
calc1.set_computer(self.computer)
calc1.set_resources({
"num_machines": 1,
"num_mpiprocs_per_machine": 1
})
calc1.label = calc1_label
calc1.store()
calc1._set_state(u'RETRIEVING')
# Export the first job calculation
filename1 = os.path.join(export_file_tmp_folder, "export1.tar.gz")
export([calc1.dbnode], outfile=filename1, silent=True)
# Reset the database
self.clean_db()
self.insert_data()
# Set the computer name to the same name as before
self.computer.set_name(comp1_name)
# Store a second calculation
calc2_label = "calc2"
calc2 = JobCalculation()
calc2.set_computer(self.computer)
calc2.set_resources({
"num_machines": 2,
"num_mpiprocs_per_machine": 2
})
calc2.label = calc2_label
calc2.store()
calc2._set_state(u'RETRIEVING')
# Export the second job calculation
filename2 = os.path.join(export_file_tmp_folder, "export2.tar.gz")
export([calc2.dbnode], outfile=filename2, silent=True)
# Reset the database
self.clean_db()
self.insert_data()
# Set the computer name to the same name as before
self.computer.set_name(comp1_name)
# Store a third calculation
calc3_label = "calc3"
calc3 = JobCalculation()
calc3.set_computer(self.computer)
calc3.set_resources({
"num_machines": 2,
"num_mpiprocs_per_machine": 2
})
calc3.label = calc3_label
calc3.store()
calc3._set_state(u'RETRIEVING')
# Export the third job calculation
filename3 = os.path.join(export_file_tmp_folder, "export3.tar.gz")
export([calc3.dbnode], outfile=filename3, silent=True)
# Clean the local database
self.clean_db()
# Check that there are no computers
qb = QueryBuilder()
qb.append(Computer, project=['*'])
self.assertEqual(
qb.count(), 0, "There should not be any computers"
"in the database at this point.")
# Check that there are no calculations
qb = QueryBuilder()
qb.append(JobCalculation, project=['*'])
self.assertEqual(
qb.count(), 0, "There should not be any "
"calculations in the database at "
"this point.")
# Import all the calculations
import_data(filename1, silent=True)
import_data(filename2, silent=True)
import_data(filename3, silent=True)
# Retrieve the calculation-computer pairs
qb = QueryBuilder()
qb.append(JobCalculation, project=['label'], tag='jcalc')
qb.append(Computer, project=['name'], computer_of='jcalc')
self.assertEqual(qb.count(), 3, "Three combinations expected.")
res = qb.all()
self.assertIn([calc1_label, comp1_name], res,
"Calc-Computer combination not found.")
self.assertIn(
[calc2_label, comp1_name + COMP_DUPL_SUFFIX.format(0)], res,
"Calc-Computer combination not found.")
self.assertIn(
[calc3_label, comp1_name + COMP_DUPL_SUFFIX.format(1)], res,
"Calc-Computer combination not found.")
finally:
# Deleting the created temporary folders
shutil.rmtree(export_file_tmp_folder, ignore_errors=True)
shutil.rmtree(unpack_tmp_folder, ignore_errors=True)
def test_same_computer_different_name_import(self):
"""
This test checks that if the computer is re-imported with a different
name to the same database, then the original computer will not be
renamed. It also checks that the names were correctly imported (without
any change since there is no computer name collision)
"""
import os
import shutil
import tempfile
from aiida.orm.importexport import export
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.computer import Computer
from aiida.orm.calculation.job import JobCalculation
# Creating a folder for the import/export files
export_file_tmp_folder = tempfile.mkdtemp()
unpack_tmp_folder = tempfile.mkdtemp()
try:
# Store a calculation
calc1_label = "calc1"
calc1 = JobCalculation()
calc1.set_computer(self.computer)
calc1.set_resources({
"num_machines": 1,
"num_mpiprocs_per_machine": 1
})
calc1.label = calc1_label
calc1.store()
calc1._set_state(u'RETRIEVING')
# Store locally the computer name
comp1_name = unicode(self.computer.name)
# Export the first job calculation
filename1 = os.path.join(export_file_tmp_folder, "export1.tar.gz")
export([calc1.dbnode], outfile=filename1, silent=True)
# Rename the computer
self.computer.set_name(comp1_name + "_updated")
# Store a second calculation
calc2_label = "calc2"
calc2 = JobCalculation()
calc2.set_computer(self.computer)
calc2.set_resources({
"num_machines": 2,
"num_mpiprocs_per_machine": 2
})
calc2.label = calc2_label
calc2.store()
calc2._set_state(u'RETRIEVING')
# Export the second job calculation
filename2 = os.path.join(export_file_tmp_folder, "export2.tar.gz")
export([calc2.dbnode], outfile=filename2, silent=True)
# Clean the local database
self.clean_db()
# Check that there are no computers
qb = QueryBuilder()
qb.append(Computer, project=['*'])
self.assertEqual(
qb.count(), 0, "There should not be any computers"
"in the database at this point.")
# Check that there are no calculations
qb = QueryBuilder()
qb.append(JobCalculation, project=['*'])
self.assertEqual(
qb.count(), 0, "There should not be any "
"calculations in the database at "
"this point.")
# Import the first calculation
import_data(filename1, silent=True)
# Check that the calculation computer is imported correctly.
qb = QueryBuilder()
qb.append(JobCalculation, project=['label'])
self.assertEqual(qb.count(), 1, "Only one calculation should be "
"found.")
self.assertEqual(unicode(qb.first()[0]), calc1_label,
"The calculation label is not correct.")
# Check that the referenced computer is imported correctly.
qb = QueryBuilder()
qb.append(Computer, project=['name', 'uuid', 'id'])
self.assertEqual(qb.count(), 1, "Only one computer should be "
"found.")
self.assertEqual(unicode(qb.first()[0]), comp1_name,
"The computer name is not correct.")
# Import the second calculation
import_data(filename2, silent=True)
# Check that the number of computers remains the same and its data
# did not change.
qb = QueryBuilder()
qb.append(Computer, project=['name'])
self.assertEqual(qb.count(), 1, "Only one computer should be "
"found.")
self.assertEqual(unicode(qb.first()[0]), comp1_name,
"The computer name is not correct.")
finally:
# Deleting the created temporary folders
shutil.rmtree(export_file_tmp_folder, ignore_errors=True)
shutil.rmtree(unpack_tmp_folder, ignore_errors=True)
def _add_dblink_from(self,
src,
label=None,
link_type=LinkType.UNSPECIFIED):
from aiida.backends.sqlalchemy import get_scoped_session
from aiida.orm.querybuilder import QueryBuilder
session = get_scoped_session()
if not isinstance(src, Node):
raise ValueError("src must be a Node instance")
if self.uuid == src.uuid:
raise ValueError("Cannot link to itself")
if self._to_be_stored:
raise ModificationNotAllowed(
"Cannot call the internal _add_dblink_from if the "
"destination node is not stored")
if src._to_be_stored:
raise ModificationNotAllowed(
"Cannot call the internal _add_dblink_from if the "
"source node is not stored")
# Check for cycles. This works if the transitive closure is enabled; if
# it isn't, this test will never fail, but then having a circular link
# is not meaningful but does not pose a huge threat
#
# I am linking src->self; a loop would be created if a DbPath exists
# already in the TC table from self to src
if link_type is LinkType.CREATE or link_type is LinkType.INPUT:
if QueryBuilder().append(
Node, filters={
'id': self.pk
}, tag='parent').append(
Node,
filters={
'id': src.pk
},
tag='child',
descendant_of='parent').count() > 0:
raise ValueError(
"The link you are attempting to create would generate a loop"
)
if label is None:
autolabel_idx = 1
existing_from_autolabels = session.query(DbLink.label).filter(
DbLink.output_id == self.dbnode.id, DbLink.label.like("link%"))
while "link_{}".format(autolabel_idx) in existing_from_autolabels:
autolabel_idx += 1
safety_counter = 0
while True:
safety_counter += 1
if safety_counter > 100:
# Well, if you have more than 100 concurrent addings
# to the same node, you are clearly doing something wrong...
raise InternalError(
"Hey! We found more than 100 concurrent"
" adds of links "
"to the same nodes! Are you really doing that??")
try:
self._do_create_link(src, "link_{}".format(autolabel_idx),
link_type)
break
except UniquenessError:
# Retry loop until you find a new loop
autolabel_idx += 1
else:
self._do_create_link(src, label, link_type)
def run(self, *args):
load_dbenv()
import argparse
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm import Group, Node, Computer
from aiida.orm.importexport import export, export_zip
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description='Export data from the DB.')
parser.add_argument('-c',
'--computers',
nargs='+',
type=int,
metavar="PK",
help="Export the given computers")
parser.add_argument('-n',
'--nodes',
nargs='+',
type=int,
metavar="PK",
help="Export the given nodes")
parser.add_argument('-g',
'--groups',
nargs='+',
metavar="GROUPNAME",
help="Export all nodes in the given group(s), "
"identified by name.",
type=str)
parser.add_argument('-G',
'--group_pks',
nargs='+',
metavar="PK",
help="Export all nodes in the given group(s), "
"identified by pk.",
type=str)
parser.add_argument('-P',
'--no-parents',
dest='no_parents',
action='store_true',
help="Store only the nodes that are explicitly "
"given, without exporting the parents")
parser.set_defaults(no_parents=False)
parser.add_argument('-O',
'--no-calc-outputs',
dest='no_calc_outputs',
action='store_true',
help="If a calculation is included in the list of "
"nodes to export, do not export its outputs")
parser.set_defaults(no_calc_outputs=False)
parser.add_argument('-y',
'--overwrite',
dest='overwrite',
action='store_true',
help="Overwrite the output file, if it exists")
parser.set_defaults(overwrite=False)
zipsubgroup = parser.add_mutually_exclusive_group()
zipsubgroup.add_argument(
'-z',
'--zipfile-compressed',
dest='zipfilec',
action='store_true',
help="Store as zip file (experimental, should be "
"faster")
zipsubgroup.add_argument('-Z',
'--zipfile-uncompressed',
dest='zipfileu',
action='store_true',
help="Store as uncompressed zip file "
"(experimental, should be faster")
parser.set_defaults(zipfilec=False)
parser.set_defaults(zipfileu=False)
parser.add_argument('output_file',
type=str,
help='The output file name for the export file')
parsed_args = parser.parse_args(args)
if parsed_args.nodes is None:
node_id_set = set()
else:
node_id_set = set(parsed_args.nodes)
group_dict = dict()
if parsed_args.groups is not None:
qb = QueryBuilder()
qb.append(Group,
tag='group',
project=['*'],
filters={'name': {
'in': parsed_args.groups
}})
qb.append(Node, tag='node', member_of='group', project=['id'])
res = qb.dict()
group_dict.update(
{_['group']['*'].name: _['group']['*'].dbgroup
for _ in res})
node_id_set.update([_['node']['id'] for _ in res])
if parsed_args.group_pks is not None:
qb = QueryBuilder()
qb.append(Group,
tag='group',
project=['*'],
filters={'id': {
'in': parsed_args.group_pks
}})
qb.append(Node, tag='node', member_of='group', project=['id'])
res = qb.dict()
group_dict.update(
{_['group']['*'].name: _['group']['*'].dbgroup
for _ in res})
node_id_set.update([_['node']['id'] for _ in res])
# The db_groups that correspond to what was searched above
dbgroups_list = group_dict.values()
# Getting the nodes that correspond to the ids that were found above
if len(node_id_set) > 0:
qb = QueryBuilder()
qb.append(Node,
tag='node',
project=['*'],
filters={'id': {
'in': node_id_set
}})
node_list = [_[0] for _ in qb.all()]
else:
node_list = list()
# Check if any of the nodes wasn't found in the database.
missing_nodes = node_id_set.difference(_.id for _ in node_list)
for id in missing_nodes:
print >> sys.stderr, ("WARNING! Node with pk= {} "
"not found, skipping.".format(id))
# The dbnodes of the above node list
dbnode_list = [_.dbnode for _ in node_list]
if parsed_args.computers is not None:
qb = QueryBuilder()
qb.append(Computer,
tag='comp',
project=['*'],
filters={'id': {
'in': set(parsed_args.computers)
}})
computer_list = [_[0] for _ in qb.all()]
missing_computers = set(parsed_args.computers).difference(
_.id for _ in computer_list)
for id in missing_computers:
print >> sys.stderr, ("WARNING! Computer with pk= {} "
"not found, skipping.".format(id))
else:
computer_list = []
# The dbcomputers of the above computer list
dbcomputer_list = [_.dbcomputer for _ in computer_list]
what_list = dbnode_list + dbcomputer_list + dbgroups_list
export_function = export
additional_kwargs = {}
if parsed_args.zipfileu:
export_function = export_zip
additional_kwargs.update({"use_compression": False})
elif parsed_args.zipfilec:
export_function = export_zip
additional_kwargs.update({"use_compression": True})
try:
export_function(what=what_list,
also_parents=not parsed_args.no_parents,
also_calc_outputs=not parsed_args.no_calc_outputs,
outfile=parsed_args.output_file,
overwrite=parsed_args.overwrite,
**additional_kwargs)
except IOError as e:
print >> sys.stderr, "IOError: {}".format(e.message)
sys.exit(1)
def get_bands_and_parents_structure(self, args):
"""
Search for bands and return bands and the closest structure that is a parent of the instance.
This is the backend independent way, can be overriden for performance reason
:returns:
A list of sublists, each latter containing (in order):
pk as string, formula as string, creation date, bandsdata-label
"""
import datetime
from aiida.utils import timezone
from aiida.orm.querybuilder import QueryBuilder
from aiida.backends.utils import get_automatic_user
from aiida.orm.implementation import User
from aiida.orm.implementation import Group
from aiida.orm.data.structure import (get_formula, get_symbols_string)
from aiida.orm.data.array.bands import BandsData
from aiida.orm.data.structure import StructureData
qb = QueryBuilder()
if args.all_users is False:
au = get_automatic_user()
user = User(dbuser=au)
qb.append(User, tag="creator", filters={"email": user.email})
else:
qb.append(User, tag="creator")
bdata_filters = {}
if args.past_days is not None:
now = timezone.now()
n_days_ago = now - datetime.timedelta(days=args.past_days)
bdata_filters.update({"ctime": {'>=': n_days_ago}})
qb.append(BandsData,
tag="bdata",
created_by="creator",
filters=bdata_filters,
project=["id", "label", "ctime"])
group_filters = {}
if args.group_name is not None:
group_filters.update({"name": {"in": args.group_name}})
if args.group_pk is not None:
group_filters.update({"id": {"in": args.group_pk}})
if group_filters:
qb.append(Group,
tag="group",
filters=group_filters,
group_of="bdata")
qb.append(
StructureData,
tag="sdata",
ancestor_of="bdata",
# We don't care about the creator of StructureData
project=["id", "attributes.kinds", "attributes.sites"])
qb.order_by({StructureData: {'ctime': 'desc'}})
list_data = qb.distinct()
entry_list = []
already_visited_bdata = set()
for [bid, blabel, bdate, sid, akinds, asites] in list_data.all():
# We process only one StructureData per BandsData.
# We want to process the closest StructureData to
# every BandsData.
# We hope that the StructureData with the latest
# creation time is the closest one.
# This will be updated when the QueryBuilder supports
# order_by by the distance of two nodes.
if already_visited_bdata.__contains__(bid):
continue
already_visited_bdata.add(bid)
if args.element is not None:
all_symbols = [_["symbols"][0] for _ in akinds]
if not any([s in args.element for s in all_symbols]):
continue
if args.element_only is not None:
all_symbols = [_["symbols"][0] for _ in akinds]
if not all([s in all_symbols for s in args.element_only]):
continue
# We want only the StructureData that have attributes
if akinds is None or asites is None:
continue
symbol_dict = {}
for k in akinds:
symbols = k['symbols']
weights = k['weights']
symbol_dict[k['name']] = get_symbols_string(symbols, weights)
try:
symbol_list = []
for s in asites:
symbol_list.append(symbol_dict[s['kind_name']])
formula = get_formula(symbol_list, mode=args.formulamode)
# If for some reason there is no kind with the name
# referenced by the site
except KeyError:
formula = "<<UNKNOWN>>"
entry_list.append(
[str(bid),
str(formula),
bdate.strftime('%d %b %Y'), blabel])
return entry_list
def code_list(self, *args):
"""
List available codes
"""
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description='List the codes in the database.')
# The default states are those that are shown if no option is given
parser.add_argument(
'-c',
'--computer',
help="Filter only codes on a given computer",
)
parser.add_argument(
'-p',
'--plugin',
help="Filter only calculation with a given plugin",
)
parser.add_argument(
'-A',
'--all-users',
dest='all_users',
action='store_true',
help="Show codes of all users",
)
parser.add_argument(
'-o',
'--show-owner',
dest='show_owner',
action='store_true',
help="Show also the owner of the code",
)
parser.add_argument(
'-a',
'--all-codes',
action='store_true',
help="Show also hidden codes",
)
parser.set_defaults(all_users=False, hidden=False)
parsed_args = parser.parse_args(args)
computer_filter = parsed_args.computer
plugin_filter = parsed_args.plugin
all_users = parsed_args.all_users
show_owner = parsed_args.show_owner
reveal_filter = parsed_args.all_codes
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.code import Code
from aiida.orm.computer import Computer
from aiida.orm.user import User
from aiida.backends.utils import get_automatic_user
qb_user_filters = dict()
if not all_users:
user = User(dbuser=get_automatic_user())
qb_user_filters['email'] = user.email
qb_computer_filters = dict()
if computer_filter is not None:
qb_computer_filters['name'] = computer_filter
qb_code_filters = dict()
if plugin_filter is not None:
qb_code_filters['attributes.input_plugin'] = plugin_filter
if not reveal_filter:
qb_code_filters['attributes.hidden'] = {"~==": True}
print "# List of configured codes:"
print "# (use 'verdi code show CODEID' to see the details)"
if computer_filter is not None:
qb = QueryBuilder()
qb.append(Code,
tag="code",
filters=qb_code_filters,
project=["id", "label"])
# We have a user assigned to the code so we can ask for the
# presence of a user even if there is no user filter
qb.append(User,
creator_of="code",
project=["email"],
filters=qb_user_filters)
# We also add the filter on computer. This will automatically
# return codes that have a computer (and of course satisfy the
# other filters). The codes that have a computer attached are the
# remote codes.
qb.append(Computer,
computer_of="code",
project=["name"],
filters=qb_computer_filters)
self.print_list_res(qb, show_owner)
# If there is no filter on computers
else:
# Print all codes that have a computer assigned to them
# (these are the remote codes)
qb = QueryBuilder()
qb.append(Code,
tag="code",
filters=qb_code_filters,
project=["id", "label"])
# We have a user assigned to the code so we can ask for the
# presence of a user even if there is no user filter
qb.append(User,
creator_of="code",
project=["email"],
filters=qb_user_filters)
qb.append(Computer, computer_of="code", project=["name"])
self.print_list_res(qb, show_owner)
# Now print all the local codes. To get the local codes we ask
# the dbcomputer_id variable to be None.
qb = QueryBuilder()
comp_non_existence = {"dbcomputer_id": {"==": None}}
if not qb_code_filters:
qb_code_filters = comp_non_existence
else:
new_qb_code_filters = {
"and": [qb_code_filters, comp_non_existence]
}
qb_code_filters = new_qb_code_filters
qb.append(Code,
tag="code",
filters=qb_code_filters,
project=["id", "label"])
# We have a user assigned to the code so we can ask for the
# presence of a user even if there is no user filter
qb.append(User,
creator_of="code",
project=["email"],
filters=qb_user_filters)
self.print_list_res(qb, show_owner)
def computer_configure(self, *args):
"""
Configure the authentication information for a given computer
"""
if not is_dbenv_loaded():
load_dbenv()
import readline
import inspect
from django.core.exceptions import ObjectDoesNotExist
from aiida.common.exceptions import (NotExistent, ValidationError)
from aiida.backends.utils import get_automatic_user
from aiida.common.utils import get_configured_user_email
from aiida.backends.settings import BACKEND
from aiida.backends.profile import BACKEND_SQLA, BACKEND_DJANGO
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description='Configure a computer for a given AiiDA user.')
# The default states are those that are shown if no option is given
parser.add_argument(
'-u',
'--user',
type=str,
metavar='EMAIL',
help=
"Configure the computer for the given AiiDA user (otherwise, configure the current default user)",
)
parser.add_argument(
'computer',
type=str,
help="The name of the computer that you want to configure")
parsed_args = parser.parse_args(args)
user_email = parsed_args.user
computername = parsed_args.computer
try:
computer = self.get_computer(name=computername)
except NotExistent:
print >> sys.stderr, "No computer exists with name '{}'".format(
computername)
sys.exit(1)
if user_email is None:
user = get_automatic_user()
else:
from aiida.orm.querybuilder import QueryBuilder
qb = QueryBuilder()
qb.append(type="user", filters={'email': user_email})
user = qb.first()
if user is None:
print >> sys.stderr, ("No user with email '{}' in the "
"database.".format(user_email))
sys.exit(1)
if BACKEND == BACKEND_DJANGO:
from aiida.backends.djsite.db.models import DbAuthInfo
try:
authinfo = DbAuthInfo.objects.get(
dbcomputer=computer.dbcomputer, aiidauser=user)
old_authparams = authinfo.get_auth_params()
except ObjectDoesNotExist:
authinfo = DbAuthInfo(dbcomputer=computer.dbcomputer,
aiidauser=user)
old_authparams = {}
elif BACKEND == BACKEND_SQLA:
from aiida.backends.sqlalchemy.models.authinfo import DbAuthInfo
from aiida.backends.sqlalchemy import session
authinfo = session.query(DbAuthInfo).filter(
DbAuthInfo.dbcomputer == computer.dbcomputer).filter(
DbAuthInfo.aiidauser == user).first()
if authinfo is None:
authinfo = DbAuthInfo(dbcomputer=computer.dbcomputer,
aiidauser=user)
old_authparams = {}
else:
old_authparams = authinfo.get_auth_params()
else:
raise Exception("Unknown backend {}".format(BACKEND))
Transport = computer.get_transport_class()
print("Configuring computer '{}' for the AiiDA user '{}'".format(
computername, user.email))
print "Computer {} has transport of type {}".format(
computername, computer.get_transport_type())
if user.email != get_configured_user_email():
print "*" * 72
print "** {:66s} **".format("WARNING!")
print "** {:66s} **".format(
" You are configuring a different user.")
print "** {:66s} **".format(
" Note that the default suggestions are taken from your")
print "** {:66s} **".format(
" local configuration files, so they may be incorrect.")
print "*" * 72
valid_keys = Transport.get_valid_auth_params()
default_authparams = {}
for k in valid_keys:
if k in old_authparams:
default_authparams[k] = old_authparams.pop(k)
if old_authparams:
print(
"WARNING: the following keys were previously in the "
"authorization parameters,")
print "but have not been recognized and have been deleted:"
print ", ".join(old_authparams.keys())
if not valid_keys:
print "There are no special keys to be configured. Configuration completed."
authinfo.set_auth_params({})
authinfo.save()
return
print ""
print "Note: to leave a field unconfigured, leave it empty and press [Enter]"
# I strip out the old auth_params that are not among the valid keys
new_authparams = {}
for k in valid_keys:
key_set = False
while not key_set:
try:
converter_name = '_convert_{}_fromstring'.format(k)
try:
converter = dict(
inspect.getmembers(Transport))[converter_name]
except KeyError:
print >> sys.stderr, (
"Internal error! "
"No {} defined in Transport {}".format(
converter_name, computer.get_transport_type()))
sys.exit(1)
if k in default_authparams:
readline.set_startup_hook(lambda: readline.insert_text(
str(default_authparams[k])))
else:
# Use suggestion only if parameters were not already set
suggester_name = '_get_{}_suggestion_string'.format(k)
try:
suggester = dict(
inspect.getmembers(Transport))[suggester_name]
suggestion = suggester(computer)
readline.set_startup_hook(
lambda: readline.insert_text(suggestion))
except KeyError:
readline.set_startup_hook()
txtval = raw_input("=> {} = ".format(k))
if txtval:
new_authparams[k] = converter(txtval)
key_set = True
except ValidationError as e:
print "Error in the inserted value: {}".format(e.message)
authinfo.set_auth_params(new_authparams)
authinfo.save()
print "Configuration stored for your user on computer '{}'.".format(
computername)
def test_creation_and_deletion(self):
from aiida.backends.djsite.db.models import DbLink # Direct links
from aiida.orm.querybuilder import QueryBuilder
from aiida.common.links import LinkType
n1 = Node().store()
n2 = Node().store()
n3 = Node().store()
n4 = Node().store()
n5 = Node().store()
n6 = Node().store()
n7 = Node().store()
n8 = Node().store()
n9 = Node().store()
# I create a strange graph, inserting links in a order
# such that I often have to create the transitive closure
# between two graphs
n3.add_link_from(n2, link_type=LinkType.CREATE)
n2.add_link_from(n1, link_type=LinkType.CREATE)
n5.add_link_from(n3, link_type=LinkType.CREATE)
n5.add_link_from(n4, link_type=LinkType.CREATE)
n4.add_link_from(n2, link_type=LinkType.CREATE)
n7.add_link_from(n6, link_type=LinkType.CREATE)
n8.add_link_from(n7, link_type=LinkType.CREATE)
# Yet, no links from 1 to 8
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 0)
n6.add_link_from(n5, link_type=LinkType.INPUT)
# Yet, now 2 links from 1 to 8
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 2)
n7.add_link_from(n9, link_type=LinkType.INPUT)
# Still two links...
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 2)
n9.add_link_from(n6, link_type=LinkType.INPUT)
# And now there should be 4 nodes
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 4)
### I start deleting now
# I cut one branch below: I should loose 2 links
DbLink.objects.filter(input=n6, output=n9).delete()
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 2)
DbLink.objects.filter(input=n2, output=n4).delete()
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 1)
#~ self.assertEquals(
#~ len(DbPath.objects.filter(parent=n1, child=n8).distinct()), 1)
# Another cut should delete all links
DbLink.objects.filter(input=n3, output=n5).delete()
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 0)
#~ self.assertEquals(
#~ len(DbPath.objects.filter(parent=n1, child=n8).distinct()), 0)
# But I did not delete everything! For instance, I can check
# the following links
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n4.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 1)
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n5.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n7.pk
}).count(), 1)
#~ self.assertEquals(
#~ len(DbPath.objects.filter(parent=n4, child=n8).distinct()), 1)
#~ self.assertEquals(
#~ len(DbPath.objects.filter(parent=n5, child=n7).distinct()), 1)
# Finally, I reconnect in a different way the two graphs and
# check that 1 and 8 are again connected
n4.add_link_from(n3, link_type=LinkType.INPUT)
self.assertEquals(
QueryBuilder().append(Node, filters={
'id': n1.pk
}, tag='anc').append(Node,
descendant_of='anc',
filters={
'id': n8.pk
}).count(), 1)
def calculation_cleanworkdir(self, *args):
"""
Clean all the content of all the output remote folders of calculations,
passed as a list of pks, or identified by modification time.
If a list of calculation PKs is not passed through -c option, one of
the option -p or -u has to be specified (if both are given, a logical
AND is done between the 2 - you clean out calculations modified AFTER
[-p option] days from now but BEFORE [-o option] days from now).
If you also pass the -f option, no confirmation will be asked.
"""
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description="Clean work directory (i.e. remote folder) of AiiDA "
"calculations.")
parser.add_argument("-k",
"--pk",
metavar="PK",
type=int,
nargs="+",
help="The principal key (PK) of the calculations "
"to clean the workdir of",
dest="pk")
parser.add_argument("-f",
"--force",
action="store_true",
help="Force the cleaning (no prompt)")
parser.add_argument("-p",
"--past-days",
metavar="N",
help="Add a filter to clean workdir of "
"calculations modified during the past N "
"days",
type=int,
action="store",
dest="past_days")
parser.add_argument("-o",
"--older-than",
metavar="N",
help="Add a filter to clean workdir of "
"calculations that have been modified on a "
"date before N days ago",
type=int,
action="store",
dest="older_than")
parser.add_argument("-c",
"--computers",
metavar="label",
nargs="+",
help="Add a filter to clean workdir of "
"calculations on this computer(s) only",
type=str,
action="store",
dest="computer")
if not is_dbenv_loaded():
load_dbenv()
from aiida.backends.utils import get_automatic_user
from aiida.backends.utils import get_authinfo
from aiida.common.utils import query_yes_no
from aiida.orm.computer import Computer as OrmComputer
from aiida.orm.user import User as OrmUser
from aiida.orm.calculation import Calculation as OrmCalculation
from aiida.orm.querybuilder import QueryBuilder
from aiida.utils import timezone
import datetime
parsed_args = parser.parse_args(args)
# If a pk is given then the -o & -p options should not be specified
if parsed_args.pk is not None:
if ((parsed_args.past_days is not None)
or (parsed_args.older_than is not None)):
print(
"You cannot specify both a list of calculation pks and "
"the -p or -o options")
return
# If no pk is given then at least one of the -o & -p options should be
# specified
else:
if ((parsed_args.past_days is None)
and (parsed_args.older_than is None)):
print(
"You should specify at least a list of calculations or "
"the -p, -o options")
return
# At this point we know that either the pk or the -p -o options are
# specified
# We also check that not both -o & -p options are specified
if ((parsed_args.past_days is not None)
and (parsed_args.older_than is not None)):
print(
"Not both of the -p, -o options can be specified in the "
"same time")
return
qb_user_filters = dict()
user = OrmUser(dbuser=get_automatic_user())
qb_user_filters["email"] = user.email
qb_computer_filters = dict()
if parsed_args.computer is not None:
qb_computer_filters["name"] = {"in": parsed_args.computer}
qb_calc_filters = dict()
if parsed_args.past_days is not None:
pd_ts = timezone.now() - datetime.timedelta(
days=parsed_args.past_days)
qb_calc_filters["mtime"] = {">": pd_ts}
if parsed_args.older_than is not None:
ot_ts = timezone.now() - datetime.timedelta(
days=parsed_args.older_than)
qb_calc_filters["mtime"] = {"<": ot_ts}
if parsed_args.pk is not None:
print("parsed_args.pk: ", parsed_args.pk)
qb_calc_filters["id"] = {"in": parsed_args.pk}
qb = QueryBuilder()
qb.append(OrmCalculation,
tag="calc",
filters=qb_calc_filters,
project=["id", "uuid", "attributes.remote_workdir"])
qb.append(OrmComputer,
computer_of="calc",
project=["*"],
filters=qb_computer_filters)
qb.append(OrmUser,
creator_of="calc",
project=["*"],
filters=qb_user_filters)
no_of_calcs = qb.count()
if no_of_calcs == 0:
print("No calculations found with the given criteria.")
return
print("Found {} calculations with the given criteria.".format(
no_of_calcs))
if not parsed_args.force:
if not query_yes_no(
"Are you sure you want to clean the work "
"directory?", "no"):
return
# get the uuids of all calculations matching the filters
calc_list_data = qb.dict()
# get all computers associated to the calc uuids above, and load them
# we group them by uuid to avoid computer duplicates
comp_uuid_to_computers = {
_["computer"]["*"].uuid: _["computer"]["*"]
for _ in calc_list_data
}
# now build a dictionary with the info of folders to delete
remotes = {}
for computer in comp_uuid_to_computers.values():
# initialize a key of info for a given computer
remotes[computer.name] = {
'transport':
get_authinfo(computer=computer,
aiidauser=user._dbuser).get_transport(),
'computer':
computer,
}
# select the calc pks done on this computer
this_calc_pks = [
_["calc"]["id"] for _ in calc_list_data
if _["computer"]["*"].id == computer.id
]
this_calc_uuids = [
unicode(_["calc"]["uuid"]) for _ in calc_list_data
if _["computer"]["*"].id == computer.id
]
remote_workdirs = [
_["calc"]["attributes.remote_workdir"] for _ in calc_list_data
if _["calc"]["id"] in this_calc_pks
if _["calc"]["attributes.remote_workdir"] is not None
]
remotes[computer.name]['remotes'] = remote_workdirs
remotes[computer.name]['uuids'] = this_calc_uuids
# now proceed to cleaning
for computer, dic in remotes.iteritems():
print(
"Cleaning the work directory on computer {}.".format(computer))
counter = 0
t = dic['transport']
with t:
remote_user = remote_user = t.whoami()
aiida_workdir = dic['computer'].get_workdir().format(
username=remote_user)
t.chdir(aiida_workdir)
# Hardcoding the sharding equal to 3 parts!
existing_folders = t.glob('*/*/*')
folders_to_delete = [
i for i in existing_folders
if i.replace("/", "") in dic['uuids']
]
for folder in folders_to_delete:
t.rmtree(folder)
counter += 1
if counter % 20 == 0 and counter > 0:
print("Deleted work directories: {}".format(counter))
print("{} remote folder(s) cleaned.".format(counter))
def query_jobcalculations_by_computer_user_state(
self,
state,
computer=None,
user=None,
only_computer_user_pairs=False,
only_enabled=True,
limit=None):
"""
Filter all calculations with a given state.
Issue a warning if the state is not in the list of valid states.
:param string state: The state to be used to filter (should be a string among
those defined in aiida.common.datastructures.calc_states)
:param computer: a Django DbComputer entry, or a Computer object, of a
computer in the DbComputer table.
A string for the hostname is also valid.
:param user: a Django entry (or its pk) of a user in the DbUser table;
if present, the results are restricted to calculations of that
specific user
:param bool only_computer_user_pairs: if False (default) return a queryset
where each element is a suitable instance of Node (it should
be an instance of Calculation, if everything goes right!)
If True, return only a list of tuples, where each tuple is
in the format
('dbcomputer__id', 'user__id')
[where the IDs are the IDs of the respective tables]
:param int limit: Limit the number of rows returned
:return: a list of calculation objects matching the filters.
"""
# I assume that calc_states are strings. If this changes in the future,
# update the filter below from dbattributes__tval to the correct field.
from aiida.orm.computer import Computer
from aiida.orm.calculation.job import JobCalculation
from aiida.orm.user import User
from aiida.orm.querybuilder import QueryBuilder
from aiida.common.exceptions import InputValidationError
from aiida.common.datastructures import calc_states
if state not in calc_states:
raise InputValidationError(
"querying for calculation state='{}', but it "
"is not a valid calculation state".format(state))
calcfilter = {'state': {'==': state}}
computerfilter = {"enabled": {'==': True}}
userfilter = {}
if computer is None:
pass
elif isinstance(computer, int):
computerfilter.update({'id': {'==': computer}})
elif isinstance(computer, Computer):
computerfilter.update({'id': {'==': computer.pk}})
else:
try:
computerfilter.update({'id': {'==': computer.id}})
except AttributeError as e:
raise Exception("{} is not a valid computer\n{}".format(
computer, e))
if user is None:
pass
elif isinstance(user, int):
userfilter.update({'id': {'==': user}})
else:
try:
userfilter.update({'id': {'==': int(user.id)}})
# Is that safe?
except:
raise Exception("{} is not a valid user".format(user))
qb = QueryBuilder()
qb.append(type="computer", tag='computer', filters=computerfilter)
qb.append(JobCalculation,
filters=calcfilter,
tag='calc',
has_computer='computer')
qb.append(type="user",
tag='user',
filters=userfilter,
creator_of="calc")
if only_computer_user_pairs:
qb.add_projection("computer", "*")
qb.add_projection("user", "*")
returnresult = qb.distinct().all()
else:
qb.add_projection("calc", "*")
if limit is not None:
qb.limit(limit)
returnresult = qb.all()
returnresult = zip(*returnresult)[0]
return returnresult
def calculation_cleanworkdir(self, *args):
"""
Clean the working directory of calculations by removing all the content of the
associated RemoteFolder node. Calculations can be identified by pk with the -k flag
or by specifying limits on the modification times with -p/-o flags
"""
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description="""
Clean all content of all output remote folders of calculations,
passed as a list of pks, or identified by modification time.
If a list of calculation PKs is not passed with the -k option, one or both
of the -p and -o options has to be specified. If both are specified, a logical
AND is done between the two, i.e. the calculations that will be cleaned have been
modified AFTER [-p option] days from now but BEFORE [-o option] days from now.
Passing the -f option will prevent the confirmation dialog from being prompted.
"""
)
parser.add_argument(
'-k', '--pk', metavar='PK', type=int, nargs='+', dest='pk',
help='The principal key (PK) of the calculations of which to clean the work directory'
)
parser.add_argument(
'-f', '--force', action='store_true',
help='Force the cleaning (no prompt)'
)
parser.add_argument(
'-p', '--past-days', metavar='N', type=int, action='store', dest='past_days',
help='Include calculations that have been modified within the last N days',
)
parser.add_argument(
'-o', '--older-than', metavar='N', type=int, action='store', dest='older_than',
help='Include calculations that have been modified more than N days ago',
)
parser.add_argument(
'-c', '--computers', metavar='label', nargs='+', type=str, action='store', dest='computer',
help='Include only calculations that were ran on these computers'
)
if not is_dbenv_loaded():
load_dbenv()
from aiida.backends.utils import get_automatic_user
from aiida.backends.utils import get_authinfo
from aiida.common.utils import query_yes_no
from aiida.orm.computer import Computer as OrmComputer
from aiida.orm.user import User as OrmUser
from aiida.orm.calculation import Calculation as OrmCalculation
from aiida.orm.querybuilder import QueryBuilder
from aiida.utils import timezone
import datetime
parsed_args = parser.parse_args(args)
# If a pk is given then the -o & -p options should not be specified
if parsed_args.pk is not None:
if (parsed_args.past_days is not None or parsed_args.older_than is not None):
print("You cannot specify both a list of calculation pks and the -p or -o options")
return
# If no pk is given then at least one of the -o & -p options should be specified
else:
if (parsed_args.past_days is None and parsed_args.older_than is None):
print("You should specify at least a list of calculations or the -p, -o options")
return
qb_user_filters = dict()
user = OrmUser(dbuser=get_automatic_user())
qb_user_filters["email"] = user.email
qb_computer_filters = dict()
if parsed_args.computer is not None:
qb_computer_filters["name"] = {"in": parsed_args.computer}
qb_calc_filters = dict()
if parsed_args.past_days is not None:
pd_ts = timezone.now() - datetime.timedelta(days=parsed_args.past_days)
qb_calc_filters["mtime"] = {">": pd_ts}
if parsed_args.older_than is not None:
ot_ts = timezone.now() - datetime.timedelta(days=parsed_args.older_than)
qb_calc_filters["mtime"] = {"<": ot_ts}
if parsed_args.pk is not None:
print("parsed_args.pk: ", parsed_args.pk)
qb_calc_filters["id"] = {"in": parsed_args.pk}
qb = QueryBuilder()
qb.append(OrmCalculation, tag="calc",
filters=qb_calc_filters,
project=["id", "uuid", "attributes.remote_workdir"])
qb.append(OrmComputer, computer_of="calc", tag="computer",
project=["*"],
filters=qb_computer_filters)
qb.append(OrmUser, creator_of="calc", tag="user",
project=["*"],
filters=qb_user_filters)
no_of_calcs = qb.count()
if no_of_calcs == 0:
print("No calculations found with the given criteria.")
return
print("Found {} calculations with the given criteria.".format(
no_of_calcs))
if not parsed_args.force:
if not query_yes_no("Are you sure you want to clean the work "
"directory?", "no"):
return
# get the uuids of all calculations matching the filters
calc_list_data = qb.dict()
# get all computers associated to the calc uuids above, and load them
# we group them by uuid to avoid computer duplicates
comp_uuid_to_computers = {_["computer"]["*"].uuid: _["computer"]["*"] for _ in calc_list_data}
# now build a dictionary with the info of folders to delete
remotes = {}
for computer in comp_uuid_to_computers.values():
# initialize a key of info for a given computer
remotes[computer.name] = {'transport': get_authinfo(
computer=computer, aiidauser=user._dbuser).get_transport(),
'computer': computer,
}
# select the calc pks done on this computer
this_calc_pks = [_["calc"]["id"] for _ in calc_list_data
if _["computer"]["*"].id == computer.id]
this_calc_uuids = [unicode(_["calc"]["uuid"])
for _ in calc_list_data
if _["computer"]["*"].id == computer.id]
remote_workdirs = [_["calc"]["attributes.remote_workdir"]
for _ in calc_list_data
if _["calc"]["id"] in this_calc_pks
if _["calc"]["attributes.remote_workdir"]
is not None]
remotes[computer.name]['remotes'] = remote_workdirs
remotes[computer.name]['uuids'] = this_calc_uuids
# now proceed to cleaning
for computer, dic in remotes.iteritems():
print("Cleaning the work directory on computer {}.".format(computer))
counter = 0
t = dic['transport']
with t:
remote_user = remote_user = t.whoami()
aiida_workdir = dic['computer'].get_workdir().format(
username=remote_user)
t.chdir(aiida_workdir)
# Hardcoding the sharding equal to 3 parts!
existing_folders = t.glob('*/*/*')
folders_to_delete = [i for i in existing_folders if
i.replace("/", "") in dic['uuids']]
for folder in folders_to_delete:
t.rmtree(folder)
counter += 1
if counter % 20 == 0 and counter > 0:
print("Deleted work directories: {}".format(counter))
print("{} remote folder(s) cleaned.".format(counter))
def create(outfile, computers, groups, nodes, group_names, no_parents,
no_calc_outputs, overwrite, archive_format):
"""
Export nodes and groups of nodes to an archive file for backup or sharing purposes
"""
import sys
from aiida.backends.utils import load_dbenv
load_dbenv()
from aiida.orm import Group, Node, Computer
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.importexport import export, export_zip
node_id_set = set(nodes)
group_dict = dict()
if group_names:
qb = QueryBuilder()
qb.append(Group,
tag='group',
project=['*'],
filters={'name': {
'in': group_names
}})
qb.append(Node, tag='node', member_of='group', project=['id'])
res = qb.dict()
group_dict.update({
group['group']['*'].name: group['group']['*'].dbgroup
for group in res
})
node_id_set.update([node['node']['id'] for node in res])
if groups:
qb = QueryBuilder()
qb.append(Group,
tag='group',
project=['*'],
filters={'id': {
'in': groups
}})
qb.append(Node, tag='node', member_of='group', project=['id'])
res = qb.dict()
group_dict.update({
group['group']['*'].name: group['group']['*'].dbgroup
for group in res
})
node_id_set.update([node['node']['id'] for node in res])
# The db_groups that correspond to what was searched above
dbgroups_list = group_dict.values()
# Getting the nodes that correspond to the ids that were found above
if len(node_id_set) > 0:
qb = QueryBuilder()
qb.append(Node,
tag='node',
project=['*'],
filters={'id': {
'in': node_id_set
}})
node_list = [node[0] for node in qb.all()]
else:
node_list = list()
# Check if any of the nodes wasn't found in the database.
missing_nodes = node_id_set.difference(node.id for node in node_list)
for node_id in missing_nodes:
print >> sys.stderr, (
'WARNING! Node with pk={} not found, skipping'.format(node_id))
# The dbnodes of the above node list
dbnode_list = [node.dbnode for node in node_list]
if computers:
qb = QueryBuilder()
qb.append(Computer,
tag='comp',
project=['*'],
filters={'id': {
'in': set(computers)
}})
computer_list = [computer[0] for computer in qb.all()]
missing_computers = set(computers).difference(
computer.id for computer in computer_list)
for computer_id in missing_computers:
print >> sys.stderr, (
'WARNING! Computer with pk={} not found, skipping'.format(
computer_id))
else:
computer_list = []
# The dbcomputers of the above computer list
dbcomputer_list = [computer.dbcomputer for computer in computer_list]
what_list = dbnode_list + dbcomputer_list + dbgroups_list
additional_kwargs = dict()
if archive_format == 'zip':
export_function = export_zip
additional_kwargs.update({'use_compression': True})
elif archive_format == 'zip-uncompressed':
export_function = export_zip
additional_kwargs.update({'use_compression': False})
elif archive_format == 'tar.gz':
export_function = export
else:
print >> sys.stderr, 'invalid --archive-format value {}'.format(
archive_format)
sys.exit(1)
try:
export_function(what=what_list,
also_parents=not no_parents,
also_calc_outputs=not no_calc_outputs,
outfile=outfile,
overwrite=overwrite,
**additional_kwargs)
except IOError as e:
print >> sys.stderr, 'IOError: {}'.format(e.message)
sys.exit(1)
def test_simple_query_django_1(self):
"""
Testing a simple query
"""
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.calculation.job import JobCalculation
from aiida.orm import Node, Data, Calculation
from datetime import datetime
from aiida.common.links import LinkType
n1 = Data()
n1.label = 'node1'
n1._set_attr('foo', ['hello', 'goodbye'])
n1.store()
n2 = Calculation()
n2.label = 'node2'
n2._set_attr('foo', 1)
n2.store()
n3 = Data()
n3.label = 'node3'
n3._set_attr('foo', 1.0000) # Stored as fval
n3.store()
n4 = Calculation()
n4.label = 'node4'
n4._set_attr('foo', 'bar')
n4.store()
n5 = Data()
n5.label = 'node5'
n5._set_attr('foo', None)
n5.store()
n2.add_link_from(n1, link_type=LinkType.INPUT)
n3.add_link_from(n2, link_type=LinkType.CREATE)
n4.add_link_from(n3, link_type=LinkType.INPUT)
n5.add_link_from(n4, link_type=LinkType.CREATE)
qb1 = QueryBuilder()
qb1.append(Node, filters={'attributes.foo': 1.000})
self.assertEqual(len(qb1.all()), 2)
qb2 = QueryBuilder()
qb2.append(Data)
self.assertEqual(qb2.count(), 3)
qb2 = QueryBuilder()
qb2.append(type='data.Data.')
self.assertEqual(qb2.count(), 3)
qb3 = QueryBuilder()
qb3.append(Node, project='label', tag='node1')
qb3.append(Node, project='label', tag='node2')
self.assertEqual(qb3.count(), 4)
qb4 = QueryBuilder()
qb4.append(Calculation, tag='node1')
qb4.append(Data, tag='node2')
self.assertEqual(qb4.count(), 2)
qb5 = QueryBuilder()
qb5.append(Data, tag='node1')
qb5.append(Calculation, tag='node2')
self.assertEqual(qb5.count(), 2)
qb6 = QueryBuilder()
qb6.append(Data, tag='node1')
qb6.append(Data, tag='node2')
self.assertEqual(qb6.count(), 0)
