diff --git a/requirements.txt b/requirements.txt index 56555fc6..a34bcc40 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ cachetools==4.2.2 configparser==5.0.2 discord.py iso8601==0.1.16 -psycopg +psycopg[pool] pytz==2021.1 topggpy psutil diff --git a/src/analytics/snapshot.py b/src/analytics/snapshot.py index 565280e1..789439eb 100644 --- a/src/analytics/snapshot.py +++ b/src/analytics/snapshot.py @@ -22,7 +22,7 @@ async def shard_snapshot(): snap = ShardSnapshot( guild_count=len(bot.guilds), voice_count=sum(len(channel.members) for guild in bot.guilds for channel in guild.voice_channels), - member_count=sum(len(guild.members) for guild in bot.guilds), + member_count=sum(guild.member_count for guild in bot.guilds), user_count=len(set(m.id for guild in bot.guilds for m in guild.members)) ) return snap diff --git a/src/bot.py b/src/bot.py index 1e47a9e4..8fda839d 100644 --- a/src/bot.py +++ b/src/bot.py @@ -69,7 +69,8 @@ async def main(): shard_count=sharding.shard_count, help_command=None, proxy=conf.bot.get('proxy', None), - translator=translator + translator=translator, + chunk_guilds_at_startup=False, ) as lionbot: ctx_bot.set(lionbot) try: diff --git a/src/meta/LionBot.py b/src/meta/LionBot.py index f316c24e..cb15a9a2 100644 --- a/src/meta/LionBot.py +++ b/src/meta/LionBot.py @@ -240,7 +240,7 @@ class LionBot(Bot): details['guild'] = f"`{ctx.guild.id}` -- `{ctx.guild.name}`" details['my_guild_perms'] = f"`{ctx.guild.me.guild_permissions.value}`" if ctx.author: - ownerstr = ' (owner)' if ctx.author == ctx.guild.owner else '' + ownerstr = ' (owner)' if ctx.author.id == ctx.guild.owner_id else '' details['author_guild_perms'] = f"`{ctx.author.guild_permissions.value}{ownerstr}`" if ctx.channel.type is discord.enums.ChannelType.private: details['channel'] = "`Direct Message`" @@ -281,3 +281,29 @@ class LionBot(Bot): def add_command(self, command): if not hasattr(command, '_placeholder_group_'): super().add_command(command) + + def request_chunking_for(self, guild): + if not guild.chunked: + return asyncio.create_task( + self._connection.chunk_guild(guild, wait=False, cache=True), + name=f"Background chunkreq for {guild.id}" + ) + + async def on_interaction(self, interaction: discord.Interaction): + """ + Adds the interaction author to guild cache if appropriate. + + This gets run a little bit late, so it is possible the interaction gets handled + without the author being in case. + """ + guild = interaction.guild + user = interaction.user + if guild is not None and user is not None and isinstance(user, discord.Member): + if not guild.get_member(user.id): + guild._add_member(user) + if guild is not None and not guild.chunked: + # Getting an interaction in the guild is a good enough reason to request chunking + logger.info( + f"Unchunked guild requesting chunking after interaction." + ) + self.request_chunking_for(guild) diff --git a/src/meta/LionTree.py b/src/meta/LionTree.py index 4461ef2e..d8275e32 100644 --- a/src/meta/LionTree.py +++ b/src/meta/LionTree.py @@ -82,7 +82,7 @@ class LionTree(CommandTree): details['guild'] = f"`{interaction.guild.id}` -- `{interaction.guild.name}`" details['my_guild_perms'] = f"`{interaction.guild.me.guild_permissions.value}`" if interaction.user: - ownerstr = ' (owner)' if interaction.user == interaction.guild.owner else '' + ownerstr = ' (owner)' if interaction.user.id == interaction.guild.owner_id else '' details['user_guild_perms'] = f"`{interaction.user.guild_permissions.value}{ownerstr}`" if interaction.channel.type is discord.enums.ChannelType.private: details['channel'] = "`Direct Message`" diff --git a/src/modules/economy/cog.py b/src/modules/economy/cog.py index aef2b8b9..9e72e4ce 100644 --- a/src/modules/economy/cog.py +++ b/src/modules/economy/cog.py @@ -190,7 +190,7 @@ class Economy(LionCog): # First fetch the members which currently exist query = self.bot.core.data.Member.table.select_where(guildid=ctx.guild.id) query.select('userid').with_no_adapter() - if 2 * len(targets) < len(ctx.guild.members): + if 2 * len(targets) < ctx.guild.member_count: # More efficient to fetch the targets explicitly query.where(userid=list(targetids)) existent_rows = await query diff --git a/src/modules/member_admin/cog.py b/src/modules/member_admin/cog.py index 6d57f222..6887ed6d 100644 --- a/src/modules/member_admin/cog.py +++ b/src/modules/member_admin/cog.py @@ -181,15 +181,17 @@ class MemberAdminCog(LionCog): finally: self._adding_roles.discard((member.guild.id, member.id)) - @LionCog.listener('on_member_remove') + @LionCog.listener('on_raw_member_remove') @log_wrap(action="Farewell") - async def admin_member_farewell(self, member: discord.Member): + async def admin_member_farewell(self, payload: discord.RawMemberRemoveEvent): # Ignore members that just joined - if (member.guild.id, member.id) in self._adding_roles: + guildid = payload.guild_id + userid = payload.user.id + if (guildid, userid) in self._adding_roles: return # Set lion last_left, creating the lion_member if needed - lion = await self.bot.core.lions.fetch_member(member.guild.id, member.id) + lion = await self.bot.core.lions.fetch_member(guildid, userid) await lion.data.update(last_left=utc_now()) # Save member roles @@ -197,18 +199,21 @@ class MemberAdminCog(LionCog): self.bot.db.conn = conn async with conn.transaction(): await self.data.past_roles.delete_where( - guildid=member.guild.id, - userid=member.id + guildid=guildid, + userid=userid ) # Insert current member roles - if member.roles: + print(type(payload.user)) + if isinstance(payload.user, discord.Member) and payload.user.roles: + member = payload.user await self.data.past_roles.insert_many( ('guildid', 'userid', 'roleid'), - *((member.guild.id, member.id, role.id) for role in member.roles) + *((guildid, userid, role.id) for role in member.roles) ) logger.debug( - f"Stored persisting roles for member in ." + f"Stored persisting roles for member in ." ) + # TODO: Event log, and include info about unchunked members @LionCog.listener('on_guild_join') async def admin_init_guild(self, guild: discord.Guild): diff --git a/src/modules/member_admin/settings.py b/src/modules/member_admin/settings.py index d175c275..deaba142 100644 --- a/src/modules/member_admin/settings.py +++ b/src/modules/member_admin/settings.py @@ -173,7 +173,7 @@ class MemberAdminSettings(SettingGroup): '{guild_name}': guild.name, '{guild_icon}': guild.icon.url if guild.icon else member.default_avatar.url, '{studying_count}': str(active), - '{member_count}': len(guild.members), + '{member_count}': guild.member_count, } recurse_map( @@ -297,7 +297,7 @@ class MemberAdminSettings(SettingGroup): '{guild_name}': guild.name, '{guild_icon}': guild.icon.url if guild.icon else member.default_avatar.url, '{studying_count}': str(active), - '{member_count}': str(len(guild.members)), + '{member_count}': str(guild.member_count), '{last_time}': str(last_seen or member.joined_at.timestamp()), } diff --git a/src/modules/ranks/cog.py b/src/modules/ranks/cog.py index 8e245cf6..5f7e9866 100644 --- a/src/modules/ranks/cog.py +++ b/src/modules/ranks/cog.py @@ -503,7 +503,15 @@ class RankCog(LionCog): # Ensure guild is chunked if not guild.chunked: - members = await guild.chunk() + try: + members = await asyncio.wait_for(guild.chunk(), timeout=60) + except asyncio.TimeoutError: + error = t(_p( + 'rank_refresh|error:cannot_chunk|desc', + "Could not retrieve member list from Discord. Please try again later." + )) + await ui.set_error(error) + return else: members = guild.members ui.stage_members = True diff --git a/src/modules/schedule/core/session.py b/src/modules/schedule/core/session.py index cd884c72..da266a9b 100644 --- a/src/modules/schedule/core/session.py +++ b/src/modules/schedule/core/session.py @@ -253,6 +253,12 @@ class ScheduledSession: overwrites = room.overwrites for member in members: mobj = guild.get_member(member.userid) + if not mobj and not guild.chunked: + self.bot.request_chunking_for(guild) + try: + mobj = await guild.fetch_member(member.userid) + except discord.HTTPException: + mobj = None if mobj: overwrites[mobj] = discord.PermissionOverwrite(connect=True, view_channel=True) try: @@ -297,6 +303,13 @@ class ScheduledSession: } for member in members: mobj = guild.get_member(member.userid) + if not mobj and not guild.chunked: + self.bot.request_chunking_for(guild) + try: + mobj = await guild.fetch_member(member.userid) + except discord.HTTPException: + mobj = None + if mobj: overwrites[mobj] = discord.PermissionOverwrite(connect=True, view_channel=True) try: diff --git a/src/modules/schedule/core/timeslot.py b/src/modules/schedule/core/timeslot.py index 1669b780..200c3e1b 100644 --- a/src/modules/schedule/core/timeslot.py +++ b/src/modules/schedule/core/timeslot.py @@ -440,7 +440,7 @@ class TimeSlot: ) def launch(self) -> asyncio.Task: - self.run_task = asyncio.create_task(self.run()) + self.run_task = asyncio.create_task(self.run(), name=f"TimeSlot {self.slotid}") return self.run_task @log_wrap(action="TimeSlot Run") diff --git a/src/modules/statistics/cog.py b/src/modules/statistics/cog.py index d572f5e2..1ef512d5 100644 --- a/src/modules/statistics/cog.py +++ b/src/modules/statistics/cog.py @@ -1,3 +1,4 @@ +import asyncio import logging from typing import Optional @@ -91,7 +92,11 @@ class StatsCog(LionCog): timestamp=utc_now(), ) await ctx.interaction.response.send_message(embed=waiting_embed) - await ctx.guild.chunk() + try: + await asyncio.wait_for(ctx.guild.chunk(), timeout=10) + pass + except asyncio.TimeoutError: + pass else: await ctx.interaction.response.defer(thinking=True) ui = LeaderboardUI(self.bot, ctx.author, ctx.guild) diff --git a/src/modules/statistics/ui/leaderboard.py b/src/modules/statistics/ui/leaderboard.py index 4dadbde8..bf2f9206 100644 --- a/src/modules/statistics/ui/leaderboard.py +++ b/src/modules/statistics/ui/leaderboard.py @@ -75,6 +75,8 @@ class LeaderboardUI(StatsUI): # (type, period) -> (pagen -> Optional[Future[Card]]) self.cache = {} + self.was_chunked: bool = guild.chunked + async def run(self, interaction: discord.Interaction): self._original = interaction @@ -136,6 +138,7 @@ class LeaderboardUI(StatsUI): # Filter out members which are not in the server and unranked roles and bots # Usually hits cache + self.was_chunked = self.guild.chunked unranked_setting = await self.bot.get_cog('StatsCog').settings.UnrankedRoles.get(self.guild.id) unranked_roleids = set(unranked_setting.data) true_leaderboard = [] @@ -435,12 +438,19 @@ class LeaderboardUI(StatsUI): Generate UI message arguments from stored data """ t = self.bot.translator.t + chunk_warning = t(_p( + 'ui:leaderboard|chunk_warning', + "**Note:** Could not retrieve member list from Discord, so some members may be missing. " + "Try again in a minute!" + )) if self.card is not None: period_start = self.period_starts[self.current_period] header = t(_p( 'ui:leaderboard|since', "Counting statistics since {timestamp}" )).format(timestamp=discord.utils.format_dt(period_start)) + if not self.was_chunked: + header = '\n'.join((header, chunk_warning)) args = MessageArgs( embed=None, content=header, @@ -473,7 +483,11 @@ class LeaderboardUI(StatsUI): )), description=empty_description ) - args = MessageArgs(content=None, embed=embed, files=[]) + args = MessageArgs( + content=chunk_warning if not self.was_chunked else None, + embed=embed, + files=[] + ) return args async def refresh_components(self): diff --git a/src/modules/sysadmin/guild_log.py b/src/modules/sysadmin/guild_log.py index a12d24e2..2e38248a 100644 --- a/src/modules/sysadmin/guild_log.py +++ b/src/modules/sysadmin/guild_log.py @@ -1,3 +1,4 @@ +import asyncio import datetime import discord @@ -22,8 +23,8 @@ class GuildLog(LionCog): embed.set_author(name="Left guild!") # Add more specific information about the guild - embed.add_field(name="Owner", value="{0.name} (ID: {0.id})".format(guild.owner), inline=False) - embed.add_field(name="Members (cached)", value="{}".format(len(guild.members)), inline=False) + embed.add_field(name="Owner", value="<@{}>".format(guild.owner_id), inline=False) + embed.add_field(name="Members", value="{}".format(guild.member_count), inline=False) embed.add_field(name="Now studying in", value="{} guilds".format(len(self.bot.guilds)), inline=False) # Retrieve the guild log channel and log the event @@ -35,39 +36,51 @@ class GuildLog(LionCog): @LionCog.listener('on_guild_join') @log_wrap(action="Log Guild Join") async def log_join_guild(self, guild: discord.Guild): - owner = guild.owner + try: + await asyncio.wait_for(guild.chunk(), timeout=60) + except asyncio.TimeoutError: + pass - bots = 0 - known = 0 - unknown = 0 - other_members = set(mem.id for mem in self.bot.get_all_members() if mem.guild != guild) + # TODO: Add info about when we last joined this guild etc once we have it. - for member in guild.members: - if member.bot: - bots += 1 - elif member.id in other_members: - known += 1 - else: - unknown += 1 + if guild.chunked: + bots = 0 + known = 0 + unknown = 0 + other_members = set(mem.id for mem in self.bot.get_all_members() if mem.guild != guild) + + for member in guild.members: + if member.bot: + bots += 1 + elif member.id in other_members: + known += 1 + else: + unknown += 1 + + mem1 = "people I know" if known != 1 else "person I know" + mem2 = "new friends" if unknown != 1 else "new friend" + mem3 = "bots" if bots != 1 else "bot" + mem4 = "total members" + known = "`{}`".format(known) + unknown = "`{}`".format(unknown) + bots = "`{}`".format(bots) + total = "`{}`".format(guild.member_count) + mem_str = "{0:<5}\t{4},\n{1:<5}\t{5},\n{2:<5}\t{6}, and\n{3:<5}\t{7}.".format( + known, + unknown, + bots, + total, + mem1, + mem2, + mem3, + mem4 + ) + else: + mem_str = ( + "`{count}` total members.\n" + "(Could not chunk guild within `60` seconds.)" + ).format(count=guild.member_count) - mem1 = "people I know" if known != 1 else "person I know" - mem2 = "new friends" if unknown != 1 else "new friend" - mem3 = "bots" if bots != 1 else "bot" - mem4 = "total members" - known = "`{}`".format(known) - unknown = "`{}`".format(unknown) - bots = "`{}`".format(bots) - total = "`{}`".format(guild.member_count) - mem_str = "{0:<5}\t{4},\n{1:<5}\t{5},\n{2:<5}\t{6}, and\n{3:<5}\t{7}.".format( - known, - unknown, - bots, - total, - mem1, - mem2, - mem3, - mem4 - ) created = "".format(int(guild.created_at.timestamp())) embed = discord.Embed( @@ -77,7 +90,7 @@ class GuildLog(LionCog): ) embed.set_author(name="Joined guild!") - embed.add_field(name="Owner", value="{0} (ID: {0.id})".format(owner), inline=False) + embed.add_field(name="Owner", value="<@{}>".format(guild.owner_id), inline=False) embed.add_field(name="Created at", value=created, inline=False) embed.add_field(name="Members", value=mem_str, inline=False) embed.add_field(name="Now studying in", value="{} guilds".format(len(self.bot.guilds)), inline=False) diff --git a/src/wards.py b/src/wards.py index 2ac290e1..5db46eb9 100644 --- a/src/wards.py +++ b/src/wards.py @@ -32,7 +32,7 @@ async def high_management(bot: LionBot, member: discord.Member, guild: discord.G async def low_management(bot: LionBot, member: discord.Member, guild: discord.Guild): if not guild: return True - if await high_management(bot, member): + if await high_management(bot, member, guild): return True return member.guild_permissions.manage_guild @@ -196,7 +196,7 @@ async def equippable_role(bot: LionBot, target_role: discord.Role, actor: discor "You need the `MANAGE_ROLES` permission before you can configure roles!" )).format(role=target_role.mention) ) - elif actor.top_role <= target_role and not actor == guild.owner: + elif actor.top_role <= target_role and not actor.id == guild.owner_id: raise UserInputError( t(_p( 'ward:equippable_role|error:actor_top_role',