Abstract db announcement methods.

diff --git a/db.py b/db.py
index d38b972..07e8930 100644 (file)
--- a/db.py
+++ b/db.py
@@ -74,9 +74,9 @@ class DBConnection(object):
     self.dbh.commit()
     self.close_db()
 
-  def get_all_announcements(self, client, servers = []):
+  def get_all_from_table(self, client, table, servers = []):
     params = [client.user.id]
-    sql = 'select * from announcements where bot_id=?'
+    sql = 'select * from {} where bot_id=?'.format(table)
     if len(servers):
       sql += ' and server_id in ({})'.format(', '.join(['?'] * len(servers)))
       for server in servers:
@@ -89,15 +89,32 @@ class DBConnection(object):
       yield dict(row)
     self.close_db()
 
-  def get_announcement(self, client, id):
-    cursor = self.query('select * from announcements where bot_id=? and id=?', [client.user.id, id])
+  def get_from_table(self, client, table, id):
+    cursor = self.query('select * from {} where bot_id=? and id=?'.format(table), [client.user.id, id])
     row = cursor.fetchone()
     if row:
-      announcement = dict(row)
+      ret = dict(row)
     else:
-      announcement = None
+      ret = None
     self.close_db()
-    return announcement
+    return ret
+
+  def delete_from_table(self, client, table, id):
+    cursor = self.query('delete from {} where bot_id=? and id=?'.format(table), [client.user.id, id])
+    if cursor.rowcount:
+      self.dbh.commit()
+      ret = True
+    else:
+      log.error('No such entry {} in {} table for {}'.format(id, table, client.user.name))
+      ret = False
+    self.close_db()
+    return ret
+
+  def get_all_announcements(self, client, servers = []):
+    yield from self.get_all_from_table(client, 'announcements', servers)
+
+  def get_announcement(self, client, id):
+    return self.get_from_table(client, 'announcements', id)
 
   def get_announcements(self, client):
     now = self.now()
@@ -126,7 +143,7 @@ class DBConnection(object):
     self.close_db()
     return ret
 
-  def update_announcement(self, client, id, **args):
+  def update_table(self, client, table, id, **args):
     values = []
     params = []
     for k, v in args.items():
@@ -134,25 +151,20 @@ class DBConnection(object):
       params.append(v)
     params.append(client.user.id)
     params.append(id)
-    sql = 'update announcements set {} where bot_id=? and id=?'.format(', '.join(['{}=?'.format(v) for v in values]))
+    sql = 'update {} set {} where bot_id=? and id=?'.format(table, ', '.join(['{}=?'.format(v) for v in values]))
     cursor = self.query(sql, params)
     if cursor.rowcount:
       self.dbh.commit()
       ret = True
     else:
-      log.error('No such announcement {} for {}'.format(id, client.user.name))
+      log.error('No such entry {} in {} table for {}'.format(id, table, client.user.name))
       self.dbh.rollback()
       ret = False
     self.close_db()
     return ret
 
+  def update_announcement(self, client, id, **args):
+    return self.update_table(client, 'announcements', id, **args)
+
   def delete_announcement(self, client, id):
-    cursor = self.query('delete from announcements where bot_id=? and id=?', [client.user.id, id])
-    if cursor.rowcount:
-      self.dbh.commit()
-      ret = True
-    else:
-      log.error('No such announcement {} for {}'.format(id, client.user.name))
-      ret = False
-    self.close_db()
-    return ret
+    return self.delete_from_table(client, 'announcements', id)