From 6e272a3f715e60f63129c072d192eeea2e44b771 Mon Sep 17 00:00:00 2001 From: Chris Sewell Date: Sun, 24 Oct 2021 21:34:09 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=91=8C=20IMPROVE:=20Add=20`AuthInfo`=20jo?= =?UTF-8?q?ins=20to=20`QueryBuilder`=20(#5195)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Allow for `with_user` and `with_computer` when querying `AuthInfo` --- aiida/orm/computers.py | 4 +- aiida/orm/implementation/querybuilder.py | 2 +- .../sqlalchemy/querybuilder/joiner.py | 377 ++++++++++-------- tests/orm/test_querybuilder.py | 20 + 4 files changed, 225 insertions(+), 178 deletions(-) diff --git a/aiida/orm/computers.py b/aiida/orm/computers.py index 4c77d96f41..ffeba6ad8f 100644 --- a/aiida/orm/computers.py +++ b/aiida/orm/computers.py @@ -84,7 +84,7 @@ def objects(cls) -> ComputerCollection: # pylint: disable=no-self-argument def __init__( # pylint: disable=too-many-arguments self, label: str = None, - hostname: str = None, + hostname: str = '', description: str = '', transport_type: str = '', scheduler_type: str = '', @@ -137,7 +137,7 @@ def _hostname_validator(cls, hostname: str) -> None: """ Validates the hostname. """ - if not hostname.strip(): + if not (hostname or hostname.strip()): raise exceptions.ValidationError('No hostname specified') @classmethod diff --git a/aiida/orm/implementation/querybuilder.py b/aiida/orm/implementation/querybuilder.py index 6e2c0de340..cb55bdc0b1 100644 --- a/aiida/orm/implementation/querybuilder.py +++ b/aiida/orm/implementation/querybuilder.py @@ -41,7 +41,7 @@ class EntityTypes(Enum): EntityRelationships: Dict[str, Set[str]] = { - 'authinfo': set(), + 'authinfo': {'with_computer', 'with_user'}, 'comment': {'with_node', 'with_user'}, 'computer': {'with_node'}, 'group': {'with_node', 'with_user'}, diff --git a/aiida/orm/implementation/sqlalchemy/querybuilder/joiner.py b/aiida/orm/implementation/sqlalchemy/querybuilder/joiner.py index 08dd41ee58..b0e407fcb8 100644 --- a/aiida/orm/implementation/sqlalchemy/querybuilder/joiner.py +++ b/aiida/orm/implementation/sqlalchemy/querybuilder/joiner.py @@ -33,6 +33,10 @@ class _EntityMapper(Protocol): # pylint: disable=invalid-name + @property + def AuthInfo(self) -> Type[Model]: + ... + @property def Node(self) -> Type[Model]: ... @@ -97,41 +101,213 @@ def _entity_join_map(self) -> Dict[str, Dict[str, JoinFuncType]]: and the second defines the relationship with respect to a given tag. """ mapping = { - 'node': { - 'with_log': self._join_log_node, - 'with_comment': self._join_comment_node, - 'with_incoming': self._join_outputs, - 'with_outgoing': self._join_inputs, - 'with_descendants': self._join_ancestors_recursive, - 'with_ancestors': self._join_descendants_recursive, - 'with_computer': self._join_to_computer_used, - 'with_user': self._join_created_by, - 'with_group': self._join_group_members, + 'authinfo': { + 'with_computer': self._join_computer_authinfo, + 'with_user': self._join_user_authinfo, }, - 'computer': { - 'with_node': self._join_computer, + 'comment': { + 'with_node': self._join_node_comment, + 'with_user': self._join_user_comment, }, - 'user': { - 'with_comment': self._join_comment_user, - 'with_node': self._join_creator_of, - 'with_group': self._join_group_user, + 'computer': { + 'with_node': self._join_node_computer, }, 'group': { - 'with_node': self._join_groups, + 'with_node': self._join_node_group, 'with_user': self._join_user_group, }, - 'comment': { - 'with_node': self._join_node_comment, - 'with_user': self._join_user_comment, - }, 'log': { 'with_node': self._join_node_log, - } + }, + 'node': { + 'with_log': self._join_log_node, + 'with_comment': self._join_comment_node, + 'with_incoming': self._join_node_outputs, + 'with_outgoing': self._join_node_inputs, + 'with_descendants': self._join_node_ancestors_recursive, + 'with_ancestors': self._join_node_descendants_recursive, + 'with_computer': self._join_computer_node, + 'with_user': self._join_user_node, + 'with_group': self._join_group_node, + }, + 'user': { + 'with_comment': self._join_comment_user, + 'with_node': self._join_node_user, + 'with_group': self._join_group_user, + }, } return mapping # type: ignore - def _join_outputs(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + def _join_computer_authinfo(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: the aliased user you want to join to + :param entity_to_join: the (aliased) node or group in the DB to join with + """ + _check_dbentities((joined_entity, self._entities.Computer), (entity_to_join, self._entities.AuthInfo), + 'with_computer') + new_query = query.join(entity_to_join, entity_to_join.dbcomputer_id == joined_entity.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_user_authinfo(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: the aliased user you want to join to + :param entity_to_join: the (aliased) node or group in the DB to join with + """ + _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.AuthInfo), 'with_user') + new_query = query.join(entity_to_join, entity_to_join.aiidauser_id == joined_entity.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_group_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: + The (aliased) ORMclass that is + a group in the database + :param entity_to_join: + The (aliased) ORMClass that is a node and member of the group + + **joined_entity** and **entity_to_join** + are joined via the table_groups_nodes table. + from **joined_entity** as group to **enitity_to_join** as node. + (**enitity_to_join** is *with_group* **joined_entity**) + """ + _check_dbentities((joined_entity, self._entities.Group), (entity_to_join, self._entities.Node), 'with_group') + aliased_group_nodes = aliased(self._entities.table_groups_nodes) + new_query = query.join(aliased_group_nodes, aliased_group_nodes.c.dbgroup_id == joined_entity.id).join( + entity_to_join, entity_to_join.id == aliased_group_nodes.c.dbnode_id, isouter=isouterjoin + ) + return JoinReturn(new_query, aliased_group_nodes) + + def _join_node_group(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: The (aliased) node in the database + :param entity_to_join: The (aliased) Group + + **joined_entity** and **entity_to_join** are + joined via the table_groups_nodes table. + from **joined_entity** as node to **enitity_to_join** as group. + (**enitity_to_join** is a group *with_node* **joined_entity**) + """ + _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Group), 'with_node') + aliased_group_nodes = aliased(self._entities.table_groups_nodes) + new_query = query.join(aliased_group_nodes, aliased_group_nodes.c.dbnode_id == joined_entity.id).join( + entity_to_join, entity_to_join.id == aliased_group_nodes.c.dbgroup_id, isouter=isouterjoin + ) + return JoinReturn(new_query, aliased_group_nodes) + + def _join_node_user(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: the aliased node + :param entity_to_join: the aliased user to join to that node + """ + _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.User), 'with_node') + new_query = query.join(entity_to_join, entity_to_join.id == joined_entity.user_id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_user_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: the aliased user you want to join to + :param entity_to_join: the (aliased) node or group in the DB to join with + """ + _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.Node), 'with_user') + new_query = query.join(entity_to_join, entity_to_join.user_id == joined_entity.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_computer_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: the (aliased) computer entity + :param entity_to_join: the (aliased) node entity + + """ + _check_dbentities((joined_entity, self._entities.Computer), (entity_to_join, self._entities.Node), + 'with_computer') + new_query = query.join(entity_to_join, entity_to_join.dbcomputer_id == joined_entity.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_node_computer(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An entity that can use a computer (eg a node) + :param entity_to_join: aliased dbcomputer entity + """ + _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Computer), 'with_node') + new_query = query.join(entity_to_join, joined_entity.dbcomputer_id == entity_to_join.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_group_user(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased dbgroup + :param entity_to_join: aliased dbuser + """ + _check_dbentities((joined_entity, self._entities.Group), (entity_to_join, self._entities.User), 'with_group') + new_query = query.join(entity_to_join, joined_entity.user_id == entity_to_join.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_user_group(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased user + :param entity_to_join: aliased group + """ + _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.Group), 'with_user') + new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.user_id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_node_comment(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased node + :param entity_to_join: aliased comment + """ + _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Comment), 'with_node') + new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.dbnode_id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_comment_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased comment + :param entity_to_join: aliased node + """ + _check_dbentities((joined_entity, self._entities.Comment), (entity_to_join, self._entities.Node), + 'with_comment') + new_query = query.join(entity_to_join, joined_entity.dbnode_id == entity_to_join.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_node_log(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased node + :param entity_to_join: aliased log + """ + _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Log), 'with_node') + new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.dbnode_id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_log_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased log + :param entity_to_join: aliased node + """ + _check_dbentities((joined_entity, self._entities.Log), (entity_to_join, self._entities.Node), 'with_log') + new_query = query.join(entity_to_join, joined_entity.dbnode_id == entity_to_join.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_user_comment(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased user + :param entity_to_join: aliased comment + """ + _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.Comment), 'with_user') + new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.user_id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_comment_user(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + """ + :param joined_entity: An aliased comment + :param entity_to_join: aliased user + """ + _check_dbentities((joined_entity, self._entities.Comment), (entity_to_join, self._entities.User), + 'with_comment') + new_query = query.join(entity_to_join, joined_entity.user_id == entity_to_join.id, isouter=isouterjoin) + return JoinReturn(new_query) + + def _join_node_outputs(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): """ :param joined_entity: The (aliased) ORMclass that is an input :param entity_to_join: The (aliased) ORMClass that is an output. @@ -147,7 +323,7 @@ def _join_outputs(self, query: Query, joined_entity, entity_to_join, isouterjoin ).join(entity_to_join, aliased_edge.output_id == entity_to_join.id, isouter=isouterjoin) return JoinReturn(new_query, aliased_edge) - def _join_inputs(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): + def _join_node_inputs(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): """ :param joined_entity: The (aliased) ORMclass that is an output :param entity_to_join: The (aliased) ORMClass that is an input. @@ -166,7 +342,7 @@ def _join_inputs(self, query: Query, joined_entity, entity_to_join, isouterjoin: ).join(entity_to_join, aliased_edge.input_id == entity_to_join.id, isouter=isouterjoin) return JoinReturn(new_query, aliased_edge) - def _join_descendants_recursive( + def _join_node_descendants_recursive( self, query: Query, joined_entity, @@ -231,7 +407,7 @@ def _join_descendants_recursive( ) return JoinReturn(new_query, descendants_recursive.c) - def _join_ancestors_recursive( + def _join_node_ancestors_recursive( self, query: Query, joined_entity, @@ -297,155 +473,6 @@ def _join_ancestors_recursive( ) return JoinReturn(new_query, ancestors_recursive.c) - def _join_group_members(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: - The (aliased) ORMclass that is - a group in the database - :param entity_to_join: - The (aliased) ORMClass that is a node and member of the group - - **joined_entity** and **entity_to_join** - are joined via the table_groups_nodes table. - from **joined_entity** as group to **enitity_to_join** as node. - (**enitity_to_join** is *with_group* **joined_entity**) - """ - _check_dbentities((joined_entity, self._entities.Group), (entity_to_join, self._entities.Node), 'with_group') - aliased_group_nodes = aliased(self._entities.table_groups_nodes) - new_query = query.join(aliased_group_nodes, aliased_group_nodes.c.dbgroup_id == joined_entity.id).join( - entity_to_join, entity_to_join.id == aliased_group_nodes.c.dbnode_id, isouter=isouterjoin - ) - return JoinReturn(new_query, aliased_group_nodes) - - def _join_groups(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: The (aliased) node in the database - :param entity_to_join: The (aliased) Group - - **joined_entity** and **entity_to_join** are - joined via the table_groups_nodes table. - from **joined_entity** as node to **enitity_to_join** as group. - (**enitity_to_join** is a group *with_node* **joined_entity**) - """ - _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Group), 'with_node') - aliased_group_nodes = aliased(self._entities.table_groups_nodes) - new_query = query.join(aliased_group_nodes, aliased_group_nodes.c.dbnode_id == joined_entity.id).join( - entity_to_join, entity_to_join.id == aliased_group_nodes.c.dbgroup_id, isouter=isouterjoin - ) - return JoinReturn(new_query, aliased_group_nodes) - - def _join_creator_of(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: the aliased node - :param entity_to_join: the aliased user to join to that node - """ - _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.User), 'with_node') - new_query = query.join(entity_to_join, entity_to_join.id == joined_entity.user_id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_created_by(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: the aliased user you want to join to - :param entity_to_join: the (aliased) node or group in the DB to join with - """ - _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.Node), 'with_user') - new_query = query.join(entity_to_join, entity_to_join.user_id == joined_entity.id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_to_computer_used(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: the (aliased) computer entity - :param entity_to_join: the (aliased) node entity - - """ - _check_dbentities((joined_entity, self._entities.Computer), (entity_to_join, self._entities.Node), - 'with_computer') - new_query = query.join(entity_to_join, entity_to_join.dbcomputer_id == joined_entity.id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_computer(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An entity that can use a computer (eg a node) - :param entity_to_join: aliased dbcomputer entity - """ - _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Computer), 'with_node') - new_query = query.join(entity_to_join, joined_entity.dbcomputer_id == entity_to_join.id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_group_user(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased dbgroup - :param entity_to_join: aliased dbuser - """ - _check_dbentities((joined_entity, self._entities.Group), (entity_to_join, self._entities.User), 'with_group') - new_query = query.join(entity_to_join, joined_entity.user_id == entity_to_join.id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_user_group(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased user - :param entity_to_join: aliased group - """ - _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.Group), 'with_user') - new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.user_id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_node_comment(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased node - :param entity_to_join: aliased comment - """ - _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Comment), 'with_node') - new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.dbnode_id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_comment_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased comment - :param entity_to_join: aliased node - """ - _check_dbentities((joined_entity, self._entities.Comment), (entity_to_join, self._entities.Node), - 'with_comment') - new_query = query.join(entity_to_join, joined_entity.dbnode_id == entity_to_join.id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_node_log(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased node - :param entity_to_join: aliased log - """ - _check_dbentities((joined_entity, self._entities.Node), (entity_to_join, self._entities.Log), 'with_node') - new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.dbnode_id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_log_node(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased log - :param entity_to_join: aliased node - """ - _check_dbentities((joined_entity, self._entities.Log), (entity_to_join, self._entities.Node), 'with_log') - new_query = query.join(entity_to_join, joined_entity.dbnode_id == entity_to_join.id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_user_comment(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased user - :param entity_to_join: aliased comment - """ - _check_dbentities((joined_entity, self._entities.User), (entity_to_join, self._entities.Comment), 'with_user') - new_query = query.join(entity_to_join, joined_entity.id == entity_to_join.user_id, isouter=isouterjoin) - return JoinReturn(new_query) - - def _join_comment_user(self, query: Query, joined_entity, entity_to_join, isouterjoin: bool, **_kw): - """ - :param joined_entity: An aliased comment - :param entity_to_join: aliased user - """ - _check_dbentities((joined_entity, self._entities.Comment), (entity_to_join, self._entities.User), - 'with_comment') - new_query = query.join(entity_to_join, joined_entity.user_id == entity_to_join.id, isouter=isouterjoin) - return JoinReturn(new_query) - def _check_dbentities(entities_cls_joined, entities_cls_to_join, relationship: str): """Type check for entities diff --git a/tests/orm/test_querybuilder.py b/tests/orm/test_querybuilder.py index e6e366dba0..75f626ecd1 100644 --- a/tests/orm/test_querybuilder.py +++ b/tests/orm/test_querybuilder.py @@ -1084,6 +1084,26 @@ def test_joins3_user_group(self): assert qb.count() == 1, 'The expected user that owns the selected group was not found.' + def test_joins_authinfo(self): + """Test querying for AuthInfo with specific computer/user.""" + user = orm.User(email='email@new.com').store() + computer = orm.Computer( + label='new', hostname='localhost', transport_type='core.local', scheduler_type='core.direct' + ).store() + authinfo = computer.configure(user) + + # Search for the user of the authinfo + qb = orm.QueryBuilder() + qb.append(orm.User, tag='user', filters={'id': {'==': user.id}}) + qb.append(orm.AuthInfo, with_user='user', filters={'id': {'==': authinfo.id}}) + assert qb.count() == 1, 'The expected user that owns the selected authinfo was not found.' + + # Search for the computer of the authinfo + qb = orm.QueryBuilder() + qb.append(orm.Computer, tag='computer', filters={'id': {'==': computer.id}}) + qb.append(orm.AuthInfo, with_computer='computer', filters={'id': {'==': authinfo.id}}) + assert qb.count() == 1, 'The expected computer that owns the selected authinfo was not found.' + def test_joins_group_node(self): """ This test checks that the querying for the nodes that belong to a group works correctly (using QueryBuilder).