Add database plugin I copied from my other project

1 files changed, 90 insertions, 0 deletions

diff --git a/src/acit/db.py b/src/acit/db.py
new file mode 100644
index 0000000..d11518d
--- /dev/null
+++ b/src/acit/db.py
@@ -0,0 +1,90 @@
+
+import cherrypy
+import mariadb
+from cherrypy.process import plugins
+import os.getenv
+from time import sleep
+from threading import Event
+
+if not mariadb.threadsafety:
+	print("Mariadb connector not threadsafe. Failing.")
+	exit(1)
+
+class DBPoolManager(plugins.SimplePlugin):
+	def __init__(self, bus):
+		plugins.SimplePlugin.__init__(self, bus)
+		self.pool:mariadb.ConnectionPool=None
+		self.poolStartedEvent=Event()
+
+	def start(self):
+		# get amount of concurrent connections
+		conns=os.getenv("DB_CONNECTIONS",4)
+		self.maxconns=conns
+		if not type(conns)==int and not conns.isdigit():
+			self._log("DB_CONNECTIONS does not have a valid value:",conns,", using 4 instead.")
+			conns=4
+		conns=int(conns)
+
+		# get host/port from MYSQL_HOST
+		HOST=os.getenv("MYSQL_HOST","database")
+		if ":" in HOST:
+			HOST,PORT=HOST.split(":",2)
+		else:
+			PORT=3306
+
+		self._log(f"Connecting to database at {HOST}:{PORT}")
+		
+		# fill connection pool
+		attempts=0
+		while True:
+			try:
+				self.pool=mariadb.ConnectionPool(
+						pool_name="main", pool_size=conns,
+						user=os.getenv("MYSQL_USER","root"),
+						password=os.getenv("MYSQL_PASSWORD"),
+						host=HOST,
+						port=PORT,
+						database=os.getenv("MYSQL_DATABASE","db")
+						)
+				self._log("Database connection pool set up in %d attempts"%attempts)
+				break
+
+			except mariadb.Error as e:
+				self._log(f"Error connecting to database: {e}")
+				attempts+=1
+				if attempts>=5:
+					self._log("ERROR: Limit for connecting to database reached, stopping cherrypy!")
+					cherrypy.engine.exit()
+					return
+
+				self._log("Attempt %d to connect to database failed (see above error), retrying in 10 seconds..."%attempts)
+				sleep(10)
+
+		cherrypy.engine.publish("db-started")
+		self.poolStartedEvent.set()
+
+	def stop(self):
+		if not self.pool:
+			self._log("Nothing to stop.")
+			return
+		self.pool.close()
+		self._log("Stopped database pool.")
+
+	def _log(self, *msg):
+		"""
+		Just a simple wrapper around cherrypy.log adding a context
+		"""
+		cherrypy.log(
+				context="DBCONN",
+				msg=" ".join(( str(i) for i in msg )),
+				)
+	
+	def get_connection(self):
+		if self.pool==None:
+			self.poolStartedEvent.wait(cherrypy)
+
+		conn=self.pool.get_connection()
+		conn.rollback() # reset any previous actions and refresh
+		return conn
+
+