def test_get_or_create_invalid_prefix(self):
"""Test the ``get_or_create_group`` method of ``Autogroup`` when there is already a group
with the same prefix, but followed by other non-underscore characters."""
label_prefix = 'new_test_prefix_TestAutogroup'
# I create a group with the same prefix, but followed by non-underscore
# characters. These should be ignored in the logic.
AutoGroup(label=f'{label_prefix}xx').store()
# Check that there are no groups to begin with
queryb = QueryBuilder().append(AutoGroup, filters={'label': label_prefix})
assert not list(queryb.all())
queryb = QueryBuilder().append(AutoGroup, filters={'label': {'like': r'{}\_%'.format(label_prefix)}})
assert not list(queryb.all())
# First group (no existing one)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix
self.assertEqual(
group.label, expected_label,
f"The auto-group should be labelled '{expected_label}', it is instead '{group.label}'"
)
# Second group (only one with no suffix existing)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = f'{label_prefix}_1'
self.assertEqual(
group.label, expected_label,
f"The auto-group should be labelled '{expected_label}', it is instead '{group.label}'"
)
def setup_groups(clear_database_before_test):
"""Create some groups to test the `GroupParamType` parameter type for the command line infrastructure.
We create an initial group with a random name and then on purpose create two groups with a name that matches exactly
the ID and UUID, respectively, of the first one. This allows us to test the rules implemented to solve ambiguities
that arise when determing the identifier type.
"""
entity_01 = Group(label='group_01').store()
entity_02 = AutoGroup(label=str(entity_01.pk)).store()
entity_03 = ImportGroup(label=str(entity_01.uuid)).store()
return entity_01, entity_02, entity_03
def get_or_create_group(self):
"""Return the current `AutoGroup`, or create one if None has been set yet.
This function implements a somewhat complex logic that is however needed
to make sure that, even if `verdi run` is called at the same time multiple
times, e.g. in a for loop in bash, there is never the risk that two ``verdi run``
Unix processes try to create the same group, with the same label, ending
up in a crash of the code (see PR #3650).
Here, instead, we make sure that if this concurrency issue happens,
one of the two will get a IntegrityError from the DB, and then recover
trying to create a group with a different label (with a numeric suffix appended),
until it manages to create it.
"""
from aiida.orm import QueryBuilder
# When this function is called, if it is the first time, just generate
# a new group name (later on, after this ``if`` block`).
# In that case, we will later cache in ``self._group_label`` the group label,
# So the group with the same name can be returned quickly in future
# calls of this method.
if self._group_label is not None:
builder = QueryBuilder().append(
AutoGroup, filters={'label': self._group_label})
results = [res[0] for res in builder.iterall()]
if results:
# If it is not empty, it should have only one result due to the uniqueness constraints
assert len(
results
) == 1, 'I got more than one autogroup with the same label!'
return results[0]
# There are no results: probably the group has been deleted.
# I continue as if it was not cached
self._group_label = None
label_prefix = self.get_group_label_prefix()
# Try to do a preliminary QB query to avoid to do too many try/except
# if many of the prefix_NUMBER groups already exist
queryb = QueryBuilder().append(
AutoGroup,
filters={
'or': [{
'label': {
'==': label_prefix
}
}, {
'label': {
'like': escape_for_sql_like(label_prefix + '_') + '%'
}
}]
},
project='label')
existing_group_labels = [
res[0][len(label_prefix):] for res in queryb.all()
]
existing_group_ints = []
for label in existing_group_labels:
if label == '':
# This is just the prefix without name - corresponds to counter = 0
existing_group_ints.append(0)
elif label.startswith('_'):
try:
existing_group_ints.append(int(label[1:]))
except ValueError:
# It's not an integer, so it will never collide - just ignore it
pass
if not existing_group_ints:
counter = 0
else:
counter = max(existing_group_ints) + 1
while True:
try:
label = label_prefix if counter == 0 else '{}_{}'.format(
label_prefix, counter)
group = AutoGroup(label=label).store()
self._group_label = group.label
except exceptions.IntegrityError:
counter += 1
else:
break
return group
def test_get_or_create(self):
"""Test the ``get_or_create_group`` method of ``Autogroup``."""
label_prefix = 'test_prefix_TestAutogroup'
# Check that there are no groups to begin with
queryb = QueryBuilder().append(AutoGroup,
filters={'label': label_prefix})
assert not list(queryb.all())
queryb = QueryBuilder().append(
AutoGroup,
filters={'label': {
'like': r'{}\_%'.format(label_prefix)
}})
assert not list(queryb.all())
# First group (no existing one)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))
# Second group (only one with no suffix existing)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix + '_1'
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))
# Second group (only one suffix _1 existing)
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix + '_2'
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))
# I create a group with a large integer suffix (9)
AutoGroup(label='{}_9'.format(label_prefix)).store()
# The next autogroup should become number 10
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix + '_10'
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))
# I create a group with a non-integer suffix (15a), it should be ignored
AutoGroup(label='{}_15b'.format(label_prefix)).store()
# The next autogroup should become number 11
autogroup = Autogroup()
autogroup.set_group_label_prefix(label_prefix)
group = autogroup.get_or_create_group()
expected_label = label_prefix + '_11'
self.assertEqual(
group.label, expected_label,
"The auto-group should be labelled '{}', it is instead '{}'".
format(expected_label, group.label))