DB Connection is now created within the same event loop as client

Referencing pool connection from client/self.bot
pull/8/head
sgoudham 4 years ago
parent 55e87b3fc3
commit d1e1eb2fc8

@ -197,7 +197,7 @@ class Guild(Cog):
else: else:
# Set up the modlogs channel within the guild # Set up the modlogs channel within the guild
mod_log_setup = True mod_log_setup = True
await storage_modlog_for_guild(ctx, channelID, mod_log_setup) await storage_modlog_for_guild(self.bot.db, ctx, channelID, mod_log_setup)
@modlogs.command(name="update") @modlogs.command(name="update")
@has_permissions(manage_guild=True) @has_permissions(manage_guild=True)
@ -235,7 +235,7 @@ class Guild(Cog):
else: else:
# Update the modlog channel within the database and cache # Update the modlog channel within the database and cache
mod_log_setup = False mod_log_setup = False
await storage_modlog_for_guild(ctx, channelID, mod_log_setup) await storage_modlog_for_guild(self.bot.db, ctx, channelID, mod_log_setup)
@modlogs.command("delete") @modlogs.command("delete")
@has_permissions(manage_guild=True) @has_permissions(manage_guild=True)

@ -104,7 +104,7 @@ async def ummute_members(self, message, targets, reason):
await message.channel.send(embed=embed) await message.channel.send(embed=embed)
async def mute_members(message, targets, reason, muted): async def mute_members(pool, message, targets, reason, muted):
""" """
Method to allow members to be muted Method to allow members to be muted
@ -119,7 +119,7 @@ async def mute_members(message, targets, reason, muted):
and not target.guild_permissions.administrator): and not target.guild_permissions.administrator):
# Store the current roles of the user within database # Store the current roles of the user within database
await storeRoles(target=target, ctx=message, member=target) await storeRoles(pool=pool, target=target, ctx=message, member=target)
# Give the user the muted role # Give the user the muted role
await target.edit(roles=[muted], reason=reason) await target.edit(roles=[muted], reason=reason)
@ -350,10 +350,10 @@ class Moderation(Cog):
await channel.set_permissions(muted, send_messages=False, read_messages=True) await channel.set_permissions(muted, send_messages=False, read_messages=True)
# Send embed of the kicked member # Send embed of the kicked member
await mute_members(ctx.message, members, reason, muted) await mute_members(self.bot.db, ctx.message, members, reason, muted)
else: else:
# Send embed of the kicked member # Send embed of the kicked member
await mute_members(ctx.message, members, reason, role) await mute_members(self.bot.db, ctx.message, members, reason, role)
@command(name="unmute", aliases=["Unmute"], usage="`<member>...` `[reason]`") @command(name="unmute", aliases=["Unmute"], usage="`<member>...` `[reason]`")
@has_guild_permissions(manage_roles=True) @has_guild_permissions(manage_roles=True)

@ -56,25 +56,35 @@ async def create_connection():
loop=client.loop) loop=client.loop)
async def startup_cache_log():
"""Store the modlogs/prefixes in cache from the database on startup"""
# Setup pool
pool = client.db
# Setup up pool connection and cursor
async with pool.acquire() as conn:
async with conn.cursor() as cur:
# Grab the prefix of the server from the database
select_query = """SELECT * FROM guilds"""
# Execute the query
await cur.execute(select_query)
results = await cur.fetchall()
# Store the guildID's, modlog channels and prefixes within cache
for row in results:
cache(guildid=row[0], channel=row[2], prefix=row[1])
# Make sure the connection is setup before the bot is ready # Make sure the connection is setup before the bot is ready
client.loop.run_until_complete(create_connection()) client.loop.run_until_complete(create_connection())
client.loop.run_until_complete(startup_cache_log())
if __name__ == '__main__': if __name__ == '__main__':
for ext in settings.extensions(): for ext in settings.extensions():
client.load_extension(ext) client.load_extension(ext)
@client.event
async def on_message(message):
"""Make sure bot messages are not tracked"""
if message.author.bot:
return
# Processing the message
await client.process_commands(message)
@tasks.loop(seconds=120, reconnect=True) @tasks.loop(seconds=120, reconnect=True)
async def change_status(): async def change_status():
"""Creating Custom Statuses as a Background Task""" """Creating Custom Statuses as a Background Task"""
@ -113,6 +123,17 @@ async def change_status():
change_status.start() change_status.start()
@client.event
async def on_message(message):
"""Make sure bot messages are not tracked"""
if message.author.bot:
return
# Processing the message
await client.process_commands(message)
@client.event @client.event
async def on_ready(): async def on_ready():
"""Displaying if Bot is Ready""" """Displaying if Bot is Ready"""
@ -204,7 +225,7 @@ async def change_prefix(ctx, new: Optional[str] = None):
# As long as a new prefix has been given and is less than 5 characters # As long as a new prefix has been given and is less than 5 characters
if new and len(new) <= 5: if new and len(new) <= 5:
# Store the new prefix in the dictionary and update the database # Store the new prefix in the dictionary and update the database
await storage_prefix_for_guild(ctx, new) await storage_prefix_for_guild(client.db, ctx, new)
# Making sure that errors are handled if prefix is above 5 characters # Making sure that errors are handled if prefix is above 5 characters
elif new and len(new) > 5: elif new and len(new) > 5:

@ -1,12 +1,9 @@
import asyncio
import random import random
from discord import Colour from discord import Colour
import db
from db import connection
# Defining a list of colours # Defining a list of colours
colors = { colors = {
'DEFAULT': 0x000000, 'DEFAULT': 0x000000,
'WHITE': 0xFFFFFF, 'WHITE': 0xFFFFFF,
@ -63,27 +60,6 @@ def rndColour():
enso_cache = {} enso_cache = {}
async def startup_cache_log():
"""Store the modlogs/prefixes in cache from the database on startup"""
# Setup pool
pool = await connection(db.loop)
# Setup up pool connection and cursor
async with pool.acquire() as conn:
async with conn.cursor() as cur:
# Grab the prefix of the server from the database
select_query = """SELECT * FROM guilds"""
# Execute the query
await cur.execute(select_query)
results = await cur.fetchall()
# Store the guildID's, modlog channels and prefixes within cache
for row in results:
cache(guildid=row[0], channel=row[2], prefix=row[1])
# Method to store the guildID, channel # Method to store the guildID, channel
def cache(guildid, channel, prefix): def cache(guildid, channel, prefix):
"""Storing GuildID, Modlogs Channel and Prefix in Cache""" """Storing GuildID, Modlogs Channel and Prefix in Cache"""
@ -105,12 +81,9 @@ def del_cache(guildid):
# --------------------------------------------!ModLogs Section!--------------------------------------------------------- # --------------------------------------------!ModLogs Section!---------------------------------------------------------
# Updating the modlog within the dict and database when the method is called # Updating the modlog within the dict and database when the method is called
async def storage_modlog_for_guild(ctx, channelID, setup): async def storage_modlog_for_guild(pool, ctx, channelID, setup):
enso_cache[str(ctx.guild.id)]["Modlogs"] = channelID enso_cache[str(ctx.guild.id)]["Modlogs"] = channelID
# Setup pool
pool = await connection(db.loop)
# Setup up pool connection and cursor # Setup up pool connection and cursor
async with pool.acquire() as conn: async with pool.acquire() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
@ -167,12 +140,9 @@ def get_modlog_for_guild(guildid):
# --------------------------------------------!Prefixes Section!-------------------------------------------------------- # --------------------------------------------!Prefixes Section!--------------------------------------------------------
# Updating the prefix within the dict and database when the method is called # Updating the prefix within the dict and database when the method is called
async def storage_prefix_for_guild(ctx, prefix): async def storage_prefix_for_guild(pool, ctx, prefix):
enso_cache[str(ctx.guild.id)]["Prefix"] = prefix enso_cache[str(ctx.guild.id)]["Prefix"] = prefix
# Setup pool
pool = await connection(db.loop)
# Setup up pool connection and cursor # Setup up pool connection and cursor
async with pool.acquire() as conn: async with pool.acquire() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
@ -238,13 +208,10 @@ def extensions():
return ext return ext
async def storeRoles(target, ctx, member): async def storeRoles(pool, target, ctx, member):
"""Storing User Roles within Database""" """Storing User Roles within Database"""
role_ids = ", ".join([str(r.id) for r in target.roles]) role_ids = ", ".join([str(r.id) for r in target.roles])
# Setup pool
pool = await connection(db.loop)
# Setup up pool connection and cursor # Setup up pool connection and cursor
async with pool.acquire() as conn: async with pool.acquire() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
@ -273,9 +240,4 @@ async def clearRoles(ctx, member, pool):
await conn.commit() await conn.commit()
print(cur.rowcount, f"Roles Cleared For User {member} in {ctx.guild.name}") print(cur.rowcount, f"Roles Cleared For User {member} in {ctx.guild.name}")
# Run the async function to store everything in cache
loop = asyncio.get_event_loop()
loop.run_until_complete(startup_cache_log())
# --------------------------------------------!End Cogs/Set Values Section!--------------------------------------------- # --------------------------------------------!End Cogs/Set Values Section!---------------------------------------------

Loading…
Cancel
Save