diff --git a/bot/LionModule.py b/bot/LionModule.py index c1fd7256..9c2a4671 100644 --- a/bot/LionModule.py +++ b/bot/LionModule.py @@ -81,17 +81,17 @@ class LionModule(Module): pass raise SafeCancellation(details="Module '{}' is not ready.".format(self.name)) - # Check that the channel and guild still exists - if not ctx.client.get_guild(ctx.guild.id) or not ctx.guild.get_channel(ctx.ch.id): - raise SafeCancellation(details='Command channel is no longer reachable.') - # Check global user blacklist - if ctx.author.id in ctx.client.objects['blacklisted_users']: + if ctx.author.id in ctx.client.user_blacklist(): raise SafeCancellation(details='User is blacklisted.') if ctx.guild: + # Check that the channel and guild still exists + if not ctx.client.get_guild(ctx.guild.id) or not ctx.guild.get_channel(ctx.ch.id): + raise SafeCancellation(details='Command channel is no longer reachable.') + # Check global guild blacklist - if ctx.guild.id in ctx.client.objects['blacklisted_guilds']: + if ctx.guild.id in ctx.client.guild_blacklist(): raise SafeCancellation(details='Guild is blacklisted.') # Check guild's own member blacklist diff --git a/bot/constants.py b/bot/constants.py index 132ce93d..e5eb789a 100644 --- a/bot/constants.py +++ b/bot/constants.py @@ -1,2 +1,2 @@ CONFIG_FILE = "config/bot.conf" -DATA_VERSION = 5 +DATA_VERSION = 6 diff --git a/bot/core/blacklists.py b/bot/core/blacklists.py index 942bd012..1ca5bd9c 100644 --- a/bot/core/blacklists.py +++ b/bot/core/blacklists.py @@ -1,9 +1,8 @@ """ Guild, user, and member blacklists. - -NOTE: The pre-loading methods are not shard-optimised. """ from collections import defaultdict +import cachetools.func from data import tables from meta import client @@ -11,32 +10,22 @@ from meta import client from .module import module -@module.init_task -def load_guild_blacklist(client): +@cachetools.func.ttl_cache(ttl=300) +def guild_blacklist(): """ - Load the blacklisted guilds. + Get the guild blacklist """ rows = tables.global_guild_blacklist.select_where() - client.objects['blacklisted_guilds'] = set(row['guildid'] for row in rows) - if rows: - client.log( - "Loaded {} blacklisted guilds.".format(len(rows)), - context="GUILD_BLACKLIST" - ) + return set(row['guildid'] for row in rows) -@module.init_task -def load_user_blacklist(client): +@cachetools.func.ttl_cache(ttl=300) +def user_blacklist(): """ - Load the blacklisted users. + Get the global user blacklist. """ rows = tables.global_user_blacklist.select_where() - client.objects['blacklisted_users'] = set(row['userid'] for row in rows) - if rows: - client.log( - "Loaded {} globally blacklisted users.".format(len(rows)), - context="USER_BLACKLIST" - ) + return set(row['userid'] for row in rows) @module.init_task @@ -62,18 +51,20 @@ def load_ignored_members(client): ) +@module.init_task +def attach_client_blacklists(client): + client.guild_blacklist = guild_blacklist + client.user_blacklist = user_blacklist + + @module.launch_task async def leave_blacklisted_guilds(client): """ Launch task to leave any blacklisted guilds we are in. - Assumes that the blacklisted guild list has been initialised. """ - # Cache to avoic repeated lookups - blacklisted = client.objects['blacklisted_guilds'] - to_leave = [ guild for guild in client.guilds - if guild.id in blacklisted + if guild.id in guild_blacklist() ] for guild in to_leave: @@ -92,7 +83,8 @@ async def check_guild_blacklist(client, guild): Guild join event handler to check whether the guild is blacklisted. If so, leaves the guild. """ - if guild.id in client.objects['blacklisted_guilds']: + # First refresh the blacklist cache + if guild.id in guild_blacklist(): await guild.leave() client.log( "Automatically left blacklisted guild '{}' (gid:{}) upon join.".format(guild.name, guild.id), diff --git a/bot/core/data.py b/bot/core/data.py index 992e9ab5..475ba5a3 100644 --- a/bot/core/data.py +++ b/bot/core/data.py @@ -20,44 +20,13 @@ user_config = RowTable( ) -@user_config.save_query -def add_pending(pending): - """ - pending: - List of tuples of the form `(userid, pending_coins, pending_time)`. - """ - with lions.conn: - cursor = lions.conn.cursor() - data = execute_values( - cursor, - """ - UPDATE members - SET - coins = coins + t.coin_diff, - tracked_time = tracked_time + t.time_diff - FROM - (VALUES %s) - AS - t (guildid, userid, coin_diff, time_diff) - WHERE - members.guildid = t.guildid - AND - members.userid = t.userid - RETURNING * - """, - pending, - fetch=True - ) - return lions._make_rows(*data) - - guild_config = RowTable( 'guild_config', ('guildid', 'admin_role', 'mod_role', 'event_log_channel', 'mod_log_channel', 'alert_channel', - 'studyban_role', + 'studyban_role', 'max_study_bans', 'min_workout_length', 'workout_reward', 'max_tasks', 'task_reward', 'task_reward_limit', - 'study_hourly_reward', 'study_hourly_live_bonus', + 'study_hourly_reward', 'study_hourly_live_bonus', 'daily_study_cap', 'renting_price', 'renting_category', 'renting_cap', 'renting_role', 'renting_sync_perms', 'accountability_category', 'accountability_lobby', 'accountability_bonus', 'accountability_reward', 'accountability_price', @@ -88,9 +57,66 @@ lions = RowTable( attach_as='lions' ) + +@lions.save_query +def add_pending(pending): + """ + pending: + List of tuples of the form `(guildid, userid, pending_coins)`. + """ + with lions.conn: + cursor = lions.conn.cursor() + data = execute_values( + cursor, + """ + UPDATE members + SET + coins = coins + t.coin_diff + FROM + (VALUES %s) + AS + t (guildid, userid, coin_diff) + WHERE + members.guildid = t.guildid + AND + members.userid = t.userid + RETURNING * + """, + pending, + fetch=True + ) + return lions._make_rows(*data) + + lion_ranks = Table('member_ranks', attach_as='lion_ranks') +@lions.save_query +def get_member_rank(guildid, userid, untracked): + """ + Get the time and coin ranking for the given member, ignoring the provided untracked members. + """ + with lions.conn as conn: + with conn.cursor() as curs: + curs.execute( + """ + SELECT + time_rank, coin_rank + FROM ( + SELECT + userid, + row_number() OVER (ORDER BY total_tracked_time DESC, userid ASC) AS time_rank, + row_number() OVER (ORDER BY total_coins DESC, userid ASC) AS coin_rank + FROM members_totals + WHERE + guildid=%s AND userid NOT IN %s + ) AS guild_ranks WHERE userid=%s + """, + (guildid, tuple(untracked), userid) + ) + return curs.fetchone() or (None, None) + + global_guild_blacklist = Table('global_guild_blacklist') global_user_blacklist = Table('global_user_blacklist') ignored_members = Table('ignored_members') diff --git a/bot/core/lion.py b/bot/core/lion.py index b9b10092..c803cff8 100644 --- a/bot/core/lion.py +++ b/bot/core/lion.py @@ -1,4 +1,5 @@ import pytz +from datetime import datetime, timedelta from meta import client from data import tables as tb @@ -11,7 +12,7 @@ class Lion: Mostly acts as a transparent interface to the corresponding Row, but also adds some transaction caching logic to `coins` and `tracked_time`. """ - __slots__ = ('guildid', 'userid', '_pending_coins', '_pending_time', '_member') + __slots__ = ('guildid', 'userid', '_pending_coins', '_member') # Members with pending transactions _pending = {} # userid -> User @@ -24,7 +25,6 @@ class Lion: self.userid = userid self._pending_coins = 0 - self._pending_time = 0 self._member = None @@ -41,6 +41,7 @@ class Lion: if key in cls._lions: return cls._lions[key] else: + # TODO: Debug log lion = tb.lions.fetch(key) if not lion: tb.lions.create_row( @@ -77,23 +78,103 @@ class Lion: @property def settings(self): """ - The UserSettings object for this user. + The UserSettings interface for this member. """ return UserSettings(self.userid) + @property + def guild_settings(self): + """ + The GuildSettings interface for this member. + """ + return GuildSettings(self.guildid) + @property def time(self): """ - Amount of time the user has spent studying, accounting for pending values. + Amount of time the user has spent studying, accounting for a current session. """ - return int(self.data.tracked_time + self._pending_time) + # Base time from cached member data + time = self.data.tracked_time + + # Add current session time if it exists + if session := self.session: + time += session.duration + + return int(time) @property def coins(self): """ - Number of coins the user has, accounting for the pending value. + Number of coins the user has, accounting for the pending value and current session. """ - return int(self.data.coins + self._pending_coins) + # Base coin amount from cached member data + coins = self.data.coins + + # Add pending coin amount + coins += self._pending_coins + + # Add current session coins if applicable + if session := self.session: + coins += session.coins_earned + + return int(coins) + + @property + def session(self): + """ + The current study session the user is in, if any. + """ + if 'sessions' not in client.objects: + raise ValueError("Cannot retrieve session before Study module is initialised!") + return client.objects['sessions'][self.guildid].get(self.userid, None) + + @property + def timezone(self): + """ + The user's configured timezone. + Shortcut to `Lion.settings.timezone.value`. + """ + return self.settings.timezone.value + + @property + def day_start(self): + """ + A timezone aware datetime representing the start of the user's day (in their configured timezone). + NOTE: This might not be accurate over DST boundaries. + """ + now = datetime.now(tz=self.timezone) + return now.replace(hour=0, minute=0, second=0, microsecond=0) + + @property + def remaining_in_day(self): + return ((self.day_start + timedelta(days=1)) - datetime.now(self.timezone)).total_seconds() + + @property + def studied_today(self): + """ + The amount of time, in seconds, that the member has studied today. + Extracted from the session history. + """ + return tb.session_history.queries.study_time_since(self.guildid, self.userid, self.day_start) + + @property + def remaining_study_today(self): + """ + Maximum remaining time (in seconds) this member can study today. + + May not account for DST boundaries and leap seconds. + """ + studied_today = self.studied_today + study_cap = self.guild_settings.daily_study_cap.value + + remaining_in_day = self.remaining_in_day + if remaining_in_day >= (study_cap - studied_today): + remaining = study_cap - studied_today + else: + remaining = remaining_in_day + study_cap + + return remaining def localize(self, naive_utc_dt): """ @@ -111,15 +192,6 @@ class Lion: if flush: self.flush() - def addTime(self, amount, flush=True): - """ - Add time to a user (in seconds), optionally storing the transaction in pending. - """ - self._pending_time += amount - self._pending[self.key] = self - if flush: - self.flush() - def flush(self): """ Flush any pending transactions to the database. @@ -137,7 +209,7 @@ class Lion: if lions: # Build userid to pending coin map pending = [ - (lion.guildid, lion.userid, int(lion._pending_coins), int(lion._pending_time)) + (lion.guildid, lion.userid, int(lion._pending_coins)) for lion in lions ] @@ -147,5 +219,4 @@ class Lion: # Cleanup pending users for lion in lions: lion._pending_coins -= int(lion._pending_coins) - lion._pending_time -= int(lion._pending_time) cls._pending.pop(lion.key, None) diff --git a/bot/core/module.py b/bot/core/module.py index daaa4bc7..2f51408a 100644 --- a/bot/core/module.py +++ b/bot/core/module.py @@ -61,16 +61,18 @@ async def preload_studying_members(client): """ userids = list(set(member.id for guild in client.guilds for ch in guild.voice_channels for member in ch.members)) if userids: - rows = client.data.lions.fetch_rows_where(userid=userids) + users = client.data.user_config.fetch_rows_where(userid=userids) + members = client.data.lions.fetch_rows_where(userid=userids) client.log( - "Preloaded member data for {} members.".format(len(rows)), + "Preloaded data for {} user with {} members.".format(len(users), len(members)), context="CORE_LOADING" ) -@module.launch_task -async def launch_lion_sync_loop(client): - asyncio.create_task(_lion_sync_loop()) +# Removing the sync loop in favour of the studybadge sync. +# @module.launch_task +# async def launch_lion_sync_loop(client): +# asyncio.create_task(_lion_sync_loop()) @module.unload_task diff --git a/bot/data/__init__.py b/bot/data/__init__.py index f048ce37..2deecc48 100644 --- a/bot/data/__init__.py +++ b/bot/data/__init__.py @@ -1,5 +1,5 @@ +from .conditions import Condition, NOT, Constant, NULL, NOTNULL # noqa from .connection import conn # noqa from .formatters import UpdateValue, UpdateValueAdd # noqa from .interfaces import Table, RowTable, Row, tables # noqa from .queries import insert, insert_many, select_where, update_where, upsert, delete_where # noqa -from .conditions import Condition, NOT, Constant, NULL, NOTNULL # noqa diff --git a/bot/data/conditions.py b/bot/data/conditions.py index 4687a929..a314616e 100644 --- a/bot/data/conditions.py +++ b/bot/data/conditions.py @@ -1,5 +1,7 @@ from .connection import _replace_char +from meta import sharding + class Condition: """ @@ -70,5 +72,21 @@ class Constant(Condition): conditions.append("{} {}".format(key, self.value)) +class SHARDID(Condition): + __slots__ = ('shardid', 'shard_count') + + def __init__(self, shardid, shard_count): + self.shardid = shardid + self.shard_count = shard_count + + def apply(self, key, values, conditions): + if self.shard_count > 1: + conditions.append("({} >> 22) %% {} = {}".format(key, self.shard_count, _replace_char)) + values.append(self.shardid) + + +THIS_SHARD = SHARDID(sharding.shard_number, sharding.shard_count) + + NULL = Constant('IS NULL') NOTNULL = Constant('IS NOT NULL') diff --git a/bot/data/interfaces.py b/bot/data/interfaces.py index 42810e72..7673b0e0 100644 --- a/bot/data/interfaces.py +++ b/bot/data/interfaces.py @@ -45,10 +45,10 @@ class Table: Intended to be subclassed to provide more derivative access for specific tables. """ conn = conn - queries = DotDict() def __init__(self, name, attach_as=None): self.name = name + self.queries = DotDict() tables[attach_as or name] = self @_connection_guard diff --git a/bot/main.py b/bot/main.py index ac818e36..066bf86e 100644 --- a/bot/main.py +++ b/bot/main.py @@ -1,4 +1,4 @@ -from meta import client, conf, log +from meta import client, conf, log, sharding from data import tables @@ -7,7 +7,12 @@ import core # noqa import modules # noqa # Load and attach app specific data -client.appdata = core.data.meta.fetch_or_create(conf.bot['data_appid']) +if sharding.sharded: + appname = f"{conf.bot['data_appid']}_{sharding.shard_count}_{sharding.shard_number}" +else: + appname = conf.bot['data_appid'] +client.appdata = core.data.meta.fetch_or_create(appname) + client.data = tables # Initialise all modules diff --git a/bot/meta/__init__.py b/bot/meta/__init__.py index dd852d4f..eab9c7b8 100644 --- a/bot/meta/__init__.py +++ b/bot/meta/__init__.py @@ -1,3 +1,5 @@ +from .logger import log, logger from .client import client from .config import conf -from .logger import log, logger +from .args import args +from . import sharding diff --git a/bot/meta/args.py b/bot/meta/args.py new file mode 100644 index 00000000..c2dd70d6 --- /dev/null +++ b/bot/meta/args.py @@ -0,0 +1,19 @@ +import argparse + +from constants import CONFIG_FILE + +# ------------------------------ +# Parsed commandline arguments +# ------------------------------ +parser = argparse.ArgumentParser() +parser.add_argument('--conf', + dest='config', + default=CONFIG_FILE, + help="Path to configuration file.") +parser.add_argument('--shard', + dest='shard', + default=None, + type=int, + help="Shard number to run, if applicable.") + +args = parser.parse_args() diff --git a/bot/meta/client.py b/bot/meta/client.py index 5310171d..50414aa8 100644 --- a/bot/meta/client.py +++ b/bot/meta/client.py @@ -1,16 +1,19 @@ from discord import Intents from cmdClient.cmdClient import cmdClient -from .config import Conf +from .config import conf +from .sharding import shard_number, shard_count -from constants import CONFIG_FILE - -# Initialise config -conf = Conf(CONFIG_FILE) # Initialise client owners = [int(owner) for owner in conf.bot.getlist('owners')] intents = Intents.all() intents.presences = False -client = cmdClient(prefix=conf.bot['prefix'], owners=owners, intents=intents) +client = cmdClient( + prefix=conf.bot['prefix'], + owners=owners, + intents=intents, + shard_id=shard_number, + shard_count=shard_count +) client.conf = conf diff --git a/bot/meta/config.py b/bot/meta/config.py index a94d2b1a..ca779924 100644 --- a/bot/meta/config.py +++ b/bot/meta/config.py @@ -1,9 +1,6 @@ import configparser as cfgp - -conf = None # type: Conf - -CONF_FILE = "bot/bot.conf" +from .args import args class Conf: @@ -57,3 +54,6 @@ class Conf: def write(self): with open(self.configfile, 'w') as conffile: self.config.write(conffile) + + +conf = Conf(args.config) diff --git a/bot/meta/logger.py b/bot/meta/logger.py index 858b1292..a95500e4 100644 --- a/bot/meta/logger.py +++ b/bot/meta/logger.py @@ -9,11 +9,18 @@ from utils.lib import mail, split_text from .client import client from .config import conf +from . import sharding # Setup the logger logger = logging.getLogger() -log_fmt = logging.Formatter(fmt='[{asctime}][{levelname:^8}] {message}', datefmt='%d/%m | %H:%M:%S', style='{') +log_fmt = logging.Formatter( + fmt=('[{asctime}][{levelname:^8}]' + + '[SHARD {}]'.format(sharding.shard_number) + + ' {message}'), + datefmt='%d/%m | %H:%M:%S', + style='{' +) # term_handler = logging.StreamHandler(sys.stdout) # term_handler.setFormatter(log_fmt) # logger.addHandler(term_handler) @@ -77,7 +84,11 @@ async def live_log(message, context, level): log_chid = conf.bot.getint('log_channel') # Generate the log messages - header = "[{}][{}]".format(logging.getLevelName(level), str(context)) + if sharding.sharded: + header = f"[{logging.getLevelName(level)}][SHARD {sharding.shard_number}][{context}]" + else: + header = f"[{logging.getLevelName(level)}][{context}]" + if len(message) > 1900: blocks = split_text(message, blocksize=1900, code=False) else: diff --git a/bot/meta/sharding.py b/bot/meta/sharding.py new file mode 100644 index 00000000..ffe86a89 --- /dev/null +++ b/bot/meta/sharding.py @@ -0,0 +1,9 @@ +from .args import args +from .config import conf + + +shard_number = args.shard or 0 + +shard_count = conf.bot.getint('shard_count', 1) + +sharded = (shard_count > 0) diff --git a/bot/modules/accountability/TimeSlot.py b/bot/modules/accountability/TimeSlot.py index 9464e4e5..81cbe38d 100644 --- a/bot/modules/accountability/TimeSlot.py +++ b/bot/modules/accountability/TimeSlot.py @@ -90,7 +90,6 @@ class TimeSlot: @property def open_embed(self): - # TODO Consider adding hint to footer timestamp = int(self.start_time.timestamp()) embed = discord.Embed( @@ -247,6 +246,34 @@ class TimeSlot: return self + async def _reload_members(self, memberids=None): + """ + Reload the timeslot members from the provided list, or data. + Also updates the channel overwrites if required. + To be used before the session has started. + """ + if self.data: + if memberids is None: + member_rows = accountability_members.fetch_rows_where(slotid=self.data.slotid) + memberids = [row.userid for row in member_rows] + + self.members = members = { + memberid: SlotMember(self.data.slotid, memberid, self.guild) + for memberid in memberids + } + + if self.channel: + # Check and potentially update overwrites + current_overwrites = self.channel.overwrites + overwrites = { + mem.member: self._member_overwrite + for mem in members.values() + if mem.member + } + overwrites[self.guild.default_role] = self._everyone_overwrite + if current_overwrites != overwrites: + await self.channel.edit(overwrites=overwrites) + def _refresh(self): """ Refresh the stored data row and reload. diff --git a/bot/modules/accountability/commands.py b/bot/modules/accountability/commands.py index 6180d4bd..a0d19509 100644 --- a/bot/modules/accountability/commands.py +++ b/bot/modules/accountability/commands.py @@ -39,6 +39,7 @@ def time_format(time): time.timestamp() + 3600, ) + user_locks = {} # Map userid -> ctx @@ -229,7 +230,10 @@ async def cmd_rooms(ctx): start_time + datetime.timedelta(hours=n) for n in range(1, 25) ) - times = [time for time in times if time not in already_joined_times] + times = [ + time for time in times + if time not in already_joined_times and (time - utc_now()).total_seconds() > 660 + ] lines = [ "`[{num:>2}]` | `{count:>{count_pad}}` attending | {time}".format( num=i, @@ -255,7 +259,7 @@ async def cmd_rooms(ctx): await ctx.cancellable( out_msg, cancel_message="Booking menu cancelled, no sessions were booked.", - timeout=70 + timeout=60 ) def check(msg): @@ -265,7 +269,7 @@ async def cmd_rooms(ctx): with ensure_exclusive(ctx): try: - message = await ctx.client.wait_for('message', check=check, timeout=60) + message = await ctx.client.wait_for('message', check=check, timeout=30) except asyncio.TimeoutError: try: await out_msg.edit( @@ -325,6 +329,7 @@ async def cmd_rooms(ctx): ) # Handle case where the slot has already opened + # TODO: Fix this, doesn't always work aguild = AGuild.cache.get(ctx.guild.id, None) if aguild: if aguild.upcoming_slot and aguild.upcoming_slot.start_time in to_book: diff --git a/bot/modules/accountability/tracker.py b/bot/modules/accountability/tracker.py index 24e1dc94..faa82867 100644 --- a/bot/modules/accountability/tracker.py +++ b/bot/modules/accountability/tracker.py @@ -10,7 +10,7 @@ from discord.utils import sleep_until from meta import client from utils.interactive import discord_shield from data import NULL, NOTNULL, tables -from data.conditions import LEQ +from data.conditions import LEQ, THIS_SHARD from settings import GuildSettings from .TimeSlot import TimeSlot @@ -67,7 +67,8 @@ async def open_next(start_time): """ # Pre-fetch the new slot data, also populating the table caches room_data = accountability_rooms.fetch_rows_where( - start_at=start_time + start_at=start_time, + guildid=THIS_SHARD ) guild_rows = {row.guildid: row for row in room_data} member_data = accountability_members.fetch_rows_where( @@ -193,11 +194,30 @@ async def turnover(): # TODO: (FUTURE) with high volume, we might want to start the sessions before moving the members. # We could break up the session starting? - # Move members of the next session over to the session channel + # ---------- Start next session ---------- current_slots = [ aguild.current_slot for aguild in AccountabilityGuild.cache.values() if aguild.current_slot is not None ] + slotmap = {slot.data.slotid: slot for slot in current_slots if slot.data} + + # Reload the slot members in case they cancelled from another shard + member_data = accountability_members.fetch_rows_where( + slotid=list(slotmap.keys()) + ) if slotmap else [] + slot_memberids = {slotid: [] for slotid in slotmap} + for row in member_data: + slot_memberids[row.slotid].append(row.userid) + reload_tasks = ( + slot._reload_members(memberids=slot_memberids[slotid]) + for slotid, slot in slotmap.items() + ) + await asyncio.gather( + *reload_tasks, + return_exceptions=True + ) + + # Move members of the next session over to the session channel movement_tasks = ( mem.member.edit( voice_channel=slot.channel, @@ -335,6 +355,7 @@ async def _accountability_system_resume(): open_room_data = accountability_rooms.fetch_rows_where( closed_at=NULL, start_at=LEQ(now), + guildid=THIS_SHARD, _extra="ORDER BY start_at ASC" ) @@ -450,8 +471,10 @@ async def launch_accountability_system(client): """ # Load the AccountabilityGuild cache guilds = tables.guild_config.fetch_rows_where( - accountability_category=NOTNULL + accountability_category=NOTNULL, + guildid=THIS_SHARD ) + # Further filter out any guilds that we aren't in [AccountabilityGuild(guild.guildid) for guild in guilds if client.get_guild(guild.guildid)] await _accountability_system_resume() asyncio.create_task(_accountability_loop()) diff --git a/bot/modules/economy/cointop_cmd.py b/bot/modules/economy/cointop_cmd.py index cd90b537..81bdbad9 100644 --- a/bot/modules/economy/cointop_cmd.py +++ b/bot/modules/economy/cointop_cmd.py @@ -43,22 +43,18 @@ async def cmd_topcoin(ctx): # Fetch the leaderboard exclude = set(m.id for m in ctx.guild_settings.unranked_roles.members) - exclude.update(ctx.client.objects['blacklisted_users']) + exclude.update(ctx.client.user_blacklist()) exclude.update(ctx.client.objects['ignored_members'][ctx.guild.id]) + args = { + 'guildid': ctx.guild.id, + 'select_columns': ('userid', 'total_coins::INTEGER'), + '_extra': "AND total_coins > 0 ORDER BY total_coins DESC " + ("LIMIT 100" if top100 else "") + } if exclude: - user_data = tables.lions.select_where( - guildid=ctx.guild.id, - userid=data.NOT(list(exclude)), - select_columns=('userid', 'coins'), - _extra="AND coins > 0 ORDER BY coins DESC " + ("LIMIT 100" if top100 else "") - ) - else: - user_data = tables.lions.select_where( - guildid=ctx.guild.id, - select_columns=('userid', 'coins'), - _extra="AND coins > 0 ORDER BY coins DESC " + ("LIMIT 100" if top100 else "") - ) + args['userid'] = data.NOT(list(exclude)) + + user_data = tables.members_totals.select_where(**args) # Quit early if the leaderboard is empty if not user_data: diff --git a/bot/modules/guild_admin/__init__.py b/bot/modules/guild_admin/__init__.py index 74d5644d..2c4f3922 100644 --- a/bot/modules/guild_admin/__init__.py +++ b/bot/modules/guild_admin/__init__.py @@ -4,3 +4,4 @@ from . import guild_config from . import statreset from . import new_members from . import reaction_roles +from . import economy diff --git a/bot/modules/guild_admin/economy/__init__.py b/bot/modules/guild_admin/economy/__init__.py new file mode 100644 index 00000000..2714bcdb --- /dev/null +++ b/bot/modules/guild_admin/economy/__init__.py @@ -0,0 +1,3 @@ +from ..module import module + +from . import set_coins \ No newline at end of file diff --git a/bot/modules/guild_admin/economy/set_coins.py b/bot/modules/guild_admin/economy/set_coins.py new file mode 100644 index 00000000..c0744a13 --- /dev/null +++ b/bot/modules/guild_admin/economy/set_coins.py @@ -0,0 +1,104 @@ +import discord +import datetime +from wards import guild_admin + +from settings import GuildSettings +from core import Lion + +from ..module import module + +POSTGRES_INT_MAX = 2147483647 + +@module.cmd( + "set_coins", + group="Guild Admin", + desc="Set coins on a member." +) +@guild_admin() +async def cmd_set(ctx): + """ + Usage``: + {prefix}set_coins + Description: + Sets the given number of coins on the mentioned user. + If a number greater than 0 is mentioned, will add coins. + If a number less than 0 is mentioned, will remove coins. + Note: LionCoins on a member cannot be negative. + Example: + {prefix}set_coins {ctx.author.mention} 100 + {prefix}set_coins {ctx.author.mention} -100 + """ + # Extract target and amount + # Handle a slightly more flexible input than stated + splits = ctx.args.split() + digits = [isNumber(split) for split in splits[:2]] + mentions = ctx.msg.mentions + if len(splits) < 2 or not any(digits) or not (all(digits) or mentions): + return await _send_usage(ctx) + + if all(digits): + # Both are digits, hopefully one is a member id, and one is an amount. + target, amount = ctx.guild.get_member(int(splits[0])), int(splits[1]) + if not target: + amount, target = int(splits[0]), ctx.guild.get_member(int(splits[1])) + if not target: + return await _send_usage(ctx) + elif digits[0]: + amount, target = int(splits[0]), mentions[0] + elif digits[1]: + target, amount = mentions[0], int(splits[1]) + + # Fetch the associated lion + target_lion = Lion.fetch(ctx.guild.id, target.id) + + # Check sanity conditions + if target == ctx.client.user: + return await ctx.embed_reply("Thanks, but Ari looks after all my needs!") + if target.bot: + return await ctx.embed_reply("We are still waiting for {} to open an account.".format(target.mention)) + + # Finally, send the amount and the ack message + # Postgres `coins` column is `integer`, sanity check postgres int limits - which are smalled than python int range + target_coins_to_set = target_lion.coins + amount + if target_coins_to_set >= 0 and target_coins_to_set <= POSTGRES_INT_MAX: + target_lion.addCoins(amount) + elif target_coins_to_set < 0: + target_coins_to_set = -target_lion.coins # Coins cannot go -ve, cap to 0 + target_lion.addCoins(target_coins_to_set) + target_coins_to_set = 0 + else: + return await ctx.embed_reply("Member coins cannot be more than {}".format(POSTGRES_INT_MAX)) + + embed = discord.Embed( + title="Funds Set", + description="You have set LionCoins on {} to **{}**!".format(target.mention,target_coins_to_set), + colour=discord.Colour.orange(), + timestamp=datetime.datetime.utcnow() + ).set_footer(text=str(ctx.author), icon_url=ctx.author.avatar_url) + + await ctx.reply(embed=embed, reference=ctx.msg) + GuildSettings(ctx.guild.id).event_log.log( + "{} set {}'s LionCoins to`{}`.".format( + ctx.author.mention, + target.mention, + target_coins_to_set + ), + title="Funds Set" + ) + +def isNumber(var): + try: + return isinstance(int(var), int) + except: + return False + +async def _send_usage(ctx): + return await ctx.error_reply( + "**Usage:** `{prefix}set_coins `\n" + "**Example:**\n" + " {prefix}set_coins {ctx.author.mention} 100\n" + " {prefix}set_coins {ctx.author.mention} -100".format( + prefix=ctx.best_prefix, + ctx=ctx + ) + ) diff --git a/bot/modules/guild_admin/guild_config.py b/bot/modules/guild_admin/guild_config.py index 608786fb..c79d62bb 100644 --- a/bot/modules/guild_admin/guild_config.py +++ b/bot/modules/guild_admin/guild_config.py @@ -1,3 +1,4 @@ +import difflib import discord from cmdClient.lib import SafeCancellation @@ -121,9 +122,15 @@ async def cmd_config(ctx, flags): name = parts[0] setting = setting_displaynames.get(name.lower(), None) if setting is None: + matches = difflib.get_close_matches(name, setting_displaynames.keys(), n=2) + match = "`{}`".format('` or `'.join(matches)) if matches else None return await ctx.error_reply( - "Server setting `{}` doesn't exist! Use `{}config` to see all server settings".format( - name, ctx.best_prefix + "Couldn't find a setting called `{}`!\n" + "{}" + "Use `{}config info` to see all the server settings.".format( + name, + "Maybe you meant {}?\n".format(match) if match else "", + ctx.best_prefix ) ) diff --git a/bot/modules/guild_admin/reaction_roles/tracker.py b/bot/modules/guild_admin/reaction_roles/tracker.py index f18e3c34..17a64960 100644 --- a/bot/modules/guild_admin/reaction_roles/tracker.py +++ b/bot/modules/guild_admin/reaction_roles/tracker.py @@ -12,6 +12,7 @@ from discord import PartialEmoji from meta import client from core import Lion from data import Row +from data.conditions import THIS_SHARD from utils.lib import utc_now from settings import GuildSettings @@ -584,5 +585,5 @@ def load_reaction_roles(client): """ Load the ReactionRoleMessages. """ - rows = reaction_role_messages.fetch_rows_where() + rows = reaction_role_messages.fetch_rows_where(guildid=THIS_SHARD) ReactionRoleMessage._messages = {row.messageid: ReactionRoleMessage(row.messageid) for row in rows} diff --git a/bot/modules/moderation/tickets/Ticket.py b/bot/modules/moderation/tickets/Ticket.py index 4d7ec5ec..afea1eef 100644 --- a/bot/modules/moderation/tickets/Ticket.py +++ b/bot/modules/moderation/tickets/Ticket.py @@ -6,6 +6,7 @@ import datetime import discord from meta import client +from data.conditions import THIS_SHARD from settings import GuildSettings from utils.lib import FieldEnum, strfdelta, utc_now @@ -283,7 +284,8 @@ class Ticket: # Get all expiring tickets expiring_rows = data.tickets.select_where( - ticket_state=TicketState.EXPIRING + ticket_state=TicketState.EXPIRING, + guildid=THIS_SHARD ) # Create new expiry tasks diff --git a/bot/modules/moderation/video/admin.py b/bot/modules/moderation/video/admin.py index 1f0ddaab..8d658c60 100644 --- a/bot/modules/moderation/video/admin.py +++ b/bot/modules/moderation/video/admin.py @@ -56,10 +56,10 @@ class video_channels(settings.ChannelList, settings.ListData, settings.Setting): if any(channel.members for channel in guild.voice_channels) ] if active_guildids: + cache = {guildid: [] for guildid in active_guildids} rows = cls._table_interface.select_where( guildid=active_guildids ) - cache = defaultdict(list) for row in rows: cache[row['guildid']].append(row['channelid']) cls._cache.update(cache) diff --git a/bot/modules/reminders/commands.py b/bot/modules/reminders/commands.py index 0bb98d60..c8637a04 100644 --- a/bot/modules/reminders/commands.py +++ b/bot/modules/reminders/commands.py @@ -3,6 +3,7 @@ import asyncio import datetime import discord +from meta import sharding from utils.lib import parse_dur, parse_ranges, multiselect_regex from .module import module @@ -55,7 +56,7 @@ async def cmd_remindme(ctx, flags): if not rows: return await ctx.reply("You have no reminders to remove!") - live = Reminder.fetch(*(row.reminderid for row in rows)) + live = [Reminder(row.reminderid) for row in rows] if not ctx.args: lines = [] @@ -209,7 +210,8 @@ async def cmd_remindme(ctx, flags): ) # Schedule reminder - reminder.schedule() + if sharding.shard_number == 0: + reminder.schedule() # Ack embed = discord.Embed( @@ -231,7 +233,7 @@ async def cmd_remindme(ctx, flags): if not rows: return await ctx.reply("You have no reminders!") - live = Reminder.fetch(*(row.reminderid for row in rows)) + live = [Reminder(row.reminderid) for row in rows] lines = [] num_field = len(str(len(live) - 1)) diff --git a/bot/modules/reminders/reminder.py b/bot/modules/reminders/reminder.py index d3e4f764..67956a1d 100644 --- a/bot/modules/reminders/reminder.py +++ b/bot/modules/reminders/reminder.py @@ -1,8 +1,9 @@ import asyncio import datetime +import logging import discord -from meta import client +from meta import client, sharding from utils.lib import strfdur from .data import reminders @@ -46,7 +47,10 @@ class Reminder: cls._live_reminders[reminderid].cancel() # Remove from data - reminders.delete_where(reminderid=reminderids) + if reminderids: + return reminders.delete_where(reminderid=reminderids) + else: + return [] @property def data(self): @@ -134,10 +138,16 @@ class Reminder: """ Execute the reminder. """ - if self.data.userid in client.objects['blacklisted_users']: + if not self.data: + # Reminder deleted elsewhere + return + + if self.data.userid in client.user_blacklist(): self.delete(self.reminderid) return + userid = self.data.userid + # Build the message embed embed = discord.Embed( title="You asked me to remind you!", @@ -155,8 +165,26 @@ class Reminder: ) ) + # Update the reminder data, and reschedule if required + if self.data.interval: + next_time = self.data.remind_at + datetime.timedelta(seconds=self.data.interval) + rows = reminders.update_where( + {'remind_at': next_time}, + reminderid=self.reminderid + ) + self.schedule() + else: + rows = self.delete(self.reminderid) + if not rows: + # Reminder deleted elsewhere + return + # Send the message, if possible - user = self.user + if not (user := client.get_user(userid)): + try: + user = await client.fetch_user(userid) + except discord.HTTPException: + pass if user: try: await user.send(embed=embed) @@ -164,21 +192,38 @@ class Reminder: # Nothing we can really do here. Maybe tell the user about their reminder next time? pass - # Update the reminder data, and reschedule if required - if self.data.interval: - next_time = self.data.remind_at + datetime.timedelta(seconds=self.data.interval) - reminders.update_where({'remind_at': next_time}, reminderid=self.reminderid) - self.schedule() - else: - self.delete(self.reminderid) + +async def reminder_poll(client): + """ + One client/shard must continually poll for new or deleted reminders. + """ + # TODO: Clean this up with database signals or IPC + while True: + await asyncio.sleep(60) + + client.log( + "Running new reminder poll.", + context="REMINDERS", + level=logging.DEBUG + ) + + rids = {row.reminderid for row in reminders.fetch_rows_where()} + + to_delete = (rid for rid in Reminder._live_reminders if rid not in rids) + Reminder.delete(*to_delete) + + [Reminder(rid).schedule() for rid in rids if rid not in Reminder._live_reminders] @module.launch_task async def schedule_reminders(client): - rows = reminders.fetch_rows_where() - for row in rows: - Reminder(row.reminderid).schedule() - client.log( - "Scheduled {} reminders.".format(len(rows)), - context="LAUNCH_REMINDERS" - ) + if sharding.shard_number == 0: + rows = reminders.fetch_rows_where() + for row in rows: + Reminder(row.reminderid).schedule() + client.log( + "Scheduled {} reminders.".format(len(rows)), + context="LAUNCH_REMINDERS" + ) + if sharding.sharded: + asyncio.create_task(reminder_poll(client)) diff --git a/bot/modules/renting/commands.py b/bot/modules/renting/commands.py index b94c01f5..0683d90f 100644 --- a/bot/modules/renting/commands.py +++ b/bot/modules/renting/commands.py @@ -54,9 +54,13 @@ async def cmd_rent(ctx): # Extract members to remove current_memberids = set(room.memberids) + if ctx.author in ctx.msg.mentions: + return await ctx.error_reply( + "You can't remove yourself from your own room!" + ) to_remove = ( member for member in ctx.msg.mentions - if member.id in current_memberids + if member.id in current_memberids and member.id != ctx.author.id ) to_remove = list(set(to_remove)) # Remove duplicates @@ -86,7 +90,7 @@ async def cmd_rent(ctx): current_memberids = set(room.memberids) to_add = ( member for member in ctx.msg.mentions - if member.id not in current_memberids and member.id != ctx.author + if member.id not in current_memberids and member.id != ctx.author.id ) to_add = list(set(to_add)) # Remove duplicates diff --git a/bot/modules/renting/rooms.py b/bot/modules/renting/rooms.py index 9e79d5b0..a8c29876 100644 --- a/bot/modules/renting/rooms.py +++ b/bot/modules/renting/rooms.py @@ -5,6 +5,7 @@ import datetime from cmdClient.lib import SafeCancellation from meta import client +from data.conditions import THIS_SHARD from settings import GuildSettings from .data import rented, rented_members @@ -187,14 +188,14 @@ class Room: except discord.HTTPException: pass - # Delete the room from data (cascades to member deletion) - self.delete() - guild_settings.event_log.log( title="Private study room expired!", description="<@{}>'s private study room expired.".format(self.data.ownerid) ) + # Delete the room from data (cascades to member deletion) + self.delete() + async def add_members(self, *members): guild_settings = GuildSettings(self.data.guildid) @@ -276,7 +277,7 @@ class Room: @module.launch_task async def load_rented_rooms(client): - rows = rented.fetch_rows_where() + rows = rented.fetch_rows_where(guildid=THIS_SHARD) for row in rows: Room(row.channelid).schedule() client.log( diff --git a/bot/modules/study/__init__.py b/bot/modules/study/__init__.py index eec16e1d..30f59149 100644 --- a/bot/modules/study/__init__.py +++ b/bot/modules/study/__init__.py @@ -1,9 +1,8 @@ from .module import module -from . import data -from . import admin -from . import badge_tracker -from . import time_tracker +from . import badges +from . import timers +from . import tracking + from . import top_cmd -from . import studybadge_cmd from . import stats_cmd diff --git a/bot/modules/study/badges/__init__.py b/bot/modules/study/badges/__init__.py new file mode 100644 index 00000000..8db92c34 --- /dev/null +++ b/bot/modules/study/badges/__init__.py @@ -0,0 +1,2 @@ +from . import badge_tracker +from . import studybadge_cmd diff --git a/bot/modules/study/badge_tracker.py b/bot/modules/study/badges/badge_tracker.py similarity index 95% rename from bot/modules/study/badge_tracker.py rename to bot/modules/study/badges/badge_tracker.py index cf69057a..721f3962 100644 --- a/bot/modules/study/badge_tracker.py +++ b/bot/modules/study/badges/badge_tracker.py @@ -6,14 +6,13 @@ import contextlib import discord -from meta import client -from data.conditions import GEQ -from core import Lion +from meta import client, sharding +from data.conditions import GEQ, THIS_SHARD from core.data import lions from utils.lib import strfdur from settings import GuildSettings -from .module import module +from ..module import module from .data import new_study_badges, study_badges @@ -55,11 +54,16 @@ async def update_study_badges(full=False): # Retrieve member rows with out of date study badges if not full and client.appdata.last_study_badge_scan is not None: + # TODO: _extra here is a hack to cover for inflexible conditionals update_rows = new_study_badges.select_where( - _timestamp=GEQ(client.appdata.last_study_badge_scan or 0) + guildid=THIS_SHARD, + _timestamp=GEQ(client.appdata.last_study_badge_scan or 0), + _extra="OR session_start IS NOT NULL AND (guildid >> 22) %% {} = {}".format( + sharding.shard_count, sharding.shard_number + ) ) else: - update_rows = new_study_badges.select_where() + update_rows = new_study_badges.select_where(guildid=THIS_SHARD) if not update_rows: client.appdata.last_study_badge_scan = datetime.datetime.utcnow() @@ -303,11 +307,10 @@ async def study_badge_tracker(): await asyncio.sleep(60) -async def _update_member_studybadge(member): +async def update_member_studybadge(member): """ Checks and (if required) updates the study badge for a single member. """ - Lion.fetch(member.guild.id, member.id).flush() update_rows = new_study_badges.select_where( guildid=member.guild.id, userid=member.id @@ -331,16 +334,6 @@ async def _update_member_studybadge(member): await _update_guild_badges(member.guild, update_rows) -@client.add_after_event("voice_state_update") -async def voice_studybadge_updater(client, member, before, after): - if not client.is_ready(): - # The poll loop will pick it up - return - - if before.channel and not after.channel: - await _update_member_studybadge(member) - - @module.launch_task async def launch_study_badge_tracker(client): asyncio.create_task(study_badge_tracker()) diff --git a/bot/modules/study/data.py b/bot/modules/study/badges/data.py similarity index 92% rename from bot/modules/study/data.py rename to bot/modules/study/badges/data.py index 005765bb..eca5f220 100644 --- a/bot/modules/study/data.py +++ b/bot/modules/study/badges/data.py @@ -2,8 +2,6 @@ from cachetools import cached from data import Table, RowTable -untracked_channels = Table('untracked_channels') - study_badges = RowTable( 'study_badges', ('badgeid', 'guildid', 'roleid', 'required_time'), diff --git a/bot/modules/study/studybadge_cmd.py b/bot/modules/study/badges/studybadge_cmd.py similarity index 99% rename from bot/modules/study/studybadge_cmd.py rename to bot/modules/study/badges/studybadge_cmd.py index 99e28fec..1d9aea9c 100644 --- a/bot/modules/study/studybadge_cmd.py +++ b/bot/modules/study/badges/studybadge_cmd.py @@ -12,7 +12,7 @@ from wards import is_guild_admin from core.data import lions from settings import GuildSettings -from .module import module +from ..module import module from .data import study_badges, guild_role_cache, new_study_badges from .badge_tracker import _update_guild_badges diff --git a/bot/modules/study/stats_cmd.py b/bot/modules/study/stats_cmd.py index 73604658..88bc8be5 100644 --- a/bot/modules/study/stats_cmd.py +++ b/bot/modules/study/stats_cmd.py @@ -1,18 +1,22 @@ -import datetime +from datetime import datetime, timedelta import discord from cmdClient.checks import in_guild -from utils.lib import strfdur +from utils.lib import prop_tabulate, utc_now from data import tables +from data.conditions import LEQ from core import Lion +from .tracking.data import session_history + from .module import module @module.cmd( "stats", group="Statistics", - desc="View a summary of your study statistics!", + desc="View your personal server study statistics!", + aliases=('profile',), allow_before_ready=True ) @in_guild() @@ -24,6 +28,7 @@ async def cmd_stats(ctx): Description: View the study statistics for yourself or the mentioned user. """ + # Identify the target if ctx.args: if not ctx.msg.mentions: return await ctx.error_reply("Please mention a user to view their statistics!") @@ -31,54 +36,231 @@ async def cmd_stats(ctx): else: target = ctx.author - # Collect the required target data + # System sync + Lion.sync() + + # Fetch the required data lion = Lion.fetch(ctx.guild.id, target.id) - rank_data = tables.lion_ranks.select_one_where( + + history = session_history.select_where( + guildid=ctx.guild.id, userid=target.id, - guildid=ctx.guild.id + select_columns=( + "start_time", + "(start_time + duration * interval '1 second') AS end_time" + ), + _extra="ORDER BY start_time DESC" ) - # Extract and format data - time = strfdur(lion.time) + # Current economy balance (accounting for current session) coins = lion.coins - workouts = lion.data.workout_count - if lion.data.last_study_badgeid: - badge_row = tables.study_badges.fetch(lion.data.last_study_badgeid) - league = "<@&{}>".format(badge_row.roleid) + season_time = lion.time + workout_total = lion.data.workout_count + + # Leaderboard ranks + exclude = set(m.id for m in ctx.guild_settings.unranked_roles.members) + exclude.update(ctx.client.user_blacklist()) + exclude.update(ctx.client.objects['ignored_members'][ctx.guild.id]) + if target.id in exclude: + time_rank = None + coin_rank = None else: - league = "No league yet!" + time_rank, coin_rank = tables.lions.queries.get_member_rank(ctx.guild.id, target.id, list(exclude or [0])) - time_lb_pos = rank_data['time_rank'] - coin_lb_pos = rank_data['coin_rank'] + # Study time + # First get the all/month/week/day timestamps + day_start = lion.day_start + period_timestamps = ( + datetime(1970, 1, 1), + day_start.replace(day=1), + day_start - timedelta(days=day_start.weekday()), + day_start + ) + study_times = [0, 0, 0, 0] + for i, timestamp in enumerate(period_timestamps): + study_time = tables.session_history.queries.study_time_since(ctx.guild.id, target.id, timestamp) + if not study_time: + # So we don't make unecessary database calls + break + study_times[i] = study_time - # Build embed + # Streak statistics + streak = 0 + current_streak = None + max_streak = 0 + + day_attended = True if 'sessions' in ctx.client.objects and lion.session else None + date = day_start + daydiff = timedelta(days=1) + + periods = [(row['start_time'], row['end_time']) for row in history] + + i = 0 + while i < len(periods): + row = periods[i] + i += 1 + if row[1] > date: + # They attended this day + day_attended = True + continue + elif day_attended is None: + # Didn't attend today, but don't break streak + day_attended = False + date -= daydiff + i -= 1 + continue + elif not day_attended: + # Didn't attend the day, streak broken + date -= daydiff + i -= 1 + pass + else: + # Attended the day + streak += 1 + + # Move window to the previous day and try the row again + day_attended = False + prev_date = date + date -= daydiff + i -= 1 + + # Special case, when the last session started in the previous day + # Then the day is already attended + if i > 1 and date < periods[i-2][0] <= prev_date: + day_attended = True + + continue + + max_streak = max(max_streak, streak) + if current_streak is None: + current_streak = streak + streak = 0 + + # Handle loop exit state, i.e. the last streak + if day_attended: + streak += 1 + max_streak = max(max_streak, streak) + if current_streak is None: + current_streak = streak + + # Accountability stats + accountability = tables.accountability_member_info.select_where( + userid=target.id, + start_at=LEQ(utc_now()), + select_columns=("*", "(duration > 0 OR last_joined_at IS NOT NULL) AS attended"), + _extra="ORDER BY start_at DESC" + ) + if len(accountability): + acc_duration = sum(row['duration'] for row in accountability) + + acc_attended = sum(row['attended'] for row in accountability) + acc_total = len(accountability) + acc_rate = (acc_attended * 100) / acc_total + else: + acc_duration = 0 + acc_rate = 0 + + # Study League + guild_badges = tables.study_badges.fetch_rows_where(guildid=ctx.guild.id) + if lion.data.last_study_badgeid: + current_badge = tables.study_badges.fetch(lion.data.last_study_badgeid) + else: + current_badge = None + + next_badge = min( + (badge for badge in guild_badges + if badge.required_time > (current_badge.required_time if current_badge else 0)), + key=lambda badge: badge.required_time, + default=None + ) + + # We have all the data + # Now start building the embed embed = discord.Embed( - colour=discord.Colour.blue(), - timestamp=datetime.datetime.utcnow(), - title="Revision Statistics" - ).set_footer(text=str(target), icon_url=target.avatar_url).set_thumbnail(url=target.avatar_url) - embed.add_field( - name="📚 Study Time", - value=time + colour=discord.Colour.orange(), + title="Study Profile for {}".format(str(target)) ) - embed.add_field( - name="🦁 Revision League", - value=league + embed.set_thumbnail(url=target.avatar_url) + + # Add studying since if they have studied + if history: + embed.set_footer(text="Studying Since") + embed.timestamp = history[-1]['start_time'] + + # Set the description based on season time and server rank + if season_time: + time_str = "**{}:{:02}**".format( + season_time // 3600, + (season_time // 60) % 60 + ) + if time_rank is None: + rank_str = None + elif time_rank == 1: + rank_str = "1st" + elif time_rank == 2: + rank_str = "2nd" + elif time_rank == 3: + rank_str = "3rd" + else: + rank_str = "{}th".format(time_rank) + + embed.description = "{} has studied for **{}**{}{}".format( + target.mention, + time_str, + " this season" if study_times[0] - season_time > 60 else "", + ", and is ranked **{}** in the server!".format(rank_str) if rank_str else "." + ) + else: + embed.description = "{} hasn't studied in this server yet!".format(target.mention) + + # Build the stats table + stats = {} + + stats['Coins Earned'] = "**{}** LC".format( + coins, + # "Rank `{}`".format(coin_rank) if coins and coin_rank else "Unranked" ) - embed.add_field( - name="🦁 LionCoins", - value=coins + if workout_total: + stats['Workouts'] = "**{}** sessions".format(workout_total) + if acc_duration: + stats['Accountability'] = "**{}** hours (`{:.0f}%` attended)".format( + acc_duration // 3600, + acc_rate + ) + stats['Study Streak'] = "**{}** days{}".format( + current_streak, + " (longest **{}** days)".format(max_streak) if max_streak else '' ) - embed.add_field( - name="🏆 Leaderboard Position", - value="Time: {}\n LC: {}".format(time_lb_pos, coin_lb_pos) - ) - embed.add_field( - name="💪 Workouts", - value=workouts - ) - embed.add_field( - name="📋 Attendence", - value="TBD" + + stats_table = prop_tabulate(*zip(*stats.items())) + + # Build the time table + time_table = prop_tabulate( + ('Daily', 'Weekly', 'Monthly', 'All Time'), + ["{:02}:{:02}".format(t // 3600, (t // 60) % 60) for t in reversed(study_times)] ) + + # Populate the embed + embed.add_field(name="Study Time", value=time_table) + embed.add_field(name="Statistics", value=stats_table) + + # Add the study league field + if current_badge or next_badge: + current_str = ( + "You are currently in <@&{}>!".format(current_badge.roleid) if current_badge else "No league yet!" + ) + if next_badge: + needed = max(next_badge.required_time - season_time, 0) + next_str = "Study for **{:02}:{:02}** more to achieve <@&{}>.".format( + needed // 3600, + (needed // 60) % 60, + next_badge.roleid + ) + else: + next_str = "You have reached the highest league! Congratulations!" + embed.add_field( + name="Study League", + value="{}\n{}".format(current_str, next_str), + inline=False + ) await ctx.reply(embed=embed) diff --git a/bot/modules/study/timers/__init__.py b/bot/modules/study/timers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/modules/study/top_cmd.py b/bot/modules/study/top_cmd.py index 774f0409..79564c1f 100644 --- a/bot/modules/study/top_cmd.py +++ b/bot/modules/study/top_cmd.py @@ -38,27 +38,20 @@ async def cmd_top(ctx): ) top100 = (ctx.args == "100" or ctx.alias == "top100") - # Flush any pending coin transactions - Lion.sync() - # Fetch the leaderboard exclude = set(m.id for m in ctx.guild_settings.unranked_roles.members) - exclude.update(ctx.client.objects['blacklisted_users']) + exclude.update(ctx.client.user_blacklist()) exclude.update(ctx.client.objects['ignored_members'][ctx.guild.id]) + args = { + 'guildid': ctx.guild.id, + 'select_columns': ('userid', 'total_tracked_time::INTEGER'), + '_extra': "AND total_tracked_time > 0 ORDER BY total_tracked_time DESC " + ("LIMIT 100" if top100 else "") + } if exclude: - user_data = tables.lions.select_where( - guildid=ctx.guild.id, - userid=data.NOT(list(exclude)), - select_columns=('userid', 'tracked_time'), - _extra="AND tracked_time > 0 ORDER BY tracked_time DESC " + ("LIMIT 100" if top100 else "") - ) - else: - user_data = tables.lions.select_where( - guildid=ctx.guild.id, - select_columns=('userid', 'tracked_time'), - _extra="AND tracked_time > 0 ORDER BY tracked_time DESC " + ("LIMIT 100" if top100 else "") - ) + args['userid'] = data.NOT(list(exclude)) + + user_data = tables.members_totals.select_where(**args) # Quit early if the leaderboard is empty if not user_data: diff --git a/bot/modules/study/tracking/__init__.py b/bot/modules/study/tracking/__init__.py new file mode 100644 index 00000000..ba8de231 --- /dev/null +++ b/bot/modules/study/tracking/__init__.py @@ -0,0 +1,3 @@ +from . import data +from . import settings +from . import session_tracker diff --git a/bot/modules/study/tracking/data.py b/bot/modules/study/tracking/data.py new file mode 100644 index 00000000..d9dcae38 --- /dev/null +++ b/bot/modules/study/tracking/data.py @@ -0,0 +1,62 @@ +from data import Table, RowTable, tables +from utils.lib import FieldEnum + + +untracked_channels = Table('untracked_channels') + + +class SessionChannelType(FieldEnum): + """ + The possible session channel types. + """ + # NOTE: "None" stands for Unknown, and the STANDARD description should be replaced with the channel name + STANDARD = 'STANDARD', "Standard" + ACCOUNTABILITY = 'ACCOUNTABILITY', "Accountability Room" + RENTED = 'RENTED', "Private Room" + EXTERNAL = 'EXTERNAL', "Unknown" + + +session_history = Table('session_history') +current_sessions = RowTable( + 'current_sessions', + ('guildid', 'userid', 'channelid', 'channel_type', + 'start_time', + 'live_duration', 'live_start', + 'stream_duration', 'stream_start', + 'video_duration', 'video_start', + 'hourly_coins', 'hourly_live_coins'), + ('guildid', 'userid'), + cache={} # Keep all current sessions in cache +) + + +@current_sessions.save_query +def close_study_session(guildid, userid): + """ + Close a member's current session if it exists and update the member cache. + """ + # Execute the `close_study_session` database function + with current_sessions.conn as conn: + cursor = conn.cursor() + cursor.callproc('close_study_session', (guildid, userid)) + rows = cursor.fetchall() + # The row has been deleted, remove the from current sessions cache + current_sessions.row_cache.pop((guildid, userid), None) + # Use the function output to update the member cache + tables.lions._make_rows(*rows) + + +@session_history.save_query +def study_time_since(guildid, userid, timestamp): + """ + Retrieve the total member study time (in seconds) since the given timestamp. + Includes the current session, if it exists. + """ + with session_history.conn as conn: + cursor = conn.cursor() + cursor.callproc('study_time_since', (guildid, userid, timestamp)) + rows = cursor.fetchall() + return (rows[0][0] if rows else None) or 0 + + +members_totals = Table('members_totals') diff --git a/bot/modules/study/tracking/session_tracker.py b/bot/modules/study/tracking/session_tracker.py new file mode 100644 index 00000000..8158f96a --- /dev/null +++ b/bot/modules/study/tracking/session_tracker.py @@ -0,0 +1,500 @@ +import asyncio +import discord +import logging +import traceback +from typing import Dict +from collections import defaultdict + +from utils.lib import utc_now +from data import tables +from data.conditions import THIS_SHARD +from core import Lion +from meta import client + +from ..module import module +from .data import current_sessions, SessionChannelType +from .settings import untracked_channels, hourly_reward, hourly_live_bonus + + +class Session: + """ + A `Session` describes an ongoing study session by a single guild member. + A member is counted as studying when they are in a tracked voice channel. + + This class acts as an opaque interface to the corresponding `sessions` data row. + """ + __slots__ = ( + 'guildid', + 'userid', + '_expiry_task' + ) + # Global cache of ongoing sessions + sessions: Dict[int, Dict[int, 'Session']] = defaultdict(dict) + + # Global cache of members pending session start (waiting for daily cap reset) + members_pending: Dict[int, Dict[int, asyncio.Task]] = defaultdict(dict) + + def __init__(self, guildid, userid): + self.guildid = guildid + self.userid = userid + + self._expiry_task: asyncio.Task = None + + @classmethod + def get(cls, guildid, userid): + """ + Fetch the current session for the provided member. + If there is no current session, returns `None`. + """ + return cls.sessions[guildid].get(userid, None) + + @classmethod + def start(cls, member: discord.Member, state: discord.VoiceState): + """ + Start a new study session for the provided member. + """ + guildid = member.guild.id + userid = member.id + now = utc_now() + + if userid in cls.sessions[guildid]: + raise ValueError("A session for this member already exists!") + + # If the user is study capped, schedule the session start for the next day + if (lion := Lion.fetch(guildid, userid)).remaining_study_today <= 0: + if pending := cls.members_pending[guildid].pop(userid, None): + pending.cancel() + task = asyncio.create_task(cls._delayed_start(guildid, userid, member, state)) + cls.members_pending[guildid][userid] = task + client.log( + "Member (uid:{}) in (gid:{}) is study capped, " + "delaying session start for {} seconds until start of next day.".format( + userid, guildid, lion.remaining_in_day + ), + context="SESSION_TRACKER", + level=logging.DEBUG + ) + return + + # TODO: More reliable channel type determination + if state.channel.id in tables.rented.row_cache: + channel_type = SessionChannelType.RENTED + elif state.channel.category and state.channel.category.id == lion.guild_settings.accountability_category.data: + channel_type = SessionChannelType.ACCOUNTABILITY + else: + channel_type = SessionChannelType.STANDARD + + current_sessions.create_row( + guildid=guildid, + userid=userid, + channelid=state.channel.id, + channel_type=channel_type, + start_time=now, + live_start=now if (state.self_video or state.self_stream) else None, + stream_start=now if state.self_stream else None, + video_start=now if state.self_video else None, + hourly_coins=hourly_reward.get(guildid).value, + hourly_live_coins=hourly_live_bonus.get(guildid).value + ) + session = cls(guildid, userid).activate() + client.log( + "Started session: {}".format(session.data), + context="SESSION_TRACKER", + level=logging.DEBUG, + ) + + @classmethod + async def _delayed_start(cls, guildid, userid, *args): + delay = Lion.fetch(guildid, userid).remaining_in_day + try: + await asyncio.sleep(delay) + except asyncio.CancelledError: + pass + else: + cls.start(*args) + + @property + def key(self): + """ + RowTable Session identification key. + """ + return (self.guildid, self.userid) + + @property + def lion(self): + """ + The Lion member object associated with this member. + """ + return Lion.fetch(self.guildid, self.userid) + + @property + def data(self): + """ + Row of the `current_sessions` table corresponding to this session. + """ + return current_sessions.fetch(self.key) + + @property + def duration(self): + """ + Current duration of the session. + """ + return (utc_now() - self.data.start_time).total_seconds() + + @property + def coins_earned(self): + """ + Number of coins earned so far. + """ + data = self.data + + coins = self.duration * data.hourly_coins + coins += data.live_duration * data.hourly_live_coins + if data.live_start: + coins += (utc_now() - data.live_start).total_seconds() * data.hourly_live_coins + return coins // 3600 + + def activate(self): + """ + Activate the study session. + This adds the session to the studying members cache, + and schedules the session expiry, based on the daily study cap. + """ + # Add to the active cache + self.sessions[self.guildid][self.userid] = self + + # Schedule the session expiry + self.schedule_expiry() + + # Return self for easy chaining + return self + + def schedule_expiry(self): + """ + Schedule session termination when the user reaches the maximum daily study time. + """ + asyncio.create_task(self._schedule_expiry()) + + async def _schedule_expiry(self): + # Cancel any existing expiry + if self._expiry_task and not self._expiry_task.done(): + self._expiry_task.cancel() + + # Wait for the maximum session length + try: + self._expiry_task = await asyncio.sleep(self.lion.remaining_study_today) + except asyncio.CancelledError: + pass + else: + if self.lion.remaining_study_today <= 0: + # End the session + # Note that the user will not automatically start a new session when the day starts + # TODO: Notify user? Disconnect them? + client.log( + "Session for (uid:{}) in (gid:{}) reached daily guild study cap.\n{}".format( + self.userid, self.guildid, self.data + ), + context="SESSION_TRACKER" + ) + self.finish() + else: + # It's possible the expiry time was pushed forwards while waiting + # If so, reschedule + self.schedule_expiry() + + def finish(self): + """ + Close the study session. + """ + # Note that save_live_status doesn't need to be called here + # The database saving procedure will account for the values. + current_sessions.queries.close_study_session(*self.key) + + # Remove session from active cache + self.sessions[self.guildid].pop(self.userid, None) + + # Cancel any existing expiry task + if self._expiry_task and not self._expiry_task.done(): + self._expiry_task.cancel() + + def save_live_status(self, state: discord.VoiceState): + """ + Update the saved live status of the member. + """ + has_video = state.self_video + has_stream = state.self_stream + is_live = has_video or has_stream + + now = utc_now() + data = self.data + + with data.batch_update(): + # Update video session stats + if data.video_start: + data.video_duration += (now - data.video_start).total_seconds() + data.video_start = now if has_video else None + + # Update stream session stats + if data.stream_start: + data.stream_duration += (now - data.stream_start).total_seconds() + data.stream_start = now if has_stream else None + + # Update overall live session stats + if data.live_start: + data.live_duration += (now - data.live_start).total_seconds() + data.live_start = now if is_live else None + + +async def session_voice_tracker(client, member, before, after): + """ + Voice update event dispatcher for study session tracking. + """ + guild = member.guild + Lion.fetch(guild.id, member.id) + session = Session.get(guild.id, member.id) + + if before.channel == after.channel: + # Voice state change without moving channel + if session and ((before.self_video != after.self_video) or (before.self_stream != after.self_stream)): + # Live status has changed! + session.save_live_status(after) + else: + # Member changed channel + # End the current session and start a new one, if applicable + if session: + if (scid := session.data.channelid) and (not before.channel or scid != before.channel.id): + client.log( + "The previous voice state for " + "member {member.name} (uid:{member.id}) in {guild.name} (gid:{guild.id}) " + "does not match their current study session!\n" + "Session channel is (cid:{scid}), but the previous channel is {previous}.".format( + member=member, + guild=member.guild, + scid=scid, + previous="{0.name} (cid:{0.id})".format(before.channel) if before.channel else "None" + ), + context="SESSION_TRACKER", + level=logging.ERROR + ) + client.log( + "Ending study session for {member.name} (uid:{member.id}) " + "in {member.guild.id} (gid:{member.guild.id}) since they left the voice channel.\n{session}".format( + member=member, + session=session.data + ), + context="SESSION_TRACKER", + post=False + ) + # End the current session + session.finish() + elif pending := Session.members_pending[guild.id].pop(member.id, None): + client.log( + "Cancelling pending study session for {member.name} (uid:{member.id}) " + "in {member.guild.name} (gid:{member.guild.id}) since they left the voice channel.".format( + member=member + ), + context="SESSION_TRACKER", + post=False + ) + pending.cancel() + + if after.channel: + blacklist = client.user_blacklist() + guild_blacklist = client.objects['ignored_members'][guild.id] + untracked = untracked_channels.get(guild.id).data + start_session = ( + (after.channel.id not in untracked) + and (member.id not in blacklist) + and (member.id not in guild_blacklist) + ) + if start_session: + # Start a new session for the member + client.log( + "Starting a new voice channel study session for {member.name} (uid:{member.id}) " + "in {member.guild.name} (gid:{member.guild.id}).".format( + member=member, + ), + context="SESSION_TRACKER", + post=False + ) + session = Session.start(member, after) + + +async def leave_guild_sessions(client, guild): + """ + `guild_leave` hook. + Close all sessions in the guild when we leave. + """ + sessions = list(Session.sessions[guild.id].values()) + for session in sessions: + session.finish() + client.log( + "Left {} (gid:{}) and closed {} ongoing study sessions.".format(guild.name, guild.id, len(sessions)), + context="SESSION_TRACKER" + ) + + +async def join_guild_sessions(client, guild): + """ + `guild_join` hook. + Refresh all sessions for the guild when we rejoin. + """ + # Delete existing current sessions, which should have been closed when we left + # It is possible we were removed from the guild during an outage + current_sessions.delete_where(guildid=guild.id) + + untracked = untracked_channels.get(guild.id).data + members = [ + member + for channel in guild.voice_channels + for member in channel.members + if channel.members and channel.id not in untracked + ] + for member in members: + client.log( + "Starting new session for '{}' (uid: {}) in '{}' (cid: {}) of '{}' (gid: {})".format( + member.name, + member.id, + member.voice.channel.name, + member.voice.channel.id, + member.guild.name, + member.guild.id + ), + context="SESSION_TRACKER", + level=logging.INFO, + post=False + ) + Session.start(member, member.voice) + + # Log newly started sessions + client.log( + "Joined {} (gid:{}) and started {} new study sessions from current voice channel members.".format( + guild.name, + guild.id, + len(members) + ), + context="SESSION_TRACKER", + ) + + +async def _init_session_tracker(client): + """ + Load ongoing saved study sessions into the session cache, + update them depending on the current voice states, + and attach the voice event handler. + """ + # Ensure the client caches are ready and guilds are chunked + await client.wait_until_ready() + + # Pre-cache the untracked channels + await untracked_channels.launch_task(client) + + # Log init start and define logging counters + client.log( + "Loading ongoing study sessions.", + context="SESSION_INIT", + level=logging.DEBUG + ) + resumed = 0 + ended = 0 + + # Grab all ongoing sessions from data + rows = current_sessions.fetch_rows_where(guildid=THIS_SHARD) + + # Iterate through, resume or end as needed + for row in rows: + if (guild := client.get_guild(row.guildid)) is not None and row.channelid is not None: + try: + # Load the Session + session = Session(row.guildid, row.userid) + + # Find the channel and member voice state + voice = None + if channel := guild.get_channel(row.channelid): + voice = next((member.voice for member in channel.members if member.id == row.userid), None) + + # Resume or end as required + if voice and voice.channel: + client.log( + "Resuming ongoing session: {}".format(row), + context="SESSION_INIT", + level=logging.DEBUG + ) + session.activate() + session.save_live_status(voice) + resumed += 1 + else: + client.log( + "Ending already completed session: {}".format(row), + context="SESSION_INIT", + level=logging.DEBUG + ) + session.finish() + ended += 1 + except Exception: + # Fatal error + client.log( + "Fatal error occurred initialising session: {}\n{}".format(row, traceback.format_exc()), + context="SESSION_INIT", + level=logging.CRITICAL + ) + module.ready = False + return + + # Log resumed sessions + client.log( + "Resumed {} ongoing study sessions, and ended {}.".format(resumed, ended), + context="SESSION_INIT", + level=logging.INFO + ) + + # Now iterate through members of all tracked voice channels + # Start sessions if they don't already exist + tracked_channels = [ + channel + for guild in client.guilds + for channel in guild.voice_channels + if channel.members and channel.id not in untracked_channels.get(guild.id).data + ] + new_members = [ + member + for channel in tracked_channels + for member in channel.members + if not Session.get(member.guild.id, member.id) + ] + for member in new_members: + client.log( + "Starting new session for '{}' (uid: {}) in '{}' (cid: {}) of '{}' (gid: {})".format( + member.name, + member.id, + member.voice.channel.name, + member.voice.channel.id, + member.guild.name, + member.guild.id + ), + context="SESSION_INIT", + level=logging.DEBUG + ) + Session.start(member, member.voice) + + # Log newly started sessions + client.log( + "Started {} new study sessions from current voice channel members.".format(len(new_members)), + context="SESSION_INIT", + level=logging.INFO + ) + + # Now that we are in a valid initial state, attach the session event handler + client.add_after_event("voice_state_update", session_voice_tracker) + client.add_after_event("guild_remove", leave_guild_sessions) + client.add_after_event("guild_join", join_guild_sessions) + + +@module.launch_task +async def launch_session_tracker(client): + """ + Launch the study session initialiser. + Doesn't block on the client being ready. + """ + client.objects['sessions'] = Session.sessions + asyncio.create_task(_init_session_tracker(client)) diff --git a/bot/modules/study/admin.py b/bot/modules/study/tracking/settings.py similarity index 74% rename from bot/modules/study/admin.py rename to bot/modules/study/tracking/settings.py index 5861ab95..53f7ecf8 100644 --- a/bot/modules/study/admin.py +++ b/bot/modules/study/tracking/settings.py @@ -1,5 +1,3 @@ -from collections import defaultdict - import settings from settings import GuildSettings from wards import guild_admin @@ -52,10 +50,10 @@ class untracked_channels(settings.ChannelList, settings.ListData, settings.Setti if any(channel.members for channel in guild.voice_channels) ] if active_guildids: + cache = {guildid: [] for guildid in active_guildids} rows = cls._table_interface.select_where( guildid=active_guildids ) - cache = defaultdict(list) for row in rows: cache[row['guildid']].append(row['channelid']) cls._cache.update(cache) @@ -111,3 +109,33 @@ class hourly_live_bonus(settings.Integer, settings.GuildSetting): @property def success_response(self): return "Members will be rewarded an extra `{}` LionCoins per hour if they stream.".format(self.formatted) + + +@GuildSettings.attach_setting +class daily_study_cap(settings.Duration, settings.GuildSetting): + category = "Study Tracking" + + attr_name = "daily_study_cap" + _data_column = "daily_study_cap" + + display_name = "daily_study_cap" + desc = "Maximum amount of recorded study time per member per day." + + _default = 16 * 60 * 60 + _default_multiplier = 60 * 60 + + _max = 25 * 60 * 60 + + long_desc = ( + "The maximum amount of study time that can be recorded for a member per day, " + "intended to remove system encouragement for unhealthy or obsessive behaviour.\n" + "The member may study for longer, but their sessions will not be tracked. " + "The start and end of the day are determined by the member's configured timezone." + ) + + @property + def success_response(self): + # Refresh expiry for all sessions in the guild + [session.schedule_expiry() for session in self.client.objects['sessions'][self.id].values()] + + return "The maximum tracked daily study time is now {}.".format(self.formatted) diff --git a/bot/modules/study/time_tracker.py b/bot/modules/study/tracking/time_tracker.py similarity index 84% rename from bot/modules/study/time_tracker.py rename to bot/modules/study/tracking/time_tracker.py index 26afddd6..46f88ec7 100644 --- a/bot/modules/study/time_tracker.py +++ b/bot/modules/study/tracking/time_tracker.py @@ -7,8 +7,8 @@ from time import time from meta import client from core import Lion -from .module import module -from . import admin +from ..module import module +from .settings import untracked_channels, hourly_reward, hourly_live_bonus last_scan = {} # guildid -> timestamp @@ -36,9 +36,9 @@ def _scan(guild): if interval > 60 * 20: return - untracked = admin.untracked_channels.get(guild.id).data - hourly_reward = admin.hourly_reward.get(guild.id).data - hourly_live_bonus = admin.hourly_live_bonus.get(guild.id).data + untracked = untracked_channels.get(guild.id).data + guild_hourly_reward = hourly_reward.get(guild.id).data + guild_hourly_live_bonus = hourly_live_bonus.get(guild.id).data channel_members = ( channel.members for channel in guild.voice_channels if channel.id not in untracked @@ -47,7 +47,7 @@ def _scan(guild): members = itertools.chain(*channel_members) # TODO filter out blacklisted users - blacklist = client.objects['blacklisted_users'] + blacklist = client.user_blacklist() guild_blacklist = client.objects['ignored_members'][guild.id] for member in members: @@ -61,9 +61,9 @@ def _scan(guild): lion.addTime(interval, flush=False) # Add coins - hour_reward = hourly_reward + hour_reward = guild_hourly_reward if member.voice.self_stream or member.voice.self_video: - hour_reward += hourly_live_bonus + hour_reward += guild_hourly_live_bonus lion.addCoins(hour_reward * interval / (3600), flush=False) @@ -102,7 +102,7 @@ async def _study_tracker(): @module.launch_task async def launch_study_tracker(client): # First pre-load the untracked channels - await admin.untracked_channels.launch_task(client) + await untracked_channels.launch_task(client) asyncio.create_task(_study_tracker()) diff --git a/bot/modules/sysadmin/blacklist.py b/bot/modules/sysadmin/blacklist.py index 12a2ed9b..90202407 100644 --- a/bot/modules/sysadmin/blacklist.py +++ b/bot/modules/sysadmin/blacklist.py @@ -7,6 +7,8 @@ import discord from cmdClient.checks import is_owner from cmdClient.lib import ResponseTimedOut +from meta.sharding import sharded + from .module import module @@ -26,14 +28,14 @@ async def cmd_guildblacklist(ctx, flags): Description: View, add, or remove guilds from the blacklist. """ - blacklist = ctx.client.objects['blacklisted_guilds'] + blacklist = ctx.client.guild_blacklist() if ctx.args: # guildid parsing items = [item.strip() for item in ctx.args.split(',')] if any(not item.isdigit() for item in items): return await ctx.error_reply( - "Please provide guilds as comma seprated guild ids." + "Please provide guilds as comma separated guild ids." ) guildids = set(int(item) for item in items) @@ -80,9 +82,18 @@ async def cmd_guildblacklist(ctx, flags): insert_keys=('guildid', 'ownerid', 'reason') ) - # Check if we are in any of these guilds - to_leave = (ctx.client.get_guild(guildid) for guildid in to_add) - to_leave = [guild for guild in to_leave if guild is not None] + # Leave freshly blacklisted guilds, accounting for shards + to_leave = [] + for guildid in to_add: + guild = ctx.client.get_guild(guildid) + if not guild and sharded: + try: + guild = await ctx.client.fetch_guild(guildid) + except discord.HTTPException: + pass + if guild: + to_leave.append(guild) + for guild in to_leave: await guild.leave() @@ -102,9 +113,8 @@ async def cmd_guildblacklist(ctx, flags): ) # Refresh the cached blacklist after modification - ctx.client.objects['blacklisted_guilds'] = set( - row['guildid'] for row in ctx.client.data.global_guild_blacklist.select_where() - ) + ctx.client.guild_blacklist.cache_clear() + ctx.client.guild_blacklist() else: # Display the current blacklist # First fetch the full blacklist data @@ -183,7 +193,7 @@ async def cmd_userblacklist(ctx, flags): Description: View, add, or remove users from the blacklist. """ - blacklist = ctx.client.objects['blacklisted_users'] + blacklist = ctx.client.user_blacklist() if ctx.args: # userid parsing @@ -245,9 +255,8 @@ async def cmd_userblacklist(ctx, flags): ) # Refresh the cached blacklist after modification - ctx.client.objects['blacklisted_users'] = set( - row['userid'] for row in ctx.client.data.global_user_blacklist.select_where() - ) + ctx.client.user_blacklist.cache_clear() + ctx.client.user_blacklist() else: # Display the current blacklist # First fetch the full blacklist data diff --git a/bot/modules/sysadmin/status.py b/bot/modules/sysadmin/status.py index 83f02cc8..853f6410 100644 --- a/bot/modules/sysadmin/status.py +++ b/bot/modules/sysadmin/status.py @@ -13,19 +13,13 @@ async def update_status(): # TODO: Make globally configurable and saveable global _last_update - if time.time() - _last_update < 30: + if time.time() - _last_update < 60: return _last_update = time.time() - student_count = sum( - len(ch.members) - for guild in client.guilds - for ch in guild.voice_channels - ) - room_count = sum( - len([vc for vc in guild.voice_channels if vc.members]) - for guild in client.guilds + student_count, room_count = client.data.current_sessions.select_one_where( + select_columns=("COUNT(*) AS studying_count", "COUNT(DISTINCT(channelid)) AS channel_count"), ) status = "{} students in {} study rooms!".format(student_count, room_count) diff --git a/bot/modules/workout/tracker.py b/bot/modules/workout/tracker.py index 90eea397..be3438df 100644 --- a/bot/modules/workout/tracker.py +++ b/bot/modules/workout/tracker.py @@ -7,6 +7,7 @@ from core import Lion from settings import GuildSettings from meta import client from data import NULL, tables +from data.conditions import THIS_SHARD from .module import module from .data import workout_sessions @@ -170,7 +171,7 @@ async def workout_voice_tracker(client, member, before, after): if member.bot: return - if member.id in client.objects['blacklisted_users']: + if member.id in client.user_blacklist(): return if member.id in client.objects['ignored_members'][member.guild.id]: return @@ -226,7 +227,8 @@ async def load_workouts(client): client.objects['current_workouts'] = {} # (guildid, userid) -> Row # Process any incomplete workouts workouts = workout_sessions.fetch_rows_where( - duration=NULL + duration=NULL, + guildid=THIS_SHARD ) count = 0 for workout in workouts: diff --git a/config/example-bot.conf b/config/example-bot.conf index b2fc7a48..d2ec5dd1 100644 --- a/config/example-bot.conf +++ b/config/example-bot.conf @@ -1,6 +1,7 @@ [DEFAULT] log_file = bot.log log_channel = +error_channel = guild_log_channel = prefix = ! @@ -10,4 +11,6 @@ owners = 413668234269818890, 389399222400712714 database = dbname=lionbot data_appid = LionBot +shard_count = 1 + lion_sync_period = 60 diff --git a/data/migration/v5-v6/migration.sql b/data/migration/v5-v6/migration.sql new file mode 100644 index 00000000..eae1f88b --- /dev/null +++ b/data/migration/v5-v6/migration.sql @@ -0,0 +1,180 @@ +-- DROP TYPE IF EXISTS SessionChannelType CASCADE; +-- DROP TABLE IF EXISTS session_history CASCADE; +-- DROP TABLE IF EXISTS current_sessions CASCADE; +-- DROP FUNCTION IF EXISTS close_study_session(_guildid BIGINT, _userid BIGINT); +-- DROP FUNCTION IF EXISTS study_time_since(_guildid BIGINT, _userid BIGINT, _timestamp TIMESTAMPTZ) + +-- DROP VIEW IF EXISTS current_sessions_totals CASCADE; + +DROP VIEW IF EXISTS member_totals CASCADE; +DROP VIEW IF EXISTS member_ranks CASCADE; +DROP VIEW IF EXISTS current_study_badges CASCADE; +DROP VIEW IF EXISTS new_study_badges CASCADE; + + +CREATE TYPE SessionChannelType AS ENUM ( + 'STANDARD', + 'ACCOUNTABILITY', + 'RENTED', + 'EXTERNAL' +); + +CREATE TABLE session_history( + sessionid SERIAL PRIMARY KEY, + guildid BIGINT NOT NULL, + userid BIGINT NOT NULL, + channelid BIGINT, + channel_type SessionChannelType, + start_time TIMESTAMPTZ NOT NULL, + duration INTEGER NOT NULL, + coins_earned INTEGER NOT NULL, + live_duration INTEGER DEFAULT 0, + stream_duration INTEGER DEFAULT 0, + video_duration INTEGER DEFAULT 0, + FOREIGN KEY (guildid, userid) REFERENCES members (guildid, userid) ON DELETE CASCADE +); +CREATE INDEX session_history_members ON session_history (guildid, userid, start_time); + +CREATE TABLE current_sessions( + guildid BIGINT NOT NULL, + userid BIGINT NOT NULL, + channelid BIGINT, + channel_type SessionChannelType, + start_time TIMESTAMPTZ DEFAULT now(), + live_duration INTEGER DEFAULT 0, + live_start TIMESTAMPTZ, + stream_duration INTEGER DEFAULT 0, + stream_start TIMESTAMPTZ, + video_duration INTEGER DEFAULT 0, + video_start TIMESTAMPTZ, + hourly_coins INTEGER NOT NULL, + hourly_live_coins INTEGER NOT NULL, + FOREIGN KEY (guildid, userid) REFERENCES members (guildid, userid) ON DELETE CASCADE +); +CREATE UNIQUE INDEX current_session_members ON current_sessions (guildid, userid); + + +CREATE FUNCTION close_study_session(_guildid BIGINT, _userid BIGINT) + RETURNS SETOF members +AS $$ + BEGIN + RETURN QUERY + WITH + current_sesh AS ( + DELETE FROM current_sessions + WHERE guildid=_guildid AND userid=_userid + RETURNING + *, + EXTRACT(EPOCH FROM (NOW() - start_time)) AS total_duration, + stream_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - stream_start)), 0) AS total_stream_duration, + video_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - video_start)), 0) AS total_video_duration, + live_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - live_start)), 0) AS total_live_duration + ), saved_sesh AS ( + INSERT INTO session_history ( + guildid, userid, channelid, channel_type, start_time, + duration, stream_duration, video_duration, live_duration, + coins_earned + ) SELECT + guildid, userid, channelid, channel_type, start_time, + total_duration, total_stream_duration, total_video_duration, total_live_duration, + (total_duration * hourly_coins + live_duration * hourly_live_coins) / 3600 + FROM current_sesh + RETURNING * + ) + UPDATE members + SET + tracked_time=(tracked_time + saved_sesh.duration), + coins=(coins + saved_sesh.coins_earned) + FROM saved_sesh + WHERE members.guildid=saved_sesh.guildid AND members.userid=saved_sesh.userid + RETURNING members.*; + END; +$$ LANGUAGE PLPGSQL; + + + +CREATE VIEW current_sessions_totals AS + SELECT + *, + EXTRACT(EPOCH FROM (NOW() - start_time)) AS total_duration, + stream_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - stream_start)), 0) AS total_stream_duration, + video_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - video_start)), 0) AS total_video_duration, + live_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - live_start)), 0) AS total_live_duration + FROM current_sessions; + + +CREATE VIEW members_totals AS + SELECT + *, + sesh.start_time AS session_start, + tracked_time + COALESCE(sesh.total_duration, 0) AS total_tracked_time, + coins + COALESCE((sesh.total_duration * sesh.hourly_coins + sesh.live_duration * sesh.hourly_live_coins) / 3600, 0) AS total_coins + FROM members + LEFT JOIN current_sessions_totals sesh USING (guildid, userid); + + +CREATE VIEW member_ranks AS + SELECT + *, + row_number() OVER (PARTITION BY guildid ORDER BY total_tracked_time DESC, userid ASC) AS time_rank, + row_number() OVER (PARTITION BY guildid ORDER BY total_coins DESC, userid ASC) AS coin_rank + FROM members_totals; + +CREATE VIEW current_study_badges AS + SELECT + *, + (SELECT r.badgeid + FROM study_badges r + WHERE r.guildid = members_totals.guildid AND members_totals.total_tracked_time > r.required_time + ORDER BY r.required_time DESC + LIMIT 1) AS current_study_badgeid + FROM members_totals; + +CREATE VIEW new_study_badges AS + SELECT + current_study_badges.* + FROM current_study_badges + WHERE + last_study_badgeid IS DISTINCT FROM current_study_badgeid + ORDER BY guildid; + + +CREATE FUNCTION study_time_since(_guildid BIGINT, _userid BIGINT, _timestamp TIMESTAMPTZ) + RETURNS INTEGER +AS $$ + BEGIN + RETURN ( + SELECT + SUM( + CASE + WHEN start_time >= _timestamp THEN duration + ELSE EXTRACT(EPOCH FROM (end_time - _timestamp)) + END + ) + FROM ( + SELECT + start_time, + duration, + (start_time + duration * interval '1 second') AS end_time + FROM session_history + WHERE + guildid=_guildid + AND userid=_userid + AND (start_time + duration * interval '1 second') >= _timestamp + UNION + SELECT + start_time, + EXTRACT(EPOCH FROM (NOW() - start_time)) AS duration, + NOW() AS end_time + FROM current_sessions + WHERE + guildid=_guildid + AND userid=_userid + ) AS sessions + ); + END; +$$ LANGUAGE PLPGSQL; + +ALTER TABLE guild_config ADD COLUMN daily_study_cap INTEGER; + +INSERT INTO VersionHistory (version, author) VALUES (6, 'v5-v6 Migration'); diff --git a/data/schema.sql b/data/schema.sql index fe0158be..a7a4af31 100644 --- a/data/schema.sql +++ b/data/schema.sql @@ -4,7 +4,7 @@ CREATE TABLE VersionHistory( time TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP NOT NULL, author TEXT ); -INSERT INTO VersionHistory (version, author) VALUES (5, 'Initial Creation'); +INSERT INTO VersionHistory (version, author) VALUES (6, 'Initial Creation'); CREATE OR REPLACE FUNCTION update_timestamp_column() @@ -77,7 +77,8 @@ CREATE TABLE guild_config( greeting_message TEXT, returning_message TEXT, starting_funds INTEGER, - persist_roles BOOLEAN + persist_roles BOOLEAN, + daily_study_cap INTEGER ); CREATE TABLE ignored_members( @@ -407,24 +408,165 @@ CREATE INDEX member_timestamps ON members (_timestamp); CREATE TRIGGER update_members_timstamp BEFORE UPDATE ON members FOR EACH ROW EXECUTE PROCEDURE update_timestamp_column(); +-- }}} + +-- Study Session Data {{{ +CREATE TYPE SessionChannelType AS ENUM ( + 'STANDARD', + 'ACCOUNTABILITY', + 'RENTED', + 'EXTERNAL', +); + + +CREATE TABLE session_history( + sessionid SERIAL PRIMARY KEY, + guildid BIGINT NOT NULL, + userid BIGINT NOT NULL, + channelid BIGINT, + channel_type SessionChannelType, + start_time TIMESTAMPTZ NOT NULL, + duration INTEGER NOT NULL, + coins_earned INTEGER NOT NULL, + live_duration INTEGER DEFAULT 0, + stream_duration INTEGER DEFAULT 0, + video_duration INTEGER DEFAULT 0, + FOREIGN KEY (guildid, userid) REFERENCES members (guildid, userid) ON DELETE CASCADE +); +CREATE INDEX session_history_members ON session_history (guildid, userid, start_time); + +CREATE TABLE current_sessions( + guildid BIGINT NOT NULL, + userid BIGINT NOT NULL, + channelid BIGINT, + channel_type SessionChannelType, + start_time TIMESTAMPTZ DEFAULT now(), + live_duration INTEGER DEFAULT 0, + live_start TIMESTAMPTZ, + stream_duration INTEGER DEFAULT 0, + stream_start TIMESTAMPTZ, + video_duration INTEGER DEFAULT 0, + video_start TIMESTAMPTZ, + hourly_coins INTEGER NOT NULL, + hourly_live_coins INTEGER NOT NULL, + FOREIGN KEY (guildid, userid) REFERENCES members (guildid, userid) ON DELETE CASCADE +); +CREATE UNIQUE INDEX current_session_members ON current_sessions (guildid, userid); + + +CREATE FUNCTION study_time_since(_guildid BIGINT, _userid BIGINT, _timestamp TIMESTAMPTZ) + RETURNS INTEGER +AS $$ + BEGIN + RETURN ( + SELECT + SUM( + CASE + WHEN start_time >= _timestamp THEN duration + ELSE EXTRACT(EPOCH FROM (end_time - _timestamp)) + END + ) + FROM ( + SELECT + start_time, + duration, + (start_time + duration * interval '1 second') AS end_time + FROM session_history + WHERE + guildid=_guildid + AND userid=_userid + AND (start_time + duration * interval '1 second') >= _timestamp + UNION + SELECT + start_time, + EXTRACT(EPOCH FROM (NOW() - start_time)) AS duration, + NOW() AS end_time + FROM current_sessions + WHERE + guildid=_guildid + AND userid=_userid + ) AS sessions + ); + END; +$$ LANGUAGE PLPGSQL; + + +CREATE FUNCTION close_study_session(_guildid BIGINT, _userid BIGINT) + RETURNS SETOF members +AS $$ + BEGIN + RETURN QUERY + WITH + current_sesh AS ( + DELETE FROM current_sessions + WHERE guildid=_guildid AND userid=_userid + RETURNING + *, + EXTRACT(EPOCH FROM (NOW() - start_time)) AS total_duration, + stream_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - stream_start)), 0) AS total_stream_duration, + video_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - video_start)), 0) AS total_video_duration, + live_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - live_start)), 0) AS total_live_duration + ), saved_sesh AS ( + INSERT INTO session_history ( + guildid, userid, channelid, channel_type, start_time, + duration, stream_duration, video_duration, live_duration, + coins_earned + ) SELECT + guildid, userid, channelid, channel_type, start_time, + total_duration, total_stream_duration, total_video_duration, total_live_duration, + (total_duration * hourly_coins + live_duration * hourly_live_coins) / 3600 + FROM current_sesh + RETURNING * + ) + UPDATE members + SET + tracked_time=(tracked_time + saved_sesh.duration), + coins=(coins + saved_sesh.coins_earned) + FROM saved_sesh + WHERE members.guildid=saved_sesh.guildid AND members.userid=saved_sesh.userid + RETURNING members.*; + END; +$$ LANGUAGE PLPGSQL; + + +CREATE VIEW current_sessions_totals AS + SELECT + *, + EXTRACT(EPOCH FROM (NOW() - start_time)) AS total_duration, + stream_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - stream_start)), 0) AS total_stream_duration, + video_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - video_start)), 0) AS total_video_duration, + live_duration + COALESCE(EXTRACT(EPOCH FROM (NOW() - live_start)), 0) AS total_live_duration + FROM current_sessions; + + +CREATE VIEW members_totals AS + SELECT + *, + sesh.start_time AS session_start, + tracked_time + COALESCE(sesh.total_duration, 0) AS total_tracked_time, + coins + COALESCE((sesh.total_duration * sesh.hourly_coins + sesh.live_duration * sesh.hourly_live_coins) / 3600, 0) AS total_coins + FROM members + LEFT JOIN current_sessions_totals sesh USING (guildid, userid); + CREATE VIEW member_ranks AS SELECT *, - row_number() OVER (PARTITION BY guildid ORDER BY tracked_time DESC, userid ASC) AS time_rank, - row_number() OVER (PARTITION BY guildid ORDER BY coins DESC, userid ASC) AS coin_rank - FROM members; - + row_number() OVER (PARTITION BY guildid ORDER BY total_tracked_time DESC, userid ASC) AS time_rank, + row_number() OVER (PARTITION BY guildid ORDER BY total_coins DESC, userid ASC) AS coin_rank + FROM members_totals; +-- }}} +-- Study Badge Data {{{ CREATE VIEW current_study_badges AS SELECT *, (SELECT r.badgeid FROM study_badges r - WHERE r.guildid = members.guildid AND members.tracked_time > r.required_time + WHERE r.guildid = members_totals.guildid AND members_totals.total_tracked_time > r.required_time ORDER BY r.required_time DESC LIMIT 1) AS current_study_badgeid - FROM members; + FROM members_totals; CREATE VIEW new_study_badges AS SELECT @@ -527,6 +669,7 @@ CREATE TABLE reaction_role_expiring( reactionid INTEGER REFERENCES reaction_role_reactions (reactionid) ON DELETE SET NULL ); CREATE UNIQUE INDEX reaction_role_expiry_members ON reaction_role_expiring (guildid, userid, roleid); +-- }}} -- Member Role Data {{{ CREATE TABLE past_member_roles( @@ -538,4 +681,5 @@ CREATE TABLE past_member_roles( ); CREATE INDEX member_role_persistence_members ON past_member_roles (guildid, userid); -- }}} + -- vim: set fdm=marker: