Skip to content

Commit

Permalink
👌 IMPROVE: Move default user caching to StorageBackend
Browse files Browse the repository at this point in the history
Currently it is cached on the `UserCollection`,
but this means that, on profile switching,
it would return a user associated with the wrong backend instance.
  • Loading branch information
chrisjsewell committed Apr 25, 2022
1 parent 6b31e77 commit d0b5e06
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 32 deletions.
17 changes: 17 additions & 0 deletions aiida/orm/implementation/storage_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
BackendQueryBuilder,
BackendUserCollection,
)
from aiida.orm.users import User
from aiida.repository.backend.abstract import AbstractRepositoryBackend

__all__ = ('StorageBackend',)
Expand Down Expand Up @@ -84,6 +85,7 @@ def __init__(self, profile: 'Profile') -> None:
"""
from aiida.orm.autogroup import AutogroupManager
self._profile = profile
self._default_user: Optional['User'] = None
self._autogroup = AutogroupManager(self)

@abc.abstractmethod
Expand Down Expand Up @@ -161,6 +163,21 @@ def nodes(self) -> 'BackendNodeCollection':
def users(self) -> 'BackendUserCollection':
"""Return the collection of users"""

@property
def default_user(self) -> Optional['User']:
"""Return the default user for the profile, if it has been created.
This is cached, since it is a frequently used operation, for creating other entities.
"""
from aiida.orm import QueryBuilder, User

if self._default_user is None and self.profile.default_user_email:
query = QueryBuilder().append(User, filters={'email': self.profile.default_user_email})
results = query.all(flat=True)
if results:
self._default_user = results[0]
return self._default_user

@abc.abstractmethod
def query(self) -> 'BackendQueryBuilder':
"""Return an instance of a query builder implementation for this backend"""
Expand Down
19 changes: 1 addition & 18 deletions aiida/orm/users.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ def _entity_base_cls() -> Type['User']:

def __init__(self, entity_class: Type['User'], backend: Optional['StorageBackend'] = None) -> None:
super().__init__(entity_class=entity_class, backend=backend)
self._default_user: Optional[User] = None

def get_or_create(self, email: str, **kwargs) -> Tuple[bool, 'User']:
"""Get the existing user with a given email address or create an unstored one
Expand All @@ -48,23 +47,7 @@ def get_or_create(self, email: str, **kwargs) -> Tuple[bool, 'User']:

def get_default(self) -> Optional['User']:
"""Get the current default user"""
if self._default_user is None:
email = self.backend.profile.default_user_email
if not email:
self._default_user = None

try:
self._default_user = self.get(email=email)
except (exceptions.MultipleObjectsError, exceptions.NotExistent):
self._default_user = None

return self._default_user

def reset(self) -> None:
"""
Reset internal caches (default user).
"""
self._default_user = None
return self.backend.default_user


class User(entities.Entity['BackendUser']):
Expand Down
21 changes: 7 additions & 14 deletions aiida/storage/psql_dos/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@

from aiida.common.exceptions import ClosedStorage, IntegrityError
from aiida.manage.configuration.profile import Profile
from aiida.orm import User
from aiida.orm.entities import EntityTypes
from aiida.orm.implementation import BackendEntity, StorageBackend
from aiida.storage.log import STORAGE_LOGGER
Expand Down Expand Up @@ -113,8 +112,6 @@ def get_session(self) -> Session:
def close(self) -> None:
if self._session_factory is None:
return # the instance is already closed, and so this is a no-op
# reset the cached default user instance, since it will now have no associated session
User.collection(self).reset()
# close the connection
# pylint: disable=no-member
engine = self._session_factory.bind
Expand All @@ -137,15 +134,13 @@ def _clear(self, recreate_user: bool = True) -> None:

# save the default user
default_user_kwargs = None
if recreate_user:
default_user = User.collection(self).get_default()
if default_user is not None:
default_user_kwargs = {
'email': default_user.email,
'first_name': default_user.first_name,
'last_name': default_user.last_name,
'institution': default_user.institution,
}
if recreate_user and self.default_user is not None:
default_user_kwargs = {
'email': self.default_user.email,
'first_name': self.default_user.first_name,
'last_name': self.default_user.last_name,
'institution': self.default_user.institution,
}

# now clear the database
for table_name in (
Expand All @@ -158,8 +153,6 @@ def _clear(self, recreate_user: bool = True) -> None:
# restore the default user
if recreate_user and default_user_kwargs:
session.add(DbUser(**default_user_kwargs))
# clear aiida's cache of the default user
User.collection(self).reset()

# Clear the repository and reset the repository UUID
container = Container(self.profile.repository_path / 'container')
Expand Down

0 comments on commit d0b5e06

Please sign in to comment.