hendrik/pollmaster-docker
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

fix crash, improve !cmd help and errors, add database locks for high traffic servers, cache prefixes, add production/development setting

Browse Source
This commit is contained in:
matnad 2019-02-20 16:12:50 +01:00
parent 32a70bb01d
commit ac1f3051ef
6 changed files with 139 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
cogs/config.py
View File

@ -29,6 +29,7 @@ class Config:
f'If you would like to add a trailing whitespace to the prefix, use `{pre}prefix {pre}\w`.' f'If you would like to add a trailing whitespace to the prefix, use `{pre}prefix {pre}\w`.'
await self.bot.db.config.update_one({'_id': str(server.id)}, {'$set': {'prefix': str(pre)}}, upsert=True) await self.bot.db.config.update_one({'_id': str(server.id)}, {'$set': {'prefix': str(pre)}}, upsert=True)
self.bot.pre[str(server.id)] = str(pre)
await self.bot.say(msg) await self.bot.say(msg)
@commands.command(pass_context=True) @commands.command(pass_context=True)

3
cogs/db_api.py
View File

@ -20,7 +20,8 @@ class DiscordBotsOrgAPI:
while True: while True:
logger.info('attempting to post server count') logger.info('attempting to post server count')
try: try:
await self.dblpy.post_server_count() if SETTINGS.mode == 'production':
await self.dblpy.post_server_count()
logger.info('posted server count ({})'.format(len(self.bot.servers))) logger.info('posted server count ({})'.format(len(self.bot.servers)))
except Exception as e: except Exception as e:
logger.exception('Failed to post server count\n{}: {}'.format(type(e).__name__, e)) logger.exception('Failed to post server count\n{}: {}'.format(type(e).__name__, e))

217
cogs/poll_controls.py
View File

@ -1,4 +1,5 @@
import argparse import argparse
import asyncio
import copy import copy
import datetime import datetime
import json import json
@ -24,6 +25,11 @@ class PollControls:
self.bot = bot self.bot = bot
# General Methods # General Methods
def get_lock(self, server_id):
if not self.bot.locks.get(str(server_id)):
self.bot.locks[server_id] = asyncio.Lock()
return self.bot.locks.get(str(server_id))
async def is_admin_or_creator(self, ctx, server, owner_id, error_msg=None): async def is_admin_or_creator(self, ctx, server, owner_id, error_msg=None):
member = server.get_member(ctx.message.author.id) member = server.get_member(ctx.message.author.id)
if member.id == owner_id: if member.id == owner_id:
@ -288,11 +294,23 @@ class PollControls:
return return
try: try:
args = parser.parse_args(cmds) args, unknown_args = parser.parse_known_args(cmds)
except SystemExit: except SystemExit:
await self.say_error(ctx, error_text=helpstring) await self.say_error(ctx, error_text=helpstring)
return return
if unknown_args:
error_text = f'**There was an error reading the command line options!**.\n' \
f'Most likely this is because you didn\'t surround the arguments with double quotes like this: ' \
f'`{pre}cmd -q "question of the poll" -o "yes, no, maybe"`' \
f'\n\nHere are the arguments I could not understand:\n'
error_text += '`'+'\n'.join(unknown_args)+'`'
error_text += f'\n\nHere are the arguments which are ok:\n'
error_text += '`' + '\n'.join([f'{k}: {v}' for k, v in vars(args).items()]) + '`'
await self.say_error(ctx, error_text=error_text, footer_text=f'type `{pre}cmd help` for details.')
return
# pass arguments to the wizard # pass arguments to the wizard
async def route(poll): async def route(poll):
await poll.set_name(force=args.question) await poll.set_name(force=args.question)
@ -306,7 +324,7 @@ class PollControls:
poll = await self.wizard(ctx, route, server) poll = await self.wizard(ctx, route, server)
if poll: if poll:
await poll.post_embed() await poll.post_embed(destination=poll.channel)
@commands.command(pass_context=True) @commands.command(pass_context=True)
@ -328,7 +346,7 @@ class PollControls:
poll = await self.wizard(ctx, route, server) poll = await self.wizard(ctx, route, server)
if poll: if poll:
await poll.post_embed() await poll.post_embed(destination=poll.channel)
@commands.command(pass_context=True) @commands.command(pass_context=True)
async def prepare(self, ctx, *, cmd=None): async def prepare(self, ctx, *, cmd=None):
@ -371,7 +389,7 @@ class PollControls:
poll = await self.wizard(ctx, route, server) poll = await self.wizard(ctx, route, server)
if poll: if poll:
await poll.post_embed() await poll.post_embed(destination=poll.channel)
# The Wizard! # The Wizard!
async def wizard(self, ctx, route, server): async def wizard(self, ctx, route, server):
@ -460,14 +478,17 @@ class PollControls:
user_msg = copy.deepcopy(message) user_msg = copy.deepcopy(message)
user_msg.author = user user_msg.author = user
server = await ask_for_server(self.bot, user_msg, label) server = await ask_for_server(self.bot, user_msg, label)
p = await Poll.load_from_db(self.bot, server.id, label)
if not isinstance(p, Poll):
return
if not p.anonymous: # this is exclusive
# for anonymous polls we can't unvote because we need to hide reactions async with self.get_lock(server.id):
member = server.get_member(user_id) p = await Poll.load_from_db(self.bot, server.id, label)
await p.unvote(member, emoji, message) if not isinstance(p, Poll):
return
if not p.anonymous:
# for anonymous polls we can't unvote because we need to hide reactions
member = server.get_member(user_id)
await p.unvote(member, emoji, message)
async def do_on_reaction_add(self, data): async def do_on_reaction_add(self, data):
@ -510,103 +531,107 @@ class PollControls:
user_msg = copy.deepcopy(message) user_msg = copy.deepcopy(message)
user_msg.author = user user_msg.author = user
server = await ask_for_server(self.bot, user_msg, label) server = await ask_for_server(self.bot, user_msg, label)
p = await Poll.load_from_db(self.bot, server.id, label)
if not isinstance(p, Poll):
return
# export # this is exclusive to keep database access sequential
if emoji == '📎': # hopefully it will scale well enough or I need a different solution
# sending file async with self.get_lock(server.id):
file = await p.export() p = await Poll.load_from_db(self.bot, server.id, label)
if file is not None: if not isinstance(p, Poll):
await self.bot.send_file( return
user,
file,
content='Sending you the requested export of "{}".'.format(p.short)
)
return
# info # export
member = server.get_member(user_id) if emoji == '📎':
if emoji == '❔': # sending file
is_open = await p.is_open() file = await p.export()
embed = discord.Embed(title=f"Info for the {'CLOSED ' if not is_open else ''}poll \"{p.name}\"", if file is not None:
description='', color=SETTINGS.color) await self.bot.send_file(
embed.set_author(name=f" >> {p.short}", icon_url=SETTINGS.author_icon) user,
file,
content='Sending you the requested export of "{}".'.format(p.short)
)
return
# vote rights # info
vote_rights = await p.has_required_role(member) member = server.get_member(user_id)
embed.add_field(name=f'{"Can you vote?" if is_open else "Could you vote?"}', if emoji == '❔':
value=f'{"✅" if vote_rights else "❎"}', inline=False) is_open = await p.is_open()
embed = discord.Embed(title=f"Info for the {'CLOSED ' if not is_open else ''}poll \"{p.name}\"",
description='', color=SETTINGS.color)
embed.set_author(name=f" >> {p.short}", icon_url=SETTINGS.author_icon)
# edit rights # vote rights
edit_rights = False vote_rights = await p.has_required_role(member)
if str(member.id) == str(p.author): embed.add_field(name=f'{"Can you vote?" if is_open else "Could you vote?"}',
edit_rights = True value=f'{"✅" if vote_rights else "❎"}', inline=False)
elif member.server_permissions.manage_server:
edit_rights = True # edit rights
else: edit_rights = False
result = await self.bot.db.config.find_one({'_id': str(server.id)}) if str(member.id) == str(p.author):
if result and result.get('admin_role') in [r.name for r in member.roles]:
edit_rights = True edit_rights = True
embed.add_field(name='Can you manage the poll?', value=f'{"✅" if edit_rights else "❎"}', inline=False) elif member.server_permissions.manage_server:
edit_rights = True
else:
result = await self.bot.db.config.find_one({'_id': str(server.id)})
if result and result.get('admin_role') in [r.name for r in member.roles]:
edit_rights = True
embed.add_field(name='Can you manage the poll?', value=f'{"✅" if edit_rights else "❎"}', inline=False)
# choices # choices
choices = 'You have not voted yet.' if vote_rights else 'You can\'t vote in this poll.' choices = 'You have not voted yet.' if vote_rights else 'You can\'t vote in this poll.'
if user.id in p.votes: if user.id in p.votes:
if p.votes[user.id]['choices'].__len__() > 0: if p.votes[user.id]['choices'].__len__() > 0:
choices = ', '.join([p.options_reaction[c] for c in p.votes[user.id]['choices']]) choices = ', '.join([p.options_reaction[c] for c in p.votes[user.id]['choices']])
embed.add_field(name=f'{"Your current votes (can be changed as long as the poll is open):" if is_open else "Your final votes:"}', embed.add_field(name=f'{"Your current votes (can be changed as long as the poll is open):" if is_open else "Your final votes:"}',
value=choices, inline=False) value=choices, inline=False)
# weight # weight
if vote_rights: if vote_rights:
weight = 1 weight = 1
if p.weights_roles.__len__() > 0: if p.weights_roles.__len__() > 0:
valid_weights = [p.weights_numbers[p.weights_roles.index(r)] for r in valid_weights = [p.weights_numbers[p.weights_roles.index(r)] for r in
list(set([n.name for n in member.roles]).intersection(set(p.weights_roles)))] list(set([n.name for n in member.roles]).intersection(set(p.weights_roles)))]
if valid_weights.__len__() > 0: if valid_weights.__len__() > 0:
weight = max(valid_weights) weight = max(valid_weights)
else: else:
weight = 'You can\'t vote in this poll.' weight = 'You can\'t vote in this poll.'
embed.add_field(name='Weight of your votes:', value=weight, inline=False) embed.add_field(name='Weight of your votes:', value=weight, inline=False)
# time left # time left
deadline = p.get_duration_with_tz() deadline = p.get_duration_with_tz()
if not is_open: if not is_open:
time_left = 'This poll is closed.' time_left = 'This poll is closed.'
elif deadline == 0: elif deadline == 0:
time_left = 'Until manually closed.' time_left = 'Until manually closed.'
else: else:
time_left = str(deadline-datetime.datetime.utcnow().replace(tzinfo=pytz.utc)).split('.', 2)[0] time_left = str(deadline-datetime.datetime.utcnow().replace(tzinfo=pytz.utc)).split('.', 2)[0]
embed.add_field(name='Time left in the poll:', value=time_left, inline=False) embed.add_field(name='Time left in the poll:', value=time_left, inline=False)
await self.bot.send_message(user, embed=embed) await self.bot.send_message(user, embed=embed)
return return
# Assume: User wants to vote with reaction # Assume: User wants to vote with reaction
# no rights, terminate function # no rights, terminate function
if not await p.has_required_role(member): if not await p.has_required_role(member):
await self.bot.remove_reaction(message, emoji, user)
await self.bot.send_message(user, f'You are not allowed to vote in this poll. Only users with '
f'at least one of these roles can vote:\n{", ".join(p.roles)}')
return
# order here is crucial since we can't determine if a reaction was removed by the bot or user
# update database with vote
await p.vote(member, emoji, message)
#
# check if we need to remove reactions (this will trigger on_reaction_remove)
if str(channel.type) != 'private':
if p.anonymous:
# immediately remove reaction and to be safe, remove all reactions
await self.bot.remove_reaction(message, emoji, user) await self.bot.remove_reaction(message, emoji, user)
elif p.multiple_choice == 1: await self.bot.send_message(user, f'You are not allowed to vote in this poll. Only users with '
# remove all other reactions f'at least one of these roles can vote:\n{", ".join(p.roles)}')
for r in message.reactions: return
if r.emoji and r.emoji != emoji:
await self.bot.remove_reaction(message, r.emoji, user) # order here is crucial since we can't determine if a reaction was removed by the bot or user
# update database with vote
await p.vote(member, emoji, message)
#
# check if we need to remove reactions (this will trigger on_reaction_remove)
if str(channel.type) != 'private':
if p.anonymous:
# immediately remove reaction and to be safe, remove all reactions
await self.bot.remove_reaction(message, emoji, user)
elif p.multiple_choice == 1:
# remove all other reactions
for r in message.reactions:
if r.emoji and r.emoji != emoji:
await self.bot.remove_reaction(message, r.emoji, user)
def setup(bot): def setup(bot):
global logger global logger

11
essentials/multi_server.py
View File

@ -11,7 +11,7 @@ async def get_pre(bot, message):
if str(message.channel.type) == 'private': if str(message.channel.type) == 'private':
shared_server_list = await get_servers(bot, message) shared_server_list = await get_servers(bot, message)
if shared_server_list.__len__() == 0: if shared_server_list.__len__() == 0:
return 'pm!!' return 'pm!'
elif shared_server_list.__len__() == 1: elif shared_server_list.__len__() == 1:
return await get_server_pre(bot, shared_server_list[0]) return await get_server_pre(bot, shared_server_list[0])
else: else:
@ -24,12 +24,13 @@ async def get_pre(bot, message):
async def get_server_pre(bot, server): async def get_server_pre(bot, server):
'''Gets the prefix for a server.''' '''Gets the prefix for a server.'''
try: try:
result = await bot.db.config.find_one({'_id': str(server.id)}) #result = await bot.db.config.find_one({'_id': str(server.id)})
result = bot.pre[str(server.id)]
except AttributeError: except AttributeError:
return 'pm!' return 'pm!'
if not result or not result.get('prefix'): if not result: #or not result.get('prefix'):
return 'pm!!' return 'pm!'
return result.get('prefix') return result #result.get('prefix')
async def get_servers(bot, message, short=None): async def get_servers(bot, message, short=None):

1
essentials/settings.py
View File

@ -22,6 +22,7 @@ class Settings:
self.dbl_token = SECRETS.dbl_token self.dbl_token = SECRETS.dbl_token
self.mongo_db = SECRETS.mongo_db self.mongo_db = SECRETS.mongo_db
self.bot_token = SECRETS.bot_token self.bot_token = SECRETS.bot_token
self.mode = SECRETS.mode
SETTINGS = Settings() SETTINGS = Settings()

8
pollmaster.py
View File

@ -1,3 +1,4 @@
import asyncio
import traceback import traceback
import logging import logging
import aiohttp import aiohttp
@ -74,6 +75,10 @@ async def on_ready():
except: except:
print("Problem verifying servers.") print("Problem verifying servers.")
# cache prefixes
bot.pre = {entry['_id']:entry['prefix'] async for entry in bot.db.config.find({}, {'_id', 'prefix'})}
bot.locks = {sid: asyncio.Lock() for sid in [s.id for s in bot.servers]}
print("Servers verified. Bot running.") print("Servers verified. Bot running.")
@ -97,6 +102,8 @@ async def on_command_error(e, ctx):
# log error # log error
logger.error(f'{type(e).__name__}: {e}\n{"".join(traceback.format_tb(e.__traceback__))}') logger.error(f'{type(e).__name__}: {e}\n{"".join(traceback.format_tb(e.__traceback__))}')
if SETTINGS.mode == 'development':
raise e
if SETTINGS.msg_errors: if SETTINGS.msg_errors:
# send discord message for unexpected errors # send discord message for unexpected errors
@ -119,6 +126,7 @@ async def on_server_join(server):
{'$set': {'prefix': 'pm!', 'admin_role': 'polladmin', 'user_role': 'polluser'}}, {'$set': {'prefix': 'pm!', 'admin_role': 'polladmin', 'user_role': 'polluser'}},
upsert=True upsert=True
) )
bot.pre[str(server.id)] = 'pm!'
bot.run(SETTINGS.bot_token, reconnect=True) bot.run(SETTINGS.bot_token, reconnect=True)