def test_no_node_export(self, temp_dir):
"""Test migration of export file that has no Nodes"""
input_file = get_archive_file('export_v0.3_no_Nodes.aiida', **self.external_archive)
output_file = os.path.join(temp_dir, 'output_file.aiida')
# Known entities
computer_uuids = [self.computer.uuid] # pylint: disable=no-member
user_emails = [orm.User.objects.get_default().email]
# Known export file content used for checks
node_count = 0
computer_count = 1 + 1 # localhost is always present
computer_uuids.append('4f33c6fd-b624-47df-9ffb-a58f05d323af')
user_emails.append('aiida@localhost')
# Perform the migration
migrate_archive(input_file, output_file)
# Load the migrated file
import_data(output_file, silent=True)
# Check known number of entities is present
self.assertEqual(orm.QueryBuilder().append(orm.Node).count(), node_count)
self.assertEqual(orm.QueryBuilder().append(orm.Computer).count(), computer_count)
# Check unique identifiers
computers = orm.QueryBuilder().append(orm.Computer, project=['uuid']).all()[0][0]
users = orm.QueryBuilder().append(orm.User, project=['email']).all()[0][0]
self.assertIn(computers, computer_uuids)
self.assertIn(users, user_emails)
def test_delete_many_ids(self):
"""Test `delete_many` method filtering on both `id` and `uuid`"""
comment1 = self.create_comment()
comment2 = self.create_comment()
comment3 = self.create_comment()
comment_uuids = []
for comment in [comment1, comment2, comment3]:
comment.store()
comment_uuids.append(str(comment.uuid))
# Make sure they exist
count_comments_found = orm.QueryBuilder().append(orm.Comment, filters={'uuid': {'in': comment_uuids}}).count()
self.assertEqual(
count_comments_found,
len(comment_uuids),
msg='There should be {} Comments, instead {} Comment(s) was/were found'.format(
len(comment_uuids), count_comments_found
)
)
# Delete last two comments (comment2, comment3)
filters = {'or': [{'id': comment2.id}, {'uuid': str(comment3.uuid)}]}
self.backend.comments.delete_many(filters=filters)
# Check they were deleted
builder = orm.QueryBuilder().append(orm.Comment, filters={'uuid': {'in': comment_uuids}}, project='uuid').all()
found_comments_uuid = [_[0] for _ in builder]
self.assertEqual([comment_uuids[0]], found_comments_uuid)
def test_delete_many_user_id(self):
"""Test `delete_many` method filtering on `user_id`"""
# Create comments and separate user
user_two = self.backend.users.create(
email='tester_two@localhost').store()
comment1 = self.create_comment(user=user_two)
comment2 = self.create_comment()
comment3 = self.create_comment()
comment_uuids = []
for comment in [comment1, comment2, comment3]:
comment.store()
comment_uuids.append(str(comment.uuid))
# Make sure they exist
builder = orm.QueryBuilder().append(orm.Comment, project='uuid')
self.assertGreater(builder.count(), 0)
found_comments_uuid = [_[0] for _ in builder.all()]
for comment_uuid in comment_uuids:
self.assertIn(comment_uuid, found_comments_uuid)
# Delete last comments for `self.user`
filters = {'user_id': self.user.id}
self.backend.comments.delete_many(filters=filters)
# Check they were deleted
builder = orm.QueryBuilder().append(orm.Comment, project='uuid')
found_comments_uuid = [_[0] for _ in builder.all()]
self.assertGreater(builder.count(), 0)
for comment_uuid in comment_uuids[1:]:
self.assertNotIn(comment_uuid, found_comments_uuid)
# Make sure the first comment (comment1) was not deleted
self.assertIn(comment_uuids[0], found_comments_uuid)
def delete_database_proportion(fraction=0.0):
"""Description pending"""
from random import shuffle
from aiida import orm
from aiida.tools import delete_nodes
query = orm.QueryBuilder()
query.append(orm.Node, project=['id'])
initial_nodecount = query.count()
nodepk_list = [nodepk for nodepk, in query.all()]
shuffle(nodepk_list)
numto_delete = int(fraction * len(nodepk_list))
numto_delete = min(numto_delete, initial_nodecount)
delete_nodes(nodepk_list[0:numto_delete], dry_run=False)
query = orm.QueryBuilder()
query.append(orm.Node)
final_nodecount = query.count()
return {
'initial_nodecount': initial_nodecount,
'final_nodecount': final_nodecount
}
def test_group_import_existing(self, temp_dir):
"""
Testing what happens when I try to import a group that already exists in the
database. This should raise an appropriate exception
"""
grouplabel = 'node_group_existing'
# Create another user
new_email = '[email protected]'
user = orm.User(email=new_email)
user.store()
# Create a structure data node
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd'
sd1.store()
# Create a group and add the data inside
group = orm.Group(label=grouplabel)
group.store()
group.add_nodes([sd1])
# At this point we export the generated data
filename = os.path.join(temp_dir, 'export1.tar.gz')
export([group], outfile=filename, silent=True)
self.clean_db()
self.insert_data()
# Creating a group of the same name
group = orm.Group(label='node_group_existing')
group.store()
import_data(filename, silent=True)
# The import should have created a new group with a suffix
# I check for this:
builder = orm.QueryBuilder().append(
orm.Group, filters={'label': {
'like': grouplabel + '%'
}})
self.assertEqual(builder.count(), 2)
# Now I check for the group having one member, and whether the name is different:
builder = orm.QueryBuilder()
builder.append(orm.Group,
filters={'label': {
'like': grouplabel + '%'
}},
tag='g',
project='label')
builder.append(orm.StructureData, with_group='g')
self.assertEqual(builder.count(), 1)
# I check that the group name was changed:
self.assertTrue(builder.all()[0][0] != grouplabel)
# I import another name, the group should not be imported again
import_data(filename, silent=True)
builder = orm.QueryBuilder()
builder.append(orm.Group,
filters={'label': {
'like': grouplabel + '%'
}})
self.assertEqual(builder.count(), 2)
def test_exclude_logs_flag(self, temp_dir):
"""Test that the `include_logs` argument for `export` works."""
log_msg = 'Testing logging of critical failure'
# Create node
calc = orm.CalculationNode()
calc.store()
calc.seal()
# Create log message
calc.logger.critical(log_msg)
# Save uuids prior to export
calc_uuid = calc.uuid
# Export, excluding logs
export_file = os.path.join(temp_dir, 'export.tar.gz')
export([calc], outfile=export_file, silent=True, include_logs=False)
# Clean database and reimport exported data
self.reset_database()
import_data(export_file, silent=True)
# Finding all the log messages
import_calcs = orm.QueryBuilder().append(orm.CalculationNode,
project=['uuid']).all()
import_logs = orm.QueryBuilder().append(orm.Log,
project=['uuid']).all()
# There should be exactly: 1 orm.CalculationNode, 0 Logs
self.assertEqual(len(import_calcs), 1)
self.assertEqual(len(import_logs), 0)
# Check it's the correct node
self.assertEqual(str(import_calcs[0][0]), calc_uuid)
def test_delete_many_ids(self):
"""Test `delete_many` method filtering on both `id` and `uuid`"""
# Create logs
log1 = self.create_log()
log2 = self.create_log()
log3 = self.create_log()
log_uuids = []
for log in [log1, log2, log3]:
log.store()
log_uuids.append(str(log.uuid))
# Make sure they exist
count_logs_found = orm.QueryBuilder().append(orm.Log, filters={'uuid': {'in': log_uuids}}).count()
self.assertEqual(
count_logs_found,
len(log_uuids),
msg=f'There should be {len(log_uuids)} Logs, instead {count_logs_found} Log(s) was/were found'
)
# Delete last two logs (log2, log3)
filters = {'or': [{'id': log2.id}, {'uuid': str(log3.uuid)}]}
self.backend.logs.delete_many(filters=filters)
# Check they were deleted
builder = orm.QueryBuilder().append(orm.Log, filters={'uuid': {'in': log_uuids}}, project='uuid').all()
found_logs_uuid = [_[0] for _ in builder]
self.assertEqual([log_uuids[0]], found_logs_uuid)
def test_no_node_migration(self, tmp_path, external_archive):
"""Test migration of archive file that has no Node entities."""
input_file = get_archive_file('export_v0.3_no_Nodes.aiida',
**external_archive)
output_file = tmp_path / 'output_file.aiida'
migrator_cls = get_migrator(detect_archive_type(input_file))
migrator = migrator_cls(input_file)
# Perform the migration
migrator.migrate(newest_version, output_file)
# Load the migrated file
import_data(output_file)
# Check known entities
assert orm.QueryBuilder().append(orm.Node).count() == 0
computer_query = orm.QueryBuilder().append(orm.Computer,
project=['uuid'])
assert computer_query.all(flat=True) == [
'4f33c6fd-b624-47df-9ffb-a58f05d323af'
]
user_query = orm.QueryBuilder().append(orm.User, project=['email'])
assert set(user_query.all(flat=True)) == {
orm.User.objects.get_default().email, 'aiida@localhost'
}
def test_querying():
"""Test querying for groups with and without subclassing."""
orm.Group(label='group').store()
orm.AutoGroup(label='auto-group').store()
# Fake a subclass by manually setting the type string
group = orm.Group(label='custom-group')
group.backend_entity.dbmodel.type_string = 'custom.group'
group.store()
assert orm.QueryBuilder().append(orm.AutoGroup).count() == 1
assert orm.QueryBuilder().append(orm.AutoGroup,
subclassing=False).count() == 1
assert orm.QueryBuilder().append(orm.Group,
subclassing=False).count() == 1
assert orm.QueryBuilder().append(orm.Group).count() == 3
assert orm.QueryBuilder().append(orm.Group,
filters={
'type_string': 'custom.group'
}).count() == 1
# Removing it as other methods might get a warning instead
group_pk = group.pk
del group
orm.Group.objects.delete(id=group_pk)
def get_input_node(cls, value):
"""Return a `Node` of a given class and given value.
If a `Node` of the given type and value already exists, that will be returned, otherwise a new one will be created,
stored and returned.
:param cls: the `Node` class
:param value: the value of the `Node`
"""
from aiida import orm
if cls in (orm.Bool, orm.Float, orm.Int, orm.Str):
result = orm.QueryBuilder().append(cls, filters={'attributes.value': value}).first()
if result is None:
node = cls(value).store()
else:
node = result[0]
elif cls is orm.Dict:
result = orm.QueryBuilder().append(cls, filters={'attributes': {'==': value}}).first()
if result is None:
node = cls(dict=value).store()
else:
node = result[0]
else:
raise NotImplementedError
return node
def test_delete_many_dbnode_id(self):
"""Test `delete_many` method filtering on `dbnode_id`"""
# Create logs and separate node
calc = self.backend.nodes.create(
node_type='', user=self.user, computer=self.computer, label='label', description='description'
).store()
log1 = self.create_log(dbnode_id=calc.id)
log2 = self.create_log()
log3 = self.create_log()
log_uuids = []
for log in [log1, log2, log3]:
log.store()
log_uuids.append(str(log.uuid))
# Make sure they exist
count_logs_found = orm.QueryBuilder().append(orm.Log, filters={'uuid': {'in': log_uuids}}).count()
self.assertEqual(
count_logs_found,
len(log_uuids),
msg=f'There should be {len(log_uuids)} Logs, instead {count_logs_found} Log(s) was/were found'
)
# Delete logs for self.node
filters = {'dbnode_id': self.node.id}
self.backend.logs.delete_many(filters=filters)
# Check they were deleted
builder = orm.QueryBuilder().append(orm.Log, filters={'uuid': {'in': log_uuids}}, project='uuid').all()
found_logs_uuid = [_[0] for _ in builder]
self.assertEqual([log_uuids[0]], found_logs_uuid)
def test_group_uuid_hashing_for_querybuidler(self):
"""QueryBuilder results should be reusable and shouldn't brake hashing."""
group = orm.Group(label='test_group')
group.store()
# Search for the UUID of the stored group
builder = orm.QueryBuilder()
builder.append(orm.Group,
project=['uuid'],
filters={'label': {
'==': 'test_group'
}})
[uuid] = builder.first()
# Look the node with the previously returned UUID
builder = orm.QueryBuilder()
builder.append(orm.Group,
project=['id'],
filters={'uuid': {
'==': uuid
}})
# Check that the query doesn't fail
builder.all()
# And that the results are correct
self.assertEqual(builder.count(), 1)
self.assertEqual(builder.first()[0], group.id)
def test_delete_many_dbnode_id(self):
"""Test `delete_many` method filtering on `dbnode_id`"""
# Create comments and separate node
calc = self.backend.nodes.create(
node_type='', user=self.user, computer=self.computer, label='label', description='description'
).store()
comment1 = self.create_comment(node=calc)
comment2 = self.create_comment()
comment3 = self.create_comment()
comment_uuids = []
for comment in [comment1, comment2, comment3]:
comment.store()
comment_uuids.append(str(comment.uuid))
# Make sure they exist
count_comments_found = orm.QueryBuilder().append(orm.Comment, filters={'uuid': {'in': comment_uuids}}).count()
self.assertEqual(
count_comments_found,
len(comment_uuids),
msg='There should be {} Comments, instead {} Comment(s) was/were found'.format(
len(comment_uuids), count_comments_found
)
)
# Delete comments for self.node
filters = {'dbnode_id': self.node.id}
self.backend.comments.delete_many(filters=filters)
# Check they were deleted
builder = orm.QueryBuilder().append(orm.Comment, filters={'uuid': {'in': comment_uuids}}, project='uuid').all()
found_comments_uuid = [_[0] for _ in builder]
self.assertEqual([comment_uuids[0]], found_comments_uuid)
def test_dangling_link_to_existing_db_node(self, temp_dir):
"""A dangling link that references a Node that is not included in the archive should `not` be importable"""
struct = orm.StructureData()
struct.store()
struct_uuid = struct.uuid
calc = orm.CalculationNode()
calc.add_incoming(struct, LinkType.INPUT_CALC, 'input')
calc.store()
calc.seal()
calc_uuid = calc.uuid
filename = os.path.join(temp_dir, 'export.aiida')
export([struct], filename=filename, file_format='tar.gz')
unpack = SandboxFolder()
with tarfile.open(filename, 'r:gz', format=tarfile.PAX_FORMAT) as tar:
tar.extractall(unpack.abspath)
with open(unpack.get_abs_path('data.json'), 'r',
encoding='utf8') as fhandle:
data = json.load(fhandle)
data['links_uuid'].append({
'output': calc.uuid,
'input': struct.uuid,
'label': 'input',
'type': LinkType.INPUT_CALC.value
})
with open(unpack.get_abs_path('data.json'), 'wb') as fhandle:
json.dump(data, fhandle)
with tarfile.open(filename, 'w:gz', format=tarfile.PAX_FORMAT) as tar:
tar.add(unpack.abspath, arcname='')
# Make sure the CalculationNode is still in the database
builder = orm.QueryBuilder().append(orm.CalculationNode,
project='uuid')
self.assertEqual(
builder.count(),
1,
msg=
f'There should be a single CalculationNode, instead {builder.count()} has been found'
)
self.assertEqual(builder.all()[0][0], calc_uuid)
with self.assertRaises(DanglingLinkError):
import_data(filename)
# Using the flag `ignore_unknown_nodes` should import it without problems
import_data(filename, ignore_unknown_nodes=True)
builder = orm.QueryBuilder().append(orm.StructureData, project='uuid')
self.assertEqual(
builder.count(),
1,
msg=
f'There should be a single StructureData, instead {builder.count()} has been found'
)
self.assertEqual(builder.all()[0][0], struct_uuid)
def test_full_type_backwards_compatibility(node_class, restapi_server,
server_url):
"""Functionality test for the compatibility with old process_type entries.
This will only check the integrity of the shape of the tree, there is no checking on how the
data should be represented internally in term of full types, labels, etc. The only thing that
must work correctly is the internal consistency of `full_type` as it is interpreted by the
get_full_type_filters and the querybuilder.
"""
node_empty = node_class()
node_empty.process_type = ''
node_empty.store()
node_nones = node_class()
node_nones.process_type = None
node_nones.store()
server = restapi_server()
server_thread = Thread(target=server.serve_forever)
try:
server_thread.start()
type_count_response = requests.get(f'{server_url}/nodes/full_types',
timeout=10)
finally:
server.shutdown()
# All nodes (contains either a process branch or data branch)
# The main branch for all nodes does not currently return a queryable full_type
current_namespace = type_count_response.json()['data']
assert len(current_namespace['subspaces']) == 1
# All subnodes (contains a workflow, calculation or data_type branch)
current_namespace = current_namespace['subspaces'][0]
query_all = orm.QueryBuilder().append(orm.Node,
filters=get_full_type_filters(
current_namespace['full_type']))
assert len(current_namespace['subspaces']) == 1
assert query_all.count() == 2
# If this is a process node, there is an extra branch before the leaf
# (calcfunction, workfunction, calcjob or workchain)
if issubclass(node_class, orm.ProcessNode):
current_namespace = current_namespace['subspaces'][0]
query_all = orm.QueryBuilder().append(
orm.Node,
filters=get_full_type_filters(current_namespace['full_type']))
assert len(current_namespace['subspaces']) == 1
assert query_all.count() == 2
# This will be the last leaf: the specific data type or the empty process_type
current_namespace = current_namespace['subspaces'][0]
query_all = orm.QueryBuilder().append(orm.Node,
filters=get_full_type_filters(
current_namespace['full_type']))
assert len(current_namespace['subspaces']) == 0
assert query_all.count() == 2
def test_edges(self):
"""
Testing how the links are stored during traversal of the graph.
"""
nodes = self._create_basic_graph()
# Forward traversal (check all nodes and all links)
basket = Basket(nodes=[nodes['data_i'].id])
queryb = orm.QueryBuilder()
queryb.append(orm.Node, tag='nodes_in_set')
queryb.append(orm.Node, with_incoming='nodes_in_set')
uprule = UpdateRule(queryb, max_iterations=2, track_edges=True)
uprule_result = uprule.run(basket.copy())
obtained = uprule_result['nodes'].keyset
expected = set(anode.id for _, anode in nodes.items())
self.assertEqual(obtained, expected)
obtained = set()
for data in uprule_result['nodes_nodes'].keyset:
obtained.add((data[0], data[1]))
expected = {
(nodes['data_i'].id, nodes['calc_0'].id),
(nodes['data_i'].id, nodes['work_1'].id),
(nodes['data_i'].id, nodes['work_2'].id),
(nodes['calc_0'].id, nodes['data_o'].id),
(nodes['work_1'].id, nodes['data_o'].id),
(nodes['work_2'].id, nodes['data_o'].id),
(nodes['work_2'].id, nodes['work_1'].id),
(nodes['work_1'].id, nodes['calc_0'].id),
}
self.assertEqual(obtained, expected)
# Backwards traversal (check partial traversal and link direction)
basket = Basket(nodes=[nodes['data_o'].id])
queryb = orm.QueryBuilder()
queryb.append(orm.Node, tag='nodes_in_set')
queryb.append(orm.Node, with_outgoing='nodes_in_set')
uprule = UpdateRule(queryb, max_iterations=1, track_edges=True)
uprule_result = uprule.run(basket.copy())
obtained = uprule_result['nodes'].keyset
expected = set(anode.id for _, anode in nodes.items())
expected = expected.difference(set([nodes['data_i'].id]))
self.assertEqual(obtained, expected)
obtained = set()
for data in uprule_result['nodes_nodes'].keyset:
obtained.add((data[0], data[1]))
expected = {
(nodes['calc_0'].id, nodes['data_o'].id),
(nodes['work_1'].id, nodes['data_o'].id),
(nodes['work_2'].id, nodes['data_o'].id),
}
self.assertEqual(obtained, expected)
def test_empty_repo_folder_export(self, temp_dir):
"""Check a Node's empty repository folder is exported properly"""
from aiida.common.folders import Folder
from aiida.tools.importexport.dbexport import export_tree
node = orm.Dict().store()
node_uuid = node.uuid
node_repo = RepositoryFolder(section=Repository._section_name, uuid=node_uuid) # pylint: disable=protected-access
self.assertTrue(
node_repo.exists(), msg='Newly created and stored Node should have had an existing repository folder'
)
for filename, is_file in node_repo.get_content_list(only_paths=False):
abspath_filename = os.path.join(node_repo.abspath, filename)
if is_file:
os.remove(abspath_filename)
else:
shutil.rmtree(abspath_filename, ignore_errors=False)
self.assertFalse(
node_repo.get_content_list(),
msg='Repository folder should be empty, instead the following was found: {}'.format(
node_repo.get_content_list()
)
)
archive_variants = {
'archive folder': os.path.join(temp_dir, 'export_tree'),
'tar archive': os.path.join(temp_dir, 'export.tar.gz'),
'zip archive': os.path.join(temp_dir, 'export.zip')
}
export_tree([node], folder=Folder(archive_variants['archive folder']), silent=True)
export([node], filename=archive_variants['tar archive'], file_format='tar.gz', silent=True)
export([node], filename=archive_variants['zip archive'], file_format='zip', silent=True)
for variant, filename in archive_variants.items():
self.reset_database()
node_count = orm.QueryBuilder().append(orm.Dict, project='uuid').count()
self.assertEqual(node_count, 0, msg='After DB reset {} Dict Nodes was (wrongly) found'.format(node_count))
import_data(filename, silent=True)
builder = orm.QueryBuilder().append(orm.Dict, project='uuid')
imported_node_count = builder.count()
self.assertEqual(
imported_node_count,
1,
msg='After {} import a single Dict Node should have been found, '
'instead {} was/were found'.format(variant, imported_node_count)
)
imported_node_uuid = builder.all()[0][0]
self.assertEqual(
imported_node_uuid,
node_uuid,
msg='The wrong UUID was found for the imported {}: '
'{}. It should have been: {}'.format(variant, imported_node_uuid, node_uuid)
)
def cmd_completion_scf(max_atoms, number_species):
"""Determine the completion rate of the reconnaissance SCF step."""
from aiida import orm
filters_structure = {}
if max_atoms is not None:
filters_structure['attributes.sites'] = {'shorter': max_atoms + 1}
if number_species is not None:
filters_structure['attributes.kinds'] = {'of_length': number_species}
query = orm.QueryBuilder()
query.append(orm.Group, filters={'label': 'structure/unique'}, tag='group')
query.append(orm.Node, with_group='group', filters=filters_structure)
nstructures = query.count()
query = orm.QueryBuilder()
query.append(orm.Group, filters={'label': 'workchain/scf'}, tag='group')
query.append(orm.WorkChainNode,
with_group='group',
tag='workchain',
project='attributes.exit_status')
query.append(orm.StructureData,
with_outgoing='workchain',
filters=filters_structure,
project='id')
active = []
failed = []
success = []
submitted = []
for exit_status, pk in query.iterall():
if exit_status == 0:
success.append(pk)
elif exit_status is None:
active.append(pk)
else:
failed.append(pk)
submitted.append(pk)
submitted = set(submitted)
table = [
['Structures', nstructures],
['Submitted', len(submitted)],
['Submitted success', len(success)],
['Submitted failed', len(failed)],
['Submitted active', len(active)],
['Submittable', nstructures - len(submitted)],
]
click.echo(tabulate.tabulate(table, tablefmt='plain'))
def cif_clean(ctx, database, cif_filter, cif_select, concurrent, interval, dry_run):
"""Clean the CIF files imported from an external database.
Run the `aiida-codtools workflow launch cif-clean` CLI script to clean the imported CIFs of the given database.
"""
from datetime import datetime
from time import sleep
from aiida import orm
from aiida_codtools.cli.workflows.cif_clean import launch_cif_clean
now = datetime.utcnow().isoformat
group_cif_raw = orm.Group.get(label='{database}/cif/raw'.format(database=database))
group_cif_clean = orm.Group.get(label='{database}/cif/clean'.format(database=database))
group_structure = orm.Group.get(label='{database}/structure/primitive'.format(database=database))
group_workchain = orm.Group.get(label='{database}/workchain/clean'.format(database=database))
while(True):
filters = {'attributes.process_state': {'or': [{'==': 'excepted'}, {'==': 'killed'}]}}
builder = orm.QueryBuilder().append(orm.ProcessNode, filters=filters)
if builder.count() > 0:
echo.echo_critical('found {} excepted or killed processes, exiting'.format(builder.count()))
filters = {'attributes.process_state': {'or': [{'==': 'waiting'}, {'==': 'running'}, {'==': 'created'}]}}
builder = orm.QueryBuilder().append(orm.WorkChainNode, filters=filters)
current = builder.count()
max_entries = concurrent - current
if current < concurrent:
echo.echo('{} | currently {} running workchains, submitting {} more'.format(now(), current, max_entries))
inputs = {
'cif_filter': cif_filter,
'cif_select': cif_select,
'group_cif_raw': group_cif_raw,
'group_cif_clean': group_cif_clean,
'group_structure': group_structure,
'group_workchain': group_workchain,
'max_entries': max_entries,
'skip_check': False,
'parse_engine': 'ase',
'daemon': True,
}
if dry_run:
from pprint import pprint
echo.echo(pprint(inputs))
return
else:
ctx.invoke(launch_cif_clean, **inputs)
else:
echo.echo('{} | currently {} running workchains, nothing to submit'.format(now(), current))
echo.echo('{} | sleeping {} seconds'.format(now(), interval))
sleep(interval)
def test_mtime_of_imported_comments(self, temp_dir):
"""
Test mtime does not change for imported comments
This is related to correct usage of `comment_mode` when importing.
"""
# Get user
user = orm.User.objects.get_default()
comment_content = 'You get what you give'
# Create node
calc = orm.CalculationNode().store()
calc.seal()
# Create comment
orm.Comment(calc, user, comment_content).store()
calc.store()
# Save UUIDs and mtime
calc_uuid = calc.uuid
builder = orm.QueryBuilder().append(orm.Comment, project=['uuid', 'mtime']).all()
comment_uuid = str(builder[0][0])
comment_mtime = builder[0][1]
builder = orm.QueryBuilder().append(orm.CalculationNode, project=['uuid', 'mtime']).all()
calc_uuid = str(builder[0][0])
calc_mtime = builder[0][1]
# Export, reset database and reimport
export_file = os.path.join(temp_dir, 'export.aiida')
export([calc], filename=export_file, silent=True)
self.reset_database()
import_data(export_file, silent=True)
# Retrieve node and comment
builder = orm.QueryBuilder().append(orm.CalculationNode, tag='calc', project=['uuid', 'mtime'])
builder.append(orm.Comment, with_node='calc', project=['uuid', 'mtime'])
import_entities = builder.all()[0]
self.assertEqual(len(import_entities), 4) # Check we have the correct amount of returned values
import_calc_uuid = str(import_entities[0])
import_calc_mtime = import_entities[1]
import_comment_uuid = str(import_entities[2])
import_comment_mtime = import_entities[3]
# Check we have the correct UUIDs
self.assertEqual(import_calc_uuid, calc_uuid)
self.assertEqual(import_comment_uuid, comment_uuid)
# Make sure the mtime is the same after import as it was before export
self.assertEqual(import_comment_mtime, comment_mtime)
self.assertEqual(import_calc_mtime, calc_mtime)
def test_delete_many_ctime_mtime(self):
"""Test `delete_many` method filtering on `ctime` and `mtime`"""
from datetime import timedelta
# Initialization
comment_uuids = []
found_comments_ctime = []
found_comments_mtime = []
found_comments_uuid = []
now = timezone.now()
two_days_ago = now - timedelta(days=2)
one_day_ago = now - timedelta(days=1)
comment_times = [now, one_day_ago, two_days_ago]
# Create comments
comment1 = self.create_comment(ctime=now, mtime=now)
comment2 = self.create_comment(ctime=one_day_ago, mtime=now)
comment3 = self.create_comment(ctime=two_days_ago, mtime=one_day_ago)
for comment in [comment1, comment2, comment3]:
comment.store()
comment_uuids.append(str(comment.uuid))
# Make sure they exist with the correct times
builder = orm.QueryBuilder().append(orm.Comment, project=['ctime', 'mtime', 'uuid'])
self.assertGreater(builder.count(), 0)
for comment in builder.all():
found_comments_ctime.append(comment[0])
found_comments_mtime.append(comment[1])
found_comments_uuid.append(comment[2])
for time, uuid in zip(comment_times, comment_uuids):
self.assertIn(time, found_comments_ctime)
self.assertIn(uuid, found_comments_uuid)
if time != two_days_ago:
self.assertIn(time, found_comments_mtime)
# Delete comments that are created more than 1 hour ago,
# unless they have been modified within 5 hours
ctime_turning_point = now - timedelta(seconds=60 * 60)
mtime_turning_point = now - timedelta(seconds=60 * 60 * 5)
filters = {'and': [{'ctime': {'<': ctime_turning_point}}, {'mtime': {'<': mtime_turning_point}}]}
self.backend.comments.delete_many(filters=filters)
# Check only the most stale comment (comment3) was deleted
builder = orm.QueryBuilder().append(orm.Comment, project='uuid')
self.assertGreater(builder.count(), 1) # There should still be at least 2
found_comments_uuid = [_[0] for _ in builder.all()]
self.assertNotIn(comment_uuids[2], found_comments_uuid)
# Make sure the other comments were not deleted
for comment_uuid in comment_uuids[:-1]:
self.assertIn(comment_uuid, found_comments_uuid)
def test_querying():
"""Test querying for groups with and without subclassing."""
orm.Group(label='group').store()
orm.AutoGroup(label='auto-group').store()
# Fake a subclass by manually setting the type string
group = orm.Group(label='custom-group')
group.backend_entity.dbmodel.type_string = 'custom.group'
group.store()
assert orm.QueryBuilder().append(orm.AutoGroup).count() == 1
assert orm.QueryBuilder().append(orm.AutoGroup, subclassing=False).count() == 1
assert orm.QueryBuilder().append(orm.Group, subclassing=False).count() == 1
assert orm.QueryBuilder().append(orm.Group).count() == 3
assert orm.QueryBuilder().append(orm.Group, filters={'type_string': 'custom.group'}).count() == 1
def test_delete_many_time(self):
"""Test `delete_many` method filtering on `time`"""
from datetime import timedelta
# Initialization
log_uuids = []
found_logs_time = []
found_logs_uuid = []
now = timezone.now()
two_days_ago = now - timedelta(days=2)
one_day_ago = now - timedelta(days=1)
log_times = [now, one_day_ago, two_days_ago]
# Create logs
log1 = self.create_log(time=now)
log2 = self.create_log(time=one_day_ago)
log3 = self.create_log(time=two_days_ago)
for log in [log1, log2, log3]:
log.store()
log_uuids.append(str(log.uuid))
# Make sure they exist with the correct times
builder = orm.QueryBuilder().append(orm.Log, project=['time', 'uuid'])
self.assertGreater(builder.count(), 0)
for log in builder.all():
found_logs_time.append(log[0])
found_logs_uuid.append(log[1])
for log_time in log_times:
self.assertIn(log_time, found_logs_time)
for log_uuid in log_uuids:
self.assertIn(log_uuid, found_logs_uuid)
# Delete logs that are older than 1 hour
turning_point = now - timedelta(seconds=60 * 60)
filters = {'time': {'<': turning_point}}
self.backend.logs.delete_many(filters=filters)
# Check they were deleted
builder = orm.QueryBuilder().append(orm.Log, project='uuid')
self.assertGreater(builder.count(),
0) # There should still be at least 1
found_logs_uuid = [_[0] for _ in builder.all()]
for log_uuid in log_uuids[1:]:
self.assertNotIn(log_uuid, found_logs_uuid)
# Make sure the newest log (log1) was not deleted
self.assertIn(log_uuids[0], found_logs_uuid)
def link_flags_import_helper(self, test_data):
"""Helper function"""
for test, (export_file, _, expected_nodes) in test_data.items():
self.clean_db()
import_data(export_file)
nodes_util = {
'work': orm.WorkflowNode,
'calc': orm.CalculationNode,
'data': orm.Data
}
for node_type, node_cls in nodes_util.items():
if node_type in expected_nodes:
builder = orm.QueryBuilder().append(node_cls,
project='uuid')
self.assertEqual(
builder.count(),
len(expected_nodes[node_type]),
msg='Expected {} {} node(s), but got {}. Test: "{}"'.
format(len(expected_nodes[node_type]), node_type,
builder.count(), test))
for node_uuid in builder.iterall():
self.assertIn(node_uuid[0],
expected_nodes[node_type],
msg=f'Failed for test: "{test}"')
def test_complex_workflow_graph_export_sets(self, temp_dir):
"""Test ex-/import of individual nodes in complex graph"""
for export_conf in range(0, 9):
_, (export_node,
export_target) = self.construct_complex_graph(export_conf)
export_target_uuids = set(_.uuid for _ in export_target)
export_file = os.path.join(temp_dir, 'export.aiida')
export([export_node], filename=export_file, overwrite=True)
export_node_str = str(export_node)
self.clean_db()
import_data(export_file)
# Get all the nodes of the database
builder = orm.QueryBuilder()
builder.append(orm.Node, project='uuid')
imported_node_uuids = set(_[0] for _ in builder.all())
self.assertSetEqual(
export_target_uuids, imported_node_uuids,
'Problem in comparison of export node: ' + export_node_str +
'\n' + 'Expected set: ' + str(export_target_uuids) + '\n' +
'Imported set: ' + str(imported_node_uuids) + '\n' +
'Difference: ' + str(
export_target_uuids.symmetric_difference(
imported_node_uuids)))
def children(self):
# type: () -> Iterable[GroupPath]
"""Iterate through all (direct) children of this path."""
query = orm.QueryBuilder()
filters = {}
if self.path:
filters['label'] = {'like': self.path + self.delimiter + '%'}
query.append(self.cls,
subclassing=False,
filters=filters,
project='label')
if query.count() == 0 and self.is_virtual:
raise NoGroupsInPathError(self)
yielded = []
for (label, ) in query.iterall():
path = label.split(self._delimiter)
if len(path) <= len(self._path_list):
continue
path_string = self._delimiter.join(path[:len(self._path_list) + 1])
if (path_string not in yielded
and path[:len(self._path_list)] == self._path_list):
yielded.append(path_string)
try:
yield GroupPath(
path=path_string,
cls=self.cls,
warn_invalid_child=self._warn_invalid_child)
except InvalidPath:
if self._warn_invalid_child:
warnings.warn(
'invalid path encountered: {}'.format(path_string)) # pylint: disable=no-member
def test_explicit_type_string():
"""Test that passing explicit `type_string` to `Group` constructor is still possible despite being deprecated.
Both constructing a group while passing explicit `type_string` as well as loading a group with unregistered
type string should emit a warning, but it should be possible.
"""
type_string = 'data.potcar' # An unregistered custom type string
with pytest.warns(UserWarning):
group = orm.Group(label='group', type_string=type_string)
group.store()
assert group.type_string == type_string
with pytest.warns(UserWarning):
loaded = orm.Group.get(label=group.label, type_string=type_string)
assert isinstance(loaded, orm.Group)
assert loaded.pk == group.pk
assert loaded.type_string == group.type_string
queried = orm.QueryBuilder().append(orm.Group,
filters={
'id': group.pk,
'type_string': type_string
}).one()[0]
assert isinstance(queried, orm.Group)
assert queried.pk == group.pk
assert queried.type_string == group.type_string
def _find(backend: orm.implementation.Backend, entity_type: orm.Entity,
**kwargs) -> orm.QueryBuilder:
for key in kwargs:
if key not in {
"filters", "order_by", "limit", "project", "offset"
}:
raise ValueError(
f"You supplied key {key}. _find() only takes the keys: "
'"filters", "order_by", "limit", "project", "offset"')
filters = kwargs.get("filters", {})
order_by = kwargs.get("order_by", None)
order_by = {
entity_type: order_by
} if order_by else {
entity_type: {
"id": "asc"
}
}
limit = kwargs.get("limit", None)
offset = kwargs.get("offset", None)
project = kwargs.get("project", [])
query = orm.QueryBuilder(backend=backend, limit=limit, offset=offset)
query.append(entity_type, project=project, filters=filters)
query.order_by(order_by)
return query
def walk_nodes(self, filters=None, node_class=None, query_batch=None):
# type: () -> Iterable[WalkNodeResult]
"""Recursively iterate through all nodes of this path and its children.
:param filters: filters to apply to the node query
:param node_class: return only nodes of a certain class (or list of classes)
:param int batch_size: The size of the batches to ask the backend to batch results in subcollections.
You can optimize the speed of the query by tuning this parameter.
Be aware though that is only safe if no commit will take place during this transaction.
"""
query = orm.QueryBuilder()
group_filters = {}
if self.path:
group_filters['label'] = {
'or': [{
'==': self.path
}, {
'like': self.path + self.delimiter + '%'
}]
}
query.append(self.cls,
subclassing=False,
filters=group_filters,
project='label',
tag='group')
query.append(
orm.Node if node_class is None else node_class,
with_group='group',
filters=filters,
project=['*'],
)
for (label, node
) in query.iterall(query_batch) if query_batch else query.all():
yield WalkNodeResult(GroupPath(label, cls=self.cls), node)
def test_import_of_django_sqla_export_file(self):
"""Check that sqla import manages to import the django export file correctly"""
from tests.utils.archives import import_archive
for archive in ['django.aiida', 'sqlalchemy.aiida']:
# Clean the database
self.reset_database()
# Import the needed data
import_archive(archive, filepath='export/compare')
# The expected metadata
comp1_metadata = {'workdir': '/tmp/aiida'}
# Check that we got the correct metadata
# Make sure to exclude the default computer
builder = orm.QueryBuilder()
builder.append(orm.Computer,
project=['metadata'],
tag='comp',
filters={'name': {
'!==': self.computer.label
}})
self.assertEqual(builder.count(), 1, 'Expected only one computer')
res = builder.dict()[0]
self.assertEqual(res['comp']['metadata'], comp1_metadata)
