From 249e5d94ebac267a5d3dc1f2e9fc1de719dbb1c6 Mon Sep 17 00:00:00 2001 From: Dustin Pianalto Date: Mon, 21 May 2018 18:31:06 -0800 Subject: [PATCH] Switch to Asyncpg for db con --- exts/admin.py | 92 +++++++++++++++++++++--------------------- exts/events.py | 6 ++- exts/fun.py | 5 ++- exts/imports/checks.py | 16 ++++---- exts/imports/utils.py | 2 +- exts/patreon.py | 49 +++++++++++----------- exts/rcon.py | 79 +++++++++++++++++++----------------- exts/utils.py | 47 ++++++++++----------- geeksbot.py | 18 ++++----- misc.py | 11 ++++- 10 files changed, 168 insertions(+), 157 deletions(-) diff --git a/exts/admin.py b/exts/admin.py index ec84081..d790d3f 100644 --- a/exts/admin.py +++ b/exts/admin.py @@ -51,6 +51,7 @@ class Admin: await ctx.send('Geeksbot is restarting.') with open(f'{config_dir}reboot', 'w') as f: f.write(f'1\n{ctx.channel.id}') + # noinspection PyProtectedMember os._exit(1) @commands.command(hidden=True) @@ -71,20 +72,19 @@ class Admin: emoji_code = f'' else: emoji_code = f'<:{emoji.name}:{emoji.id}>' - if self.bot.con.all('select id from geeksbot_emojis where id = %(id)s', {'id': emoji.id}): - self.bot.con.run("update geeksbot_emojis set id = %(id)s, name = %(name)s, code = %(emoji_code)s " - "where name = %(name)s", - {'name': emoji.name, 'id': emoji.id, 'emoji_code': emoji_code}) + if self.bot.db_con.fetchall('select id from geeksbot_emojis where id = $1', emoji.id): + self.bot.db_con.execute("update geeksbot_emojis set id = $2, name = $1, code = $3 where name = $1", + emoji.name, emoji.id, emoji_code) else: - self.bot.con.run("insert into geeksbot_emojis(id,name,code) values (%(id)s,%(name)s,%(emoji_code)s)", - {'name': emoji.name, 'id': emoji.id, 'emoji_code': emoji_code}) + self.bot.db_con.execute("insert into geeksbot_emojis(id,name,code) values ($2,$1,$3)", + emoji.name, emoji.id, emoji_code) await ctx.message.add_reaction('✅') await ctx.send(f'Emojis have been updated in the database.') @commands.command(hidden=True) @commands.check(checks.is_guild_owner) async def get_guild_config(self, ctx): - config = self.bot.con.one('select * from guild_config where guild_id = %(id)s', {'id': ctx.guild.id}) + config = self.bot.db_con.fetchval('select * from guild_config where guild_id = $1', ctx.guild.id) configs = [str(config)[i:i+1990] for i in range(0, len(config), 1990)] await ctx.message.author.send(f'The current config for the {ctx.guild.name} guild is:\n') admin_log.info(configs) @@ -112,8 +112,8 @@ class Admin: if ctx.guild: if checks.is_admin(self.bot, ctx): if channel is not None: - self.bot.con.run('update guild_config set admin_chat = %(chan)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'chan': channel.id}) + self.bot.db_con.execute('update guild_config set admin_chat = $2 where guild_id = $1', + ctx.guild.id, channel.id) await ctx.send(f'{channel.name} is now set as the Admin Chat channel for this guild.') @set.command(name='channel_lockdown', aliases=['lockdown', 'restrict_access', 'cl']) @@ -121,18 +121,18 @@ class Admin: if ctx.guild: if checks.is_admin(self.bot, ctx): if str(config).lower() == 'true': - if self.bot.con.one('select allowed_channels from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) is []: + if self.bot.db_con.fetchval('select allowed_channels from guild_config where guild_id = $1', + ctx.guild.id) is []: await ctx.send('Please set at least one allowed channel before running this command.') else: - self.bot.con.run('update guild_config set channel_lockdown = True where guild_id = %(id)s', - {'id': ctx.guild.id}) + self.bot.db_con.execute('update guild_config set channel_lockdown = True where guild_id = $1', + ctx.guild.id) await ctx.send('Channel Lockdown is now active.') elif str(config).lower() == 'false': - if self.bot.con.one('select channel_lockdown from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}): - self.bot.con.run('update guild_config set channel_lockdown = False where guild_id = %(id)s', - {'id': ctx.guild.id}) + if self.bot.db_con.fetchval('select channel_lockdown from guild_config where guild_id = $1', + ctx.guild.id): + self.bot.db_con.execute('update guild_config set channel_lockdown = False where guild_id = $1', + ctx.guild.id) await ctx.send('Channel Lockdown has been deactivated.') else: await ctx.send('Channel Lockdown is already deactivated.') @@ -153,28 +153,27 @@ class Admin: await ctx.send(f'{channel} is not a valid text channel in this guild.') else: admin_log.info('Chan found') - if self.bot.con.one('select allowed_channels from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}): - if chnl.id in json.loads(self.bot.con.one('select allowed_channels from guild_config ' - 'where guild_id = %(id)s', - {'id': ctx.guild.id})): + if self.bot.db_con.fetchval('select allowed_channels from guild_config where guild_id = $1', + ctx.guild.id): + if chnl.id in json.loads(self.bot.db_con.fetchval('select allowed_channels ' + 'from guild_config where guild_id = $1', + ctx.guild.id)): admin_log.info('Chan found in config') await ctx.send(f'{channel} is already in the list of allowed channels. Skipping...') else: admin_log.info('Chan not found in config') - allowed_channels = json.loads(self.bot.con.one('select allowed_channels from ' - 'guild_config where guild_id = %(id)s', - {'id': ctx.guild.id})).append(chnl.id) - self.bot.con.run('update guild_config set allowed_channels = %(channels)s ' - 'where guild_id = %(id)s', - {'id': ctx.guild.id, 'channels': allowed_channels}) + allowed_channels = json.loads(self.bot.db_con.fetchval('select allowed_channels ' + 'from guild_config ' + 'where guild_id = $1', + ctx.guild.id)).append(chnl.id) + self.bot.db_con.execute('update guild_config set allowed_channels = $2 ' + 'where guild_id = $1', ctx.guild.id, allowed_channels) added = f'{added}\n{channel}' else: admin_log.info('Chan not found in config') allowed_channels = [chnl.id] - self.bot.con.run('update guild_config set allowed_channels = %(channels)s ' - 'where guild_id = %(id)s', - {'id': ctx.guild.id, 'channels': allowed_channels}) + self.bot.db_con.execute('update guild_config set allowed_channels = $2 where guild_id = $1', + ctx.guild.id, allowed_channels) added = f'{added}\n{channel}' if added != '': await ctx.send(f'The following channels have been added to the allowed channel list: {added}') @@ -196,8 +195,7 @@ class Admin: async def add_prefix(self, ctx, *, prefix=None): if ctx.guild: if checks.is_admin(self.bot, ctx): - prefixes = self.bot.con.one('select prefix from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + prefixes = self.bot.db_con.fetchval('select prefix from guild_config where guild_id = $1', ctx.guild.id) if prefix is None: await ctx.send(prefixes) return @@ -209,8 +207,8 @@ class Admin: if len(prefixes) > 10: await ctx.send(f'Only 10 prefixes are allowed per guild.\nPlease remove some before adding more.') prefixes = prefixes[:10] - self.bot.con.run('update guild_config set prefix = %(prefixes)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'prefixes': prefixes}) + self.bot.db_con.execute('update guild_config set prefix = $2 where guild_id = $1', + ctx.guild.id, prefixes) await ctx.guild.me.edit(nick=f'[{prefixes[0]}] Geeksbot') await ctx.send(f"Updated. You currently have {len(prefixes)} " f"{'prefix' if len(prefixes) == 1 else 'prefixes'} " @@ -225,8 +223,8 @@ class Admin: async def remove_prefix(self, ctx, *, prefix=None): if ctx.guild: if checks.is_admin(self.bot, ctx): - prefixes = self.bot.con.one('select prefix from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + prefixes = self.bot.db_con.fetchval('select prefix from guild_config where guild_id = $1', + ctx.guild.id) found = 0 if prefix is None: await ctx.send(prefixes) @@ -243,8 +241,8 @@ class Admin: else: await ctx.send(f'The prefix {p} is not in the config for this guild.') if found: - self.bot.con.run('update guild_config set prefix = %(prefixes)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'prefixes': prefixes}) + self.bot.db_con.execute('update guild_config set prefix = $2 where guild_id = $1', + ctx.guild.id, prefixes) await ctx.guild.me.edit(nick=f'[{prefixes[0] if len(prefixes) != 0 else self.bot.default_prefix}] ' f'Geeksbot') await ctx.send(f"Updated. You currently have {len(prefixes)} " @@ -261,14 +259,14 @@ class Admin: async def _add_admin_role(self, ctx, role=None): role = discord.utils.get(ctx.guild.roles, name=role) if role is not None: - roles = json.loads(self.bot.con.one('select admin_roles from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id})) + roles = json.loads(self.bot.db_con.fetchval('select admin_roles from guild_config where guild_id = $1', + ctx.guild.id)) if role.name in roles: await ctx.send(f'{role.name} is already registered as an admin role in this guild.') else: roles[role.name] = role.id - self.bot.con.run('update guild_config set admin_roles = %(roles)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'roles': json.dumps(roles)}) + self.bot.db_con.execute('update guild_config set admin_roles = $2 where guild_id = $1', + ctx.guild.id, json.dumps(roles)) await ctx.send(f'{role.name} has been added to the list of admin roles for this guild.') else: await ctx.send('You must include a role with this command.') @@ -279,12 +277,12 @@ class Admin: async def _remove_admin_role(self, ctx, role=None): role = discord.utils.get(ctx.guild.roles, name=role) if role is not None: - roles = json.loads(self.bot.con.one('select admin_roles from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id})) + roles = json.loads(self.bot.db_con.fetchval('select admin_roles from guild_config where guild_id = $1', + ctx.guild.id)) if role.name in roles: del roles[role.name] - self.bot.con.run('update guild_config set admin_roles = %(roles)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'roles': roles}) + self.bot.db_con.execute('update guild_config set admin_roles = $2 where guild_id = $1', + ctx.guild.id, json.dumps(roles)) await ctx.send(f'{role.name} has been removed from the list of admin roles for this guild.') else: await ctx.send(f'{role.name} is not registered as an admin role in this guild.') diff --git a/exts/events.py b/exts/events.py index 827a873..85c8ff5 100644 --- a/exts/events.py +++ b/exts/events.py @@ -55,11 +55,13 @@ class BotEvents: datetime.utcnow(), msg_id) async def on_raw_bulk_message_delete(self, msg_ids, chan_id): - sql = await self.bot.db_con.prepare('update messages set deleted_at = $1 where id = $2') + del_time = datetime.utcnow() for msg_id in msg_ids: - await sql.execute(datetime.utcnow(), msg_id) + await self.bot.db_con.execute('update messages set deleted_at = $1 where id = $2', + del_time, msg_id) async def on_message(self, ctx): + # noinspection PyBroadException try: if ctx.author in self.bot.infected: if datetime.now().timestamp() > self.bot.infected[ctx.author][1] + 300: diff --git a/exts/fun.py b/exts/fun.py index 6521079..acf8162 100644 --- a/exts/fun.py +++ b/exts/fun.py @@ -82,8 +82,8 @@ class Fun: @commands.command() @commands.cooldown(1, 5, type=commands.BucketType.user) - async def fact(self, ctx, number:int): - if number < 20001 and number > 0: + async def fact(self, ctx, number: int): + if 0 < number < 20001: n = 1990 with ctx.channel.typing(): a = await self.bot.loop.run_in_executor(None, self.get_factorial, number) @@ -141,6 +141,7 @@ class Fun: else: await ctx.send('Not connected to that voice channel.') + # noinspection PyUnusedLocal @commands.command(hidden=True) @commands.is_owner() async def volume(self, ctx, volume: float): diff --git a/exts/imports/checks.py b/exts/imports/checks.py index 30de021..1175086 100644 --- a/exts/imports/checks.py +++ b/exts/imports/checks.py @@ -6,8 +6,8 @@ owner_id = 351794468870946827 def check_admin_role(bot, ctx, member): - admin_roles = json.loads(bot.con.one(f"select admin_roles from guild_config where guild_id = %(id)s", - {'id': ctx.guild.id})) + admin_roles = json.loads(bot.db_con.fetchval(f"select admin_roles from guild_config where guild_id = $1", + ctx.guild.id)) for role in admin_roles: if discord.utils.get(ctx.guild.roles, id=admin_roles[role]) in member.roles: return True @@ -15,8 +15,8 @@ def check_admin_role(bot, ctx, member): def check_rcon_role(bot, ctx, member): - rcon_admin_roles = json.loads(bot.con.one("select rcon_admin_roles from guild_config where guild_id = %(id)s", - {'id': ctx.guild.id})) + rcon_admin_roles = json.loads(bot.db_con.fetchval("select rcon_admin_roles from guild_config where guild_id = $1", + ctx.guild.id)) for role in rcon_admin_roles: if discord.utils.get(ctx.guild.roles, id=rcon_admin_roles[role]) in member.roles: return True @@ -24,8 +24,8 @@ def check_rcon_role(bot, ctx, member): def is_admin(bot, ctx): - admin_roles = json.loads(bot.con.one("select admin_roles from guild_config where guild_id = %(id)s", - {'id': ctx.guild.id})) + admin_roles = json.loads(bot.db_con.fetchval("select admin_roles from guild_config where guild_id = $1", + ctx.guild.id)) for role in admin_roles: if discord.utils.get(ctx.guild.roles, id=admin_roles[role]) in ctx.message.author.roles: return True @@ -39,8 +39,8 @@ def is_guild_owner(ctx): def is_rcon_admin(bot, ctx): - rcon_admin_roles = json.loads(bot.con.one("select rcon_admin_roles from guild_config where guild_id = %(id)s", - {'id': ctx.guild.id})) + rcon_admin_roles = json.loads(bot.db_con.fetchval("select rcon_admin_roles from guild_config where guild_id = $1", + ctx.guild.id)) for role in rcon_admin_roles: if discord.utils.get(ctx.guild.roles, id=rcon_admin_roles[role]) in ctx.message.author.roles: return True diff --git a/exts/imports/utils.py b/exts/imports/utils.py index cb067fe..bac011b 100644 --- a/exts/imports/utils.py +++ b/exts/imports/utils.py @@ -19,7 +19,7 @@ class Capturing(list): async def mute(bot, ctx, admin=0, member_id=None): - mute_role = bot.con.one(f'select muted_role from guild_config where guild_id = {ctx.guild.id}') + mute_role = bot.db_con.fetchval(f'select muted_role from guild_config where guild_id = $1', ctx.guild.id) if mute_role: if admin or checks.is_admin(bot, ctx): if ctx.guild.me.guild_permissions.manage_roles: diff --git a/exts/patreon.py b/exts/patreon.py index 1e1e4d2..81b15de 100644 --- a/exts/patreon.py +++ b/exts/patreon.py @@ -17,10 +17,9 @@ class Patreon: @commands.cooldown(1, 5, type=commands.BucketType.user) async def get_patreon_links(self, ctx, target: discord.Member=None): """Prints Patreon information for creators on the server.""" - if self.bot.con.one('select patreon_enabled from guild_config where guild_id = %(id)s', {'id': ctx.guild.id}): - patreon_info = self.bot.con.one('select patreon_message,patreon_links\ - from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + if self.bot.db_con.fetchval('select patreon_enabled from guild_config where guild_id = $1', ctx.guild.id): + patreon_info = self.bot.db_con.fetchval('select patreon_message,patreon_links from guild_config ' + 'where guild_id = $1', ctx.guild.id) message = patreon_info[0].replace('\\n', '\n') patreon_links = json.loads(patreon_info[1]) for key in patreon_links: @@ -35,13 +34,13 @@ class Patreon: @commands.command(aliases=['patreon_message']) async def set_patreon_message(self, ctx, message): if checks.is_admin(self.bot, ctx): - patreon_message = self.bot.con.one('select patreon_message from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + patreon_message = self.bot.db_con.fetchval('select patreon_message from guild_config where guild_id = $1', + ctx.guild.id) if message == patreon_message: await ctx.send('That is already the current message for this guild.') else: - self.bot.con.run('update guild_config set patreon_message = %(message)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'message': message}) + self.bot.db_con.execute('update guild_config set patreon_message = $2 where guild_id = $1', + ctx.guild.id, message) await ctx.send(f'The patreon message for this guild has been set to:\n{message}') else: await ctx.send(f'You are not authorized to run this command.') @@ -49,8 +48,8 @@ class Patreon: @commands.command(aliases=['add_patreon', 'set_patreon']) async def add_patreon_info(self, ctx, name, url): if checks.is_admin(self.bot, ctx): - patreon_info = self.bot.con.one('select patreon_links from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + patreon_info = self.bot.db_con.fetchval('select patreon_links from guild_config where guild_id = $1', + ctx.guild.id) patreon_links = {} update = 0 if patreon_info: @@ -58,8 +57,8 @@ class Patreon: if name in patreon_links: update = 1 patreon_links[name] = url - self.bot.con.run('update guild_config set patreon_links = %(links)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'links': json.dumps(patreon_links)}) + self.bot.db_con.execute('update guild_config set patreon_links = $2 where guild_id = $1', + ctx.guild.id, json.dumps(patreon_links)) await ctx.send(f"The Patreon link for {name} has been " f"{'updated to the new url.' if update else'added to the config for this guild.'}") else: @@ -68,14 +67,14 @@ class Patreon: @commands.command(aliases=['remove_patreon']) async def remove_patreon_info(self, ctx, name): if checks.is_admin(self.bot, ctx): - patreon_info = self.bot.con.one('select patreon_links from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + patreon_info = self.bot.db_con.fetchval('select patreon_links from guild_config where guild_id = $1', + ctx.guild.id) if patreon_info: patreon_links = json.loads(patreon_info) if name in patreon_links: del patreon_links[name] - self.bot.con.run('update guild_config set patreon_links = %(links)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'links': json.dumps(patreon_links)}) + self.bot.db_con.execute('update guild_config set patreon_links = $2 where guild_id = $1', + ctx.guild.id, json.dumps(patreon_links)) await ctx.send(f'The Patreon link for {name} has been removed from the config for this guild.') return else: @@ -88,17 +87,17 @@ class Patreon: @commands.command() async def enable_patreon(self, ctx, state: bool=True): if checks.is_admin(self.bot, ctx): - patreon_status = self.bot.con.one('select patreon_enabled from guild_config where guild_id = %(id)s', - {'id': ctx.guild.id}) + patreon_status = self.bot.db_con.fetchval('select patreon_enabled from guild_config where guild_id = $1', + ctx.guild.id) if patreon_status and state: await ctx.send('Patreon is already enabled for this guild.') elif patreon_status and not state: - self.bot.con.run('update guild_config set patreon_enabled = %(state)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'state': state}) + self.bot.db_con.execute('update guild_config set patreon_enabled = $2 where guild_id = $1', + ctx.guild.id, state) await ctx.send('Patreon has been disabled for this guild.') elif not patreon_status and state: - self.bot.con.run('update guild_config set patreon_enabled = %(state)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'state': state}) + self.bot.db_con.execute('update guild_config set patreon_enabled = $2 where guild_id = $1', + ctx.guild.id, state) await ctx.send('Patreon has been enabled for this guild.') elif not patreon_status and not state: await ctx.send('Patreon is already disabled for this guild.') @@ -107,9 +106,9 @@ class Patreon: @commands.cooldown(1, 5, type=commands.BucketType.user) async def referral_links(self, ctx, target: discord.Member=None): """Prints G-Portal Referral Links.""" - if self.bot.con.one('select referral_enabled from guild_config where guild_id = %(id)s', {'id': ctx.guild.id}): - referral_info = self.bot.con.one('select referral_message,referral_links from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id}) + if self.bot.db_con.fetchval('select referral_enabled from guild_config where guild_id = $1', ctx.guild.id): + referral_info = self.bot.db_con.fetchval('select referral_message,referral_links from guild_config ' + 'where guild_id = $1', ctx.guild.id) message = referral_info[0] referral_links = json.loads(referral_info[1]) for key in referral_links: diff --git a/exts/rcon.py b/exts/rcon.py index f1af0d5..fe0bd4b 100644 --- a/exts/rcon.py +++ b/exts/rcon.py @@ -120,6 +120,7 @@ class Rcon: True) con.exec_command('ServerChatToPlayer "{0}" GeeksBot: Admin Geeks have been notified you need assistance. ' 'Please be patient.'.format(player)) + # noinspection PyProtectedMember con._sock.close() for role in admin_roles: msg = '{0} {1}'.format(msg, discord.utils.get(ctx.guild.roles, id=admin_roles[role]).mention) @@ -139,13 +140,13 @@ class Rcon: To view all the valid ARK servers for this guild see list_ark_servers.""" if checks.is_rcon_admin(self.bot, ctx): if server is not None: - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) server = server.replace('_', ' ').title() if server in rcon_connections: rcon_connections[server]["monitoring_chat"] = 1 - self.bot.con.run('update guild_config set rcon_connections = %(json)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'json': json.dumps(rcon_connections)}) + self.bot.db_con.execute('update guild_config set rcon_connections = $2 where guild_id = $1', + ctx.guild.id, json.dumps(rcon_connections)) channel = self.bot.get_channel(rcon_connections[server]['game_chat_chan_id']) await channel.send('Started monitoring on the {0} server.'.format(server)) await ctx.message.add_reaction('✅') @@ -158,6 +159,7 @@ class Rcon: True) messages = await self.bot.loop.run_in_executor(None, self.server_chat_background_process, ctx.guild.id, con) + # noinspection PyProtectedMember con._sock.close() except TimeoutError: rcon_log.error(traceback.format_exc()) @@ -180,8 +182,10 @@ class Rcon: message = func(ctx, message, rcon_connections['server']) await channel.send('{0}'.format(message)) await asyncio.sleep(1) - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections ' + 'from guild_config ' + 'where guild_id = $1', + ctx.guild.id)) await channel.send('Monitoring Stopped') else: await ctx.send(f'Server not found: {server}') @@ -197,13 +201,13 @@ class Rcon: Context is the same as monitor_chat""" if checks.is_rcon_admin(self.bot, ctx): if server is not None: - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) server = server.replace('_', ' ').title() if server in rcon_connections: rcon_connections[server]["monitoring_chat"] = 0 - self.bot.con.run('update guild_config set rcon_connections = %(json)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'json': json.dumps(rcon_connections)}) + self.bot.db_con.execute('update guild_config set rcon_connections = $2 where guild_id = $1', + ctx.guild.id, json.dumps(rcon_connections)) else: await ctx.send(f'Server not found: {server}') else: @@ -224,8 +228,8 @@ class Rcon: "first last" To view all the valid ARK servers for this guild see list_ark_servers.""" if checks.is_rcon_admin(self.bot, ctx): - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) if server is not None: server = server.replace('_', ' ').title() if server in rcon_connections: @@ -241,9 +245,10 @@ class Rcon: '"ip" port "password" if you would like to get info from it.'.format(server)) else: for server in rcon_connections: + msg = await ctx.send('Getting Data for the {0} server'.format(server.title())) + # noinspection PyBroadException try: connection_info = rcon_connections[server] - msg = await ctx.send('Getting Data for the {0} server'.format(server.title())) async with ctx.channel.typing(): message = self._listplayers(connection_info) except Exception as e: @@ -262,8 +267,8 @@ class Rcon: All strings (, , ) must be contained inside double quotes.""" if checks.is_rcon_admin(self.bot, ctx): server = server.title() - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) if server not in rcon_connections: rcon_connections[server] = { 'ip': ip, @@ -274,8 +279,8 @@ class Rcon: 'msg_chan_id': 0, 'monitoring_chat': 0 } - self.bot.con.run('update guild_config set rcon_connections = %(connections)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'connections': json.dumps(rcon_connections)}) + self.bot.db_con.execute('update guild_config set rcon_connections = $2 where guild_id = $1', + ctx.guild.id, json.dumps(rcon_connections)) await ctx.send('{0} server has been added to my configuration.'.format(server)) else: await ctx.send('This server name is already in my configuration. Please choose another.') @@ -291,12 +296,12 @@ class Rcon: All strings must be contained inside double quotes.""" if checks.is_rcon_admin(self.bot, ctx): server = server.title() - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) if server in rcon_connections: del rcon_connections[server] - self.bot.con.run('update guild_config set rcon_connections = %(connections)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'connections': json.dumps(rcon_connections)}) + self.bot.db_con.execute('update guild_config set rcon_connections = $2 where guild_id = $1', + ctx.guild.id, json.dumps(rcon_connections)) await ctx.send('{0} has been removed from my configuration.'.format(server)) else: await ctx.send('{0} is not in my configuration.'.format(server)) @@ -311,8 +316,8 @@ class Rcon: Example: 76561198024193239,76561198024193239,76561198024193239""" if checks.is_rcon_admin(self.bot, ctx): if steam_ids is not None: - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) error = 0 error_msg = '' success_msg = 'Adding to the running whitelist on all servers.' @@ -357,8 +362,8 @@ class Rcon: If a server is not specified it will default to running saveworld on all servers in the guild\'s config. Will print out "World Saved" for each server when the command completes successfully.""" if checks.is_rcon_admin(self.bot, ctx): - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) success_msg = 'Running saveworld' if server is None: success_msg += ' on all the servers:' @@ -369,7 +374,7 @@ class Rcon: await msg.edit(content=success_msg.strip()) message = await self.bot.loop.run_in_executor(None, self._saveworld, rcon_connections[server]) except Exception as e: - success_msg = '{0}\n{1}'.format(success_msg, e.strip()) + success_msg = '{0}\n{1}'.format(success_msg, e) await msg.edit(content=success_msg.strip()) else: success_msg = '{0}\n{1}'.format(success_msg, message.strip()) @@ -399,8 +404,8 @@ class Rcon: The message will be prefixed with the Discord name of the person running the command. Will print "Success" for each server once the broadcast is sent.""" if checks.is_rcon_admin(self.bot, ctx): - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) if message is not None: message = f'{ctx.author.display_name}: {message}' success_msg = f'Broadcasting "{message}" to all servers.' @@ -414,7 +419,7 @@ class Rcon: rcon_connections[server], message) except Exception as e: - success_msg = '{0}\n{1}'.format(success_msg, e.strip()) + success_msg = '{0}\n{1}'.format(success_msg, e) await msg.edit(content=success_msg.strip()) else: for mesg in messages: @@ -436,8 +441,8 @@ class Rcon: If has more than one word in it's name it will either need to be surrounded by double quotes or the words separated by _""" if checks.is_rcon_admin(self.bot, ctx): - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) if server is not None: server = server.replace('_', ' ').title() if message is not None: @@ -474,8 +479,8 @@ class Rcon: These channels will be added to the guild's rcon config and are where the server chat messages will be sent when monitor_chat is run.""" if checks.is_rcon_admin(self.bot, ctx): - rcon_connections = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + rcon_connections = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) edited = 0 category = discord.utils.get(ctx.guild.categories, name='Server Chats') if category is None: @@ -498,8 +503,8 @@ class Rcon: rcon_connections[server]['game_chat_chan_id'] = chan.id edited = 1 if edited == 1: - self.bot.con.run('update guild_config set rcon_connections = %(json)s where guild_id = %(id)s', - {'id': ctx.guild.id, 'json': json.dumps(rcon_connections)}) + self.bot.db_con.execute('update guild_config set rcon_connections = $2 where guild_id = $1', + ctx.guild.id, json.dumps(rcon_connections)) await ctx.message.add_reaction('✅') else: await ctx.send(f'You are not authorized to run this command.') @@ -509,8 +514,8 @@ class Rcon: @commands.check(checks.is_restricted_chan) async def list_ark_servers(self, ctx): """Returns a list of all the ARK servers in the current guild\'s config.""" - servers = json.loads(self.bot.con.one('select rcon_connections from guild_config\ - where guild_id = %(id)s', {'id': ctx.guild.id})) + servers = json.loads(self.bot.db_con.fetchval('select rcon_connections from guild_config ' + 'where guild_id = $1', ctx.guild.id)) em = discord.Embed(style='rich', title=f'__**There are currently {len(servers)} ARK servers in my config:**__', color=discord.Colour.green() diff --git a/exts/utils.py b/exts/utils.py index 430962a..4251a03 100644 --- a/exts/utils.py +++ b/exts/utils.py @@ -242,20 +242,20 @@ class Utils: if ctx.guild: if request_msg is not None: if len(request_msg) < 1000: - self.bot.con.run('insert into admin_requests (issuing_member_id, guild_orig, request_text,' - 'request_time) values (%(member_id)s, %(guild_id)s, %(text)s, %(time)s)', - {'member_id': ctx.author.id, 'guild_id': ctx.guild.id, 'text': request_msg, - 'time': ctx.message.created_at}) - channel = self.bot.con.one(f'select admin_chat from guild_config where guild_id = {ctx.guild.id}') + self.bot.db_con.execute('insert into admin_requests (issuing_member_id, guild_orig, request_text,' + 'request_time) values ($1, $2, $3, $4)', + ctx.author.id, ctx.guild.id, request_msg, ctx.message.created_at) + channel = self.bot.db_con.fetchval(f'select admin_chat from guild_config where guild_id = $1', + ctx.guild.id) if channel: chan = discord.utils.get(ctx.guild.channels, id=channel) msg = '' admin_roles = [] - roles = self.bot.con.one(f'select admin_roles,rcon_admin_roles from guild_config where ' - f'guild_id = %(id)s', {'id': ctx.guild.id}) - request_id = self.bot.con.one(f'select id from admin_requests where ' - f'issuing_member_id = %(member_id)s and request_time = %(time)s', - {'member_id': ctx.author.id, 'time': ctx.message.created_at}) + roles = self.bot.db_con.fetchval(f'select admin_roles,rcon_admin_roles from guild_config where ' + f'$1', ctx.guild.id) + request_id = self.bot.db_con.fetchval(f'select id from admin_requests where ' + f'issuing_member_id = $1 and request_time = $2', + ctx.author.id, ctx.message.created_at) for item in roles: i = json.loads(item) for j in i: @@ -294,8 +294,8 @@ class Utils: ) if checks.is_admin(self.bot, ctx) or checks.is_rcon_admin(self.bot, ctx): if assigned_to is None: - requests = self.bot.con.all(f'select * from admin_requests where guild_orig = %(guild_id)s ' - f'and completed_time is null', {'guild_id': ctx.guild.id}) + requests = self.bot.db_con.fetchall(f'select * from admin_requests where guild_orig = $1 ' + f'and completed_time is null', ctx.guild.id) em.title = f'Admin help requests for {ctx.guild.name}' if requests: for request in requests: @@ -314,9 +314,9 @@ class Utils: else: if checks.check_admin_role(self.bot, ctx, assigned_to)\ or checks.check_rcon_role(self.bot, ctx, assigned_to): - requests = self.bot.con.all('select * from admin_requests where assigned_to = %(admin_id)s ' - 'and guild_orig = %(guild_id)s and completed_time is null', - {'admin_id': assigned_to.id, 'guild_id': ctx.guild.id}) + requests = self.bot.db_con.fetchall('select * from admin_requests where assigned_to = $1 ' + 'and guild_orig = $2 and completed_time is null', + assigned_to.id, ctx.guild.id) em.title = f'Admin help requests assigned to {assigned_to.display_name} in {ctx.guild.name}' if requests: for request in requests: @@ -334,9 +334,9 @@ class Utils: else: em.title = f'{assigned_to.display_name} is not an admin in this guild.' else: - requests = self.bot.con.all('select * from admin_requests where issuing_member_id = %(member_id)s ' - 'and guild_orig = %(guild_id)s and completed_time is null', - {'member_id': ctx.author.id, 'guild_id': ctx.guild.id}) + requests = self.bot.db_con.fetchall('select * from admin_requests where issuing_member_id = $1 ' + 'and guild_orig = $2 and completed_time is null', + ctx.author.id, ctx.guild.id) em.title = f'Admin help requests for {ctx.author.display_name}' if requests: for request in requests: @@ -367,14 +367,12 @@ class Utils: except ValueError: await ctx.send(f'{request_id} is not a valid request id.') else: - request = self.bot.con.one(f'select * from admin_requests where id = %(request_id)s', - {'request_id': request_id}) + request = self.bot.db_con.fetchval(f'select * from admin_requests where id = $1', request_id) if request: if request[3] == ctx.guild.id: if request[6] is None: - self.bot.con.run('update admin_requests set completed_time = %(time_now)s where ' - 'id = %(request_id)s', - {'time_now': ctx.message.created_at, 'request_id': request_id}) + self.bot.db_con.execute('update admin_requests set completed_time = $1 where ' + 'id = $2', ctx.message.created_at, request_id) await ctx.send(f'Request {request_id} by ' f'{ctx.guild.get_member(request[1]).display_name}' f' has been marked complete.') @@ -456,8 +454,7 @@ class Utils: def is_me(message): if message.author == self.bot.user: return True - prefixes = self.bot.db_con.fetchval('select prefix from guild_config where guild_id = $1', - ctx.guild.id) + prefixes = self.bot.db_con.fetchval('select prefix from guild_config where guild_id = $1', ctx.guild.id) if prefixes: for prefix in prefixes: if message.content.startswith(prefix): diff --git a/geeksbot.py b/geeksbot.py index a4f9822..c0af47b 100644 --- a/geeksbot.py +++ b/geeksbot.py @@ -8,6 +8,7 @@ import aiohttp from googleapiclient.discovery import build import asyncpg from concurrent import futures +import asyncio log_format = '{asctime}.{msecs:03.0f}|{levelname:<8}|{name}::{message}' @@ -56,8 +57,15 @@ class Geeksbot(commands.Bot): self.guild_config = {} self.infected = {} self.TOKEN = self.bot_secrets['token'] + + async def connect_db(): + return await asyncpg.create_pool(host=self.bot_secrets['db_con']['host'], + database=self.bot_secrets['db_con']['db_name'], + user=self.bot_secrets['db_con']['user'], + password=self.bot_secrets['db_con']['password'], + loop=asyncio.get_event_loop()) del self.bot_secrets['token'] - self.db_con = None + self.db_con = asyncio.get_event_loop().create_task(connect_db()) self.default_prefix = 'g~' self.voice_chans = {} self.spam_list = {} @@ -72,13 +80,6 @@ class Geeksbot(commands.Bot): 'left_fist': '🤛', } - async def connect_db(self): - self.db_con = await asyncpg.create_pool(host=self.bot_secrets['db_con']['host'], - database=self.bot_secrets['db_con']['db_name'], - user=self.bot_secrets['db_con']['user'], - password=self.bot_secrets['db_con']['password'], - loop=self.loop) - @staticmethod async def get_custom_prefix(bot_inst, message): return await bot_inst.db_con.fetchval('select prefix from guild_config where guild_id = $1', @@ -153,7 +154,6 @@ async def on_message(ctx): @bot.event async def on_ready(): - if bot.db_con is None: await bot.connect_db() bot.recent_msgs = {} logging.info('Logged in as {0.name}|{0.id}'.format(bot.user)) load_list = bot.bot_config['load_list'] diff --git a/misc.py b/misc.py index bf16706..6960d6c 100644 --- a/misc.py +++ b/misc.py @@ -32,4 +32,13 @@ info = func() bot.voice_chan.play(discord.FFmpegPCMAudio(info['url'])) #async while bot.voice_chan.is_playing(): # pass -#await bot.voice_chan.disconnect() \ No newline at end of file +#await bot.voice_chan.disconnect() + +# Run event in loop after number of seconds +from functools import partial +return bot.loop.call_later(120, partial(bot.loop.create_task, ctx.send(f"{ctx.author.mention} Timer's Up"))) + +# Get the number of tasks currently in the loop +import asyncio +return len(asyncio.Task.all_tasks()) +