From 916de8dd4c503dab6d885ea1792560beff1be659 Mon Sep 17 00:00:00 2001 From: Conatum Date: Fri, 18 Nov 2022 08:49:47 +0200 Subject: [PATCH] rewrite: Extending core cog. --- bot/core/__init__.py | 5 ++ bot/core/cog.py | 71 +++++++++++++++++++ bot/core/data.py | 160 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 236 insertions(+) create mode 100644 bot/core/__init__.py create mode 100644 bot/core/cog.py create mode 100644 bot/core/data.py diff --git a/bot/core/__init__.py b/bot/core/__init__.py new file mode 100644 index 00000000..6217ee8f --- /dev/null +++ b/bot/core/__init__.py @@ -0,0 +1,5 @@ +from .cog import CoreCog + + +async def setup(bot): + await bot.add_cog(CoreCog(bot)) diff --git a/bot/core/cog.py b/bot/core/cog.py new file mode 100644 index 00000000..f561c011 --- /dev/null +++ b/bot/core/cog.py @@ -0,0 +1,71 @@ +from typing import Optional + +import discord + +from meta import LionBot, LionCog +from meta.app import shardname, appname +from meta.logger import log_wrap +from utils.lib import utc_now + +from settings.groups import SettingGroup + +from .data import CoreData + + +class CoreCog(LionCog): + def __init__(self, bot: LionBot): + self.bot = bot + self.data = CoreData() + bot.db.load_registry(self.data) + + self.app_config: Optional[CoreData.AppConfig] = None + self.bot_config: Optional[CoreData.BotConfig] = None + self.shard_data: Optional[CoreData.Shard] = None + + # Some global setting registries + # Do not use these for direct setting access + # Instead, import the setting directly or use the cog API + self.bot_setting_groups: list[SettingGroup] = [] + self.guild_setting_groups: list[SettingGroup] = [] + self.user_setting_groups: list[SettingGroup] = [] + + self.app_cmd_cache: list[discord.app_commands.AppCommand] = [] + self.cmd_name_cache: dict[str, discord.app_commands.AppCommand] = {} + + async def cog_load(self): + # Fetch (and possibly create) core data rows. + conn = await self.bot.db.get_connection() + async with conn.transaction(): + self.app_config = await self.data.AppConfig.fetch_or_create(appname) + self.bot_config = await self.data.BotConfig.fetch_or_create(appname) + self.shard_data = await self.data.Shard.fetch_or_create( + shardname, + appname=appname, + shard_id=self.bot.shard_id, + shard_count=self.bot.shard_count + ) + self.bot.add_listener(self.shard_update_guilds, name='on_guild_join') + self.bot.add_listener(self.shard_update_guilds, name='on_guild_leave') + + self.bot.core = self + + # Load the app command cache + for guildid in self.bot.testing_guilds: + self.app_cmd_cache += await self.bot.tree.fetch_commands(guild=discord.Object(guildid)) + self.app_cmd_cache += await self.bot.tree.fetch_commands() + self.cmd_name_cache = {cmd.name: cmd for cmd in self.app_cmd_cache} + + async def cog_unload(self): + self.bot.remove_listener(self.shard_update_guilds, name='on_guild_join') + self.bot.remove_listener(self.shard_update_guilds, name='on_guild_leave') + self.bot.core = None + + @LionCog.listener('on_ready') + @log_wrap(action='Touch shard data') + async def touch_shard_data(self): + # Update the last login and guild count for this shard + await self.shard_data.update(last_login=utc_now(), guild_count=len(self.bot.guilds)) + + @log_wrap(action='Update shard guilds') + async def shard_update_guilds(self, guild): + await self.shard_data.update(guild_count=len(self.bot.guilds)) diff --git a/bot/core/data.py b/bot/core/data.py new file mode 100644 index 00000000..f71eecd7 --- /dev/null +++ b/bot/core/data.py @@ -0,0 +1,160 @@ +from cachetools import TTLCache + +from data import Table, Registry, Column, RowModel +from data.columns import Integer, String, Bool, Timestamp + + +class CoreData(Registry, name="core"): + class AppConfig(RowModel): + """ + Schema + ------ + CREATE TABLE app_config( + appname TEXT PRIMARY KEY, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """ + _tablename_ = 'app_config' + + appname = String(primary=True) + created_at = Timestamp() + + class BotConfig(RowModel): + """ + Schema + ------ + CREATE TABLE bot_config( + appname TEXT PRIMARY KEY REFERENCES app_config(appname) ON DELETE CASCADE, + default_skin TEXT + ); + """ + _tablename_ = 'bot_config' + + appname = String(primary=True) + default_skin = String() + + class Shard(RowModel): + """ + Schema + ------ + CREATE TABLE shard_data( + shardname TEXT PRIMARY KEY, + appname TEXT REFERENCES bot_config(appname) ON DELETE CASCADE, + shard_id INTEGER NOT NULL, + shard_count INTEGER NOT NULL, + last_login TIMESTAMPTZ, + guild_count INTEGER + ); + """ + _tablename_ = 'shard_data' + + shardname = String(primary=True) + appname = String() + shard_id = Integer() + shard_count = Integer() + last_login = Timestamp() + guild_count = Integer() + + class User(RowModel): + """User model, representing configuration data for a single user.""" + + _tablename_ = "user_config" + _cache_: TTLCache[tuple[int], 'User'] = TTLCache(5000, ttl=60*5) + + userid = Integer(primary=True) + timezone = Column() + topgg_vote_reminder = Column() + avatar_hash = String() + gems = Integer() + + class Guild(RowModel): + """Guild model, representing configuration data for a single guild.""" + + _tablename_ = "guild_config" + _cache_: TTLCache[tuple[int], 'Guild'] = TTLCache(2500, ttl=60*5) + + guildid = Integer(primary=True) + + admin_role = Integer() + mod_role = Integer() + event_log_channel = Integer() + mod_log_channel = Integer() + alert_channel = Integer() + + studyban_role = Integer() + max_study_bans = Integer() + + min_workout_length = Integer() + workout_reward = Integer() + + max_tasks = Integer() + task_reward = Integer() + task_reward_limit = Integer() + + study_hourly_reward = Integer() + study_hourly_live_bonus = Integer() + daily_study_cap = Integer() + + renting_price = Integer() + renting_category = Integer() + renting_cap = Integer() + renting_role = Integer() + renting_sync_perms = Bool() + + accountability_category = Integer() + accountability_lobby = Integer() + accountability_bonus = Integer() + accountability_reward = Integer() + accountability_price = Integer() + + video_studyban = Bool() + video_grace_period = Integer() + + greeting_channel = Integer() + greeting_message = String() + returning_message = String() + + starting_funds = Integer() + persist_roles = Bool() + + pomodoro_channel = Integer() + + name = String() + + unranked_rows = Table('unranked_rows') + + donator_roles = Table('donator_roles') + + class Member(RowModel): + """Member model, representing configuration data for a single member.""" + + _tablename_ = 'members' + _cache_: TTLCache[tuple[int, int], 'Member'] = TTLCache(5000, ttl=60*5) + + guildid = Integer(primary=True) + userid = Integer(primary=True) + + tracked_time = Integer() + coins = Integer() + + workout_count = Integer() + last_workout_start = Column() + revision_mute_count = Integer() + last_study_badgeid = Integer() + video_warned = Bool() + display_name = String() + + _timestamp = Column() + + @classmethod + async def add_pending(cls, pending: tuple[int, int, int]) -> list['Member']: + """ + Safely add pending coins to a list of members. + + Arguments + --------- + pending: + List of tuples of the form `(guildid, userid, pending_coins)`. + """ + # TODO: Replace with copy syntax/query? + ...