Move postgres models to new file (#631)

* Move DBs to new file

.pylintrc

@@ -73,6 +73,7 @@ disable=C0103,
         R0912,
         R0913,
         R0914,
+        R0915,
         R0916,
         W0201,
         W0406,

Documentation/Extension-howto.md

@@ -10,24 +10,6 @@ async def setup(bot):
 ```
 This code is run when loading the extension, is used to add model classes and config entries.
 
-
-## Models
-
-This defines any database models used in this extension.
-
-```py
-class Model_name(bot.db.Model):
-    __tablename__ = "<table-name>"
-
-    Int_Entry = bot.db.Column(bot.db.Integer, primary_key=True)
-    Str_Entry = bot.db.Column(bot.db.String, default=None)
-    guild = bot.db.Column(bot.db.String)
-```
-This creates a table called `<table-name>` and adds the columns `Int_entry`, `Str_entry`and `guild` to it. 
-Make sure to include `guild` along with handling for it so the entry is only callable in the server a command was invoked from.
-The `default` argument is optional but preferred, since an exception will be thrown if there wasn't a value assigned but it is attempted to be accessed.
-
-
 ## Config entries
 
 ```py
@@ -62,10 +44,30 @@ NOTE: The entry might not automatically get added to the file and will have to b
 await bot.add_cog(extension-name(bot=bot))
 bot.add_extension_config("extension-name", config)
 ```
-This registers the extension, assumes the extension name is the filename if the `extension_name` argument wasn't supplied. If any custom models are defined, make sure to add the `models` argument with its value being `[Model_name_1, Model_name_2, ...]`
+This registers the extension, assumes the extension name is the filename if the `extension_name` argument wasn't supplied.
 
 The second line adds the extension to the config .json file.
 
+## Optional: Postgres
+
+This defines any database models used in this extension.
+
+```py
+class Model_name(bot.db.Model):
+    __tablename__ = "<table-name>"
+
+    Int_Entry = bot.db.Column(bot.db.Integer, primary_key=True)
+    Str_Entry = bot.db.Column(bot.db.String, default=None)
+    guild = bot.db.Column(bot.db.String)
+```
+This creates a table called `<table-name>` and adds the columns `Int_entry`, `Str_entry`and `guild` to it. 
+Make sure to include `guild` along with handling for it so the entry is only callable in the server a command was invoked from.
+The `default` argument is optional but preferred, since an exception will be thrown if there wasn't a value assigned but it is attempted to be accessed.
+
+Add this to the base/databases.py file, and register it in the bot.models variable at the bottom of the file.
+You can access it by "self.bot.models.Model_name"
+
+
 
 ## Optional: Command checks
 

techsupport_bot/base/__init__.py

@@ -3,4 +3,5 @@
 from .auxiliary import *
 from .cogs import *
 from .data import *
+from .databases import *
 from .extension import *

techsupport_bot/base/databases.py

@@ -0,0 +1,112 @@
+"""This file stores all of the postgres table declarations
+All models can be used by any extension
+"""
+
+import datetime
+
+
+def setup_models(bot):
+    """A function to setup all of the postgres tables
+    This is stored in bot.models variable
+
+    Args:
+        bot (TechSupportBot): The bot object to register the databases to
+    """
+
+    class DuckUser(bot.db.Model):
+        """The postgres table for ducks
+        Currently used in duck.py"""
+
+        __tablename__ = "duckusers"
+
+        pk = bot.db.Column(bot.db.Integer, primary_key=True, autoincrement=True)
+        author_id = bot.db.Column(bot.db.String)
+        guild_id = bot.db.Column(bot.db.String)
+        befriend_count = bot.db.Column(bot.db.Integer, default=0)
+        kill_count = bot.db.Column(bot.db.Integer, default=0)
+        updated = bot.db.Column(bot.db.DateTime, default=datetime.datetime.utcnow)
+        speed_record = bot.db.Column(bot.db.Float, default=80.0)
+
+    class Factoid(bot.db.Model):
+        """The postgres table for factoids
+        Currently used in factoid.py"""
+
+        __tablename__ = "factoids"
+
+        factoid_id = bot.db.Column(bot.db.Integer, primary_key=True)
+        name = bot.db.Column(bot.db.String)
+        guild = bot.db.Column(bot.db.String)
+        message = bot.db.Column(bot.db.String)
+        time = bot.db.Column(bot.db.DateTime, default=datetime.datetime.utcnow)
+        embed_config = bot.db.Column(bot.db.String, default=None)
+        hidden = bot.db.Column(bot.db.Boolean, default=False)
+        alias = bot.db.Column(bot.db.String, default=None)
+
+    class FactoidJob(bot.db.Model):
+        """The postgres table for factoid loops
+        Currently used in factoid.py"""
+
+        __tablename__ = "factoid_jobs"
+
+        job_id = bot.db.Column(bot.db.Integer, primary_key=True)
+        factoid = bot.db.Column(
+            bot.db.Integer, bot.db.ForeignKey("factoids.factoid_id")
+        )
+        channel = bot.db.Column(bot.db.String)
+        cron = bot.db.Column(bot.db.String)
+
+    class Grab(bot.db.Model):
+        """The postgres table for grabs
+        Currently used in grab.py"""
+
+        __tablename__ = "grabs"
+
+        pk = bot.db.Column(bot.db.Integer, primary_key=True)
+        author_id = bot.db.Column(bot.db.String)
+        channel = bot.db.Column(bot.db.String)
+        guild = bot.db.Column(bot.db.String)
+        message = bot.db.Column(bot.db.String)
+        time = bot.db.Column(bot.db.DateTime, default=datetime.datetime.utcnow)
+        nsfw = bot.db.Column(bot.db.Boolean, default=False)
+
+    class IRCChannelMapping(bot.db.Model):
+        """The postgres table for IRC->discord maps
+        Currently used in relay.py"""
+
+        __tablename__ = "ircchannelmap"
+        map_id = bot.db.Column(bot.db.Integer, primary_key=True)
+        guild_id = bot.db.Column(bot.db.String, default=None)
+        discord_channel_id = bot.db.Column(bot.db.String, default=None)
+        irc_channel_id = bot.db.Column(bot.db.String, default=None)
+
+    class UserNote(bot.db.Model):
+        """The postgres table for notes
+        Currently used in who.py"""
+
+        __tablename__ = "usernote"
+
+        pk = bot.db.Column(bot.db.Integer, primary_key=True, autoincrement=True)
+        user_id = bot.db.Column(bot.db.String)
+        guild_id = bot.db.Column(bot.db.String)
+        updated = bot.db.Column(bot.db.DateTime, default=datetime.datetime.utcnow)
+        author_id = bot.db.Column(bot.db.String)
+        body = bot.db.Column(bot.db.String)
+
+    class Warning(bot.db.Model):
+        """The postgres table for warnings
+        Currently used in protect.py and who.py"""
+
+        __tablename__ = "warnings"
+        pk = bot.db.Column(bot.db.Integer, primary_key=True)
+        user_id = bot.db.Column(bot.db.String)
+        guild_id = bot.db.Column(bot.db.String)
+        reason = bot.db.Column(bot.db.String)
+        time = bot.db.Column(bot.db.DateTime, default=datetime.datetime.utcnow)
+
+    bot.models.DuckUser = DuckUser
+    bot.models.Factoid = Factoid
+    bot.models.FactoidJob = FactoidJob
+    bot.models.Grab = Grab
+    bot.models.IRCChannelMapping = IRCChannelMapping
+    bot.models.UserNote = UserNote
+    bot.models.Warning = Warning

techsupport_bot/bot.py

@@ -7,6 +7,7 @@
 import botlogging
 import cogs as builtin_cogs
 import ircrelay
+import munch
 
 
 class TechSupportBot(base.AdvancedBot):
@@ -61,6 +62,8 @@ async def setup_hook(self):
 
         if self.db:
             await self.logger.debug("Syncing Postgres tables...")
+            self.models = munch.DefaultMunch(None)
+            base.databases.setup_models(self)
             await self.db.gino.create_all()
 
         await self.logger.debug("Loading Help commands...")

techsupport_bot/extensions/duck.py

@@ -16,19 +16,6 @@
 async def setup(bot):
     """Method to add duck into the config file"""
 
-    class DuckUser(bot.db.Model):
-        """The descripiton duck table in postgres to be used with gino"""
-
-        __tablename__ = "duckusers"
-
-        pk = bot.db.Column(bot.db.Integer, primary_key=True, autoincrement=True)
-        author_id = bot.db.Column(bot.db.String)
-        guild_id = bot.db.Column(bot.db.String)
-        befriend_count = bot.db.Column(bot.db.Integer, default=0)
-        kill_count = bot.db.Column(bot.db.Integer, default=0)
-        updated = bot.db.Column(bot.db.DateTime, default=datetime.datetime.utcnow)
-        speed_record = bot.db.Column(bot.db.Float, default=80.0)
-
     config = bot.ExtensionConfig()
     config.add(
         key="hunt_channels",
@@ -73,7 +60,7 @@ class DuckUser(bot.db.Model):
         default=50,
     )
 
-    await bot.add_cog(DuckHunt(bot=bot, models=[DuckUser], extension_name="duck"))
+    await bot.add_cog(DuckHunt(bot=bot, extension_name="duck"))
     bot.add_extension_config("duck", config)
 
 
@@ -194,7 +181,7 @@ async def handle_winner(self, winner, guild, action, raw_duration, channel):
 
         duck_user = await self.get_duck_user(winner.id, guild.id)
         if not duck_user:
-            duck_user = self.models.DuckUser(
+            duck_user = self.bot.models.DuckUser(
                 author_id=str(winner.id),
                 guild_id=str(guild.id),
                 befriend_count=0,
@@ -305,10 +292,10 @@ def message_check(self, config, channel, message):
     async def get_duck_user(self, user_id, guild_id):
         """Method to get the duck winner"""
         duck_user = (
-            await self.models.DuckUser.query.where(
-                self.models.DuckUser.author_id == str(user_id)
+            await self.bot.models.DuckUser.query.where(
+                self.bot.models.DuckUser.author_id == str(user_id)
             )
-            .where(self.models.DuckUser.guild_id == str(guild_id))
+            .where(self.bot.models.DuckUser.guild_id == str(guild_id))
             .gino.first()
         )
 
@@ -321,8 +308,8 @@ async def get_global_record(self, guild_id):
         Parametrs:
         guild_id -> The ID of the guild in question
         """
-        query = await self.models.DuckUser.query.where(
-            self.models.DuckUser.guild_id == str(guild_id)
+        query = await self.bot.models.DuckUser.query.where(
+            self.bot.models.DuckUser.guild_id == str(guild_id)
         ).gino.all()
 
         speed_records = [record.speed_record for record in query]
@@ -390,11 +377,11 @@ async def stats(self, ctx, *, user: discord.Member = None):
     async def friends(self, ctx):
         """Method for viewing top friend counts"""
         duck_users = (
-            await self.models.DuckUser.query.order_by(
-                -self.models.DuckUser.befriend_count
+            await self.bot.models.DuckUser.query.order_by(
+                -self.bot.models.DuckUser.befriend_count
             )
-            .where(self.models.DuckUser.befriend_count > 0)
-            .where(self.models.DuckUser.guild_id == str(ctx.guild.id))
+            .where(self.bot.models.DuckUser.befriend_count > 0)
+            .where(self.bot.models.DuckUser.guild_id == str(ctx.guild.id))
             .gino.all()
         )
 
@@ -456,10 +443,10 @@ async def record(self, ctx):
             )
             return
         record_user = (
-            await self.models.DuckUser.query.where(
-                self.models.DuckUser.speed_record == record_time
+            await self.bot.models.DuckUser.query.where(
+                self.bot.models.DuckUser.speed_record == record_time
             )
-            .where(self.models.DuckUser.guild_id == str(ctx.guild.id))
+            .where(self.bot.models.DuckUser.guild_id == str(ctx.guild.id))
             .gino.first()
         )
 
@@ -480,9 +467,11 @@ async def record(self, ctx):
     async def killers(self, ctx):
         """Method for viewing top killer counts"""
         duck_users = (
-            await self.models.DuckUser.query.order_by(-self.models.DuckUser.kill_count)
-            .where(self.models.DuckUser.kill_count > 0)
-            .where(self.models.DuckUser.guild_id == str(ctx.guild.id))
+            await self.bot.models.DuckUser.query.order_by(
+                -self.bot.models.DuckUser.kill_count
+            )
+            .where(self.bot.models.DuckUser.kill_count > 0)
+            .where(self.bot.models.DuckUser.guild_id == str(ctx.guild.id))
             .gino.all()
         )