Add possibility to pass keyword arguments to create_engine

The function `aiida.backends.utils.create_sqlalchemy_engine` now takes
keyword arguments that are passed straight to the SQLAlchemy function
`create_engine` which creates the engine that connects to the database.
This utility function is used by the SQLAlchemy and by both backends
for the `QueryBuilder`. By allowing these keyword arguments to be passed
and plumbing them all the way in to the `Backend` class, the parameters
of the SQLAlchemy engine and the associated queue pool can be
customized. This is not meant to be used by end users, but can be
important for important use, for testing and performance reasons.

aiida/backends/djsite/__init__.py

@@ -31,7 +31,7 @@ def reset_session():
     SESSION_FACTORY = None
 
 
-def get_scoped_session():
+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
@@ -43,7 +43,11 @@ def get_scoped_session():
     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 profile: :class:`aiida.manage.configuration.profile.Profile` for which to configure the engine.
+    :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
@@ -56,7 +60,7 @@ def get_scoped_session():
         return session
 
     if ENGINE is None:
-        ENGINE = create_sqlalchemy_engine(get_profile())
+        ENGINE = create_sqlalchemy_engine(get_profile(), **kwargs)
 
     SESSION_FACTORY = create_scoped_session_factory(ENGINE)
 

aiida/backends/djsite/manager.py

@@ -33,11 +33,20 @@ def get_settings_manager(self):
 
         return self._settings_manager
 
-    def _load_backend_environment(self):
-        """Load the backend environment."""
+    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

aiida/backends/manager.py

@@ -106,20 +106,24 @@ def get_settings_manager(self):
         :return: `SettingsManager`
         """
 
-    def load_backend_environment(self, profile, validate_schema=True):
+    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()
+        self._load_backend_environment(**kwargs)
 
         if validate_schema:
             self.validate_schema(profile)
 
     @abc.abstractmethod
-    def _load_backend_environment(self):
-        """Load the backend environment."""
+    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):

aiida/backends/sqlalchemy/__init__.py

@@ -31,11 +31,16 @@ def reset_session():
     SESSION_FACTORY = None
 
 
-def get_scoped_session():
+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
 
@@ -47,7 +52,7 @@ def get_scoped_session():
         return session
 
     if ENGINE is None:
-        ENGINE = create_sqlalchemy_engine(get_profile())
+        ENGINE = create_sqlalchemy_engine(get_profile(), **kwargs)
 
     SESSION_FACTORY = create_scoped_session_factory(ENGINE, expire_on_commit=True)
 

aiida/backends/sqlalchemy/manager.py

@@ -60,9 +60,13 @@ def get_settings_manager(self):
 
         return self._settings_manager
 
-    def _load_backend_environment(self):
-        """Load the backend environment."""
-        get_scoped_session()
+    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."""

aiida/backends/utils.py

@@ -7,13 +7,20 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
+"""Backend-agnostic utility functions"""
 from aiida.backends import BACKEND_SQLA, BACKEND_DJANGO
 from aiida.manage import configuration
 
 AIIDA_ATTRIBUTE_SEP = '.'
 
 
-def create_sqlalchemy_engine(profile):
+def create_sqlalchemy_engine(profile, **kwargs):
+    """Create SQLAlchemy engine (to be used for QueryBuilder queries)
+
+    :param kwargs: keyword arguments that will be passed on to `sqlalchemy.create_engine`.
+        See https://docs.sqlalchemy.org/en/13/core/engines.html?highlight=create_engine#sqlalchemy.create_engine for
+        more info.
+    """
     from sqlalchemy import create_engine
     from aiida.common import json
 
@@ -26,15 +33,19 @@ def create_sqlalchemy_engine(profile):
         port=profile.database_port,
         name=profile.database_name
     )
-    return create_engine(engine_url, json_serializer=json.dumps, json_deserializer=json.loads, encoding='utf-8')
+    return create_engine(
+        engine_url, json_serializer=json.dumps, json_deserializer=json.loads, encoding='utf-8', **kwargs
+    )
 
 
 def create_scoped_session_factory(engine, **kwargs):
+    """Create scoped SQLAlchemy session factory"""
     from sqlalchemy.orm import scoped_session, sessionmaker
     return scoped_session(sessionmaker(bind=engine, **kwargs))
 
 
 def delete_nodes_and_connections(pks):
+    """Backend-agnostic function to delete Nodes and connections"""
     if configuration.PROFILE.database_backend == BACKEND_DJANGO:
         from aiida.backends.djsite.utils import delete_nodes_and_connections_django as delete_nodes_backend
     elif configuration.PROFILE.database_backend == BACKEND_SQLA: