Refactor and multi-thread

mssql.py

@@ -1,8 +1,11 @@
 #!/usr/bin/env python3
 
 import znc
-import pymssql
 import re
+import multiprocessing
+import warnings
+import traceback
+
 from datetime import datetime
 from contextlib import closing
 
@@ -11,14 +14,65 @@ class mssql(znc.Module):
 	description = "Log all channels to a Microsoft SQL Server database"
 	has_args = False
 
-        sql_server = "db.spd.internal"
+	log_queue = multiprocessing.SimpleQueue()
-        sql_user = "_srv_znc"
+	internal_log = None
-        sql_password = "zncznc"
+	hook_debugging = True
-        sql_database = "znc"
 
-#       def OnLoad(self, args,  message):
+	def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network):
-#               self.PutModule("Loading MSSQL...")
+		"""
-#               return True
+        Adds the log line to database write queue.
+        """
+		regex = re.compile("\x1f|\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE)
+		message = regex.sub("", message)
+		message = self.colourstrip(message)
+		
+		self.log_queue.put({
+			'znc_user': self.GetUser().GetUserName() if self.GetUser() is not None else None,
+			'code': code,
+			'channel': channel,
+			'host': host,
+			'user': user,
+			'ident': ident,
+			'user_mode': user_mode,
+			'date': date,
+			'target_user': target_user,
+			'message': message,
+			'network': network
+		})
+
+	def OnLoad(self, args, message):
+		"""
+        This module hook is called when a module is loaded.
+        :type args: const CString &
+        :type args: CString &
+        :rtype: bool
+        :param args: The arguments for the modules.
+        :param message: A message that may be displayed to the user after loading the module.
+        :return: True if the module loaded successfully, else False.
+        """
+		self.internal_log = InternalLog(self.GetSavePath())
+		self.debug_hook()
+
+		try:
+			db = MSSQLDatabase({	
+								'host': "db.spd.internal",
+								'user': "_srv_znc",
+								'passwd': "zncznc",
+								'db': "znc"
+							})
+			
+			multiprocessing.Process(target=DatabaseThread.worker_safe,
+									args=(db, self.log_queue, self.internal_log)).start()
+			return True
+		except Exception as e:
+			message.s = str(e)
+
+			with self.internal_log.error() as target:
+				target.write('Could not initialize module caused by: {} {}\n'.format(type(e), str(e)))
+				target.write('Stack trace: ' + traceback.format_exc())
+				target.write('\n')
+
+			return False
 
 	def OnOp(self, user, target_user, channel, noChange):
 		self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
@@ -102,16 +156,6 @@ class mssql(znc.Module):
 		self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
 		return True
 
-#       def OnLoad(self, args, message):
-#               match = re.search("(.*?):(.*?)@(.*)", args)
-#               if match:
-#                       self.username = match.group(1)
-#                       self.password = match.group(2)
-#                       self.host = match.group(3)
-#                       return True
-#               else:
-#                       return False
-
 	def findMode(self, channel, user):
 		realUser = channel.FindNick(user.GetNick())
 		if not realUser:
@@ -166,25 +210,85 @@ class mssql(znc.Module):
 			data = data.replace('\x0f','')
 			return data
 
-        def insert(self, code, channel, host, user, ident, user_mode, date, target_user, message, network):
+class DatabaseThread:
-#                self.PutModule("Inserting row...")
+    @staticmethod
-                regex = re.compile("\x1f|\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE)
+    def worker_safe(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None:
+        try:
+            DatabaseThread.worker(db, log_queue, internal_log)
+        except Exception as e:
+            with internal_log.error() as target:
+                target.write('Unrecoverable exception in worker thread: {0} {1}\n'.format(type(e), str(e)))
+                target.write('Stack trace: ' + traceback.format_exc())
+                target.write('\n')
+            raise
+
+    @staticmethod
+    def worker(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None:
+        db.connect()
+
+        while True:
+            item = log_queue.get()
+            if item is None:
+                break
 
             try:
-                        with closing(pymssql.connect ("db.spd.internal", "_srv_znc", "znc", "znc")) as conn:
+                db.ensure_connected()
-                                with closing(conn.cursor()) as cursor:
+                db.insert_into('logs', item)
-#                                        sql = "INSERT INTO log (date, code, network, channel, host, [user], ident, mode, target, message) VALUES(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
-                                        sql = "EXEC dbo.add_log (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
-                                        if message is not None:
-#                                                message = message.encode('utf-8').decode('utf-8', 'replace')
-                                                message = regex.sub("", message)
-                                                message = self.colourstrip(message)
-                                        znc_user = self.GetUser().GetUserName() if self.GetUser() is not None else None
-                                        params = (znc_user, str(date), code, network, channel, host, user, ident, user_mode, target_user, message)
-#                                        cursor.execute (sql, params)
-                                        cursor.callproc("add_log", params)
-                                conn.commit()
-                                conn.close()
-
             except Exception as e:
-                        self.PutModule("Could not save {0} to database caused by: {1} - {2}".format(code, type(e), str(e)))
+                sleep_for = 10
+
+                with internal_log.error() as target:
+                    target.write('Could not save to database caused by: {0} {1}\n'.format(type(e), str(e)))
+                    if 'open' in dir(db.conn):
+                        target.write('Database handle state: {}\n'.format(db.conn.open))
+                    target.write('Stack trace: ' + traceback.format_exc())
+                    target.write('Current log: ')
+                    json.dump(item, target)
+                    target.write('\n\n')
+                    target.write('Retry in {} s\n'.format(sleep_for))
+
+                sleep(sleep_for)
+
+                with internal_log.error() as target:
+                    target.write('Retrying now.\n'.format(sleep_for))
+                    log_queue.put(item)
+
+class InternalLog:
+    def __init__(self, save_path: str):
+        self.save_path = save_path
+
+    def debug(self):
+        return self.open('debug')
+
+    def error(self):
+        return self.open('error')
+
+    def open(self, level: str):
+        target = open(os.path.join(self.save_path, level + '.log'), 'a')
+        line = 'Log opened at: {} UTC\n'.format(datetime.utcnow())
+        target.write(line)
+        target.write('=' * len(line) + '\n\n')
+        return target
+
+
+
+class Database:
+    def __init__(self, dsn: dict):
+        self.dsn = dsn
+        self.conn = None
+	
+class MSSQLDatabase(Database):
+	def connect(self) -> None:
+		import pymssql
+		self.conn = pymssql.connect (**self.dsn)
+
+	def ensure_connected(self):
+		if self.conn.status == 0:
+			self.connect()
+	    
+	def insert_into(self, table, row):
+#		self.PutModule("Inserting row...")
+		cols = ', '.join('"{}"'.format(col) for col in row.keys())
+		vals = ', '.join('%({})s'.format(col) for col in row.keys())	
+		self.conn.cursor().callproc("add_log", vals)
+		self.conn.commit()