♻️ REFACTOR: Reorganise Backend class UML

aiida/backends/__init__.py

@@ -11,20 +11,3 @@
 
 BACKEND_DJANGO = 'django'
 BACKEND_SQLA = 'sqlalchemy'
-
-
-def get_backend_manager(backend):
-    """Get an instance of the `BackendManager` for the current backend.
-
-    :param backend: the type of the database backend
-    :return: `BackendManager`
-    """
-    if backend == BACKEND_DJANGO:
-        from aiida.backends.djsite.manager import DjangoBackendManager
-        return DjangoBackendManager()
-
-    if backend == BACKEND_SQLA:
-        from aiida.backends.sqlalchemy.manager import SqlaBackendManager
-        return SqlaBackendManager()
-
-    raise Exception(f'unknown backend type `{backend}`')

aiida/backends/djsite/__init__.py

@@ -9,59 +9,3 @@
 ###########################################################################
 # pylint: disable=global-statement
 """Module with implementation of the database backend using Django."""
-from aiida.backends.utils import create_scoped_session_factory, create_sqlalchemy_engine
-
-ENGINE = None
-SESSION_FACTORY = None
-
-
-def reset_session():
-    """Reset the session which means setting the global engine and session factory instances to `None`."""
-    global ENGINE
-    global SESSION_FACTORY
-
-    if ENGINE is not None:
-        ENGINE.dispose()
-
-    if SESSION_FACTORY is not None:
-        SESSION_FACTORY.expunge_all()  # pylint: disable=no-member
-        SESSION_FACTORY.close()  # pylint: disable=no-member
-
-    ENGINE = None
-    SESSION_FACTORY = None
-
-
-def get_scoped_session(**kwargs):
-    """Return a scoped session for the given profile that is exclusively to be used for the `QueryBuilder`.
-
-    Since the `QueryBuilder` implementation uses SqlAlchemy to map the query onto the models in order to generate the
-    SQL to be sent to the database, it requires a session, which is an :class:`sqlalchemy.orm.session.Session` instance.
-    The only purpose is for SqlAlchemy to be able to connect to the database perform the query and retrieve the results.
-    Even the Django backend implementation will use SqlAlchemy for its `QueryBuilder` and so also needs an SqlA session.
-    It is important that we do not reuse the scoped session factory in the SqlAlchemy implementation, because that runs
-    the risk of cross-talk once profiles can be switched dynamically in a single python interpreter. Therefore the
-    Django implementation of the `QueryBuilder` should keep its own SqlAlchemy engine and scoped session factory
-    instances that are used to provide the query builder with a session.
-
-    :param kwargs: keyword arguments that will be passed on to :py:func:`aiida.backends.utils.create_sqlalchemy_engine`,
-        opening the possibility to change QueuePool time outs and more.
-        See https://docs.sqlalchemy.org/en/13/core/engines.html?highlight=create_engine#sqlalchemy.create_engine for
-        more info.
-
-    :return: :class:`sqlalchemy.orm.session.Session` instance with engine configured for the given profile.
-    """
-    from aiida.manage.configuration import get_profile
-
-    global ENGINE
-    global SESSION_FACTORY
-
-    if SESSION_FACTORY is not None:
-        session = SESSION_FACTORY()
-        return session
-
-    if ENGINE is None:
-        ENGINE = create_sqlalchemy_engine(get_profile(), **kwargs)
-
-    SESSION_FACTORY = create_scoped_session_factory(ENGINE)
-
-    return SESSION_FACTORY()

aiida/backends/djsite/manage.py

@@ -24,7 +24,7 @@ def main(profile, command):  # pylint: disable=unused-argument
     from aiida.manage.manager import get_manager
 
     manager = get_manager()
-    manager._load_backend(schema_check=False)  # pylint: disable=protected-access
+    manager._load_backend(validate_db=False)  # pylint: disable=protected-access
 
     # The `execute_from_command` expects a list of command line arguments where the first is the program name that one
     # would normally call directly. Since this is now replaced by our `click` command we just spoof a random name.

aiida/backends/djsite/manager.py

@@ -9,11 +9,6 @@
 ###########################################################################
 # pylint: disable=import-error,no-name-in-module
 """Utilities and configuration of the Django database schema."""
-
-import os
-
-import django
-
 from aiida.common import NotExistent
 
 from ..manager import SCHEMA_VERSION_DESCRIPTION, SCHEMA_VERSION_KEY, BackendManager, Setting, SettingsManager
@@ -35,25 +30,6 @@ def get_settings_manager(self):
 
         return self._settings_manager
 
-    def _load_backend_environment(self, **kwargs):
-        """Load the backend environment.
-
-        The scoped session is needed for the QueryBuilder only.
-
-        :param kwargs: keyword arguments that will be passed on to :py:func:`aiida.backends.djsite.get_scoped_session`.
-        """
-        os.environ['DJANGO_SETTINGS_MODULE'] = 'aiida.backends.djsite.settings'
-        django.setup()  # pylint: disable=no-member
-
-        # For QueryBuilder only
-        from . import get_scoped_session
-        get_scoped_session(**kwargs)
-
-    def reset_backend_environment(self):
-        """Reset the backend environment."""
-        from . import reset_session
-        reset_session()
-
     def is_database_schema_ahead(self):
         """Determine whether the database schema version is ahead of the code schema version.
 
@@ -85,12 +61,8 @@ def get_schema_generation_database(self):
         """
         from django.db.utils import ProgrammingError
 
-        from aiida.manage.manager import get_manager
-
-        backend = get_manager()._load_backend(schema_check=False, repository_check=False)  # pylint: disable=protected-access
-
         try:
-            result = backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'schema_generation';""")
+            result = self._backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'schema_generation';""")
         except ProgrammingError:
             # If this value does not exist, the schema has to correspond to the first generation which didn't actually
             # record its value explicitly in the database until ``aiida-core>=1.0.0``.
@@ -108,14 +80,10 @@ def get_schema_version_database(self):
         """
         from django.db.utils import ProgrammingError
 
-        from aiida.manage.manager import get_manager
-
-        backend = get_manager()._load_backend(schema_check=False, repository_check=False)  # pylint: disable=protected-access
-
         try:
-            result = backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'db|schemaversion';""")
+            result = self._backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'db|schemaversion';""")
         except ProgrammingError:
-            result = backend.execute_raw(r"""SELECT tval FROM db_dbsetting WHERE key = 'db|schemaversion';""")
+            result = self._backend.execute_raw(r"""SELECT tval FROM db_dbsetting WHERE key = 'db|schemaversion';""")
         return result[0][0]
 
     def set_schema_version_database(self, version):
@@ -131,13 +99,9 @@ def _migrate_database_generation(self):
         For Django we also have to clear the `django_migrations` table that contains a history of all applied
         migrations. After clearing it, we reinsert the name of the new initial schema .
         """
-        # pylint: disable=cyclic-import
-        from aiida.manage.manager import get_manager
         super()._migrate_database_generation()
-
-        backend = get_manager()._load_backend(schema_check=False, repository_check=False)  # pylint: disable=protected-access
-        backend.execute_raw(r"""DELETE FROM django_migrations WHERE app = 'db';""")
-        backend.execute_raw(
+        self._backend.execute_raw(r"""DELETE FROM django_migrations WHERE app = 'db';""")
+        self._backend.execute_raw(
             r"""INSERT INTO django_migrations (app, name, applied) VALUES ('db', '0001_initial', NOW());"""
         )
 

aiida/backends/manager.py

@@ -100,36 +100,17 @@ class BackendManager:
 
     _settings_manager = None
 
+    def __init__(self, backend) -> None:
+        from aiida.backends.sqlalchemy.manager import SqlaBackendManager
+        self._backend: SqlaBackendManager = backend
+
     @abc.abstractmethod
     def get_settings_manager(self):
         """Return an instance of the `SettingsManager`.
 
         :return: `SettingsManager`
         """
 
-    def load_backend_environment(self, profile, validate_schema=True, **kwargs):
-        """Load the backend environment.
-
-        :param profile: the profile whose backend environment to load
-        :param validate_schema: boolean, if True, validate the schema first before loading the environment.
-        :param kwargs: keyword arguments that will be passed on to the backend specific scoped session getter function.
-        """
-        self._load_backend_environment(**kwargs)
-
-        if validate_schema:
-            self.validate_schema(profile)
-
-    @abc.abstractmethod
-    def _load_backend_environment(self, **kwargs):
-        """Load the backend environment.
-
-        :param kwargs: keyword arguments that will be passed on to the backend specific scoped session getter function.
-        """
-
-    @abc.abstractmethod
-    def reset_backend_environment(self):
-        """Reset the backend environment."""
-
     def migrate(self):
         """Migrate the database to the latest schema generation or version."""
         try:

aiida/backends/sqlalchemy/__init__.py

@@ -9,51 +9,3 @@
 ###########################################################################
 # pylint: disable=global-statement
 """Module with implementation of the database backend using SqlAlchemy."""
-from aiida.backends.utils import create_scoped_session_factory, create_sqlalchemy_engine
-
-ENGINE = None
-SESSION_FACTORY = None
-
-
-def reset_session():
-    """Reset the session which means setting the global engine and session factory instances to `None`."""
-    global ENGINE
-    global SESSION_FACTORY
-
-    if ENGINE is not None:
-        ENGINE.dispose()
-
-    if SESSION_FACTORY is not None:
-        SESSION_FACTORY.expunge_all()  # pylint: disable=no-member
-        SESSION_FACTORY.close()  # pylint: disable=no-member
-
-    ENGINE = None
-    SESSION_FACTORY = None
-
-
-def get_scoped_session(**kwargs):
-    """Return a scoped session
-
-    According to SQLAlchemy docs, this returns always the same object within a thread, and a different object in a
-    different thread. Moreover, since we update the session class upon forking, different session objects will be used.
-
-    :param kwargs: keyword argument that will be passed on to :py:func:`aiida.backends.utils.create_sqlalchemy_engine`,
-        opening the possibility to change QueuePool time outs and more.
-        See https://docs.sqlalchemy.org/en/13/core/engines.html?highlight=create_engine#sqlalchemy.create_engine for
-        more info.
-    """
-    from aiida.manage.configuration import get_profile
-
-    global ENGINE
-    global SESSION_FACTORY
-
-    if SESSION_FACTORY is not None:
-        session = SESSION_FACTORY()
-        return session
-
-    if ENGINE is None:
-        ENGINE = create_sqlalchemy_engine(get_profile(), **kwargs)
-
-    SESSION_FACTORY = create_scoped_session_factory(ENGINE, expire_on_commit=True)
-
-    return SESSION_FACTORY()

aiida/backends/sqlalchemy/manage.py

@@ -22,11 +22,12 @@ def execute_alembic_command(command_name, **kwargs):
     :param command_name: the sub command name
     :param kwargs: parameters to pass to the command
     """
-    from aiida.backends.sqlalchemy.manager import SqlaBackendManager
+    from aiida.manage.configuration import get_profile
+    from aiida.orm.implementation.sqlalchemy.backend import SqlaBackend
 
-    manager = SqlaBackendManager()
-
-    with manager.alembic_config() as config:
+    # create a new backend which does not validate the schema version
+    backend = SqlaBackend(get_profile(), validate_db=False)
+    with backend._backend_manager.alembic_config() as config:  # pylint: disable=protected-access
         command = getattr(alembic.command, command_name)
         command(config, **kwargs)
 
@@ -40,7 +41,7 @@ def alembic_cli(profile):
 
     load_profile(profile=profile.name)
     manager = get_manager()
-    manager._load_backend(schema_check=False)  # pylint: disable=protected-access
+    manager._load_backend(validate_db=False)  # pylint: disable=protected-access
 
 
 @alembic_cli.command('revision')

aiida/backends/sqlalchemy/manager.py

@@ -14,7 +14,6 @@
 import sqlalchemy
 from sqlalchemy.orm.exc import NoResultFound
 
-from aiida.backends.sqlalchemy import get_scoped_session
 from aiida.common import NotExistent
 
 from ..manager import BackendManager, Setting, SettingsManager
@@ -28,19 +27,16 @@
 class SqlaBackendManager(BackendManager):
     """Class to manage the database schema."""
 
-    @staticmethod
     @contextlib.contextmanager
-    def alembic_config():
+    def alembic_config(self):
         """Context manager to return an instance of an Alembic configuration.
 
         The current database connection is added in the `attributes` property, through which it can then also be
         retrieved, also in the `env.py` file, which is run when the database is migrated.
         """
         from alembic.config import Config
 
-        from . import ENGINE
-
-        with ENGINE.begin() as connection:
+        with self._backend.get_session().bind.begin() as connection:
             dir_path = os.path.dirname(os.path.realpath(__file__))
             config = Config()
             config.set_main_option('script_location', os.path.join(dir_path, ALEMBIC_REL_PATH))
@@ -77,23 +73,10 @@ def get_settings_manager(self):
         :return: `SettingsManager`
         """
         if self._settings_manager is None:
-            self._settings_manager = SqlaSettingsManager()
+            self._settings_manager = SqlaSettingsManager(self._backend)
 
         return self._settings_manager
 
-    def _load_backend_environment(self, **kwargs):
-        """Load the backend environment.
-
-        :param kwargs: keyword arguments that will be passed on to
-            :py:func:`aiida.backends.sqlalchemy.get_scoped_session`.
-        """
-        get_scoped_session(**kwargs)
-
-    def reset_backend_environment(self):
-        """Reset the backend environment."""
-        from . import reset_session
-        reset_session()
-
     def is_database_schema_ahead(self):
         """Determine whether the database schema version is ahead of the code schema version.
 
@@ -147,12 +130,15 @@ class SqlaSettingsManager(SettingsManager):
 
     table_name = 'db_dbsetting'
 
+    def __init__(self, backend) -> None:
+        self._backend = backend
+
     def validate_table_existence(self):
         """Verify that the `DbSetting` table actually exists.
 
         :raises: `~aiida.common.exceptions.NotExistent` if the settings table does not exist
         """
-        inspector = sqlalchemy.inspect(get_scoped_session().bind)
+        inspector = sqlalchemy.inspect(self._backend.get_session().bind)
         if self.table_name not in inspector.get_table_names():
             raise NotExistent('the settings table does not exist')
 
@@ -167,7 +153,7 @@ def get(self, key):
         self.validate_table_existence()
 
         try:
-            setting = get_scoped_session().query(DbSetting).filter_by(key=key).one()
+            setting = self._backend.get_session().query(DbSetting).filter_by(key=key).one()
         except NoResultFound:
             raise NotExistent(f'setting `{key}` does not exist') from NoResultFound
 
@@ -190,7 +176,7 @@ def set(self, key, value, description=None):
         if description is not None:
             other_attribs['description'] = description
 
-        DbSetting.set_value(key, value, other_attribs=other_attribs)
+        DbSetting.set_value(self._backend.get_session(), key, value, other_attribs=other_attribs)
 
     def delete(self, key):
         """Delete the setting with the given key.
@@ -202,7 +188,7 @@ def delete(self, key):
         self.validate_table_existence()
 
         try:
-            setting = get_scoped_session().query(DbSetting).filter_by(key=key).one()
+            setting = self._backend.get_session().query(DbSetting).filter_by(key=key).one()
             setting.delete()
         except NoResultFound:
             raise NotExistent(f'setting `{key}` does not exist') from NoResultFound