Initial switch to gino for discordbot db

discordbot/titanembeds/bot.py

@@ -68,7 +68,7 @@ class Titan(discord.AutoShardedClient):
         await self.change_presence(status=discord.Status.online, activity=game)
 
         try:
-            self.database.connect(config["database-uri"])
+            await self.database.connect(config["database-uri"])
         except Exception:
             self.logger.error("Unable to connect to specified database!")
             traceback.print_exc()

discordbot/titanembeds/database/__init__.py

@@ -1,13 +1,8 @@
-from contextlib import contextmanager
+from gino import Gino
-import sqlalchemy as db
-from sqlalchemy.engine import Engine, create_engine
-from sqlalchemy.orm import sessionmaker, Session
-from sqlalchemy.ext.declarative import declarative_base
-
 import json
 import discord
 
-Base = declarative_base()
+db = Gino()
 
 from titanembeds.database.guilds import Guilds
 from titanembeds.database.messages import Messages
@@ -18,87 +13,47 @@ from titanembeds.database.unauthenticated_bans import UnauthenticatedBans
 from titanembeds.utils import get_message_author, get_message_mentions, get_webhooks_list, get_emojis_list, get_roles_list, get_channels_list, list_role_ids, get_attachments_list, get_embeds_list
 
 class DatabaseInterface(object):
-    # Courtesy of https://github.com/SunDwarf/Jokusoramame
     def __init__(self, bot):
         self.bot = bot
 
-        self.engine = None  # type: Engine
+    async def connect(self, dburi):
-        self._sessionmaker = None  # type: sessionmaker
+        await db.set_bind(dburi)
-
-    def connect(self, dburi):
-        self.engine = create_engine(dburi, pool_recycle=10)
-
-    @contextmanager
-    def get_session(self):
-        Session = sessionmaker(bind=self.engine)
-        session = Session()
-        try:
-            yield session
-        except:
-            session.rollback()
-        finally:
-            session.commit()
 
     async def push_message(self, message):
-        self.bot.loop.run_in_executor(None, self._push_message, message)
-    
-    def _push_message(self, message):
         if message.guild:
-            with self.get_session() as session:
             edit_ts = message.edited_at
             if not edit_ts:
                 edit_ts = None
             else:
                 edit_ts = str(edit_ts)
-
+            await Messages.create(
-                msg = Messages(
+                message_id = int(message.id),
-                    int(message.guild.id),
+                guild_id = int(message.guild.id),
-                    int(message.channel.id),
+                channel_id = int(message.channel.id),
-                    int(message.id),
+                content = message.content,
-                    message.content,
+                author = json.dumps(get_message_author(message)),
-                    json.dumps(get_message_author(message)),
+                timestamp = str(message.created_at),
-                    str(message.created_at),
+                edited_timestamp = edit_ts,
-                    edit_ts,
+                mentions = json.dumps(get_message_mentions(message.mentions)),
-                    json.dumps(get_message_mentions(message.mentions)),
+                attachments = json.dumps(get_attachments_list(message.attachments)),
-                    json.dumps(get_attachments_list(message.attachments)),
+                embeds = json.dumps(get_embeds_list(message.embeds))
-                    json.dumps(get_embeds_list(message.embeds))
             )
-                session.add(msg)
-                session.commit()
                 
     async def update_message(self, message):
-        self.bot.loop.run_in_executor(None, self._update_message, message)
-
-    def _update_message(self, message):
         if message.guild:
-            with self.get_session() as session:
+            await Messages.get(int(message.id)).update(
-                msg = session.query(Messages) \
+                content = message.content,
-                    .filter(Messages.guild_id == message.guild.id) \
+                timestamp = message.created_at,
-                    .filter(Messages.channel_id == message.channel.id) \
+                edited_timestamp = message.edited_at,
-                    .filter(Messages.message_id == message.id).first()
+                mentions = json.dumps(get_message_mentions(message.mentions)),
-                if msg:
+                attachments = json.dumps(get_attachments_list(message.attachments)),
-                    msg.content = message.content
+                embeds = json.dumps(get_embeds_list(message.embeds)),
-                    msg.timestamp = message.created_at
+                author = json.dumps(get_message_author(message))
-                    msg.edited_timestamp = message.edited_at
+            ).apply()
-                    msg.mentions = json.dumps(get_message_mentions(message.mentions))
-                    msg.attachments = json.dumps(get_attachments_list(message.attachments))
-                    msg.embeds = json.dumps(get_embeds_list(message.embeds))
-                    msg.author = json.dumps(get_message_author(message))
-                    session.commit()
 
     async def delete_message(self, message):
-        self.bot.loop.run_in_executor(None, self._delete_message, message)
-    
-    def _delete_message(self, message):
         if message.guild:
-            with self.get_session() as session:
+            await Messages.get(int(message.id)).delete()
-                msg = session.query(Messages) \
-                    .filter(Messages.guild_id == int(message.guild.id)) \
-                    .filter(Messages.channel_id == int(message.channel.id)) \
-                    .filter(Messages.message_id == int(message.id)).first()
-                if msg:
-                    session.delete(msg)
-                    session.commit()
     
     async def update_guild(self, guild):
         if guild.me.guild_permissions.manage_webhooks:
@@ -108,163 +63,124 @@ class DatabaseInterface(object):
                 server_webhooks = []
         else:
             server_webhooks = []
-        self.bot.loop.run_in_executor(None, self._update_guild, guild, server_webhooks)
+        gui = await Guilds.get(guild.id)
-
-    def _update_guild(self, guild, server_webhooks):
-        with self.get_session() as session:
-            gui = session.query(Guilds).filter(Guilds.guild_id == guild.id).first()
         if not gui:
-                gui = Guilds(
+            await Guilds.create(
-                    int(guild.id),
+                guild_id = int(guild.id),
-                    guild.name,
+                name = guild.name,
-                    json.dumps(get_roles_list(guild.roles)),
+                unauth_users = True,
-                    json.dumps(get_channels_list(guild.channels)),
+                visitor_view = False,
-                    json.dumps(get_webhooks_list(server_webhooks)),
+                webhook_messages = False,
-                    json.dumps(get_emojis_list(guild.emojis)),
+                guest_icon = None,
-                    int(guild.owner_id),
+                chat_links = True,
-                    guild.icon
+                bracket_links = True,
+                unauth_captcha = True,
+                mentions_limit = -1,
+                roles = json.dumps(get_roles_list(guild.roles)),
+                channels = json.dumps(get_channels_list(guild.channels)),
+                webhooks = json.dumps(get_webhooks_list(server_webhooks)),
+                emojis = json.dumps(get_emojis_list(guild.emojis)),
+                owner_id = int(guild.owner_id),
+                icon = guild.icon
             )
-                session.add(gui)
         else:
-                gui.name = guild.name
+            await gui.update(
-                gui.roles = json.dumps(get_roles_list(guild.roles))
+                name = guild.name,
-                gui.channels = json.dumps(get_channels_list(guild.channels))
+                roles = json.dumps(get_roles_list(guild.roles)),
-                gui.webhooks = json.dumps(get_webhooks_list(server_webhooks))
+                channels = json.dumps(get_channels_list(guild.channels)),
-                gui.emojis = json.dumps(get_emojis_list(guild.emojis))
+                webhooks = json.dumps(get_webhooks_list(server_webhooks)),
-                gui.owner_id = int(guild.owner_id)
+                emojis = json.dumps(get_emojis_list(guild.emojis)),
-                gui.icon = guild.icon
+                owner_id = int(guild.owner_id),
-            session.commit()
+                icon = guild.icon
+            ).apply()
 
     async def remove_unused_guilds(self, guilds):
-        self.bot.loop.run_in_executor(None, self._remove_unused_guilds, guilds)
+        dbguilds = await Guilds.query.gino.all()
-
-    def _remove_unused_guilds(self, guilds):
-        with self.get_session() as session:
-            dbguilds = session.query(Guilds).all()
-            changed = False
         for guild in dbguilds:
             disguild = discord.utils.get(guilds, id=guild.guild_id)
             if not disguild:
-                    changed = True
+                await Messages.delete.where(Messages.guild_id == int(guild.guild_id)).gino.status()
-                    dbmsgs = session.query(Messages).filter(Messages.guild_id == int(guild.guild_id)).all()
-                    for msg in dbmsgs:
-                        session.delete(msg)
-                    session.delete(guild)
-            if changed:
-                session.commit()
 
     async def remove_guild(self, guild):
-        self.bot.loop.run_in_executor(None, self._remove_guild, guild)
+        gui = await Guilds.get(int(guild.id))
-
-    def _remove_guild(self, guild):
-        with self.get_session() as session:
-            gui = session.query(Guilds).filter(Guilds.guild_id == int(guild.id)).first()
         if gui:
-                dbmsgs = session.query(Messages).filter(Messages.guild_id == int(guild.id)).delete()
+            await Messages.delete.where(Messages.guild_id == int(guild.id)).gino.status()
-                session.delete(gui)
+            await gui.delete()
-                session.commit()
     
     async def update_guild_member(self, member, active=True, banned=False, guild=None):
-        self.bot.loop.run_in_executor(None, self._update_guild_member, member, active, banned, guild)
-
-    def _update_guild_member(self, member, active=True, banned=False, guild=None):
-        with self.get_session() as session:
         if guild:
-                dbmember = session.query(GuildMembers) \
+            dbmember = await GuildMembers.query \
-                    .filter(GuildMembers.guild_id == int(guild.id)) \
+                .where(GuildMembers.guild_id == int(guild.id)) \
-                    .filter(GuildMembers.user_id == int(member.id)) \
+                .where(GuildMembers.user_id == int(member.id)) \
-                    .order_by(GuildMembers.id).all()
+                .order_by(GuildMembers.id).gino.all()
         else:
-                dbmember = session.query(GuildMembers) \
+            dbmember = await GuildMembers.query \
-                    .filter(GuildMembers.guild_id == int(member.guild.id)) \
+                .where(GuildMembers.guild_id == int(member.guild.id)) \
-                    .filter(GuildMembers.user_id == int(member.id)) \
+                .where(GuildMembers.user_id == int(member.id)) \
-                    .order_by(GuildMembers.id).all()
+                .order_by(GuildMembers.id).gino.all()
         if not dbmember:
-                dbmember = GuildMembers(
+            await GuildMembers.create(
-                    int(member.guild.id),
+                guild_id = int(member.guild.id),
-                    int(member.id),
+                user_id = int(member.id),
-                    member.name,
+                username = member.name,
-                    member.discriminator,
+                discriminator = member.discriminator,
-                    member.nick,
+                nickname = member.nick,
-                    member.avatar,
+                avatar = member.avatar,
-                    active,
+                active = active,
-                    banned,
+                banned = banned,
-                    json.dumps(list_role_ids(member.roles))
+                roles = json.dumps(list_role_ids(member.roles))
             )
-                session.add(dbmember)
         else:
             if len(dbmember) > 1:
                 for mem in dbmember[1:]:
-                        session.delete(mem)
+                    await mem.delete()
                 dbmember = dbmember[0]
             if dbmember.banned != banned or dbmember.active != active or dbmember.username != member.name or dbmember.discriminator != int(member.discriminator) or dbmember.nickname != member.nick or dbmember.avatar != member.avatar or set(json.loads(dbmember.roles)) != set(list_role_ids(member.roles)):
-                    dbmember.banned = banned
+                await dbmember.update(
-                    dbmember.active = active
+                    banned = banned,
-                    dbmember.username = member.name
+                    active = active,
-                    dbmember.discriminator = member.discriminator
+                    username = member.name,
-                    dbmember.nickname = member.nick
+                    discriminator = member.discriminator,
-                    dbmember.avatar = member.avatar
+                    nickname = member.nick,
-                    dbmember.roles = json.dumps(list_role_ids(member.roles))
+                    avatar = member.avatar,
-            session.commit()
+                    roles = json.dumps(list_role_ids(member.roles))
+                ).apply()
     
     async def unban_server_user(self, user, server):
-        self.bot.loop.run_in_executor(None, self._unban_server_user, user, server)
+        await GuildMembers.query \
-
+            .where(GuildMembers.guild_id == int(server.id)) \
-    def _unban_server_user(self, user, server):
+            .where(GuildMembers.user_id == int(user.id)) \
-        with self.get_session() as session:
+            .update(banned = False).apply()
-            dbmember = session.query(GuildMembers) \
-                .filter(GuildMembers.guild_id == int(server.id)) \
-                .filter(GuildMembers.user_id == int(user.id)).first()
-            if dbmember:
-                dbmember.banned = False
-                session.commit()
 
     async def flag_unactive_guild_members(self, guild_id, guild_members):
-        self.bot.loop.run_in_executor(None, self._flag_unactive_guild_members, guild_id, guild_members)
+        async with db.transaction():
-
+            async for member in GuildMembers.query \
-    def _flag_unactive_guild_members(self, guild_id, guild_members):
+                .where(GuildMembers.guild_id == int(guild_id)) \
-        with self.get_session() as session:
+                .where(GuildMembers.active == True).gino.iterate():
-            changed = False
-            dbmembers = session.query(GuildMembers) \
-                .filter(GuildMembers.guild_id == int(guild_id)) \
-                .filter(GuildMembers.active == True).all()
-            for member in dbmembers:
                 dismember = discord.utils.get(guild_members, id=member.user_id)
                 if not dismember:
-                    changed = True
+                    await member.update(active = False).apply()
-                    member.active = False
-            if changed:
-                session.commit()
 
     async def flag_unactive_bans(self, guild_id, guildbans):
-        self.bot.loop.run_in_executor(None, self._flag_unactive_bans, guild_id, guildbans)
-
-    def _flag_unactive_bans(self, guild_id, guildbans):
-        with self.get_session() as session:
-            changed = False
         for usr in guildbans:
-                dbusr = session.query(GuildMembers) \
+            dbusr = await GuildMembers.query \
-                    .filter(GuildMembers.guild_id == int(guild_id)) \
+                .where(GuildMembers.guild_id == int(guild_id)) \
-                    .filter(GuildMembers.user_id == int(usr.id)) \
+                .where(GuildMembers.user_id == int(usr.id)) \
-                    .filter(GuildMembers.active == False).first()
+                .where(GuildMembers.active == False).gino.first()
-                changed = True
             if dbusr:
-                    dbusr.banned = True
+                dbusr.update(banned=True).apply()
             else:
-                    dbusr = GuildMembers(
+                await GuildMembers.create(
-                        int(guild_id),
+                    guild_id = int(guild_id),
-                        int(usr.id),
+                    user_id = int(usr.id),
-                        usr.name,
+                    username = usr.name,
-                        usr.discriminator,
+                    discriminator = usr.discriminator,
-                        None,
+                    nickname = None,
-                        usr.avatar,
+                    avatar = usr.avatar,
-                        False,
+                    active = False,
-                        True,
+                    banned = True,
-                        "[]"
+                    roles = "[]"
                 )
-                    session.add(dbusr)
-            if changed:
-                session.commit()
                 
     async def ban_unauth_user_by_query(self, guild_id, placer_id, username, discriminator):
         self.bot.loop.run_in_executor(None, self._ban_unauth_user_by_query, guild_id, placer_id, username, discriminator)
@@ -324,9 +240,4 @@ class DatabaseInterface(object):
             return "Successfully kicked **{}#{}**!".format(dbuser.username, dbuser.discriminator)
     
     async def delete_all_messages_from_channel(self, channel_id):
-        self.bot.loop.run_in_executor(None, self._delete_all_messages_from_channel, channel_id)            
+        await Messages.delete.where(Messages.channel_id == int(channel_id)).gino.status()
-    
-    def _delete_all_messages_from_channel(self, channel_id):
-        with self.get_session() as session:
-            session.query(Messages).filter(Messages.channel_id == int(channel_id)).delete()
-            session.commit()

discordbot/titanembeds/database/guild_members.py

@@ -1,6 +1,6 @@
-from titanembeds.database import db, Base
+from titanembeds.database import db
 
-class GuildMembers(Base):
+class GuildMembers(db.Model):
     __tablename__ = "guild_members"
     id = db.Column(db.Integer, primary_key=True)    # Auto incremented id
     guild_id = db.Column(db.BigInteger)            # Discord guild id
@@ -12,17 +12,3 @@ class GuildMembers(Base):
     active = db.Column(db.Boolean())                # If the user is a member of the guild
     banned = db.Column(db.Boolean())                # If the user is banned in the guild
     roles = db.Column(db.Text())                    # Member roles
-
-    def __init__(self, guild_id, user_id, username, discriminator, nickname, avatar, active, banned, roles):
-        self.guild_id = guild_id
-        self.user_id = user_id
-        self.username = username
-        self.discriminator = discriminator
-        self.nickname = nickname
-        self.avatar = avatar
-        self.active = active
-        self.banned = banned
-        self.roles = roles
-
-    def __repr__(self):
-        return '<GuildMembers {0} {1} {2} {3} {4}>'.format(self.id, self.guild_id, self.user_id, self.username, self.discriminator)

discordbot/titanembeds/database/guilds.py

@@ -1,6 +1,6 @@
-from titanembeds.database import db, Base
+from titanembeds.database import db
 
-class Guilds(Base):
+class Guilds(db.Model):
     __tablename__ = "guilds"
     guild_id = db.Column(db.BigInteger, primary_key=True)            # Discord guild id
     name = db.Column(db.String(255))                # Name
@@ -24,24 +24,3 @@ class Guilds(Base):
     banned_words_enabled = db.Column(db.Boolean(), nullable=False, server_default="0") # If banned words are enforced
     banned_words_global_included = db.Column(db.Boolean(), nullable=False, server_default="0") # Add global banned words to the list
     banned_words = db.Column(db.Text(), nullable=False, server_default="[]")    # JSON list of strings to block from sending
-
-    def __init__(self, guild_id, name, roles, channels, webhooks, emojis, owner_id, icon):
-        self.guild_id = guild_id
-        self.name = name
-        self.unauth_users = True # defaults to true
-        self.visitor_view = False
-        self.webhook_messages = False
-        self.guest_icon = None
-        self.chat_links = True
-        self.bracket_links = True
-        self.unauth_captcha = True
-        self.mentions_limit = -1 # -1 = unlimited mentions
-        self.roles = roles
-        self.channels = channels
-        self.webhooks = webhooks
-        self.emojis = emojis
-        self.owner_id = owner_id
-        self.icon = icon
-
-    def __repr__(self):
-        return '<Guilds {0} {1}>'.format(self.id, self.guild_id)

discordbot/titanembeds/database/messages.py

@@ -1,6 +1,6 @@
-from titanembeds.database import db, Base
+from titanembeds.database import db
 
-class Messages(Base):
+class Messages(db.Model):
     __tablename__ = "messages"
     message_id = db.Column(db.BigInteger, primary_key=True)          # Message snowflake
     guild_id = db.Column(db.BigInteger)            # Discord guild id
@@ -12,18 +12,3 @@ class Messages(Base):
     mentions = db.Column(db.Text())                 # Mentions serialized
     attachments = db.Column(db.Text())              # serialized attachments
     embeds = db.Column(db.Text().with_variant(db.Text(length=4294967295), 'mysql')) # message embeds
-
-    def __init__(self, guild_id, channel_id, message_id, content, author, timestamp, edited_timestamp, mentions, attachments, embeds):
-        self.guild_id = guild_id
-        self.channel_id = channel_id
-        self.message_id = message_id
-        self.content = content
-        self.author = author
-        self.timestamp = timestamp
-        self.edited_timestamp = edited_timestamp
-        self.mentions = mentions
-        self.attachments = attachments
-        self.embeds = embeds
-
-    def __repr__(self):
-        return '<Messages {0} {1} {2} {3} {4}>'.format(self.id, self.guild_id, self.guild_id, self.channel_id, self.message_id)

discordbot/titanembeds/database/unauthenticated_bans.py

@@ -1,8 +1,8 @@
-from titanembeds.database import db, Base
+from titanembeds.database import db
 import datetime
 import time
 
-class UnauthenticatedBans(Base):
+class UnauthenticatedBans(db.Model):
     __tablename__ = "unauthenticated_bans"
     id = db.Column(db.Integer, primary_key=True)    # Auto increment id
     guild_id = db.Column(db.String(255))            # Guild pretaining to the unauthenticated user
@@ -13,16 +13,3 @@ class UnauthenticatedBans(Base):
     reason = db.Column(db.Text())                   # The reason of the ban set by the guild moderators
     lifter_id = db.Column(db.BigInteger)           # Discord Client ID of the user who lifted the ban
     placer_id = db.Column(db.BigInteger)           # The id of who placed the ban
-
-    def __init__(self, guild_id, ip_address, last_username, last_discriminator, reason, placer_id):
-        self.guild_id = guild_id
-        self.ip_address = ip_address
-        self.last_username = last_username
-        self.last_discriminator = last_discriminator
-        self.timestamp = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S')
-        self.reason = reason
-        self.lifter_id = None
-        self.placer_id = placer_id
-
-    def __repr__(self):
-        return '<UnauthenticatedBans {0} {1} {2} {3} {4} {5}'.format(self.id, self.guild_id, self.ip_address, self.last_username, self.last_discriminator, self.timestamp)

discordbot/titanembeds/database/unauthenticated_users.py

@@ -1,10 +1,6 @@
-from titanembeds.database import db, Base
+from titanembeds.database import db
-import datetime
-import time
-import random
-import string
 
-class UnauthenticatedUsers(Base):
+class UnauthenticatedUsers(db.Model):
     __tablename__ = "unauthenticated_users"
     id = db.Column(db.Integer, primary_key=True)    # Auto increment id
     guild_id = db.Column(db.BigInteger)            # Guild pretaining to the unauthenticated user
@@ -13,6 +9,3 @@ class UnauthenticatedUsers(Base):
     user_key = db.Column(db.Text())                 # The secret key used to identify the user holder
     ip_address = db.Column(db.String(255))          # The IP Address of the user
     revoked = db.Column(db.Boolean())               # If the user's key has been revoked and a new one is required to be generated
-
-    def __repr__(self):
-        return '<UnauthenticatedUsers {0} {1} {2} {3} {4} {5} {6}>'.format(self.id, self.guild_id, self.username, self.discriminator, self.user_key, self.ip_address, self.revoked)

requirements.txt

@@ -13,3 +13,4 @@ aioredis
 Flask-Babel
 patreon
 flask-redis
+gino