Add cache to autojoin.py

Author: linuxdaemon

plugins/autojoin.py

@@ -1,7 +1,8 @@
 import asyncio
+from collections import defaultdict
+from threading import RLock
 
 from sqlalchemy import PrimaryKeyConstraint, Column, String, Table, and_
-from sqlalchemy.exc import IntegrityError
 
 from cloudbot import hook
 from cloudbot.util import database
@@ -14,40 +15,54 @@
     PrimaryKeyConstraint('conn', 'chan')
 )
 
+chan_cache = defaultdict(set)
+db_lock = RLock()
+
 
 def get_channels(db, conn):
     return db.execute(table.select().where(table.c.conn == conn.name.casefold())).fetchall()
 
 
+@hook.on_start
+def load_cache(db):
+    with db_lock:
+        chan_cache.clear()
+        for row in db.execute(table.select()):
+            chan_cache[row['conn']].add(row['chan'])
+
+
 @asyncio.coroutine
 @hook.irc_raw('376')
-def do_joins(db, conn, async_call):
-    chans = yield from async_call(get_channels, db, conn)
+def do_joins(conn):
     join_throttle = conn.config.get("join_throttle", 0.4)
-    for chan in chans:
-        conn.join(chan[1])
+    for chan in chan_cache[conn.name]:
+        conn.join(chan)
         yield from asyncio.sleep(join_throttle)
 
 
 @hook.irc_raw('JOIN', singlethread=True)
 def add_chan(db, conn, chan, nick):
-    if nick.casefold() == conn.nick.casefold():
-        try:
+    chans = chan_cache[conn.name]
+    chan = chan.casefold()
+    if nick.casefold() == conn.nick.casefold() and chan not in chans:
+        with db_lock:
             db.execute(table.insert().values(conn=conn.name.casefold(), chan=chan.casefold()))
             db.commit()
-        except IntegrityError:
-            pass
+
+        load_cache(db)
 
 
 @hook.irc_raw('PART', singlethread=True)
 def on_part(db, conn, chan, nick):
     if nick.casefold() == conn.nick.casefold():
-        db.execute(table.delete().where(and_(table.c.conn == conn.name.casefold(), table.c.chan == chan.casefold())))
-        db.commit()
+        with db_lock:
+            db.execute(
+                table.delete().where(and_(table.c.conn == conn.name.casefold(), table.c.chan == chan.casefold())))
+            db.commit()
+
+        load_cache(db)
 
 
 @hook.irc_raw('KICK', singlethread=True)
 def on_kick(db, conn, chan, target):
-    if target.casefold() == conn.nick.casefold():
-        db.execute(table.delete().where(and_(table.c.conn == conn.name.casefold(), table.c.chan == chan.casefold())))
-        db.commit()
+    on_part(db, conn, chan, target)