Merge pull request CloudBotIRC#166 from linuxdaemon/gonzobot+herald-cache

Add caching to herald.py

Author: edwardslabs

plugins/herald.py

@@ -1,6 +1,7 @@
 import random
 import re
 import time
+from collections import defaultdict
 
 from sqlalchemy import Table, Column, String, PrimaryKeyConstraint, select
 
@@ -19,28 +20,37 @@
     PrimaryKeyConstraint('name', 'chan')
 )
 
+herald_cache = defaultdict(dict)
+
+
+@hook.on_start
+def load_cache(db):
+    herald_cache.clear()
+    for row in db.execute(table.select()):
+        herald_cache[row["chan"]][row["name"]] = row["quote"]
+
 
 @hook.command()
-def herald(text, nick, chan, db):
+def herald(text, nick, chan, db, reply):
     """{<message>|show|delete|remove} - adds a greeting for your nick that will be announced everytime you join the channel. Using .herald show will show your current herald and .herald delete will remove your greeting."""
     if text.lower() == "show":
-        query = select([table.c.quote]).where(table.c.name == nick.lower()).where(table.c.chan == chan.lower())
-        greeting = db.execute(query).fetchone()
-        if greeting:
-            return greeting[0]
-        else:
+        greeting = herald_cache[chan.casefold()].get(nick.casefold())
+        if greeting is None:
             return "you don't have a herald set try .herald <message> to set your greeting."
+
+        return greeting
     elif text.lower() in ["delete", "remove"]:
-        greeting = db.execute(
-            select([table.c.quote]).where(table.c.name == nick.lower()).where(table.c.chan == chan.lower())
-        ).fetchone()
+        greeting = herald_cache[chan.casefold()].get(nick.casefold())
+        if greeting is None:
+            return "no herald set, unable to delete."
+
         query = table.delete().where(table.c.name == nick.lower()).where(table.c.chan == chan.lower())
-        res = db.execute(query)
+        db.execute(query)
         db.commit()
-        if res.rowcount > 0:
-            return "greeting \'{}\' for {} has been removed".format(greeting[0], nick)
-        else:
-            return "no herald set, unable to delete."
+
+        reply("greeting \'{}\' for {} has been removed".format(greeting[0], nick))
+
+        load_cache(db)
     else:
         res = db.execute(
             table.update().where(table.c.name == nick.lower()).where(table.c.chan == chan.lower()).values(quote=text)
@@ -49,11 +59,13 @@ def herald(text, nick, chan, db):
             db.execute(table.insert().values(name=nick.lower(), chan=chan.lower(), quote=text))
 
         db.commit()
-        return "greeting successfully added"
+        reply("greeting successfully added")
+
+        load_cache(db)
 
 
 @hook.command(permissions=["botcontrol", "snoonetstaff"])
-def deleteherald(text, chan, db):
+def deleteherald(text, chan, db, reply):
     """<nickname> - Delete [nickname]'s herald."""
 
     nick = text.strip()
@@ -65,13 +77,15 @@ def deleteherald(text, chan, db):
     db.commit()
 
     if res.rowcount > 0:
-        return "greeting for {} has been removed".format(text.lower())
+        reply("greeting for {} has been removed".format(text.lower()))
     else:
-        return "{} does not have a herald".format(text.lower())
+        reply("{} does not have a herald".format(text.lower()))
+
+    load_cache(db)
 
 
 @hook.irc_raw("JOIN", singlethread=True)
-def welcome(nick, message, db, bot, chan):
+def welcome(nick, message, bot, chan):
     decoy = re.compile('[Òo○O0öøóȯôőŏᴏōο][<>＜]')
     colors_re = re.compile("\x02|\x03(?:\d{1,2}(?:,\d{1,2})?)?", re.UNICODE)
     bino_re = re.compile('b+i+n+o+', re.IGNORECASE)
@@ -85,9 +99,7 @@ def welcome(nick, message, db, bot, chan):
     else:
         floodcheck[chan] = time.time()
 
-    welcome = db.execute(
-        select([table.c.quote]).where(table.c.name == nick.lower()).where(table.c.chan == chan.lower())
-    ).fetchone()
+    welcome = herald_cache[chan.casefold()].get(nick.casefold())
     if welcome:
         greet = welcome[0]
         stripped = greet.translate(dict.fromkeys(["\u200b", " ", "\u202f", "\x02"]))