database: Create a new way to control when a database session is committed

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/src/buildservice/database.py b/src/buildservice/database.py
index 9192dc9de7ad3401fb523d2e4cd4d7a7c3ff5a09..ecf6633d7a9e4bc31ef15ccea1d868bd27874817 100644 (file)
--- a/src/buildservice/database.py
+++ b/src/buildservice/database.py
@@ -198,6 +198,29 @@ class Connection(object):
                                "Unknown migration operation: %s" % op,
                        )
 
+       # This context manager can be used to easily complete a session even though the
+       # task is not done (because sometimes we want to commit the transaction long)
+       # before the task has completed.
+       #
+       # It can be used as follows:
+       #
+       #     async with backend.db:
+       #         ...
+
+       async def __aenter__(self):
+               """
+                       Returns the current database session of the task
+               """
+               return await self.session()
+
+       async def __aexit__(self, type, exception, traceback):
+               # This method will be called when the block is being excited and it will
+               # release the database session. Usually that means that there will be a commit.
+               task = asyncio.current_task()
+
+               # Immediately release the session
+               await self.__release_session(task)
+
        async def session(self):
                """
                        Returns a session from the engine
@@ -232,14 +255,17 @@ class Connection(object):
                self.backend.run_task(self.__release_session, task)
 
        async def __release_session(self, task):
+               exception = None
+
                # Retrieve the session
                try:
                        session = self.__sessions[task]
                except KeyError:
                        return
 
-               # Fetch any exception
-               exception = task.exception()
+               # Fetch any exception if the task is done
+               if task.done():
+                       exception = task.exception()
 
                # If there is no exception, we can commit
                if exception is None: