test

sopel_modules/SpiceBot/Database.py

@@ -7,11 +7,11 @@ This is the SpiceBot Database
 # sopel imports
 from sopel.tools import Identifier
 import sopel.db
-from sopel.db import SopelDB, _deserialize, NickValues, ChannelValues
+from sopel.db import SopelDB, _deserialize, Nicknames
 
 from .Config import config as botconfig
 
-from sqlalchemy import Column, String
+from sqlalchemy import Column, String, ForeignKey, Integer
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.exc import SQLAlchemyError
 
@@ -21,12 +21,35 @@ import json
 BASE = declarative_base()
 
 
+class NickValues(BASE):
+    """
+    NickValues SQLAlchemy Class
+    """
+    __tablename__ = 'nick_values'
+    nick_id = Column(Integer, ForeignKey('nick_ids.nick_id'), primary_key=True)
+    category = Column(String(255), primary_key=True)
+    key = Column(String(255), primary_key=True)
+    value = Column(String(255))
+
+
+class ChannelValues(BASE):
+    """
+    ChannelValues SQLAlchemy Class
+    """
+    __tablename__ = 'channel_values'
+    channel = Column(String(255), primary_key=True)
+    category = Column(String(255), primary_key=True)
+    key = Column(String(255), primary_key=True)
+    value = Column(String(255))
+
+
 class PluginValues(BASE):
     """
     PluginValues SQLAlchemy Class
     """
     __tablename__ = 'plugin_values'
     plugin = Column(String(255), primary_key=True)
+    category = Column(String(255), primary_key=True)
     key = Column(String(255), primary_key=True)
     value = Column(String(255))
 
@@ -35,7 +58,54 @@ class SpiceDB(object):
 
     # NICK FUNCTIONS
 
-    def delete_nick_value(self, nick, key):
+    def set_nick_value(self, nick, key, value, category='default'):
+        """Sets the value for a given key to be associated with the nick."""
+        nick = Identifier(nick)
+        value = json.dumps(value, ensure_ascii=False)
+        nick_id = self.get_nick_id(nick)
+        session = self.ssession()
+        try:
+            result = session.query(NickValues) \
+                .filter(NickValues.nick_id == nick_id) \
+                .filter(NickValues.category == category)\
+                .filter(NickValues.key == key) \
+                .one_or_none()
+            # NickValue exists, update
+            if result:
+                result.value = value
+                session.commit()
+            # DNE - Insert
+            else:
+                new_nickvalue = NickValues(nick_id=nick_id, category=category, key=key, value=value)
+                session.add(new_nickvalue)
+                session.commit()
+        except SQLAlchemyError:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def get_nick_value(self, nick, key, category='default'):
+        """Retrieves the value for a given key associated with a nick."""
+        nick = Identifier(nick)
+        session = self.ssession()
+        try:
+            result = session.query(NickValues) \
+                .filter(Nicknames.nick_id == NickValues.nick_id) \
+                .filter(Nicknames.slug == nick.lower()) \
+                .filter(NickValues.category == category)\
+                .filter(NickValues.key == key) \
+                .one_or_none()
+            if result is not None:
+                result = result.value
+            return _deserialize(result)
+        except SQLAlchemyError:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def delete_nick_value(self, nick, key, category='default'):
         """Deletes the value for a given key to be associated with the nick."""
         nick = Identifier(nick)
         nick_id = self.get_nick_id(nick)
@@ -43,6 +113,7 @@ class SpiceDB(object):
         try:
             result = session.query(NickValues) \
                 .filter(NickValues.nick_id == nick_id) \
+                .filter(NickValues.category == category)\
                 .filter(NickValues.key == key) \
                 .one_or_none()
             # NickValue exists, delete
@@ -55,7 +126,7 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def adjust_nick_value(self, nick, key, value):
+    def adjust_nick_value(self, nick, key, value, category='default'):
         """Sets the value for a given key to be associated with the nick."""
         nick = Identifier(nick)
         value = json.dumps(value, ensure_ascii=False)
@@ -64,6 +135,7 @@ class SpiceDB(object):
         try:
             result = session.query(NickValues) \
                 .filter(NickValues.nick_id == nick_id) \
+                .filter(NickValues.category == category)\
                 .filter(NickValues.key == key) \
                 .one_or_none()
             # NickValue exists, update
@@ -72,7 +144,7 @@ class SpiceDB(object):
                 session.commit()
             # DNE - Insert
             else:
-                new_nickvalue = NickValues(nick_id=nick_id, key=key, value=float(value))
+                new_nickvalue = NickValues(nick_id=nick_id, category=category, key=key, value=float(value))
                 session.add(new_nickvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -81,7 +153,7 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def adjust_nick_list(self, nick, key, entries, adjustmentdirection):
+    def adjust_nick_list(self, nick, key, entries, adjustmentdirection, category='default'):
         """Sets the value for a given key to be associated with the nick."""
         nick = Identifier(nick)
         if not isinstance(entries, list):
@@ -92,6 +164,7 @@ class SpiceDB(object):
         try:
             result = session.query(NickValues) \
                 .filter(NickValues.nick_id == nick_id) \
+                .filter(NickValues.category == category)\
                 .filter(NickValues.key == key) \
                 .one_or_none()
             # NickValue exists, update
@@ -116,7 +189,7 @@ class SpiceDB(object):
                     for entry in entries:
                         while entry in values:
                             values.remove(entry)
-                new_nickvalue = NickValues(nick_id=nick_id, key=key, value=values)
+                new_nickvalue = NickValues(nick_id=nick_id, category=category, key=key, value=values)
                 session.add(new_nickvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -127,13 +200,59 @@ class SpiceDB(object):
 
     # CHANNEL FUNCTIONS
 
-    def delete_channel_value(self, channel, key):
+    def set_channel_value(self, channel, key, value, category='default'):
+        """Sets the value for a given key to be associated with the channel."""
+        channel = Identifier(channel).lower()
+        value = json.dumps(value, ensure_ascii=False)
+        session = self.ssession()
+        try:
+            result = session.query(ChannelValues) \
+                .filter(ChannelValues.channel == channel)\
+                .filter(ChannelValues.category == category)\
+                .filter(ChannelValues.key == key) \
+                .one_or_none()
+            # ChannelValue exists, update
+            if result:
+                result.value = value
+                session.commit()
+            # DNE - Insert
+            else:
+                new_channelvalue = ChannelValues(channel=channel, category=category, key=key, value=value)
+                session.add(new_channelvalue)
+                session.commit()
+        except SQLAlchemyError:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def get_channel_value(self, channel, key, category='default'):
+        """Retrieves the value for a given key associated with a channel."""
+        channel = Identifier(channel).lower()
+        session = self.ssession()
+        try:
+            result = session.query(ChannelValues) \
+                .filter(ChannelValues.channel == channel)\
+                .filter(ChannelValues.category == category)\
+                .filter(ChannelValues.key == key) \
+                .one_or_none()
+            if result is not None:
+                result = result.value
+            return _deserialize(result)
+        except SQLAlchemyError:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def delete_channel_value(self, channel, key, category='default'):
         """Sets the value for a given key to be associated with the channel."""
         channel = Identifier(channel).lower()
         session = self.ssession()
         try:
             result = session.query(ChannelValues) \
                 .filter(ChannelValues.channel == channel)\
+                .filter(ChannelValues.category == category)\
                 .filter(ChannelValues.key == key) \
                 .one_or_none()
             # ChannelValue exists, delete
@@ -146,7 +265,7 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def adjust_channel_value(self, channel, key, value):
+    def adjust_channel_value(self, channel, key, value, category='default'):
         """Sets the value for a given key to be associated with the channel."""
         channel = Identifier(channel).lower()
         value = json.dumps(value, ensure_ascii=False)
@@ -154,6 +273,7 @@ class SpiceDB(object):
         try:
             result = session.query(ChannelValues) \
                 .filter(ChannelValues.channel == channel)\
+                .filter(ChannelValues.category == category)\
                 .filter(ChannelValues.key == key) \
                 .one_or_none()
             # ChannelValue exists, update
@@ -162,7 +282,7 @@ class SpiceDB(object):
                 session.commit()
             # DNE - Insert
             else:
-                new_channelvalue = ChannelValues(channel=channel, key=key, value=float(value))
+                new_channelvalue = ChannelValues(channel=channel, category=category, key=key, value=float(value))
                 session.add(new_channelvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -171,7 +291,7 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def adjust_channel_list(self, channel, key, entries, adjustmentdirection):
+    def adjust_channel_list(self, channel, key, entries, adjustmentdirection, category='default'):
         """Sets the value for a given key to be associated with the channel."""
         channel = Identifier(channel).lower()
         if not isinstance(entries, list):
@@ -181,6 +301,7 @@ class SpiceDB(object):
         try:
             result = session.query(ChannelValues) \
                 .filter(ChannelValues.channel == channel)\
+                .filter(ChannelValues.category == category)\
                 .filter(ChannelValues.key == key) \
                 .one_or_none()
             # ChannelValue exists, update
@@ -205,7 +326,7 @@ class SpiceDB(object):
                     for entry in entries:
                         while entry in values:
                             values.remove(entry)
-                new_channelvalue = ChannelValues(channel=channel, key=key, value=values)
+                new_channelvalue = ChannelValues(channel=channel, category=category, key=key, value=values)
                 session.add(new_channelvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -216,7 +337,7 @@ class SpiceDB(object):
 
     # PLUGIN FUNCTIONS
 
-    def set_plugin_value(self, plugin, key, value):
+    def set_plugin_value(self, plugin, key, value, category='default'):
         """Sets the value for a given key to be associated with the plugin."""
         plugin = Identifier(plugin).lower()
         value = json.dumps(value, ensure_ascii=False)
@@ -224,6 +345,7 @@ class SpiceDB(object):
         try:
             result = session.query(PluginValues) \
                 .filter(PluginValues.plugin == plugin)\
+                .filter(ChannelValues.category == category)\
                 .filter(PluginValues.key == key) \
                 .one_or_none()
             # PluginValues exists, update
@@ -232,7 +354,7 @@ class SpiceDB(object):
                 session.commit()
             # DNE - Insert
             else:
-                new_pluginvalue = PluginValues(plugin=plugin, key=key, value=value)
+                new_pluginvalue = PluginValues(plugin=plugin, category=category, key=key, value=value)
                 session.add(new_pluginvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -241,13 +363,14 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def get_plugin_value(self, plugin, key):
+    def get_plugin_value(self, plugin, key, category='default'):
         """Retrieves the value for a given key associated with a plugin."""
         plugin = Identifier(plugin).lower()
         session = self.ssession()
         try:
             result = session.query(PluginValues) \
                 .filter(PluginValues.plugin == plugin)\
+                .filter(ChannelValues.category == category)\
                 .filter(PluginValues.key == key) \
                 .one_or_none()
             if result is not None:
@@ -259,13 +382,14 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def delete_plugin_value(self, plugin, key):
+    def delete_plugin_value(self, plugin, key, category='default'):
         """Deletes the value for a given key to be associated with the plugin."""
         plugin = Identifier(plugin).lower()
         session = self.ssession()
         try:
             result = session.query(PluginValues) \
                 .filter(PluginValues.plugin == plugin)\
+                .filter(ChannelValues.category == category)\
                 .filter(PluginValues.key == key) \
                 .one_or_none()
             # PluginValues exists, delete
@@ -278,7 +402,7 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def adjust_plugin_value(self, plugin, key, value):
+    def adjust_plugin_value(self, plugin, key, value, category='default'):
         """Sets the value for a given key to be associated with the plugin."""
         plugin = Identifier(plugin).lower()
         value = json.dumps(value, ensure_ascii=False)
@@ -286,6 +410,7 @@ class SpiceDB(object):
         try:
             result = session.query(PluginValues) \
                 .filter(PluginValues.plugin == plugin)\
+                .filter(ChannelValues.category == category)\
                 .filter(PluginValues.key == key) \
                 .one_or_none()
             # ChannelValue exists, update
@@ -294,7 +419,7 @@ class SpiceDB(object):
                 session.commit()
             # DNE - Insert
             else:
-                new_pluginvalue = PluginValues(plugin=plugin, key=key, value=float(value))
+                new_pluginvalue = PluginValues(plugin=plugin, category=category, key=key, value=float(value))
                 session.add(new_pluginvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -303,7 +428,7 @@ class SpiceDB(object):
         finally:
             session.close()
 
-    def adjust_plugin_list(self, plugin, key, entries, adjustmentdirection):
+    def adjust_plugin_list(self, plugin, key, entries, adjustmentdirection, category='default'):
         """Sets the value for a given key to be associated with the plugin."""
         plugin = Identifier(plugin).lower()
         if not isinstance(entries, list):
@@ -313,6 +438,7 @@ class SpiceDB(object):
         try:
             result = session.query(PluginValues) \
                 .filter(PluginValues.plugin == plugin)\
+                .filter(ChannelValues.category == category)\
                 .filter(PluginValues.key == key) \
                 .one_or_none()
             # ChannelValue exists, update
@@ -337,7 +463,7 @@ class SpiceDB(object):
                     for entry in entries:
                         while entry in values:
                             values.remove(entry)
-                new_pluginvalue = PluginValues(plugin=plugin, key=key, value=values)
+                new_pluginvalue = PluginValues(plugin=plugin, category=category, key=key, value=values)
                 session.add(new_pluginvalue)
                 session.commit()
         except SQLAlchemyError:
@@ -352,10 +478,16 @@ class BotDatabase():
 
     def __init__(self):
 
+        sopel.db.NickValues = NickValues
+        SopelDB.get_nick_value = SpiceDB.get_nick_value
+        SopelDB.set_nick_value = SpiceDB.set_nick_value
         SopelDB.delete_nick_value = SpiceDB.delete_nick_value
         SopelDB.adjust_nick_value = SpiceDB.adjust_nick_value
         SopelDB.adjust_nick_list = SpiceDB.adjust_nick_list
 
+        sopel.db.ChannelValues = ChannelValues
+        SopelDB.get_channel_value = SpiceDB.get_channel_value
+        SopelDB.set_channel_value = SpiceDB.set_channel_value
         SopelDB.delete_channel_value = SpiceDB.delete_channel_value
         SopelDB.adjust_channel_value = SpiceDB.adjust_channel_value
         SopelDB.adjust_channel_list = SpiceDB.adjust_channel_list