Skip to content

Commit

Permalink
Make Group sub classable through entry points
Browse files Browse the repository at this point in the history
We add the `aiida.groups` entry point group where sub classes of the
`aiida.orm.groups.Group` class can be registered. A new metaclass is
used to automatically set the `type_string` based on the entry point of
the `Group` sub class. This will make it possible to reload the correct
sub class when reloading from the database.

If the `GroupMeta` metclass fails cannot retrieve the corresponding
entry point of the subclass, a warning is issued that any instances
of this class will not be storable and the `__type_string` attribute is
set to `None`. This can be checked by the `store` method which will
make it fail. We choose to only except in the `store` method such that
it is still possible to define and instantiate subclasses of `Group`
that have not yet been registered. This is useful for testing and
experimenting.

Since the group type strings are now based on the entry point names, the
existing group type strings in the database have to be migrated:

 * `user` -> `core`
 * `data.upf.family` -> `core.upf`
 * `auto.import` -> `core.import`
 * `auto.run` -> `core.run`

When loading a `Group` instance from the database, the loader will try
to resolve the type string to the corresponding subclass through the
entry points. If this fails, a warning is issued and we fallback on the
base `Group` class.
  • Loading branch information
sphuber committed Apr 5, 2020
1 parent 5c0d86f commit c3906c6
Show file tree
Hide file tree
Showing 26 changed files with 554 additions and 233 deletions.
44 changes: 44 additions & 0 deletions aiida/backends/djsite/db/migrations/0044_dbgroup_type_string.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
# pylint: disable=invalid-name,too-few-public-methods
"""Migration after the `Group` class became pluginnable and so the group `type_string` changed."""

# pylint: disable=no-name-in-module,import-error
from django.db import migrations
from aiida.backends.djsite.db.migrations import upgrade_schema_version

REVISION = '1.0.44'
DOWN_REVISION = '1.0.43'

forward_sql = [
"""UPDATE db_dbgroup SET type_string = 'core' WHERE type_string = 'user';""",
"""UPDATE db_dbgroup SET type_string = 'core.upf' WHERE type_string = 'data.upf';""",
"""UPDATE db_dbgroup SET type_string = 'core.import' WHERE type_string = 'auto.import';""",
"""UPDATE db_dbgroup SET type_string = 'core.run' WHERE type_string = 'auto.run';""",
]

reverse_sql = [
"""UPDATE db_dbgroup SET type_string = 'user' WHERE type_string = 'core';""",
"""UPDATE db_dbgroup SET type_string = 'data.upf' WHERE type_string = 'core.upf';""",
"""UPDATE db_dbgroup SET type_string = 'auto.import' WHERE type_string = 'core.import';""",
"""UPDATE db_dbgroup SET type_string = 'auto.run' WHERE type_string = 'core.run';""",
]


class Migration(migrations.Migration):
"""Migration after the update of group `type_string`"""
dependencies = [
('db', '0043_default_link_label'),
]

operations = [
migrations.RunSQL(sql='\n'.join(forward_sql), reverse_sql='\n'.join(reverse_sql)),
upgrade_schema_version(REVISION, DOWN_REVISION),
]
2 changes: 1 addition & 1 deletion aiida/backends/djsite/db/migrations/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ class DeserializationException(AiidaException):
pass


LATEST_MIGRATION = '0043_default_link_label'
LATEST_MIGRATION = '0044_dbgroup_type_string'


def _update_schema_version(version, apps, _):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
# -*- coding: utf-8 -*-
"""Migration after the `Group` class became pluginnable and so the group `type_string` changed.
Revision ID: bf591f31dd12
Revises: 118349c10896
Create Date: 2020-03-31 10:00:52.609146
"""
# pylint: disable=no-name-in-module,import-error,invalid-name,no-member
from alembic import op
from sqlalchemy.sql import text

forward_sql = [
"""UPDATE db_dbgroup SET type_string = 'core' WHERE type_string = 'user';""",
"""UPDATE db_dbgroup SET type_string = 'core.upf' WHERE type_string = 'data.upf';""",
"""UPDATE db_dbgroup SET type_string = 'core.import' WHERE type_string = 'auto.import';""",
"""UPDATE db_dbgroup SET type_string = 'core.run' WHERE type_string = 'auto.run';""",
]

reverse_sql = [
"""UPDATE db_dbgroup SET type_string = 'user' WHERE type_string = 'core';""",
"""UPDATE db_dbgroup SET type_string = 'data.upf' WHERE type_string = 'core.upf';""",
"""UPDATE db_dbgroup SET type_string = 'auto.import' WHERE type_string = 'core.import';""",
"""UPDATE db_dbgroup SET type_string = 'auto.run' WHERE type_string = 'core.run';""",
]

# revision identifiers, used by Alembic.
revision = 'bf591f31dd12'
down_revision = '118349c10896'
branch_labels = None
depends_on = None


def upgrade():
"""Migrations for the upgrade."""
conn = op.get_bind()
statement = text('\n'.join(forward_sql))
conn.execute(statement)


def downgrade():
"""Migrations for the downgrade."""
conn = op.get_bind()
statement = text('\n'.join(reverse_sql))
conn.execute(statement)
11 changes: 1 addition & 10 deletions aiida/cmdline/commands/cmd_data/cmd_upf.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,22 +64,13 @@ def upf_listfamilies(elements, with_description):
"""
from aiida import orm
from aiida.plugins import DataFactory
from aiida.orm.nodes.data.upf import UPFGROUP_TYPE

UpfData = DataFactory('upf') # pylint: disable=invalid-name
query = orm.QueryBuilder()
query.append(UpfData, tag='upfdata')
if elements is not None:
query.add_filter(UpfData, {'attributes.element': {'in': elements}})
query.append(
orm.Group,
with_node='upfdata',
tag='group',
project=['label', 'description'],
filters={'type_string': {
'==': UPFGROUP_TYPE
}}
)
query.append(orm.UpfFamily, with_node='upfdata', tag='group', project=['label', 'description'])

query.distinct()
if query.count() > 0:
Expand Down
22 changes: 4 additions & 18 deletions aiida/cmdline/commands/cmd_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

from aiida.common.exceptions import UniquenessError
from aiida.cmdline.commands.cmd_verdi import verdi
from aiida.cmdline.params import options, arguments, types
from aiida.cmdline.params import options, arguments
from aiida.cmdline.utils import echo
from aiida.cmdline.utils.decorators import with_dbenv

Expand Down Expand Up @@ -178,18 +178,6 @@ def group_show(group, raw, limit, uuid):
echo.echo(tabulate(table, headers=header))


@with_dbenv()
def valid_group_type_strings():
from aiida.orm import GroupTypeString
return tuple(i.value for i in GroupTypeString)


@with_dbenv()
def user_defined_group():
from aiida.orm import GroupTypeString
return GroupTypeString.USER.value


@verdi_group.command('list')
@options.ALL_USERS(help='Show groups for all users, rather than only for the current user')
@click.option(
Expand All @@ -204,8 +192,7 @@ def user_defined_group():
'-t',
'--type',
'group_type',
type=types.LazyChoice(valid_group_type_strings),
default=user_defined_group,
default='core',
help='Show groups of a specific type, instead of user-defined groups. Start with semicolumn if you want to '
'specify aiida-internal type'
)
Expand Down Expand Up @@ -330,9 +317,8 @@ def group_list(
def group_create(group_label):
"""Create an empty group with a given name."""
from aiida import orm
from aiida.orm import GroupTypeString

group, created = orm.Group.objects.get_or_create(label=group_label, type_string=GroupTypeString.USER.value)
group, created = orm.Group.objects.get_or_create(label=group_label)

if created:
echo.echo_success("Group created with PK = {} and name '{}'".format(group.id, group.label))
Expand All @@ -351,7 +337,7 @@ def group_copy(source_group, destination_group):
Note that the destination group may not exist."""
from aiida import orm

dest_group, created = orm.Group.objects.get_or_create(label=destination_group, type_string=source_group.type_string)
dest_group, created = orm.Group.objects.get_or_create(label=destination_group)

# Issue warning if destination group is not empty and get user confirmation to continue
if not created and not dest_group.is_empty:
Expand Down
1 change: 1 addition & 0 deletions aiida/cmdline/commands/cmd_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,5 +150,6 @@ def run(scriptname, varargs, auto_group, auto_group_label_prefix, group_name, ex
# Re-raise the exception to have the error code properly returned at the end
raise
finally:
autogroup.current_autogroup = None
if handle:
handle.close()
4 changes: 2 additions & 2 deletions aiida/cmdline/params/types/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,12 @@ def orm_class_loader(self):

@with_dbenv()
def convert(self, value, param, ctx):
from aiida.orm import Group, GroupTypeString
from aiida.orm import Group
try:
group = super().convert(value, param, ctx)
except click.BadParameter:
if self._create_if_not_exist:
group = Group(label=value, type_string=GroupTypeString.USER.value)
group = Group(label=value)
else:
raise

Expand Down
33 changes: 12 additions & 21 deletions aiida/orm/autogroup.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,18 @@
from aiida.common import exceptions, timezone
from aiida.common.escaping import escape_for_sql_like, get_regex_pattern_from_sql
from aiida.common.warnings import AiidaDeprecationWarning
from aiida.orm import GroupTypeString, Group
from aiida.orm import AutoGroup
from aiida.plugins.entry_point import get_entry_point_string_from_class

CURRENT_AUTOGROUP = None

VERDIAUTOGROUP_TYPE = GroupTypeString.VERDIAUTOGROUP_TYPE.value


class Autogroup:
"""
An object used for the autogrouping of objects.
The autogrouping is checked by the Node.store() method.
In the store(), the Node will check if CURRENT_AUTOGROUP is != None.
If so, it will call Autogroup.is_to_be_grouped, and decide whether to put it in a group.
Such autogroups are going to be of the VERDIAUTOGROUP_TYPE.
"""Class to create a new `AutoGroup` instance that will, while active, automatically contain all nodes being stored.
The autogrouping is checked by the `Node.store()` method which, if `CURRENT_AUTOGROUP is not None` the method
`Autogroup.is_to_be_grouped` is called to decide whether to put the current node being stored in the current
`AutoGroup` instance.
The exclude/include lists are lists of strings like:
``aiida.data:int``, ``aiida.calculation:quantumespresso.pw``,
Expand Down Expand Up @@ -198,7 +195,7 @@ def clear_group_cache(self):
self._group_label = None

def get_or_create_group(self):
"""Return the current Autogroup, or create one if None has been set yet.
"""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
Expand All @@ -219,16 +216,10 @@ def get_or_create_group(self):
# So the group with the same name can be returned quickly in future
# calls of this method.
if self._group_label is not None:
results = [
res[0] for res in QueryBuilder().
append(Group, filters={
'label': self._group_label,
'type_string': VERDIAUTOGROUP_TYPE
}, project='*').iterall()
]
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
# 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.
Expand All @@ -239,7 +230,7 @@ def get_or_create_group(self):
# 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(
Group,
AutoGroup,
filters={
'or': [{
'label': {
Expand Down Expand Up @@ -274,7 +265,7 @@ def get_or_create_group(self):
while True:
try:
label = label_prefix if counter == 0 else '{}_{}'.format(label_prefix, counter)
group = Group(label=label, type_string=VERDIAUTOGROUP_TYPE).store()
group = AutoGroup(label=label).store()
self._group_label = group.label
except exceptions.IntegrityError:
counter += 1
Expand Down
5 changes: 3 additions & 2 deletions aiida/orm/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,9 @@ def _(backend_entity):

@get_orm_entity.register(BackendGroup)
def _(backend_entity):
from . import groups
return groups.Group.from_backend_entity(backend_entity)
from .groups import load_group_class
group_class = load_group_class(backend_entity.type_string)
return group_class.from_backend_entity(backend_entity)


@get_orm_entity.register(BackendComputer)
Expand Down
Loading

0 comments on commit c3906c6

Please sign in to comment.