improve threading, add modcommand, batch posts, default workers =4

This commit is contained in:
2026-10-19 14:37:59 +10:00
parent 8531294ca5
commit eef1ce9cc2
+577 -228
View File
@@ -2,273 +2,597 @@
import znc import znc
import re import re
import multiprocessing import queue
import threading
import warnings import warnings
import traceback import traceback
import inspect
import os
import pprint
import json
import uuid
from datetime import datetime from datetime import datetime
from time import sleep
from contextlib import closing from contextlib import closing
class mssql(znc.Module): class mssql(znc.Module):
module_types = [znc.CModInfo.GlobalModule] module_types = [znc.CModInfo.GlobalModule]
description = "Log all channels to a Microsoft SQL Server database" description = "Log all channels to a Microsoft SQL Server database"
has_args = False has_args = True
log_queue = multiprocessing.SimpleQueue() log_queue = None
internal_log = None worker_threads = None
hook_debugging = True internal_log = None
hook_debugging = False
debug_log_saves = False
default_workers = 4
num_workers = 4
max_workers = 32
batch_size = 50
db_config = None
next_worker_id = 0
def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network): def put_log(self, code, channel, host, user, ident, user_mode, date, target_user, message, network):
""" """
Adds the log line to database write queue. 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({ decolouriser = re.compile(r"\x1f|\x02|\x1d|\x16|\x11|\x1e|\x0f|\x03(?:\d{1,2}(?:,\d{1,2})?)?|\x04(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE)
'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): if message is not None:
""" message.encode('utf8', 'replace').decode('utf8')
This module hook is called when a module is loaded. message = decolouriser.sub("", message)
: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: if self.log_queue is None:
db = MSSQLDatabase({ return
'host': "db.spd.internal",
'user': "_srv_znc",
'passwd': "zncznc",
'db': "znc"
})
multiprocessing.Process(target=DatabaseThread.worker_safe, queue_len = self.log_queue.qsize() if self.log_queue is not None else 0
args=(db, self.log_queue, self.internal_log)).start() if queue_len > 10:
return True with self.internal_log.error() as target:
except Exception as e: target.write('Queue length high: {} entries at {}\n'.format(queue_len, datetime.utcnow()))
message.s = str(e)
with self.internal_log.error() as target: try:
target.write('Could not initialize module caused by: {} {}\n'.format(type(e), str(e))) self.log_queue.put({
target.write('Stack trace: ' + traceback.format_exc()) 'znc_user': self.GetUser().GetUserName() if self.GetUser() is not None else None,
target.write('\n') 'code': code,
'channel': channel,
'host': host,
'user': user,
'ident': ident,
'user_mode': user_mode,
'date': str(date),
'target_user': target_user,
'message': message,
'network': network,
'_id': uuid.uuid4()
})
except Exception as e:
with self.internal_log.error() as target:
target.write('Failed to enqueue log event: {0} {1}\n'.format(type(e), str(e)))
target.write('Stack trace: ' + traceback.format_exc())
target.write('\n')
return False def debug_hook(self):
"""
Dumps parent calling method name and its arguments to debug logfile.
"""
def OnOp(self, user, target_user, channel, noChange): if self.hook_debugging is not True:
self.insert("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) return
return True
def OnPrivNotice(self, user, message): frameinfo = inspect.stack()[1]
self.insert("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) argvals = frameinfo.frame.f_locals
return True
def OnDeop(self, user, target_user, channel, noChange): with self.internal_log.debug() as target:
self.insert("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) target.write('Called method: ' + frameinfo.function + '()\n')
return True for argname in argvals:
if argname == 'self':
continue
target.write(' ' + argname + ' -> ' + pprint.pformat(argvals[argname]) + '\n')
target.write('\n')
def OnVoice(self, user, target_user, channel, noChange): def _start_workers(self, count):
self.insert("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) """Start `count` more worker threads on the shared queue."""
return True for _ in range(count):
idx = self.next_worker_id
self.next_worker_id += 1
db = MSSQLDatabase(self.db_config)
worker_thread = threading.Thread(
target=DatabaseThread.worker_safe,
args=(db, self.log_queue, self.internal_log, self.debug_log_saves, self.batch_size),
daemon=True,
name=f'DatabaseWorker-{idx}'
)
worker_thread.start()
self.worker_threads.append(worker_thread)
def OnDevoice(self, user, target_user, channel, noChange): def OnModCommand(self, command):
self.insert("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork())) """
return True Commands: Status, AddWorkers <n>
"""
parts = str(command).split()
if not parts:
return
cmd = parts[0].lower()
def OnMode(self, user, channel, mode, argument, added, noChange): if cmd == 'status':
mode = chr(mode) alive = len([t for t in self.worker_threads if t.is_alive()]) if self.worker_threads else 0
code = "MODE" qsize = self.log_queue.qsize() if self.log_queue is not None else 0
if mode == "b": self.PutModule('Workers alive: {}, queue length: {}, batch size: {}'.format(alive, qsize, self.batch_size))
code = "BAN" if added else "UNBAN" elif cmd == 'addworkers':
else: try:
sign = "+" if added else "-" n = int(parts[1])
argument = "{0}{1} {2}".format(sign, mode, argument) except (IndexError, ValueError):
self.PutModule('Usage: AddWorkers <n>')
return
current = len(self.worker_threads)
n = max(0, min(n, self.max_workers - current))
if n == 0:
self.PutModule('No workers added (limit is {}, currently {})'.format(self.max_workers, current))
return
self._start_workers(n)
self.PutModule('Started {} workers, now {}'.format(n, len(self.worker_threads)))
else:
self.PutModule('Commands: Status, AddWorkers <n>')
self.insert(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork())) def OnLoad(self, args, message):
return True """
This module hook is called when a module is loaded.
"""
# Initialize internal log first for debugging
self.internal_log = InternalLog(self.GetSavePath())
def OnQuit(self, user, message, channels): # Parse num_workers from args (default 4)
for channel in channels: self.num_workers = self.default_workers
self.insert("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) args_str = ""
return True if args:
try:
args_str = str(args) if not hasattr(args, 's') else args.s
if args_str:
self.num_workers = int(args_str.strip())
if self.num_workers < 1:
self.num_workers = 1
if self.num_workers > self.max_workers:
self.num_workers = self.max_workers
except (ValueError, AttributeError, TypeError) as e:
with self.internal_log.debug() as target:
target.write('Failed to parse args: {} (args_str: "{}", type: {})\n'.format(e, args_str, type(args)))
def OnNick(self, user, new_nick, channels): with self.internal_log.debug() as target:
for channel in channels: target.write('mssql module loaded at: {}\n'.format(datetime.utcnow()))
self.insert("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork())) target.write('save_path: {}\n'.format(self.GetSavePath()))
return True target.write('raw args: {} (type: {})\n'.format(repr(args), type(args).__name__))
target.write('parsed args_str: "{}"\n'.format(args_str))
target.write('num_workers: {}\n'.format(self.num_workers))
target.write('\n')
def OnKick(self, user, target_nick, channel, message): with self.internal_log.error() as target:
self.insert("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork())) target.write('mssql module loaded at: {}\n'.format(datetime.utcnow()))
return True target.write('hook_debugging: {}\n'.format(self.hook_debugging))
target.write('num_workers: {}\n'.format(self.num_workers))
target.write('\n')
def OnJoin(self, user, channel): self.debug_hook()
self.insert("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork()))
return True
def OnPart(self, user, channel, message): try:
self.insert("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork())) self.log_queue = queue.Queue()
return True self.db_config = {
'host': "dbx.spd.internal",
'user': "_srv_znc",
'password': "znc",
'database': "znc",
'port': '1433',
'appname': 'ZNC Logger'
}
def OnUserAction (self, target, message): self.worker_threads = []
(channel, user) = self.resolveTarget(target) self.next_worker_id = 0
if channel == None: self._start_workers(self.num_workers)
return True
return self.OnChanAction(user, channel, message)
def OnPrivMsg(self, nick, message): return True
self.insert("MSG", "", nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) except Exception as e:
return True message.s = str(e)
def OnUserMsg (self, target, message): with self.internal_log.error() as target:
(channel, user) = self.resolveTarget(target) target.write('Could not initialize module caused by: {} {}\n'.format(type(e), str(e)))
if channel == None: target.write('Stack trace: ' + traceback.format_exc())
return True target.write('\n')
return self.OnChanMsg(user, channel, message)
def OnChanAction(self, user, channel, message): return False
self.insert("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True
def OnChanMsg(self, user, channel, message):
self.insert("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork()))
return True
def OnTopic(self, user, channel, message): def OnUnload(self):
self.insert("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork())) if self.log_queue is None:
return True return
def findMode(self, channel, user): with self.internal_log.debug() as target:
realUser = channel.FindNick(user.GetNick()) target.write('Unloading module, draining queue length {}\n'.format(self.log_queue.qsize()))
if not realUser:
return None
return str(realUser.GetPermChar())
def resolveTarget(self, target): # Send one sentinel value per worker thread
target = target.s if self.worker_threads:
channel = self.GetNetwork().FindChan(target) for _ in self.worker_threads:
if channel == None: self.log_queue.put(None)
return (None, None)
user = self.GetNetwork().GetIRCNick()
return (channel, user)
def colourstrip(self, data): self.log_queue.join()
find = data.find('\x03')
while find > -1:
done = False
data = data[0:find] + data[find+1:]
if len(data) <= find+1:
done = True
try:
assert not done
assert int(data[find])
while True:
assert int(data[find])
data = data[0:find] + data[find+1:]
except:
if not done:
if data[find] != ',': done = True
if (not done) and (len(data) > find+1) and (data[find] == ','): if self.worker_threads is not None:
try: for worker_thread in self.worker_threads:
assert not done worker_thread.join(timeout=30)
assert int(data[find+1])
data = data[0:find] + data[find+1:]
data = data[0:find] + data[find+1:]
except:
done = True
try:
assert not done
while True:
assert int(data[find])
data = data[0:find] + data[find+1:]
except: pass
find = data.find('\x03') with self.internal_log.debug() as target:
data = data.replace('\x02','') target.write('Unload complete, queue drained\n')
data = data.replace('\x1d','')
data = data.replace('\x1f','') def __del__(self):
data = data.replace('\x16','') # Terminate worker threads cleanly if possible.
data = data.replace('\x0f','') try:
return data if self.log_queue is not None:
# Send sentinel values for all workers
if self.worker_threads:
for _ in self.worker_threads:
self.log_queue.put_nowait(None)
except Exception:
pass
def OnIRCConnected(self):
"""
This module hook is called after a successful login to IRC.
:rtype: None
"""
self.debug_hook()
self.put_log("CONNECT", None, None, None, None, None, datetime.now(), None, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnIRCDisconnected(self):
"""
This module hook is called when a user gets disconnected from IRC.
:rtype: None
"""
self.debug_hook()
self.put_log("DISCONNECT", None, None, None, None, None, datetime.now(), None, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnOp(self, user, target_user, channel, noChange):
self.debug_hook()
self.put_log("OP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnChanNotice(self, nick, channel, message):
"""
Called when we receive a channel NOTICE message from IRC.
:type nick: CNick &
:type channel: CChan &
:type message: CString &
:rtype: EModRet
"""
self.debug_hook()
self.put_log("NOTICE", channel.GetName(), nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnPrivNotice(self, user, message):
self.debug_hook()
self.put_log("NOTICE", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnUserNotice(self, target, message):
"""
This module hook is called when a user sends a NOTICE message.
:type target: CString &
:type message: CString &
:rtype: EModRet
"""
self.debug_hook()
network = self.GetNetwork()
if network:
self.put_log("NOTICE", None, target.GetHost(), target.GetNick(), target.GetIdent(), None, datetime.now(), target.GetNick(), message, str(self.GetNetwork()))
def OnDeop(self, user, target_user, channel, noChange):
self.debug_hook()
self.put_log("DEOP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnVoice(self, user, target_user, channel, noChange):
self.debug_hook()
self.put_log("VOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnDevoice(self, user, target_user, channel, noChange):
self.debug_hook()
self.put_log("DEVOICE", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_user.GetNick(), None, str(self.GetNetwork()))
return znc.CONTINUE
def OnMode(self, user, channel, mode, argument, added, noChange):
self.debug_hook()
mode = chr(mode)
code = "MODE"
if mode == "b":
code = "BAN" if added else "UNBAN"
else:
sign = "+" if added else "-"
argument = "{0}{1} {2}".format(sign, mode, argument)
self.put_log(code, channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, argument, str(self.GetNetwork()))
return znc.CONTINUE
def OnQuit(self, user, message, channels):
self.debug_hook()
for channel in channels:
self.put_log("QUIT", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return znc.CONTINUE
def OnNick(self, user, new_nick, channels):
self.debug_hook()
for channel in channels:
self.put_log("NICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), new_nick, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnKick(self, user, target_nick, channel, message):
self.debug_hook()
self.put_log("KICK", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), target_nick, message, str(self.GetNetwork()))
return znc.CONTINUE
def OnJoin(self, user, channel):
self.debug_hook()
self.put_log("JOIN", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, None, str(self.GetNetwork()))
return znc.CONTINUE
def OnPart(self, user, channel, message):
self.debug_hook()
self.put_log("PART", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return znc.CONTINUE
def OnUserAction (self, target, message):
self.debug_hook()
(channel, user) = self.resolveTarget(target)
if channel == None:
self.put_log("ME", None, user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message, str(self.GetNetwork()))
return znc.CONTINUE
return self.OnChanAction(user, channel, message)
def OnPrivAction(self, nick, message):
"""
Called when we receive a private CTCP ACTION ("/me" in query) from IRC.
:type nick: CNick &
:type message: CString &
:rtype: EModRet
"""
self.debug_hook()
self.put_log("ME", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnPrivMsg(self, nick, message):
self.debug_hook()
self.put_log("MSG", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnPrivCTCP(self, nick, message):
self.debug_hook()
self.put_log("CTCP", None, nick.GetHost(), nick.GetNick(), nick.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnChanCTCP(self, user, channel, message):
self.debug_hook()
self.put_log("CTCP", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnUserCTCP(self, target, message):
self.debug_hook()
local_nick = self.GetNetwork().GetIRCNick()
if local_nick is not None:
self.put_log(
"CTCP",
None,
local_nick.GetHost(),
local_nick.GetNick(),
local_nick.GetIdent(),
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
else:
self.put_log(
"CTCP",
None,
None,
None,
None,
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
return znc.CONTINUE
def OnUserMsg (self, target, message):
self.debug_hook()
(channel, user) = self.resolveTarget(target)
if channel == None:
local_nick = self.GetNetwork().GetIRCNick()
if local_nick is not None:
self.put_log(
"MSG",
None,
local_nick.GetHost(),
local_nick.GetNick(),
local_nick.GetIdent(),
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
else:
self.put_log(
"MSG",
None,
None,
None,
None,
None,
datetime.now(),
target.s,
message.s,
str(self.GetNetwork()),
)
return znc.CONTINUE
return self.OnChanMsg(user, channel, message)
def OnChanAction(self, user, channel, message):
self.debug_hook()
self.put_log("ME", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return True
def OnChanMsg(self, user, channel, message):
self.debug_hook()
self.put_log("SAY", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), self.findMode(channel, user), datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def OnTopic(self, user, channel, message):
self.debug_hook()
self.put_log("TOPIC", channel.GetName(), user.GetHost(), user.GetNick(), user.GetIdent(), None, datetime.now(), None, message.s, str(self.GetNetwork()))
return znc.CONTINUE
def findMode(self, channel, user):
realUser = channel.FindNick(user.GetNick())
if not realUser:
return None
return str(realUser.GetPermChar())
def resolveTarget(self, target):
target = target.s
channel = self.GetNetwork().FindChan(target)
if channel == None:
return (None, None)
user = self.GetNetwork().GetIRCNick()
return (channel, user)
def colourstrip(self, data):
find = data.find('\x03')
while find > -1:
done = False
data = data[0:find] + data[find+1:]
if len(data) <= find+1:
done = True
try:
assert not done
assert int(data[find])
while True:
assert int(data[find])
data = data[0:find] + data[find+1:]
except:
if not done:
if data[find] != ',': done = True
if (not done) and (len(data) > find+1) and (data[find] == ','):
try:
assert not done
assert int(data[find+1])
data = data[0:find] + data[find+1:]
data = data[0:find] + data[find+1:]
except:
done = True
try:
assert not done
while True:
assert int(data[find])
data = data[0:find] + data[find+1:]
except: pass
find = data.find('\x03')
data = data.replace('\x02','')
data = data.replace('\x1d','')
data = data.replace('\x1f','')
data = data.replace('\x16','')
data = data.replace('\x0f','')
return data
class DatabaseThread: class DatabaseThread:
@staticmethod @staticmethod
def worker_safe(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None: def worker_safe(db, log_queue, internal_log, debug_log_saves=False, batch_size=50) -> None:
try: while True:
DatabaseThread.worker(db, log_queue, internal_log) try:
except Exception as e: DatabaseThread.worker(db, log_queue, internal_log, debug_log_saves, batch_size)
with internal_log.error() as target: break
target.write('Unrecoverable exception in worker thread: {0} {1}\n'.format(type(e), str(e))) except Exception as e:
target.write('Stack trace: ' + traceback.format_exc()) with internal_log.error() as target:
target.write('\n') target.write('Worker thread crashed: {0} {1}\n'.format(type(e), str(e)))
raise target.write('Stack trace: ' + traceback.format_exc())
target.write('Restarting worker thread in 5 seconds...\n\n')
sleep(5)
continue
@staticmethod @staticmethod
def worker(db, log_queue: multiprocessing.SimpleQueue, internal_log) -> None: def worker(db, log_queue, internal_log, debug_log_saves=False, batch_size=50) -> None:
db.connect() stop = False
while not stop:
first = log_queue.get()
if first is None:
log_queue.task_done() # count the sentinel too, or join() hangs
return
while True: # Collect more rows without blocking, up to batch_size
item = log_queue.get() taken = 1
if item is None: items = [first]
break while len(items) < batch_size:
try:
nxt = log_queue.get_nowait()
except queue.Empty:
break
taken += 1
if nxt is None:
stop = True
break
items.append(nxt)
try: try:
db.ensure_connected() DatabaseThread.save_batch(db, items, log_queue, internal_log, debug_log_saves)
db.insert_into('logs', item) finally:
except Exception as e: # exactly one task_done() per get(), after the batch is finished
sleep_for = 10 for _ in range(taken):
log_queue.task_done()
with internal_log.error() as target: @staticmethod
target.write('Could not save to database caused by: {0} {1}\n'.format(type(e), str(e))) def save_batch(db, items, log_queue, internal_log, debug_log_saves) -> None:
if 'open' in dir(db.conn): while True:
target.write('Database handle state: {}\n'.format(db.conn.open)) try:
target.write('Stack trace: ' + traceback.format_exc()) db.ensure_connected(internal_log, str(items[0]['_id']))
target.write('Current log: ') db.insert_batch(items)
json.dump(item, target) if debug_log_saves:
target.write('\n\n') with internal_log.debug() as target:
target.write('Retry in {} s\n'.format(sleep_for)) target.write('Saved {} log items queue_size={}\n'.format(len(items), log_queue.qsize()))
return
except Exception as e:
sleep_for = 10
sleep(sleep_for) with internal_log.error() as target:
target.write('Could not save to database caused by: {0} {1}\n'.format(type(e), str(e)))
target.write('Stack trace: ' + traceback.format_exc())
target.write('Batch of {} rows, first log: '.format(len(items)))
item_copy = items[0].copy()
item_copy['_id'] = str(item_copy['_id'])
json.dump(item_copy, target)
target.write('\n\n')
target.write('Retry in {} s\n'.format(sleep_for))
with internal_log.error() as target: sleep(sleep_for)
target.write('Retrying now.\n'.format(sleep_for))
log_queue.put(item)
class InternalLog: class InternalLog:
def __init__(self, save_path: str): def __init__(self, save_path: str):
self.save_path = save_path self.save_path = save_path
def debug(self): def debug(self):
return self.open('debug') return self.open('debug')
def error(self): def error(self):
return self.open('error') return self.open('error')
def open(self, level: str): def open(self, level: str):
target = open(os.path.join(self.save_path, level + '.log'), 'a') target = open(os.path.join(self.save_path, level + '.log'), 'a')
line = 'Log opened at: {} UTC\n'.format(datetime.utcnow()) return target
target.write(line)
target.write('=' * len(line) + '\n\n')
return target
@@ -278,17 +602,42 @@ class Database:
self.conn = None self.conn = None
class MSSQLDatabase(Database): class MSSQLDatabase(Database):
def connect(self) -> None: def connect(self) -> None:
import pymssql import pymssql
self.conn = pymssql.connect (**self.dsn) self.conn = pymssql.connect(**self.dsn)
def ensure_connected(self): def ensure_connected(self, internal_log, log_id):
if self.conn.status == 0: # No per-row SELECT 1. A dead connection is detected when an
self.connect() # insert fails; insert_batch() then drops it and the retry reconnects.
if self.conn is None:
self.connect()
def insert_into(self, table, row): def insert_batch(self, rows):
# self.PutModule("Inserting row...") """Insert all rows, then commit once. All or nothing: on any error
cols = ', '.join('"{}"'.format(col) for col in row.keys()) the transaction is rolled back and the connection is dropped."""
vals = ', '.join('%({})s'.format(col) for col in row.keys()) try:
self.conn.cursor().callproc("add_log", vals) for row in rows:
self.conn.commit() params = (row['znc_user'], row['date'], row['code'], row['network'], row['channel'], row['host'], row['user'], row['ident'], row['user_mode'], row['target_user'], row['message'])
cur = self.conn.cursor()
try:
cur.callproc("add_log", params)
finally:
try:
cur.close()
except Exception:
pass
self.conn.commit()
except Exception:
try:
self.conn.rollback()
except Exception:
pass
try:
self.conn.close()
except Exception:
pass
self.conn = None
raise
def insert_into(self, table, row):
self.insert_batch([row])