close django connection in REST API

aiida/backends/utils.py

@@ -8,11 +8,22 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Backend-agnostic utility functions"""
+from sqlalchemy.pool import NullPool
 from aiida.backends import BACKEND_SQLA, BACKEND_DJANGO
 from aiida.manage import configuration
+from aiida.common import json
 
 AIIDA_ATTRIBUTE_SEP = '.'
 
+SQLALCHEMY_ENGINE_DEFAULTS = dict(
+    json_serializer=json.dumps,
+    json_deserializer=json.loads,
+    encoding='utf-8',
+    poolclass=NullPool,
+    # pool_size=1,
+    # max_overflow=0,
+)
+
 
 def create_sqlalchemy_engine(profile, **kwargs):
     """Create SQLAlchemy engine (to be used for QueryBuilder queries)
@@ -22,7 +33,6 @@ def create_sqlalchemy_engine(profile, **kwargs):
         more info.
     """
     from sqlalchemy import create_engine
-    from aiida.common import json
 
     # The hostname may be `None`, which is a valid value in the case of peer authentication for example. In this case
     # it should be converted to an empty string, because otherwise the `None` will be converted to string literal "None"
@@ -37,9 +47,9 @@ def create_sqlalchemy_engine(profile, **kwargs):
         port=profile.database_port,
         name=profile.database_name
     )
-    return create_engine(
-        engine_url, json_serializer=json.dumps, json_deserializer=json.loads, encoding='utf-8', **kwargs
-    )
+    kwargs_new = SQLALCHEMY_ENGINE_DEFAULTS.copy()
+    kwargs_new.update(kwargs)
+    return create_engine(engine_url, **kwargs_new)
 
 
 def create_scoped_session_factory(engine, **kwargs):

aiida/restapi/common/utils.py

@@ -857,4 +857,10 @@ def close_session(wrapped, _, args, kwargs):
     try:
         return wrapped(*args, **kwargs)
     finally:
-        get_manager().get_backend().get_session().close()
+        backend = get_manager().get_backend()
+        backend.get_session().close()
+
+        # Close django connection if open (not needed by REST API)
+        if 'Django' in type(backend).__name__:  # don't want to import the backend class, gives error
+            from django.db import connection
+            connection.close()

tests/restapi/conftest.py

@@ -51,10 +51,11 @@ def server_url():
 def restrict_sqlalchemy_queuepool(aiida_profile):
     """Create special SQLAlchemy engine for use with QueryBuilder - backend-agnostic"""
     from aiida.manage.manager import get_manager
+    from sqlalchemy.pool import QueuePool
 
     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)
+    backend_manager.load_backend_environment(aiida_profile, poolclass=QueuePool, pool_timeout=1, max_overflow=0)
 
 
 @pytest.fixture