From 2eea40f6790b9132a36934b1e63593d419650c94 Mon Sep 17 00:00:00 2001 From: Conatum Date: Sun, 20 Nov 2022 08:34:18 +0200 Subject: [PATCH] rewrite: Snapshots and Lion. --- bot/analytics/cog.py | 3 + bot/core/cog.py | 12 +- bot/core/data.py | 172 +++++++++++++++++++++++++--- bot/core/lion.py | 124 ++++++++++++++++++++ bot/data/columns.py | 11 +- bot/data/models.py | 6 +- bot/meta/LionBot.py | 2 +- bot/meta/LionContext.py | 3 + data/migration/v12-13/migration.sql | 18 ++- 9 files changed, 322 insertions(+), 29 deletions(-) create mode 100644 bot/core/lion.py diff --git a/bot/analytics/cog.py b/bot/analytics/cog.py index 78c24de8..ae4a00fa 100644 --- a/bot/analytics/cog.py +++ b/bot/analytics/cog.py @@ -51,6 +51,9 @@ class Analytics(LionCog): elif before.channel and not after.channel: # Member left channel action = VoiceAction.LEFT + else: + # Member change state, we don't need to deal with that + return event = VoiceEvent( appname=appname, diff --git a/bot/core/cog.py b/bot/core/cog.py index d584d537..a0978642 100644 --- a/bot/core/cog.py +++ b/bot/core/cog.py @@ -2,7 +2,7 @@ from typing import Optional import discord -from meta import LionBot, LionCog +from meta import LionBot, LionCog, LionContext from meta.app import shardname, appname from meta.logger import log_wrap from utils.lib import utc_now @@ -10,6 +10,7 @@ from utils.lib import utc_now from settings.groups import SettingGroup from .data import CoreData +from .lion import Lions class CoreCog(LionCog): @@ -17,6 +18,7 @@ class CoreCog(LionCog): self.bot = bot self.data = CoreData() bot.db.load_registry(self.data) + self.lions = Lions(bot) self.app_config: Optional[CoreData.AppConfig] = None self.bot_config: Optional[CoreData.BotConfig] = None @@ -32,6 +34,12 @@ class CoreCog(LionCog): self.app_cmd_cache: list[discord.app_commands.AppCommand] = [] self.cmd_name_cache: dict[str, discord.app_commands.AppCommand] = {} + async def bot_check_once(self, ctx: LionContext): + lion = await self.lions.fetch(ctx.guild.id if ctx.guild else 0, ctx.author.id) + await lion.touch_discord_models(ctx.author) + ctx.alion = lion + return True + async def cog_load(self): # Fetch (and possibly create) core data rows. conn = await self.bot.db.get_connection() @@ -48,6 +56,7 @@ class CoreCog(LionCog): self.bot.add_listener(self.shard_update_guilds, name='on_guild_remove') self.bot.core = self + await self.bot.add_cog(self.lions) # Load the app command cache for guildid in self.bot.testing_guilds: @@ -56,6 +65,7 @@ class CoreCog(LionCog): self.cmd_name_cache = {cmd.name: cmd for cmd in self.app_cmd_cache} async def cog_unload(self): + await self.bot.remove_cog(self.lions.qualified_name) 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 diff --git a/bot/core/data.py b/bot/core/data.py index f71eecd7..41559b3a 100644 --- a/bot/core/data.py +++ b/bot/core/data.py @@ -1,6 +1,9 @@ +from itertools import chain +from psycopg import sql from cachetools import TTLCache from data import Table, Registry, Column, RowModel +from data.models import WeakCache from data.columns import Integer, String, Bool, Timestamp @@ -56,22 +59,86 @@ class CoreData(Registry, name="core"): guild_count = Integer() class User(RowModel): - """User model, representing configuration data for a single user.""" + """ + User model, representing configuration data for a single user. + + Schema + ------ + CREATE TABLE user_config( + userid BIGINT PRIMARY KEY, + timezone TEXT, + topgg_vote_reminder BOOLEAN, + avatar_hash TEXT, + name TEXT, + API_timestamp BIGINT, + gems INTEGER DEFAULT 0, + first_seen TIMESTAMPTZ DEFAULT now(), + last_seen TIMESTAMPTZ + ); + """ _tablename_ = "user_config" - _cache_: TTLCache[tuple[int], 'User'] = TTLCache(5000, ttl=60*5) + _cache_: WeakCache[tuple[int], 'CoreData.User'] = WeakCache(TTLCache(1000, ttl=60*5)) userid = Integer(primary=True) - timezone = Column() - topgg_vote_reminder = Column() + timezone = String() + topgg_vote_reminder = Bool() avatar_hash = String() + name = String() + API_timestamp = Integer() gems = Integer() + first_seen = Timestamp() + last_seen = Timestamp() class Guild(RowModel): - """Guild model, representing configuration data for a single guild.""" + """ + Guild model, representing configuration data for a single guild. + + Schema + ------ + CREATE TABLE guild_config( + guildid BIGINT PRIMARY KEY, + admin_role BIGINT, + mod_role BIGINT, + event_log_channel BIGINT, + mod_log_channel BIGINT, + alert_channel BIGINT, + studyban_role BIGINT, + 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, + renting_price INTEGER, + renting_category BIGINT, + renting_cap INTEGER, + renting_role BIGINT, + renting_sync_perms BOOLEAN, + accountability_category BIGINT, + accountability_lobby BIGINT, + accountability_bonus INTEGER, + accountability_reward INTEGER, + accountability_price INTEGER, + video_studyban BOOLEAN, + video_grace_period INTEGER, + greeting_channel BIGINT, + greeting_message TEXT, + returning_message TEXT, + starting_funds INTEGER, + persist_roles BOOLEAN, + daily_study_cap INTEGER, + pomodoro_channel BIGINT, + name TEXT, + first_joined_at TIMESTAMPTZ DEFAULT now(), + left_at TIMESTAMPTZ + ); + + """ _tablename_ = "guild_config" - _cache_: TTLCache[tuple[int], 'Guild'] = TTLCache(2500, ttl=60*5) + _cache_: WeakCache[tuple[int], 'CoreData.Guild'] = WeakCache(TTLCache(1000, ttl=60*5)) guildid = Integer(primary=True) @@ -121,15 +188,41 @@ class CoreData(Registry, name="core"): name = String() + first_joined_at = Timestamp() + left_at = Timestamp() + unranked_rows = Table('unranked_rows') donator_roles = Table('donator_roles') - class Member(RowModel): - """Member model, representing configuration data for a single member.""" + member_ranks = Table('member_ranks') + class Member(RowModel): + """ + Member model, representing configuration data for a single member. + + Schema + ------ + CREATE TABLE members( + guildid BIGINT, + userid BIGINT, + tracked_time INTEGER DEFAULT 0, + coins INTEGER DEFAULT 0, + workout_count INTEGER DEFAULT 0, + revision_mute_count INTEGER DEFAULT 0, + last_workout_start TIMESTAMP, + last_study_badgeid INTEGER REFERENCES study_badges ON DELETE SET NULL, + video_warned BOOLEAN DEFAULT FALSE, + display_name TEXT, + first_joined TIMESTAMPTZ DEFAULT now(), + last_left TIMESTAMPTZ, + _timestamp TIMESTAMP DEFAULT (now() at time zone 'utc'), + PRIMARY KEY(guildid, userid) + ); + CREATE INDEX member_timestamps ON members (_timestamp); + """ _tablename_ = 'members' - _cache_: TTLCache[tuple[int, int], 'Member'] = TTLCache(5000, ttl=60*5) + _cache_: WeakCache[tuple[int, int], 'CoreData.Member'] = WeakCache(TTLCache(5000, ttl=60*5)) guildid = Integer(primary=True) userid = Integer(primary=True) @@ -138,16 +231,18 @@ class CoreData(Registry, name="core"): coins = Integer() workout_count = Integer() - last_workout_start = Column() revision_mute_count = Integer() + last_workout_start = Timestamp() last_study_badgeid = Integer() video_warned = Bool() display_name = String() - _timestamp = Column() + first_joined = Timestamp() + last_left = Timestamp() + _timestamp = Timestamp() @classmethod - async def add_pending(cls, pending: tuple[int, int, int]) -> list['Member']: + async def add_pending(cls, pending: list[tuple[int, int, int]]) -> list['CoreData.Member']: """ Safely add pending coins to a list of members. @@ -156,5 +251,56 @@ class CoreData(Registry, name="core"): pending: List of tuples of the form `(guildid, userid, pending_coins)`. """ + query = sql.SQL(""" + UPDATE members + SET + coins = LEAST(coins + t.coin_diff, 2147483647) + FROM + (VALUES {}) + AS + t (guildid, userid, coin_diff) + WHERE + members.guildid = t.guildid + AND + members.userid = t.userid + RETURNING * + """).format( + sql.SQL(', ').join( + sql.SQL("({}, {}, {})").format(sql.Placeholder(), sql.Placeholder(), sql.Placeholder()) + for _ in pending + ) + ) # TODO: Replace with copy syntax/query? - ... + conn = await cls.table.connector.get_connection() + async with conn.cursor() as cursor: + await cursor.execute( + query, + tuple(chain(*pending)) + ) + rows = await cursor.fetchall() + return cls._make_rows(*rows) + + @classmethod + async def get_member_rank(cls, guildid, userid, untracked): + """ + Get the time and coin ranking for the given member, ignoring the provided untracked members. + """ + conn = await cls.table.connector.get_connection() + async with conn.cursor() as curs: + await 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 (await curs.fetchone()) or (None, None) diff --git a/bot/core/lion.py b/bot/core/lion.py new file mode 100644 index 00000000..ae790fe7 --- /dev/null +++ b/bot/core/lion.py @@ -0,0 +1,124 @@ +from typing import Optional +from cachetools import LRUCache +import discord + +from meta import LionCog, LionBot, LionContext +from settings import InteractiveSetting +from utils.lib import utc_now +from data import WeakCache + +from .data import CoreData + + +class Lion: + """ + A Lion is a high level representation of a Member in the LionBot paradigm. + + All members interacted with by the application should be available as Lions. + It primarily provides an interface to the User and Member data. + Lion also provides centralised access to various Member properties and methods, + that would normally be served by other cogs. + + Many Lion methods may only be used when the required cogs and extensions are loaded. + A Lion may exist without a Bot instance or a Member in cache, + although the functionality available will be more limited. + + There is no guarantee that a corresponding discord Member actually exists. + """ + __slots__ = ('data', 'user_data', 'guild_data', '_member', '__weakref__') + + def __init__(self, data: CoreData.Member, user_data: CoreData.User, guild_data: CoreData.Guild): + self.data = data + self.user_data = user_data + self.guild_data = guild_data + + self._member: Optional[discord.Member] = None + + # Data properties + + @property + def key(self): + return (self.data.guildid, self.data.userid) + + @property + def guildid(self): + return self.data.guildid + + @property + def userid(self): + return self.data.userid + + @classmethod + def get(cls, guildid, userid): + return cls._cache_.get((guildid, userid), None) + + # Setting interfaces + # Each of these return an initialised member setting + + @property + def timezone(self): + pass + + @property + def locale(self): + pass + + # Time utilities + @property + def now(self): + """ + Returns current time-zone aware time for the member. + """ + pass + + # Discord data cache + async def touch_discord_models(self, member: discord.Member): + """ + Update the stored discord data from the givem member. + Intended to be used when we get member data from events that may not be available in cache. + """ + # Can we do these in one query? + if member.guild and (self.guild_data.name != member.guild.name): + await self.guild_data.update(name=member.guild.name) + + avatar_key = member.avatar.key if member.avatar else None + await self.user_data.update(avatar_hash=avatar_key, name=member.name, last_seen=utc_now()) + + if member.display_name != self.data.display_name: + await self.data.update(display_name=member.display_name) + + +class Lions(LionCog): + def __init__(self, bot: LionBot): + self.bot = bot + + # Full Lions cache + # Don't expire Lions with strong references + self._cache_: WeakCache[tuple[int, int], 'Lion'] = WeakCache(LRUCache(5000)) + + self._settings_: dict[str, InteractiveSetting] = {} + + async def fetch(self, guildid, userid) -> Lion: + """ + Fetch or create the given Member. + If the guild or user row doesn't exist, also creates it. + Relies on the core cog existing, to retrieve the core data. + """ + # TODO: Find a way to reduce this to one query, while preserving cache + lion = self._cache_.get((guildid, userid)) + if lion is None: + if self.bot.core: + data = self.bot.core.data + else: + raise ValueError("Cannot fetch Lion before core module is attached.") + + guild = await data.Guild.fetch_or_create(guildid) + user = await data.User.fetch_or_create(userid) + member = await data.Member.fetch_or_create(guildid, userid) + lion = Lion(member, user, guild) + self._cache_[(guildid, userid)] = lion + return lion + + def add_model_setting(self, setting: InteractiveSetting): + self._settings_[setting.__class__.__name__] = setting + return setting diff --git a/bot/data/columns.py b/bot/data/columns.py index 10298a82..252db83b 100644 --- a/bot/data/columns.py +++ b/bot/data/columns.py @@ -119,10 +119,9 @@ class Column(ColumnExpr, Generic[T]): def __set_name__(self, owner, name): # Only allow setting the owner once - if self.owner is None: - self.name = self.name or name - self.owner = owner - self.expr = sql.Identifier(self.owner._schema_, self.owner._tablename_, self.name) + self.name = self.name or name + self.owner = owner + self.expr = sql.Identifier(self.owner._schema_, self.owner._tablename_, self.name) @overload def __get__(self: 'Column[T]', obj: None, objtype: "None | Type['RowModel']") -> 'Column[T]': @@ -136,10 +135,8 @@ class Column(ColumnExpr, Generic[T]): # Get value from row data or session if obj is None: return self - elif obj is self.owner: - return obj.data[self.name] else: - return self + return obj.data[self.name] class Integer(Column[int]): diff --git a/bot/data/models.py b/bot/data/models.py index 8fef0f10..7003699e 100644 --- a/bot/data/models.py +++ b/bot/data/models.py @@ -69,7 +69,11 @@ class RowTable(Table, Generic[RowT]): ).where(*args, **kwargs) -class WeakCache(MutableMapping): +WK = TypeVar('WK') +WV = TypeVar('WV') + + +class WeakCache(Generic[WK, WV], MutableMapping[WK, WV]): def __init__(self, ref_cache): self.ref_cache = ref_cache self.weak_cache = WeakValueDictionary() diff --git a/bot/meta/LionBot.py b/bot/meta/LionBot.py index 20dfd5da..45c328d6 100644 --- a/bot/meta/LionBot.py +++ b/bot/meta/LionBot.py @@ -110,7 +110,7 @@ class LionBot(Bot): if isinstance(ctx.command, HybridCommand) and ctx.command.app_command: cmd_str = ctx.command.app_command.to_dict() try: - raise exception from None + raise exception except (HybridCommandError, CommandInvokeError, appCommandInvokeError): try: if isinstance(exception.original, (HybridCommandError, CommandInvokeError, appCommandInvokeError)): diff --git a/bot/meta/LionContext.py b/bot/meta/LionContext.py index eb0ffaa8..cc9d3918 100644 --- a/bot/meta/LionContext.py +++ b/bot/meta/LionContext.py @@ -8,6 +8,7 @@ from discord.ext.commands import Context if TYPE_CHECKING: from .LionBot import LionBot + from core.lion import Lion logger = logging.getLogger(__name__) @@ -44,6 +45,8 @@ class LionContext(Context['LionBot']): Extends Context to add Lion-specific methods and attributes. Also adds several contextual wrapped utilities for simpler user during command invocation. """ + alion: 'Lion' + def __repr__(self): parts = {} if self.interaction is not None: diff --git a/data/migration/v12-13/migration.sql b/data/migration/v12-13/migration.sql index b828d5fc..f04af6ca 100644 --- a/data/migration/v12-13/migration.sql +++ b/data/migration/v12-13/migration.sql @@ -1,11 +1,10 @@ --- Add gem support ALTER TABLE user_config ADD COLUMN name TEXT; -INSERT INTO VersionHistory (version, author) VALUES (13, 'v12-v13 migration'); - --- Add first_joined_at to guild table --- Add left_at to guild table -ALTER TABLE guild_config ADD COLUMN first_joined_at TIMESTAMPTZ; +ALTER TABLE guild_config ADD COLUMN first_joined_at TIMESTAMPTZ DEFAULT now(); ALTER TABLE guild_config ADD COLUMN left_at TIMESTAMPTZ; +ALTER TABLE members ADD COLUMN first_joined TIMESTAMPTZ DEFAULT now(); +ALTER TABLE members ADD COLUMN last_left TIMESTAMPTZ; +ALTER TABLE user_config ADD COLUMN first_seen TIMESTAMPTZ DEFAULT now(); +ALTER TABLE user_config ADD COLUMN last_seen TIMESTAMPTZ; -- Bot config data @@ -114,3 +113,10 @@ CREATE TABLE analytics.gui_renders( cardname TEXT NOT NULL, duration INTEGER NOT NULL ) INHERITS (analytics.events); + + +-- TODO: Correct foreign keys for member table +-- TODO: Add name to user +-- TODO: Add first_joined and last_left time to member +-- TODO: Add first_seen and last_seen time to User +INSERT INTO VersionHistory (version, author) VALUES (13, 'v12-v13 migration');