Properly handle multiple requests to threaded REST API

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:

aiida/manage/tests/unittest_classes.py

@@ -82,7 +82,7 @@ def run(self, suite, backend=None, profile_name=None):
         import warnings
         from aiida.common.warnings import AiidaDeprecationWarning
         warnings.warn(  # pylint: disable=no-member
-            'Please use "pytest" for testing AiiDA plugins. Support for "unittest" be removed in `v2.0.0`',
+            'Please use "pytest" for testing AiiDA plugins. Support for "unittest" will be removed in `v2.0.0`',
             AiidaDeprecationWarning
         )
 

aiida/restapi/common/identifiers.py

@@ -31,8 +31,7 @@
     'process.calculation%.calcfunction.%|aiida.calculations:arithmetic.add'  # More than one operator in segment
 
 """
-
-import collections
+from collections.abc import MutableMapping
 
 from aiida.common.escaping import escape_for_sql_like
 
@@ -163,7 +162,7 @@ def load_entry_point_from_full_type(full_type):
     raise EntryPointError('entry point of the given full type cannot be loaded')
 
 
-class Namespace(collections.MutableMapping):
+class Namespace(MutableMapping):
     """Namespace that can be used to map the node class hierarchy."""
 
     namespace_separator = '.'

aiida/restapi/common/utils.py

@@ -8,13 +8,15 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """ Util methods """
-import urllib.parse
 from datetime import datetime, timedelta
+import urllib.parse
 
 from flask import jsonify
 from flask.json import JSONEncoder
+from wrapt import decorator
 
 from aiida.common.exceptions import InputValidationError, ValidationError
+from aiida.manage.manager import get_manager
 from aiida.restapi.common.exceptions import RestValidationError, \
     RestInputValidationError
 
@@ -845,3 +847,17 @@ def list_routes():
         output.append(line)
 
     return sorted(set(output))
+
+
+@decorator
+def close_session(wrapped, _, args, kwargs):
+    """Close AiiDA SQLAlchemy (QueryBuilder) session
+
+    This decorator can be used for router endpoints to close the SQLAlchemy global scoped session after the response
+    has been created. This is needed, since the QueryBuilder uses a SQLAlchemy global scoped session no matter the
+    profile's database backend.
+    """
+    try:
+        return wrapped(*args, **kwargs)
+    finally:
+        get_manager().get_backend().get_session().close()

aiida/restapi/resources.py

@@ -15,12 +15,11 @@
 
 from aiida.common.lang import classproperty
 from aiida.restapi.common.exceptions import RestInputValidationError
-from aiida.restapi.common.utils import Utils
+from aiida.restapi.common.utils import Utils, close_session
 
 
 class ServerInfo(Resource):
-    # pylint: disable=fixme
-    """Endpointd to return general server info"""
+    """Endpoint to return general server info"""
 
     def __init__(self, **kwargs):
         # Configure utils
@@ -97,6 +96,8 @@ class BaseResource(Resource):
     _translator_class = BaseTranslator
     _parse_pk_uuid = None  # Flag to tell the path parser whether to expect a pk or a uuid pattern
 
+    method_decorators = [close_session]  # Close SQLA session after any method call
+
     ## TODO add the caching support. I cache total count, results, and possibly
 
     def __init__(self, **kwargs):
@@ -106,11 +107,13 @@ def __init__(self, **kwargs):
         utils_conf_keys = ('PREFIX', 'PERPAGE_DEFAULT', 'LIMIT_DEFAULT')
         self.utils_confs = {k: kwargs[k] for k in utils_conf_keys if k in kwargs}
         self.utils = Utils(**self.utils_confs)
-        self.method_decorators = {'get': kwargs.get('get_decorators', [])}
+
+        # HTTP Request method decorators
+        if 'get_decorators' in kwargs and isinstance(kwargs['get_decorators'], (tuple, list, set)):
+            self.method_decorators = {'get': list(kwargs['get_decorators'])}
 
     @classproperty
-    def parse_pk_uuid(cls):
-        # pylint: disable=no-self-argument
+    def parse_pk_uuid(cls):  # pylint: disable=no-self-argument
         return cls._parse_pk_uuid
 
     def _load_and_verify(self, node_id=None):
@@ -212,17 +215,6 @@ class Node(BaseResource):
     _translator_class = NodeTranslator
     _parse_pk_uuid = 'uuid'  # Parse a uuid pattern in the URL path (not a pk)
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        from aiida.orm import Node as tNode
-        self.tclass = tNode
-
-        # Configure utils
-        utils_conf_keys = ('PREFIX', 'PERPAGE_DEFAULT', 'LIMIT_DEFAULT')
-        self.utils_confs = {k: kwargs[k] for k in utils_conf_keys if k in kwargs}
-        self.utils = Utils(**self.utils_confs)
-        self.method_decorators = {'get': kwargs.get('get_decorators', [])}
-
     def get(self, id=None, page=None):  # pylint: disable=redefined-builtin,invalid-name,unused-argument
         # pylint: disable=too-many-locals,too-many-statements,too-many-branches,fixme,unused-variable
         """

aiida/restapi/run_api.py

@@ -56,7 +56,7 @@ def run_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **kwargs)
     port = kwargs.pop('port', CLI_DEFAULTS['PORT'])
     debug = kwargs.pop('debug', APP_CONFIG['DEBUG'])
 
-    app, api = configure_api(flask_app, flask_api, **kwargs)
+    api = configure_api(flask_app, flask_api, **kwargs)
 
     if hookup:
         # Run app through built-in werkzeug server
@@ -66,7 +66,7 @@ def run_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **kwargs)
     else:
         # Return the app & api without specifying port/host to be handled by an external server (e.g. apache).
         # Some of the user-defined configuration of the app is ineffective (only affects built-in server).
-        return (app, api)
+        return api.app, api
 
 
 def configure_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **kwargs):
@@ -81,7 +81,8 @@ def configure_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **k
     :param catch_internal_server:  If true, catch and print all inter server errors
     :param wsgi_profile: use WSGI profiler middleware for finding bottlenecks in web application
 
-    :returns: tuple (app, api)
+    :returns: Flask RESTful API
+    :rtype: :py:class:`flask_restful.Api`
     """
 
     # Unpack parameters
@@ -119,6 +120,5 @@ def configure_api(flask_app=api_classes.App, flask_api=api_classes.AiidaApi, **k
         app.config['PROFILE'] = True
         app.wsgi_app = ProfilerMiddleware(app.wsgi_app, restrictions=[30])
 
-    # Instantiate an Api by associating its app
-    api = flask_api(app, **API_CONFIG)
-    return (app, api)
+    # Instantiate and return a Flask RESTful API by associating its app
+    return flask_api(app, **API_CONFIG)

docs/source/developer_guide/core/extend_restapi.rst

@@ -369,13 +369,14 @@ as confirmed by the response to the GET request.
 
 As a final remark, there might be circumstances in which you do not want to use the internal werkzeug-based server.
 For example, you might want to run the app through Apache using a wsgi script.
-In this case, simply use ``configure_api`` to return two custom objects ``app`` and ``api``:
+In this case, simply use ``configure_api`` to return a custom object ``api``:
 
 .. code-block:: python
 
-    (app, api) = configure_api(App, MycloudApi, **kwargs)
+    api = configure_api(App, MycloudApi, **kwargs)
 
 
+The ``app`` can be retrieved by ``api.app``.
 This snippet of code becomes the fundamental block of a *wsgi* file used by Apache as documented in  :ref:`restapi_apache`.
 Moreover, we recommend to consult the documentation of `mod_wsgi <https://modwsgi.readthedocs.io/>`_.
 

tests/restapi/conftest.py

@@ -0,0 +1,50 @@
+"""pytest fixtures for use with the aiida.restapi tests"""
+import pytest
+
+
+@pytest.fixture(scope='function')
+def restapi_server():
+    """Make REST API server"""
+    from werkzeug.serving import make_server
+
+    from aiida.restapi.common.config import CLI_DEFAULTS
+    from aiida.restapi.run_api import configure_api
+
+    def _restapi_server(restapi=None):
+        if restapi is None:
+            flask_restapi = configure_api()
+        else:
+            flask_restapi = configure_api(flask_api=restapi)
+
+        return make_server(
+            host=CLI_DEFAULTS['HOST_NAME'],
+            port=int(CLI_DEFAULTS['PORT']),
+            app=flask_restapi.app,
+            threaded=True,
+            processes=1,
+            request_handler=None,
+            passthrough_errors=True,
+            ssl_context=None,
+            fd=None
+        )
+
+    return _restapi_server
+
+
+@pytest.fixture
+def server_url():
+    from aiida.restapi.common.config import CLI_DEFAULTS, API_CONFIG
+
+    return 'http://{hostname}:{port}{api}'.format(
+        hostname=CLI_DEFAULTS['HOST_NAME'], port=CLI_DEFAULTS['PORT'], api=API_CONFIG['PREFIX']
+    )
+
+
+@pytest.fixture
+def restrict_sqlalchemy_queuepool(aiida_profile):
+    """Create special SQLAlchemy engine for use with QueryBuilder - backend-agnostic"""
+    from aiida.manage.manager import get_manager
+
+    backend_manager = get_manager().get_backend_manager()
+    backend_manager.reset_backend_environment()
+    backend_manager.load_backend_environment(aiida_profile, pool_timeout=1, max_overflow=0)