Initial switch to gino for discordbot db

This commit is contained in:
Jeremy Zhang 2018-07-10 07:04:19 +00:00
parent 3ca6a33be7
commit ac225c84f5
8 changed files with 153 additions and 311 deletions

View File

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

View File

@ -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 json
import discord import discord
Base = declarative_base() db = Gino()
from titanembeds.database.guilds import Guilds from titanembeds.database.guilds import Guilds
from titanembeds.database.messages import Messages 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 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): class DatabaseInterface(object):
# Courtesy of https://github.com/SunDwarf/Jokusoramame
def __init__(self, bot): def __init__(self, bot):
self.bot = 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): 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: if message.guild:
with self.get_session() as session: edit_ts = message.edited_at
edit_ts = message.edited_at if not edit_ts:
if not edit_ts: edit_ts = None
edit_ts = None else:
else: edit_ts = str(edit_ts)
edit_ts = str(edit_ts) await Messages.create(
message_id = int(message.id),
msg = Messages( guild_id = int(message.guild.id),
int(message.guild.id), channel_id = int(message.channel.id),
int(message.channel.id), content = message.content,
int(message.id), author = json.dumps(get_message_author(message)),
message.content, timestamp = str(message.created_at),
json.dumps(get_message_author(message)), edited_timestamp = edit_ts,
str(message.created_at), mentions = json.dumps(get_message_mentions(message.mentions)),
edit_ts, attachments = json.dumps(get_attachments_list(message.attachments)),
json.dumps(get_message_mentions(message.mentions)), embeds = json.dumps(get_embeds_list(message.embeds))
json.dumps(get_attachments_list(message.attachments)), )
json.dumps(get_embeds_list(message.embeds))
)
session.add(msg)
session.commit()
async def update_message(self, message): 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: 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): 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: 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): async def update_guild(self, guild):
if guild.me.guild_permissions.manage_webhooks: if guild.me.guild_permissions.manage_webhooks:
@ -108,163 +63,124 @@ class DatabaseInterface(object):
server_webhooks = [] server_webhooks = []
else: else:
server_webhooks = [] server_webhooks = []
self.bot.loop.run_in_executor(None, self._update_guild, guild, server_webhooks) gui = await Guilds.get(guild.id)
if not gui:
def _update_guild(self, guild, server_webhooks): await Guilds.create(
with self.get_session() as session: guild_id = int(guild.id),
gui = session.query(Guilds).filter(Guilds.guild_id == guild.id).first() name = guild.name,
if not gui: unauth_users = True,
gui = Guilds( visitor_view = False,
int(guild.id), webhook_messages = False,
guild.name, guest_icon = None,
json.dumps(get_roles_list(guild.roles)), chat_links = True,
json.dumps(get_channels_list(guild.channels)), bracket_links = True,
json.dumps(get_webhooks_list(server_webhooks)), unauth_captcha = True,
json.dumps(get_emojis_list(guild.emojis)), mentions_limit = -1,
int(guild.owner_id), roles = json.dumps(get_roles_list(guild.roles)),
guild.icon channels = json.dumps(get_channels_list(guild.channels)),
) webhooks = json.dumps(get_webhooks_list(server_webhooks)),
session.add(gui) emojis = json.dumps(get_emojis_list(guild.emojis)),
else: owner_id = int(guild.owner_id),
gui.name = guild.name icon = guild.icon
gui.roles = json.dumps(get_roles_list(guild.roles)) )
gui.channels = json.dumps(get_channels_list(guild.channels)) else:
gui.webhooks = json.dumps(get_webhooks_list(server_webhooks)) await gui.update(
gui.emojis = json.dumps(get_emojis_list(guild.emojis)) name = guild.name,
gui.owner_id = int(guild.owner_id) roles = json.dumps(get_roles_list(guild.roles)),
gui.icon = guild.icon channels = json.dumps(get_channels_list(guild.channels)),
session.commit() 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
).apply()
async def remove_unused_guilds(self, guilds): 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()
for guild in dbguilds:
def _remove_unused_guilds(self, guilds): disguild = discord.utils.get(guilds, id=guild.guild_id)
with self.get_session() as session: if not disguild:
dbguilds = session.query(Guilds).all() await Messages.delete.where(Messages.guild_id == int(guild.guild_id)).gino.status()
changed = False
for guild in dbguilds:
disguild = discord.utils.get(guilds, id=guild.guild_id)
if not disguild:
changed = True
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): async def remove_guild(self, guild):
self.bot.loop.run_in_executor(None, self._remove_guild, guild) gui = await Guilds.get(int(guild.id))
if gui:
def _remove_guild(self, guild): await Messages.delete.where(Messages.guild_id == int(guild.id)).gino.status()
with self.get_session() as session: await gui.delete()
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()
session.delete(gui)
session.commit()
async def update_guild_member(self, member, active=True, banned=False, guild=None): 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) if guild:
dbmember = await GuildMembers.query \
def _update_guild_member(self, member, active=True, banned=False, guild=None): .where(GuildMembers.guild_id == int(guild.id)) \
with self.get_session() as session: .where(GuildMembers.user_id == int(member.id)) \
if guild: .order_by(GuildMembers.id).gino.all()
dbmember = session.query(GuildMembers) \ else:
.filter(GuildMembers.guild_id == int(guild.id)) \ dbmember = await GuildMembers.query \
.filter(GuildMembers.user_id == int(member.id)) \ .where(GuildMembers.guild_id == int(member.guild.id)) \
.order_by(GuildMembers.id).all() .where(GuildMembers.user_id == int(member.id)) \
else: .order_by(GuildMembers.id).gino.all()
dbmember = session.query(GuildMembers) \ if not dbmember:
.filter(GuildMembers.guild_id == int(member.guild.id)) \ await GuildMembers.create(
.filter(GuildMembers.user_id == int(member.id)) \ guild_id = int(member.guild.id),
.order_by(GuildMembers.id).all() user_id = int(member.id),
if not dbmember: username = member.name,
dbmember = GuildMembers( discriminator = member.discriminator,
int(member.guild.id), nickname = member.nick,
int(member.id), avatar = member.avatar,
member.name, active = active,
member.discriminator, banned = banned,
member.nick, roles = json.dumps(list_role_ids(member.roles))
member.avatar, )
active, else:
banned, if len(dbmember) > 1:
json.dumps(list_role_ids(member.roles)) for mem in dbmember[1:]:
) await mem.delete()
session.add(dbmember)
else:
if len(dbmember) > 1:
for mem in dbmember[1:]:
session.delete(mem)
dbmember = dbmember[0] 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)): 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): 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): 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) dismember = discord.utils.get(guild_members, id=member.user_id)
if not dismember: 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): async def flag_unactive_bans(self, guild_id, guildbans):
self.bot.loop.run_in_executor(None, self._flag_unactive_bans, guild_id, guildbans) for usr in guildbans:
dbusr = await GuildMembers.query \
def _flag_unactive_bans(self, guild_id, guildbans): .where(GuildMembers.guild_id == int(guild_id)) \
with self.get_session() as session: .where(GuildMembers.user_id == int(usr.id)) \
changed = False .where(GuildMembers.active == False).gino.first()
for usr in guildbans: if dbusr:
dbusr = session.query(GuildMembers) \ dbusr.update(banned=True).apply()
.filter(GuildMembers.guild_id == int(guild_id)) \ else:
.filter(GuildMembers.user_id == int(usr.id)) \ await GuildMembers.create(
.filter(GuildMembers.active == False).first() guild_id = int(guild_id),
changed = True user_id = int(usr.id),
if dbusr: username = usr.name,
dbusr.banned = True discriminator = usr.discriminator,
else: nickname = None,
dbusr = GuildMembers( avatar = usr.avatar,
int(guild_id), active = False,
int(usr.id), banned = True,
usr.name, roles = "[]"
usr.discriminator, )
None,
usr.avatar,
False,
True,
"[]"
)
session.add(dbusr)
if changed:
session.commit()
async def ban_unauth_user_by_query(self, guild_id, placer_id, username, discriminator): 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) 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) return "Successfully kicked **{}#{}**!".format(dbuser.username, dbuser.discriminator)
async def delete_all_messages_from_channel(self, channel_id): 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()

View File

@ -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" __tablename__ = "guild_members"
id = db.Column(db.Integer, primary_key=True) # Auto incremented id id = db.Column(db.Integer, primary_key=True) # Auto incremented id
guild_id = db.Column(db.BigInteger) # Discord guild id guild_id = db.Column(db.BigInteger) # Discord guild id
@ -11,18 +11,4 @@ class GuildMembers(Base):
avatar = db.Column(db.String(255)) # The avatar str of the user avatar = db.Column(db.String(255)) # The avatar str of the user
active = db.Column(db.Boolean()) # If the user is a member of the guild 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 banned = db.Column(db.Boolean()) # If the user is banned in the guild
roles = db.Column(db.Text()) # Member roles 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)

View File

@ -1,6 +1,6 @@
from titanembeds.database import db, Base from titanembeds.database import db
class Guilds(Base): class Guilds(db.Model):
__tablename__ = "guilds" __tablename__ = "guilds"
guild_id = db.Column(db.BigInteger, primary_key=True) # Discord guild id guild_id = db.Column(db.BigInteger, primary_key=True) # Discord guild id
name = db.Column(db.String(255)) # Name name = db.Column(db.String(255)) # Name
@ -23,25 +23,4 @@ class Guilds(Base):
max_message_length = db.Column(db.Integer, nullable=False, server_default="300") # Chars length the message should be before being rejected by the server max_message_length = db.Column(db.Integer, nullable=False, server_default="300") # Chars length the message should be before being rejected by the server
banned_words_enabled = db.Column(db.Boolean(), nullable=False, server_default="0") # If banned words are enforced 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_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 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)

View File

@ -1,6 +1,6 @@
from titanembeds.database import db, Base from titanembeds.database import db
class Messages(Base): class Messages(db.Model):
__tablename__ = "messages" __tablename__ = "messages"
message_id = db.Column(db.BigInteger, primary_key=True) # Message snowflake message_id = db.Column(db.BigInteger, primary_key=True) # Message snowflake
guild_id = db.Column(db.BigInteger) # Discord guild id guild_id = db.Column(db.BigInteger) # Discord guild id
@ -11,19 +11,4 @@ class Messages(Base):
edited_timestamp = db.Column(db.TIMESTAMP) # Timestamp of when content is edited edited_timestamp = db.Column(db.TIMESTAMP) # Timestamp of when content is edited
mentions = db.Column(db.Text()) # Mentions serialized mentions = db.Column(db.Text()) # Mentions serialized
attachments = db.Column(db.Text()) # serialized attachments attachments = db.Column(db.Text()) # serialized attachments
embeds = db.Column(db.Text().with_variant(db.Text(length=4294967295), 'mysql')) # message embeds 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)

View File

@ -1,8 +1,8 @@
from titanembeds.database import db, Base from titanembeds.database import db
import datetime import datetime
import time import time
class UnauthenticatedBans(Base): class UnauthenticatedBans(db.Model):
__tablename__ = "unauthenticated_bans" __tablename__ = "unauthenticated_bans"
id = db.Column(db.Integer, primary_key=True) # Auto increment id id = db.Column(db.Integer, primary_key=True) # Auto increment id
guild_id = db.Column(db.String(255)) # Guild pretaining to the unauthenticated user guild_id = db.Column(db.String(255)) # Guild pretaining to the unauthenticated user
@ -12,17 +12,4 @@ class UnauthenticatedBans(Base):
timestamp = db.Column(db.TIMESTAMP) # The timestamp of when the user got banned timestamp = db.Column(db.TIMESTAMP) # The timestamp of when the user got banned
reason = db.Column(db.Text()) # The reason of the ban set by the guild moderators 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 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 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)

View File

@ -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" __tablename__ = "unauthenticated_users"
id = db.Column(db.Integer, primary_key=True) # Auto increment id id = db.Column(db.Integer, primary_key=True) # Auto increment id
guild_id = db.Column(db.BigInteger) # Guild pretaining to the unauthenticated user guild_id = db.Column(db.BigInteger) # Guild pretaining to the unauthenticated user
@ -12,7 +8,4 @@ class UnauthenticatedUsers(Base):
discriminator = db.Column(db.Integer) # The discriminator to distinguish unauth users with each other discriminator = db.Column(db.Integer) # The discriminator to distinguish unauth users with each other
user_key = db.Column(db.Text()) # The secret key used to identify the user holder 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 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 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)

View File

@ -12,4 +12,5 @@ redis
aioredis aioredis
Flask-Babel Flask-Babel
patreon patreon
flask-redis flask-redis
gino