Improve SqlAlchemy session handling for `QueryBuilder` #3708

sphuber · 2020-01-07T14:41:51Z

Fixes #3645 and fixes #3372

CasperWA · 2020-01-19T08:58:46Z

Concerning using this in a REST-API server, i.e., aiida-optimade, what do you suggest to do?
I have tried to retrieve the session via get_manager().get_backend().get_session() before handling a request and then calling session.close() at the end, but this still prompts a QueuePool overflow for both backends.

Edit:
Calling get_manager().get_backend_manager().reset_backend_environment() after every response handle seems to work.
However, this meant I also discovered a wrong import in your code for that function in the django backend.
I am still not sure this is the correct approach though.

CasperWA · 2020-01-19T09:11:21Z

aiida/backends/djsite/manager.py

@@ -32,6 +32,8 @@ def _load_backend_environment(self):

    def reset_backend_environment(self):
        """Reset the backend environment."""
+        from aiida.orm.implementation.django.querybuilder import reset_session


Suggested change

from aiida.orm.implementation.django.querybuilder import reset_session

from . import reset_session

I think this is still an issue?

CasperWA

Look's great. I do have some change requests though (see comments).

CasperWA · 2020-01-30T14:33:29Z

aiida/backends/sqlalchemy/manager.py

@@ -62,7 +62,7 @@ def reset_backend_environment(self):
        from aiida.backends import sqlalchemy
        if sqlalchemy.ENGINE is not None:
            sqlalchemy.ENGINE.dispose()
-        sqlalchemy.SCOPED_SESSION_CLASS = None
+        sqlalchemy.SESSION_FACTORY = None


Shouldn't this method simply call aiida.backends.sqlalchemy.reset_session()?

Sure does :)

CasperWA · 2020-01-30T14:33:46Z

aiida/backends/djsite/manager.py

@@ -32,6 +32,8 @@ def _load_backend_environment(self):

    def reset_backend_environment(self):
        """Reset the backend environment."""
+        from aiida.orm.implementation.django.querybuilder import reset_session


I think this is still an issue?

CasperWA · 2020-01-30T14:37:56Z

aiida/orm/implementation/querybuilder.py

@@ -9,6 +9,12 @@
 ###########################################################################
 """Backend query implementation classes"""


Maybe note here that the general QueryBuilder here is based on the assumption that all queries are performed with SQLAlchemy, no matter the actual backend? I.e., just copy some stuff in here from your commit messages.

CasperWA · 2020-01-30T14:41:14Z

aiida/orm/implementation/sqlalchemy/backend.py

+    def get_session():
+        """Return a database session that can be used by the `QueryBuilder` to perform its query.
+
+        If there is an exception within the context then the changes will be rolled back and the state will


Shouldn't this be under the description of the transaction method?

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.

The scoped session is an instance of SqlAlchemy's `Session` class that is used by the query builder to connect to the database, for both the SqlAlchemy database backend as well as for Django. Both database backends need to maintain their own scoped session factory which can be called to get a session instance. Certain applications need access to the session. For example, applications that run AiiDA in a threaded way, such as a REST API server need to manually close the session after the query has finished because this is not done automatically when the thread ends. The associated database connection remains open causing an eventual timeout when a new request comes in. The method `Backend.get_session` provides an official API to access the global scoped session instance which can then be closed. Additionally, a lot of code that was duplicated across the two implementations of the `QueryBuilder` for the two database backends has been moved to the abstract `BackendQueryBuilder`. Normally this code does indeed belong in the implementations but since the current implementation for both backends is based on SqlAlchemy they are both nearly identical. When in the future a new backend is implemented that does not use SqlAlchemy the current code can be factored out to a specific `SqlAlchemyQueryBuilder` that can be used for both database backends.

Presumably, the `multiprocessing.recreate_after_fork` hook on the session creation of the SqlAlchemy was added because of the design of the old daemon of the v0.* series. There the daemon was implemented using `celery` which would spawn workers, by forking the main process. Since the definition of forking process means "cloning a process with the exact same state", the forked process would have the same session which is undesirable. This is why the session has to be reset such that an independent one has to be initialized. With the new engine design, there is no forking anymore by the daemon and it is also not supported. Instead, each daemon worker is launched as a completely new process and so will not start with an existing session. Finally, the `multiprocessing.recreate_after_fork` is not even documented and according to `https://bugs.python.org/issue21372` this is for a reason, because it is not intended for external use. Combined with the fact that its cross-platform consistent operation is pulled into question, it is better to not use it.

CasperWA

Cheers @sphuber!
Note for later: We should add good documentation about threat handling with AiiDA SQLAlchemy sessions / use of QueryBuilder. This is especially important for REST API frameworks and servers, where multi-threading is commonplace.

sphuber mentioned this pull request Jan 7, 2020

[WIP] Centralize get_session() for backend QueryBuilders #3676

Closed

CasperWA mentioned this pull request Jan 8, 2020

Update AiiDA session middleware aiidateam/aiida-optimade#23

Merged

CasperWA reviewed Jan 19, 2020

View reviewed changes

CasperWA mentioned this pull request Jan 19, 2020

Update session handling per request as per new AiiDA PR 3708 aiidateam/aiida-optimade#34

Merged

sphuber force-pushed the fix_3645_backend_session_handling branch from a8b6026 to 84044e2 Compare January 22, 2020 09:50

CasperWA suggested changes Jan 30, 2020

View reviewed changes

sphuber added 3 commits January 31, 2020 11:06

sphuber force-pushed the fix_3645_backend_session_handling branch from 9bc6f85 to 3c304e9 Compare January 31, 2020 10:06

CasperWA approved these changes Jan 31, 2020

View reviewed changes

sphuber merged commit ffcef0b into aiidateam:develop Jan 31, 2020

sphuber mentioned this pull request Jan 31, 2020

Use Aldjemy to create the session for Django QueryBuilder implementation #3372

Closed

sphuber deleted the fix_3645_backend_session_handling branch January 31, 2020 11:27

sphuber mentioned this pull request Jan 31, 2020

Fix of the multithreading connection problem #3655

Closed

giovannipizzi mentioned this pull request Mar 6, 2020

Make sure threads are handled correctly for the REST API #3743

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improve SqlAlchemy session handling for `QueryBuilder` #3708

Improve SqlAlchemy session handling for `QueryBuilder` #3708

sphuber commented Jan 7, 2020 •

edited

Loading

CasperWA commented Jan 19, 2020 •

edited

Loading

CasperWA Jan 19, 2020

CasperWA Jan 30, 2020

CasperWA left a comment

CasperWA Jan 30, 2020

sphuber Jan 31, 2020

CasperWA Jan 30, 2020

CasperWA Jan 30, 2020

CasperWA Jan 30, 2020

CasperWA left a comment

	from aiida.orm.implementation.django.querybuilder import reset_session
	from . import reset_session

		@@ -9,6 +9,12 @@
		###########################################################################
		"""Backend query implementation classes"""

Improve SqlAlchemy session handling for QueryBuilder #3708

Improve SqlAlchemy session handling for QueryBuilder #3708

Conversation

sphuber commented Jan 7, 2020 • edited Loading

CasperWA commented Jan 19, 2020 • edited Loading

CasperWA Jan 19, 2020

Choose a reason for hiding this comment

CasperWA Jan 30, 2020

Choose a reason for hiding this comment

CasperWA left a comment

Choose a reason for hiding this comment

CasperWA Jan 30, 2020

Choose a reason for hiding this comment

sphuber Jan 31, 2020

Choose a reason for hiding this comment

CasperWA Jan 30, 2020

Choose a reason for hiding this comment

CasperWA Jan 30, 2020

Choose a reason for hiding this comment

CasperWA Jan 30, 2020

Choose a reason for hiding this comment

CasperWA left a comment

Choose a reason for hiding this comment

Improve SqlAlchemy session handling for `QueryBuilder` #3708

Improve SqlAlchemy session handling for `QueryBuilder` #3708

sphuber commented Jan 7, 2020 •

edited

Loading

CasperWA commented Jan 19, 2020 •

edited

Loading