Throw threadpool out of the window

This commit is contained in:
Jeremy Zhang 2018-03-25 22:16:55 +00:00
parent e0b18fb94f
commit 5fa1e43e2a
2 changed files with 223 additions and 240 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:
await self.database.connect(config["database-uri"]) 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,5 +1,4 @@
from contextlib import contextmanager from contextlib import contextmanager
from asyncio_extras import threadpool
import sqlalchemy as db import sqlalchemy as db
from sqlalchemy.engine import Engine, create_engine from sqlalchemy.engine import Engine, create_engine
from sqlalchemy.orm import sessionmaker, Session from sqlalchemy.orm import sessionmaker, Session
@ -26,26 +25,22 @@ class DatabaseInterface(object):
self.engine = None # type: Engine self.engine = None # type: Engine
self._sessionmaker = None # type: sessionmaker self._sessionmaker = None # type: sessionmaker
async def connect(self, dburi): def connect(self, dburi):
async with threadpool():
self.engine = create_engine(dburi, pool_recycle=10) self.engine = create_engine(dburi, pool_recycle=10)
self._sessionmaker = sessionmaker(bind=self.engine, expire_on_commit=False)
@contextmanager @contextmanager
def get_session(self) -> Session: def get_session(self):
session = self._sessionmaker() # type: Session Session = sessionmaker(bind=self.engine)
session = Session()
try: try:
yield session yield session
session.commit()
except: except:
session.rollback() session.rollback()
raise
finally: finally:
session.close() session.commit()
async def push_message(self, message): async def push_message(self, message):
if message.guild: if message.guild:
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
edit_ts = message.edited_at edit_ts = message.edited_at
if not edit_ts: if not edit_ts:
@ -70,7 +65,6 @@ class DatabaseInterface(object):
async def update_message(self, message): async def update_message(self, message):
if message.guild: if message.guild:
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
msg = session.query(Messages) \ msg = session.query(Messages) \
.filter(Messages.guild_id == message.guild.id) \ .filter(Messages.guild_id == message.guild.id) \
@ -88,7 +82,6 @@ class DatabaseInterface(object):
async def delete_message(self, message): async def delete_message(self, message):
if message.guild: if message.guild:
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
msg = session.query(Messages) \ msg = session.query(Messages) \
.filter(Messages.guild_id == int(message.guild.id)) \ .filter(Messages.guild_id == int(message.guild.id)) \
@ -106,7 +99,6 @@ class DatabaseInterface(object):
server_webhooks = [] server_webhooks = []
else: else:
server_webhooks = [] server_webhooks = []
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
gui = session.query(Guilds).filter(Guilds.guild_id == guild.id).first() gui = session.query(Guilds).filter(Guilds.guild_id == guild.id).first()
if not gui: if not gui:
@ -132,7 +124,6 @@ class DatabaseInterface(object):
session.commit() session.commit()
async def remove_unused_guilds(self, guilds): async def remove_unused_guilds(self, guilds):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
dbguilds = session.query(Guilds).all() dbguilds = session.query(Guilds).all()
changed = False changed = False
@ -148,7 +139,6 @@ class DatabaseInterface(object):
session.commit() session.commit()
async def remove_guild(self, guild): async def remove_guild(self, guild):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
gui = session.query(Guilds).filter(Guilds.guild_id == int(guild.id)).first() gui = session.query(Guilds).filter(Guilds.guild_id == int(guild.id)).first()
if gui: if gui:
@ -157,7 +147,6 @@ class DatabaseInterface(object):
session.commit() 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):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
if guild: if guild:
dbmember = session.query(GuildMembers) \ dbmember = session.query(GuildMembers) \
@ -198,7 +187,6 @@ class DatabaseInterface(object):
session.commit() session.commit()
async def unban_server_user(self, user, server): async def unban_server_user(self, user, server):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
dbmember = session.query(GuildMembers) \ dbmember = session.query(GuildMembers) \
.filter(GuildMembers.guild_id == int(server.id)) \ .filter(GuildMembers.guild_id == int(server.id)) \
@ -208,7 +196,6 @@ class DatabaseInterface(object):
session.commit() session.commit()
async def flag_unactive_guild_members(self, guild_id, guild_members): async def flag_unactive_guild_members(self, guild_id, guild_members):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
changed = False changed = False
dbmembers = session.query(GuildMembers) \ dbmembers = session.query(GuildMembers) \
@ -223,7 +210,6 @@ class DatabaseInterface(object):
session.commit() session.commit()
async def flag_unactive_bans(self, guild_id, guildbans): async def flag_unactive_bans(self, guild_id, guildbans):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
changed = False changed = False
for usr in guildbans: for usr in guildbans:
@ -251,7 +237,6 @@ class DatabaseInterface(object):
session.commit() 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):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
dbuser = None dbuser = None
if discriminator: if discriminator:
@ -281,7 +266,6 @@ class DatabaseInterface(object):
return "Guest user, **{}#{}**, has successfully been added to the ban list!".format(dbban.last_username, dbban.last_discriminator) return "Guest user, **{}#{}**, has successfully been added to the ban list!".format(dbban.last_username, dbban.last_discriminator)
async def revoke_unauth_user_by_query(self, guild_id, username, discriminator): async def revoke_unauth_user_by_query(self, guild_id, username, discriminator):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
dbuser = None dbuser = None
if discriminator: if discriminator:
@ -304,7 +288,6 @@ 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):
async with threadpool():
with self.get_session() as session: with self.get_session() as session:
session.query(Messages).filter(Messages.channel_id == int(channel_id)).delete() session.query(Messages).filter(Messages.channel_id == int(channel_id)).delete()
session.commit() session.commit()