From: Michael Tremer <michael.tremer@ipfire.org>
Date: Sun, 29 Jun 2025 17:50:41 +0000 (+0000)
Subject: database: Create a new way to control when a database session is committed
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=5bd896c05c5730a22878c4444bdacc2b58b8389a;p=pbs.git

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 9192dc9d..ecf6633d 100644
--- 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: