thread watcher writes to database incase of disconnect

This commit is contained in:
iou1name 2018-02-13 16:31:23 -05:00
parent 29b740cd33
commit fb2637a14e
2 changed files with 259 additions and 221 deletions

405
db.py
View File

@ -1,6 +1,3 @@
# coding=utf-8
from __future__ import unicode_literals, absolute_import, print_function, division
import json import json
import os.path import os.path
import sys import sys
@ -8,240 +5,236 @@ import sqlite3
from tools import Identifier from tools import Identifier
if sys.version_info.major >= 3:
unicode = str
basestring = str
def _deserialize(value): def _deserialize(value):
if value is None: if value is None:
return None return None
# sqlite likes to return ints for strings that look like ints, even though # sqlite likes to return ints for strings that look like ints, even though
# the column type is string. That's how you do dynamic typing wrong. # the column type is string. That's how you do dynamic typing wrong.
value = unicode(value) value = str(value)
# Just in case someone's mucking with the DB in a way we can't account for, # Just in case someone's mucking with the DB in a way we can't account for,
# ignore json parsing errors # ignore json parsing errors
try: try:
value = json.loads(value) value = json.loads(value)
except: except:
pass pass
return value return value
class SopelDB(object): class SopelDB(object):
"""*Availability: 5.0+* """*Availability: 5.0+*
This defines an interface for basic, common operations on a sqlite This defines an interface for basic, common operations on a sqlite
database. It simplifies those common operations, and allows direct access database. It simplifies those common operations, and allows direct access
to the database, wherever the user has configured it to be. to the database, wherever the user has configured it to be.
When configured with a relative filename, it is assumed to be in the same When configured with a relative filename, it is assumed to be in the same
directory as the config.""" directory as the config."""
def __init__(self, config): def __init__(self, config):
path = config.core.db_filename path = config.core.db_filename
config_dir, config_file = os.path.split(config.filename) config_dir, config_file = os.path.split(config.filename)
config_name, _ = os.path.splitext(config_file) config_name, _ = os.path.splitext(config_file)
if path is None: if path is None:
path = os.path.join(config_dir, config_name + '.db') path = os.path.join(config_dir, config_name + '.db')
path = os.path.expanduser(path) path = os.path.expanduser(path)
if not os.path.isabs(path): if not os.path.isabs(path):
path = os.path.normpath(os.path.join(config_dir, path)) path = os.path.normpath(os.path.join(config_dir, path))
self.filename = path self.filename = path
self._create() self._create()
def connect(self): def connect(self):
"""Return a raw database connection object.""" """Return a raw database connection object."""
return sqlite3.connect(self.filename, timeout=10) return sqlite3.connect(self.filename, timeout=10)
def execute(self, *args, **kwargs): def execute(self, *args, **kwargs):
"""Execute an arbitrary SQL query against the database. """Execute an arbitrary SQL query against the database.
Returns a cursor object, on which things like `.fetchall()` can be Returns a cursor object, on which things like `.fetchall()` can be
called per PEP 249.""" called per PEP 249."""
with self.connect() as conn: with self.connect() as conn:
cur = conn.cursor() cur = conn.cursor()
return cur.execute(*args, **kwargs) return cur.execute(*args, **kwargs)
def _create(self): def _create(self):
"""Create the basic database structure.""" """Create the basic database structure."""
# Do nothing if the db already exists. # Do nothing if the db already exists.
try: try:
self.execute('SELECT * FROM nick_ids;') self.execute('SELECT * FROM nick_ids;')
self.execute('SELECT * FROM nicknames;') self.execute('SELECT * FROM nicknames;')
self.execute('SELECT * FROM nick_values;') self.execute('SELECT * FROM nick_values;')
self.execute('SELECT * FROM channel_values;') self.execute('SELECT * FROM channel_values;')
except: except:
pass pass
else: else:
return return
self.execute( self.execute(
'CREATE TABLE nick_ids (nick_id INTEGER PRIMARY KEY AUTOINCREMENT)' 'CREATE TABLE nick_ids (nick_id INTEGER PRIMARY KEY AUTOINCREMENT)'
) )
self.execute( self.execute(
'CREATE TABLE nicknames ' 'CREATE TABLE nicknames '
'(nick_id INTEGER REFERENCES nick_ids, ' '(nick_id INTEGER REFERENCES nick_ids, '
'slug STRING PRIMARY KEY, canonical string)' 'slug STRING PRIMARY KEY, canonical string)'
) )
self.execute( self.execute(
'CREATE TABLE nick_values ' 'CREATE TABLE nick_values '
'(nick_id INTEGER REFERENCES nick_ids(nick_id), ' '(nick_id INTEGER REFERENCES nick_ids(nick_id), '
'key STRING, value STRING, ' 'key STRING, value STRING, '
'PRIMARY KEY (nick_id, key))' 'PRIMARY KEY (nick_id, key))'
) )
self.execute( self.execute(
'CREATE TABLE channel_values ' 'CREATE TABLE channel_values '
'(channel STRING, key STRING, value STRING, ' '(channel STRING, key STRING, value STRING, '
'PRIMARY KEY (channel, key))' 'PRIMARY KEY (channel, key))'
) )
def get_uri(self): def get_uri(self):
"""Returns a URL for the database, usable to connect with SQLAlchemy. """Returns a URL for the database, usable to connect with SQLAlchemy.
""" """
return 'sqlite://{}'.format(self.filename) return 'sqlite://{}'.format(self.filename)
# NICK FUNCTIONS # NICK FUNCTIONS
def get_nick_id(self, nick, create=True): def get_nick_id(self, nick, create=True):
"""Return the internal identifier for a given nick. """Return the internal identifier for a given nick.
This identifier is unique to a user, and shared across all of that This identifier is unique to a user, and shared across all of that
user's aliases. If create is True, a new ID will be created if one does user's aliases. If create is True, a new ID will be created if one does
not already exist""" not already exist"""
slug = nick.lower() slug = nick.lower()
nick_id = self.execute('SELECT nick_id from nicknames where slug = ?', nick_id = self.execute('SELECT nick_id from nicknames where slug = ?',
[slug]).fetchone() [slug]).fetchone()
if nick_id is None: if nick_id is None:
if not create: if not create:
raise ValueError('No ID exists for the given nick') raise ValueError('No ID exists for the given nick')
with self.connect() as conn: with self.connect() as conn:
cur = conn.cursor() cur = conn.cursor()
cur.execute('INSERT INTO nick_ids VALUES (NULL)') cur.execute('INSERT INTO nick_ids VALUES (NULL)')
nick_id = cur.execute('SELECT last_insert_rowid()').fetchone()[0] nick_id = cur.execute('SELECT last_insert_rowid()').fetchone()[0]
cur.execute( cur.execute(
'INSERT INTO nicknames (nick_id, slug, canonical) VALUES ' 'INSERT INTO nicknames (nick_id, slug, canonical) VALUES '
'(?, ?, ?)', '(?, ?, ?)',
[nick_id, slug, nick] [nick_id, slug, nick]
) )
nick_id = self.execute('SELECT nick_id from nicknames where slug = ?', nick_id = self.execute('SELECT nick_id from nicknames where slug = ?',
[slug]).fetchone() [slug]).fetchone()
return nick_id[0] return nick_id[0]
def alias_nick(self, nick, alias): def alias_nick(self, nick, alias):
"""Create an alias for a nick. """Create an alias for a nick.
Raises ValueError if the alias already exists. If nick does not already Raises ValueError if the alias already exists. If nick does not already
exist, it will be added along with the alias.""" exist, it will be added along with the alias."""
nick = Identifier(nick) nick = Identifier(nick)
alias = Identifier(alias) alias = Identifier(alias)
nick_id = self.get_nick_id(nick) nick_id = self.get_nick_id(nick)
sql = 'INSERT INTO nicknames (nick_id, slug, canonical) VALUES (?, ?, ?)' sql = 'INSERT INTO nicknames (nick_id, slug, canonical) VALUES (?, ?, ?)'
values = [nick_id, alias.lower(), alias] values = [nick_id, alias.lower(), alias]
try: try:
self.execute(sql, values) self.execute(sql, values)
except sqlite3.IntegrityError: except sqlite3.IntegrityError:
raise ValueError('Alias already exists.') raise ValueError('Alias already exists.')
def set_nick_value(self, nick, key, value): def set_nick_value(self, nick, key, value):
"""Sets the value for a given key to be associated with the nick.""" """Sets the value for a given key to be associated with the nick."""
nick = Identifier(nick) nick = Identifier(nick)
value = json.dumps(value, ensure_ascii=False) value = json.dumps(value, ensure_ascii=False)
nick_id = self.get_nick_id(nick) nick_id = self.get_nick_id(nick)
self.execute('INSERT OR REPLACE INTO nick_values VALUES (?, ?, ?)', self.execute('INSERT OR REPLACE INTO nick_values VALUES (?, ?, ?)',
[nick_id, key, value]) [nick_id, key, value])
def get_nick_value(self, nick, key): def get_nick_value(self, nick, key):
"""Retrieves the value for a given key associated with a nick.""" """Retrieves the value for a given key associated with a nick."""
nick = Identifier(nick) nick = Identifier(nick)
result = self.execute( result = self.execute(
'SELECT value FROM nicknames JOIN nick_values ' 'SELECT value FROM nicknames JOIN nick_values '
'ON nicknames.nick_id = nick_values.nick_id ' 'ON nicknames.nick_id = nick_values.nick_id '
'WHERE slug = ? AND key = ?', 'WHERE slug = ? AND key = ?',
[nick.lower(), key] [nick.lower(), key]
).fetchone() ).fetchone()
if result is not None: if result is not None:
result = result[0] result = result[0]
return _deserialize(result) return _deserialize(result)
def unalias_nick(self, alias): def unalias_nick(self, alias):
"""Removes an alias. """Removes an alias.
Raises ValueError if there is not at least one other nick in the group. Raises ValueError if there is not at least one other nick in the group.
To delete an entire group, use `delete_group`. To delete an entire group, use `delete_group`.
""" """
alias = Identifier(alias) alias = Identifier(alias)
nick_id = self.get_nick_id(alias, False) nick_id = self.get_nick_id(alias, False)
count = self.execute('SELECT COUNT(*) FROM nicknames WHERE nick_id = ?', count = self.execute('SELECT COUNT(*) FROM nicknames WHERE nick_id = ?',
[nick_id]).fetchone()[0] [nick_id]).fetchone()[0]
if count <= 1: if count <= 1:
raise ValueError('Given alias is the only entry in its group.') raise ValueError('Given alias is the only entry in its group.')
self.execute('DELETE FROM nicknames WHERE slug = ?', [alias.lower()]) self.execute('DELETE FROM nicknames WHERE slug = ?', [alias.lower()])
def delete_nick_group(self, nick): def delete_nick_group(self, nick):
"""Removes a nickname, and all associated aliases and settings. """Removes a nickname, and all associated aliases and settings.
""" """
nick = Identifier(nick) nick = Identifier(nick)
nick_id = self.get_nick_id(nick, False) nick_id = self.get_nick_id(nick, False)
self.execute('DELETE FROM nicknames WHERE nick_id = ?', [nick_id]) self.execute('DELETE FROM nicknames WHERE nick_id = ?', [nick_id])
self.execute('DELETE FROM nick_values WHERE nick_id = ?', [nick_id]) self.execute('DELETE FROM nick_values WHERE nick_id = ?', [nick_id])
def merge_nick_groups(self, first_nick, second_nick): def merge_nick_groups(self, first_nick, second_nick):
"""Merges the nick groups for the specified nicks. """Merges the nick groups for the specified nicks.
Takes two nicks, which may or may not be registered. Unregistered Takes two nicks, which may or may not be registered. Unregistered
nicks will be registered. Keys which are set for only one of the given nicks will be registered. Keys which are set for only one of the given
nicks will be preserved. Where multiple nicks have values for a given nicks will be preserved. Where multiple nicks have values for a given
key, the value set for the first nick will be used. key, the value set for the first nick will be used.
Note that merging of data only applies to the native key-value store. Note that merging of data only applies to the native key-value store.
If modules define their own tables which rely on the nick table, they If modules define their own tables which rely on the nick table, they
will need to have their merging done separately.""" will need to have their merging done separately."""
first_id = self.get_nick_id(Identifier(first_nick)) first_id = self.get_nick_id(Identifier(first_nick))
second_id = self.get_nick_id(Identifier(second_nick)) second_id = self.get_nick_id(Identifier(second_nick))
self.execute( self.execute(
'UPDATE OR IGNORE nick_values SET nick_id = ? WHERE nick_id = ?', 'UPDATE OR IGNORE nick_values SET nick_id = ? WHERE nick_id = ?',
[first_id, second_id]) [first_id, second_id])
self.execute('DELETE FROM nick_values WHERE nick_id = ?', [second_id]) self.execute('DELETE FROM nick_values WHERE nick_id = ?', [second_id])
self.execute('UPDATE nicknames SET nick_id = ? WHERE nick_id = ?', self.execute('UPDATE nicknames SET nick_id = ? WHERE nick_id = ?',
[first_id, second_id]) [first_id, second_id])
# CHANNEL FUNCTIONS # CHANNEL FUNCTIONS
def set_channel_value(self, channel, key, value): def set_channel_value(self, channel, key, value):
channel = Identifier(channel).lower() channel = Identifier(channel).lower()
value = json.dumps(value, ensure_ascii=False) value = json.dumps(value, ensure_ascii=False)
self.execute('INSERT OR REPLACE INTO channel_values VALUES (?, ?, ?)', self.execute('INSERT OR REPLACE INTO channel_values VALUES (?, ?, ?)',
[channel, key, value]) [channel, key, value])
def get_channel_value(self, channel, key): def get_channel_value(self, channel, key):
"""Retrieves the value for a given key associated with a channel.""" """Retrieves the value for a given key associated with a channel."""
channel = Identifier(channel).lower() channel = Identifier(channel).lower()
result = self.execute( result = self.execute(
'SELECT value FROM channel_values WHERE channel = ? AND key = ?', 'SELECT value FROM channel_values WHERE channel = ? AND key = ?',
[channel, key] [channel, key]
).fetchone() ).fetchone()
if result is not None: if result is not None:
result = result[0] result = result[0]
return _deserialize(result) return _deserialize(result)
# NICK AND CHANNEL FUNCTIONS # NICK AND CHANNEL FUNCTIONS
def get_nick_or_channel_value(self, name, key): def get_nick_or_channel_value(self, name, key):
"""Gets the value `key` associated to the nick or channel `name`. """Gets the value `key` associated to the nick or channel `name`.
""" """
name = Identifier(name) name = Identifier(name)
if name.is_nick(): if name.is_nick():
return self.get_nick_value(name, key) return self.get_nick_value(name, key)
else: else:
return self.get_channel_value(name, key) return self.get_channel_value(name, key)
def get_preferred_value(self, names, key): def get_preferred_value(self, names, key):
"""Gets the value for the first name which has it set. """Gets the value for the first name which has it set.
`names` is a list of channel and/or user names. Returns None if none of `names` is a list of channel and/or user names. Returns None if none of
the names have the key set.""" the names have the key set."""
for name in names: for name in names:
value = self.get_nick_or_channel_value(name, key) value = self.get_nick_or_channel_value(name, key)
if value is not None: if value is not None:
return value return value

View File

@ -15,6 +15,27 @@ def setup(bot):
""" """
bot.memory["watcher"] = {} bot.memory["watcher"] = {}
con = bot.db.connect()
cur = con.cursor()
try:
watching = cur.execute("SELECT * FROM watcher").fetchall()
except:
cur.execute("CREATE TABLE watcher("
"api_url STRING PRIMARY KEY,"
"name STRING DEFAULT 'Anonymous',"
"last_post INT,"
"time_since STRING,"
"channel STRING)")
cur.commit()
con.close()
else:
for thread in watching:
t = WatcherThread(bot, thread[0], thread[1], thread[2], thread[3],
thread[4])
t.start()
bot.memory["watcher"][get_thread_url(thread[0])] = t
con.close()
def get_time(): def get_time():
""" """
@ -64,7 +85,9 @@ def watch(bot, trigger):
A thread watcher for 4chan. A thread watcher for 4chan.
""" """
url = trigger.group(3) url = trigger.group(3)
op_name = trigger.group(4) name = trigger.group(4)
if not name:
name = "Anonymous"
if url in bot.memory["watcher"].keys(): if url in bot.memory["watcher"].keys():
return bot.say("Error: I'm already watching that thread.") return bot.say("Error: I'm already watching that thread.")
@ -75,12 +98,21 @@ def watch(bot, trigger):
return bot.say("404: thread not found") return bot.say("404: thread not found")
thread = res.json() thread = res.json()
last_post = get_last_post(thread, op_name) last_post = get_last_post(thread, name)
time_since = get_time() time_since = get_time()
t = WatcherThread(bot, api_url, op_name, last_post, time_since) t = WatcherThread(bot, api_url, name, last_post, time_since, trigger.sender)
t.start() t.start()
bot.memory["watcher"][url] = t bot.memory["watcher"][url] = t
con = bot.db.connect()
cur = con.cursor()
cur.execute("INSERT INTO watcher(api_url, name, last_post, time_since, channel)"
" VALUES(?,?,?,?,?)", (api_url, name, last_post, time_since,
trigger.sender))
con.commit()
con.close()
bot.say("[\x0304Watcher\x03] Watching thread: \x0307" + url) bot.say("[\x0304Watcher\x03] Watching thread: \x0307" + url)
@ -96,51 +128,64 @@ def unwatch(bot, trigger):
bot.memory["watcher"].pop(url) bot.memory["watcher"].pop(url)
except KeyError: except KeyError:
return bot.say("Error: I'm not watching that thread.") return bot.say("Error: I'm not watching that thread.")
return bot.say("[\x0304Watcher\x03] No longer watching: \x0307" + url)
con = bot.db.connect()
cur = con.cursor()
cur.execute("DELETE FROM watcher WHERE api_url = ?", (get_api_url(url),))
con.commit()
con.close()
bot.say("[\x0304Watcher\x03] No longer watching: \x0307" + url)
class WatcherThread(threading.Thread): class WatcherThread(threading.Thread):
def __init__(self, bot, api_url, op_name, last_post, time_since): def __init__(self, bot, api_url, name, last_post, time_since, channel):
threading.Thread.__init__(self) threading.Thread.__init__(self)
self.stop = threading.Event() self.stop = threading.Event()
self.period = 20 self.period = 20
self.bot = bot self.bot = bot
self.api_url = api_url self.api_url = api_url
self.op_name = op_name self.name = name
self.last_post = last_post self.last_post = last_post
self.time_since = time_since self.time_since = time_since
self.channel = channel
def run(self): def run(self):
while not self.stop.is_set(): while not self.stop.is_set():
self.stop.wait(self.period)
headers = {"If-Modified-Since": self.time_since} headers = {"If-Modified-Since": self.time_since}
res = requests.get(self.api_url, headers=headers, verify=True) try:
self.time_since = get_time() res = requests.get(self.api_url, headers=headers, verify=True)
self.time_since = get_time()
except urllib3.exceptions.NewConnectionError:
print(f"Watcher: Thread {self.api_url}: Connection error")
continue
if res.status_code == 404: if res.status_code == 404:
msg = "[\x0304Watcher\x03] Thread deleted: " \ msg = "[\x0304Watcher\x03] Thread deleted: " \
+ f"\x0307{get_thread_url(self.api_url)}" + f"\x0307{get_thread_url(self.api_url)}"
self.bot.say(msg) self.bot.say(msg, channel)
self.stop.set() self.stop.set()
continue continue
if res.status_code == 304: if res.status_code == 304:
self.stop.wait(self.period)
continue continue
thread = res.json() thread = res.json()
if thread["posts"][0].get("closed"): if thread["posts"][0].get("closed"):
msg = "[\x0304Watcher\x03] Thread closed: " \ msg = "[\x0304Watcher\x03] Thread closed: " \
+ f"\x0307{get_thread_url(self.api_url)}" + f"\x0307{get_thread_url(self.api_url)}"
self.bot.say(msg) self.bot.say(msg, channel)
self.stop.set() self.stop.set()
continue continue
new_last_post = get_last_post(thread, self.op_name) new_last_post = get_last_post(thread, self.name)
if new_last_post > self.last_post: if new_last_post > self.last_post:
self.last_post = new_last_post self.last_post = new_last_post
msg = "[\x0304Watcher\x03] New post from \x0308" \ msg = "[\x0304Watcher\x03] New post from \x0308" \
+ f"{self.op_name}\x03 in \x0307{get_thread_url(self.api_url)}" + f"{self.name}\x03 in \x0307{get_thread_url(self.api_url)}"
self.bot.say(msg) self.bot.say(msg, channel)
self.stop.wait(self.period)