From 92fee23afa6ea4f2ad6f514799788fea19bc9960 Mon Sep 17 00:00:00 2001 From: Interitio Date: Sun, 6 Oct 2024 15:43:49 +1000 Subject: [PATCH] feat(profiles): Add profile base and users. --- data/schema.sql | 21 ++- src/meta/LionBot.py | 10 ++ src/meta/LionCog.py | 1 + src/modules/__init__.py | 1 + src/modules/profiles/cog.py | 245 +++++++++++++++++++++--------- src/modules/profiles/community.py | 0 src/modules/profiles/data.py | 13 +- src/modules/profiles/profile.py | 124 +++++++++++++++ src/twitch/userflow.py | 2 +- 9 files changed, 332 insertions(+), 85 deletions(-) create mode 100644 src/modules/profiles/community.py create mode 100644 src/modules/profiles/profile.py diff --git a/data/schema.sql b/data/schema.sql index 01898dfc..faf668c8 100644 --- a/data/schema.sql +++ b/data/schema.sql @@ -1498,10 +1498,19 @@ CREATE INDEX voice_role_channels on voice_roles (channelid); -- }}} -- User and Community Profiles {{{ +DROP TABLE IF EXISTS community_members; +DROP TABLE IF EXISTS communities_twitch; +DROP TABLE IF EXISTS communities_discord; +DROP TABLE IF EXISTS communities; +DROP TABLE IF EXISTS profiles_twitch; +DROP TABLE IF EXISTS profiles_discord; +DROP TABLE IF EXISTS user_profiles; + CREATE TABLE user_profiles( profileid SERIAL PRIMARY KEY, nickname TEXT, + migrated INTEGER REFERENCES user_profiles (profileid) ON DELETE CASCADE ON UPDATE CASCADE, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); @@ -1511,8 +1520,8 @@ CREATE TABLE profiles_discord( userid BIGINT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); -CREATE UNIQUE INDEX profiles_discord_profileid ON profiles_discord (profileid); -CREATE INDEX profiles_discord_userid ON profiles_discord (userid); +CREATE INDEX profiles_discord_profileid ON profiles_discord (profileid); +CREATE UNIQUE INDEX profiles_discord_userid ON profiles_discord (userid); CREATE TABLE profiles_twitch( linkid SERIAL PRIMARY KEY, @@ -1520,8 +1529,8 @@ CREATE TABLE profiles_twitch( userid TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); -CREATE UNIQUE INDEX profiles_twitch_profileid ON profiles_twitch (profileid); -CREATE INDEX profiles_twitch_userid ON profiles_twitch (userid); +CREATE INDEX profiles_twitch_profileid ON profiles_twitch (profileid); +CREATE UNIQUE INDEX profiles_twitch_userid ON profiles_twitch (userid); CREATE TABLE communities( @@ -1534,14 +1543,14 @@ CREATE TABLE communities_discord( communityid INTEGER NOT NULL REFERENCES communities (communityid) ON DELETE CASCADE ON UPDATE CASCADE, linked_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); -CREATE UNIQUE INDEX communities_discord_communityid ON communities_discord (communityid); +CREATE INDEX communities_discord_communityid ON communities_discord (communityid); CREATE TABLE communities_twitch( channelid TEXT PRIMARY KEY, communityid INTEGER NOT NULL REFERENCES communities (communityid) ON DELETE CASCADE ON UPDATE CASCADE, linked_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); -CREATE UNIQUE INDEX communities_twitch_communityid ON communities_twitch (communityid); +CREATE INDEX communities_twitch_communityid ON communities_twitch (communityid); CREATE TABLE community_members( memberid SERIAL PRIMARY KEY, diff --git a/src/meta/LionBot.py b/src/meta/LionBot.py index 19904f9f..ebfc2875 100644 --- a/src/meta/LionBot.py +++ b/src/meta/LionBot.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: from meta.CrocBot import CrocBot from core.cog import CoreCog from core.config import ConfigCog + from twitch.cog import TwitchAuthCog from tracking.voice.cog import VoiceTrackerCog from tracking.text.cog import TextTrackerCog from modules.config.cog import GuildConfigCog @@ -49,6 +50,7 @@ if TYPE_CHECKING: from modules.topgg.cog import TopggCog from modules.user_config.cog import UserConfigCog from modules.video_channels.cog import VideoCog + from modules.profiles.cog import ProfileCog logger = logging.getLogger(__name__) @@ -142,6 +144,10 @@ class LionBot(Bot): # To make the type checker happy about fetching cogs by name # TODO: Move this to stubs at some point + @overload + def get_cog(self, name: Literal['ProfileCog']) -> 'ProfileCog': + ... + @overload def get_cog(self, name: Literal['CoreCog']) -> 'CoreCog': ... @@ -154,6 +160,10 @@ class LionBot(Bot): def get_cog(self, name: Literal['VoiceTrackerCog']) -> 'VoiceTrackerCog': ... + @overload + def get_cog(self, name: Literal['TwitchAuthCog']) -> 'TwitchAuthCog': + ... + @overload def get_cog(self, name: Literal['TextTrackerCog']) -> 'TextTrackerCog': ... diff --git a/src/meta/LionCog.py b/src/meta/LionCog.py index a2b1b625..f1719a6f 100644 --- a/src/meta/LionCog.py +++ b/src/meta/LionCog.py @@ -22,6 +22,7 @@ class LionCog(Cog): cls._placeholder_groups_ = set() cls._twitch_cmds_ = {} cls._twitch_events_ = {} + cls._twitch_events_loaded_ = set() for base in reversed(cls.__mro__): for elem, value in base.__dict__.items(): diff --git a/src/modules/__init__.py b/src/modules/__init__.py index 9e6bb1fd..ddafc4dd 100644 --- a/src/modules/__init__.py +++ b/src/modules/__init__.py @@ -2,6 +2,7 @@ this_package = 'modules' active_discord = [ '.sysadmin', + '.profiles', '.config', '.user_config', '.skins', diff --git a/src/modules/profiles/cog.py b/src/modules/profiles/cog.py index ed2eb756..3e0c49a0 100644 --- a/src/modules/profiles/cog.py +++ b/src/modules/profiles/cog.py @@ -1,76 +1,36 @@ import asyncio from enum import Enum -from typing import Optional +from typing import Optional, overload from datetime import timedelta import discord from discord.ext import commands as cmds import twitchio from twitchio.ext import commands +from twitchio import User from twitchAPI.object.api import TwitchUser from data.queries import ORDER -from meta import LionCog, LionBot, CrocBot +from meta import LionCog, LionBot, CrocBot, LionContext +from meta.logger import log_wrap from utils.lib import utc_now from . import logger from .data import ProfileData - - -class UserProfile: - def __init__(self, data, profile_row, *, discord_row=None, twitch_row=None): - self.data: ProfileData = data - self.profile_row: ProfileData.UserProfileRow = profile_row - - self.discord_row: Optional[ProfileData.DiscordProfileRow] = discord_row - self.twitch_row: Optional[ProfileData.TwitchProfileRow] = twitch_row - - @property - def profileid(self): - return self.profile_row.profileid - - async def attach_discord(self, user: discord.User | discord.Member): - """ - Attach a new discord user to this profile. - """ - # TODO: Attach whatever other data we want to cache here. - # Currently Lion also caches most of this data - discord_row = await self.data.DiscordProfileRow.create( - profileid=self.profileid, - userid=user.id - ) - - async def attach_twitch(self, user: TwitchUser): - """ - Attach a new Twitch user to this profile. - """ - ... - - @classmethod - async def fetch_profile( - cls, data: ProfileData, - *, - profile_id: Optional[int] = None, - profile_row: Optional[ProfileData.UserProfileRow] = None, - discord_row: Optional[ProfileData.DiscordProfileRow] = None, - twitch_row: Optional[ProfileData.TwitchProfileRow] = None, - ): - if not any((profile_id, profile_row, discord_row, twitch_row)): - raise ValueError("UserProfile needs an id or a data row to construct.") - if profile_id is None: - profile_id = (profile_row or discord_row or twitch_row).profileid - profile_row = profile_row or await data.UserProfileRow.fetch(profile_id) - discord_row = discord_row or await data.DiscordProfileRow.fetch_profile(profile_id) - twitch_row = twitch_row or await data.TwitchProfileRow.fetch_profile(profile_id) - - return cls(data, profile_row, discord_row=discord_row, twitch_row=twitch_row) +from .profile import UserProfile class ProfileCog(LionCog): def __init__(self, bot: LionBot): self.bot = bot + + assert bot.crocbot is not None + self.crocbot: CrocBot = bot.crocbot self.data = bot.db.load_registry(ProfileData()) + self._profile_migrators = {} + self._comm_migrators = {} + async def cog_load(self): await self.data.init() @@ -78,34 +38,84 @@ class ProfileCog(LionCog): return True # Profile API - async def fetch_profile_discord(self, userid: int, create=True): - """ - Fetch or create a UserProfile from the given Discord userid. - """ - # TODO: (Extension) May be context dependent - # Current model assumes profile (one->0..n) discord - discord_row = next(await self.data.DiscordProfileRow.fetch_where(userid=userid), None) - if discord_row is None: - profile_row = await self.data.UserProfileRow.create() + def add_profile_migrator(self, migrator, name=None): + name = name or migrator.__name__ + self._profile_migrators[name or migrator.__name__] = migrator - async def fetch_profile_twitch(self, userid: int, create=True): - """ - Fetch or create a UserProfile from the given Twitch userid. - """ - ... + logger.info( + f"Added user profile migrator {name}: {migrator}" + ) + return migrator - async def fetch_profile(self, profileid: int): + def del_profile_migrator(self, name: str): + migrator = self._profile_migrators.pop(name, None) + + logger.info( + f"Removed user profile migrator {name}: {migrator}" + ) + + @log_wrap(action="profile migration") + async def migrate_profile(self, source_profile, target_profile) -> list[str]: + logger.info( + f"Beginning user profile migration from {source_profile!r} to {target_profile!r}" + ) + results = [] + # Wrap this in a transaction so if something goes wrong with migration, + # we roll back safely (although this may mess up caches) + async with self.bot.db.connection() as conn: + self.bot.db.conn = conn + async with conn.transaction(): + for name, migrator in self._profile_migrators.items(): + try: + result = await migrator(source_profile, target_profile) + if result: + results.append(result) + except Exception: + logger.exception( + f"Unexpected exception running user profile migrator {name} " + f"migrating {source_profile!r} to {target_profile!r}." + ) + raise + + # Move all Discord and Twitch profile references over to the new profile + discord_rows = await self.data.DiscordProfileRow.table.update_where( + profileid=source_profile.profileid + ).set(profileid=target_profile.profileid) + results.append(f"Migrated {len(discord_rows)} attached discord account(s).") + + twitch_rows = await self.data.TwitchProfileRow.table.update_where( + profileid=source_profile.profileid + ).set(profileid=target_profile.profileid) + results.append(f"Migrated {len(twitch_rows)} attached twitch account(s).") + + # And then mark the old profile as migrated + await source_profile.update(migrate=target_profile.profileid) + results.append("Marking old profile as migrated.. finished!") + return results + + async def fetch_profile_by_id(self, profile_id: int) -> UserProfile: """ Fetch a UserProfile by the given id. """ - ... + return await UserProfile.fetch_profile(self.bot, profile_id=profile_id) - async def merge_profiles(self, sourceid: int, targetid: int): + async def fetch_profile_discord(self, user: discord.Member | discord.User) -> UserProfile: """ - Merge two UserProfiles by id. - Merges the 'sourceid' into the 'targetid'. + Fetch or create a UserProfile from the provided discord account. """ - ... + profile = await UserProfile.fetch_from_discordid(user.id) + if profile is None: + profile = await UserProfile.create_from_discord(user) + return profile + + async def fetch_profile_twitch(self, user: twitchio.User) -> UserProfile: + """ + Fetch or create a UserProfile from the provided twitch account. + """ + profile = await UserProfile.fetch_from_twitchid(user.id) + if profile is None: + profile = await UserProfile.create_from_twitch(user) + return profile async def fetch_community_discord(self, guildid: int, create=True): ... @@ -118,4 +128,95 @@ class ProfileCog(LionCog): # ----- Profile Commands ----- - # Link twitch profile + @cmds.hybrid_group( + name='profiles', + description="Base comand group for user profiles." + ) + async def profiles_grp(self, ctx: LionContext): + ... + + @profiles_grp.group( + name='link', + description="Base command group for linking profiles" + ) + async def profiles_link_grp(self, ctx: LionContext): + ... + + @profiles_link_grp.command( + name='twitch', + description="Link a twitch account to your current profile." + ) + async def profiles_link_twitch_cmd(self, ctx: LionContext): + if not ctx.interaction: + return + + await ctx.interaction.response.defer(ephemeral=True) + + # Ask the user to go through auth to get their userid + auth_cog = self.bot.get_cog('TwitchAuthCog') + flow = await auth_cog.start_auth() + message = await ctx.reply( + f"Please [click here]({flow.auth.return_auth_url()}) to link your profile " + "to Twitch." + ) + authrow = await flow.run() + await message.edit( + content="Authentication Complete! Beginning profile merge..." + ) + + results = await self.crocbot.fetch_users(ids=[authrow.userid]) + if not results: + logger.error( + f"User {authrow} obtained from Twitch authentication does not exist." + ) + await ctx.error_reply("Sorry, something went wrong. Please try again later!") + + user = results[0] + + # Retrieve author's profile if it exists + author_profile = await UserProfile.fetch_from_discordid(self.bot, ctx.author.id) + + # Check if the twitch-side user has a profile + source_profile = await UserProfile.fetch_from_twitchid(self.bot, user.id) + + if author_profile and source_profile is None: + # All we need to do is attach the twitch row + await author_profile.attach_twitch(user) + await message.edit( + content=f"Successfully added Twitch account **{user.name}**! There was no profile data to merge." + ) + elif source_profile and author_profile is None: + # Attach the discord row to the profile + await source_profile.attach_discord(ctx.author) + await message.edit( + content=f"Successfully connect to Twitch profile **{user.name}**! There was no profile data to merge." + ) + elif source_profile is None and author_profile is None: + profile = await UserProfile.create_from_discord(self.bot, ctx.author) + await profile.attach_twitch(user) + + await message.edit( + content=f"Opened a new user profile for you and linked Twitch account **{user.name}**." + ) + elif author_profile.profileid == source_profile.profileid: + await message.edit( + content=f"The Twitch account **{user.name}** is already linked to your profile!" + ) + else: + # Migrate the existing profile data to the new profiles + try: + results = await self.migrate_profile(source_profile, author_profile) + except Exception: + await ctx.error_reply( + "An issue was encountered while merging your account profiles!\n" + "Migration rolled back, no data has been lost.\n" + "The developer has been notified. Please try again later!" + ) + raise + + content = '\n'.join(( + "## Connecting Twitch account and merging profiles...", + *results, + "**Successfully linked account and merge profile data!**" + )) + await message.edit(content=content) diff --git a/src/modules/profiles/community.py b/src/modules/profiles/community.py new file mode 100644 index 00000000..e69de29b diff --git a/src/modules/profiles/data.py b/src/modules/profiles/data.py index dc85eb1e..eed48792 100644 --- a/src/modules/profiles/data.py +++ b/src/modules/profiles/data.py @@ -10,6 +10,7 @@ class ProfileData(Registry): CREATE TABLE user_profiles( profileid SERIAL PRIMARY KEY, nickname TEXT, + migrated INTEGER REFERENCES user_profiles (profileid) ON DELETE CASCADE ON UPDATE CASCADE, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); """ @@ -18,8 +19,10 @@ class ProfileData(Registry): profileid = Integer(primary=True) nickname = String() + migrated = Integer() created_at = Timestamp() + class DiscordProfileRow(RowModel): """ Schema @@ -30,8 +33,8 @@ class ProfileData(Registry): userid BIGINT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); - CREATE UNIQUE INDEX profiles_discord_profileid ON profiles_discord (profileid); - CREATE INDEX profiles_discord_userid ON profiles_discord (userid); + CREATE INDEX profiles_discord_profileid ON profiles_discord (profileid); + CREATE UNIQUE INDEX profiles_discord_userid ON profiles_discord (userid); """ _tablename_ = 'profiles_discord' _cache_ = {} @@ -57,8 +60,8 @@ class ProfileData(Registry): userid TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); - CREATE UNIQUE INDEX profiles_twitch_profileid ON profiles_twitch (profileid); - CREATE INDEX profiles_twitch_userid ON profiles_twitch (userid); + CREATE INDEX profiles_twitch_profileid ON profiles_twitch (profileid); + CREATE UNIQUE INDEX profiles_twitch_userid ON profiles_twitch (userid); """ _tablename_ = 'profiles_twitch' _cache_ = {} @@ -97,7 +100,6 @@ class ProfileData(Registry): communityid INTEGER NOT NULL REFERENCES communities (communityid) ON DELETE CASCADE ON UPDATE CASCADE, linked_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); - CREATE UNIQUE INDEX communities_discord_communityid ON communities_discord (communityid); """ _tablename_ = 'communities_discord' _cache_ = {} @@ -120,7 +122,6 @@ class ProfileData(Registry): communityid INTEGER NOT NULL REFERENCES communities (communityid) ON DELETE CASCADE ON UPDATE CASCADE, linked_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ); - CREATE UNIQUE INDEX communities_twitch_communityid ON communities_twitch (communityid); """ _tablename_ = 'communities_twitch' _cache_ = {} diff --git a/src/modules/profiles/profile.py b/src/modules/profiles/profile.py new file mode 100644 index 00000000..9e1151c5 --- /dev/null +++ b/src/modules/profiles/profile.py @@ -0,0 +1,124 @@ +from typing import Optional, Self + +import discord +import twitchio + +from meta import LionBot +from utils.lib import utc_now + +from . import logger +from .data import ProfileData + + + +class UserProfile: + def __init__(self, bot: LionBot, profile_row): + self.bot = bot + self.profile_row: ProfileData.UserProfileRow = profile_row + + @property + def cog(self): + return self.bot.get_cog('ProfileCog') + + @property + def data(self) -> ProfileData: + return self.cog.data + + @property + def profileid(self): + return self.profile_row.profileid + + def __repr__(self): + return f"" + + async def attach_discord(self, user: discord.User | discord.Member): + """ + Attach a new discord user to this profile. + Assumes the discord user does not itself have a profile. + """ + discord_row = await self.data.DiscordProfileRow.create( + profileid=self.profileid, + userid=user.id + ) + logger.info( + f"Attached discord user {user!r} to profile {self!r}" + ) + return discord_row + + async def attach_twitch(self, user: twitchio.User): + """ + Attach a new Twitch user to this profile. + """ + twitch_row = await self.data.TwitchProfileRow.create( + profileid=self.profileid, + userid=str(user.id) + ) + logger.info( + f"Attached twitch user {user!r} to profile {self!r}" + ) + return twitch_row + + async def discord_accounts(self) -> list[ProfileData.DiscordProfileRow]: + """ + Fetch the Discord accounts associated to this profile. + """ + return await self.data.DiscordProfileRow.fetch_where(profileid=self.profileid) + + async def twitch_accounts(self) -> list[ProfileData.DiscordProfileRow]: + """ + Fetch the Twitch accounts associated to this profile. + """ + return await self.data.TwitchProfileRow.fetch_where(profileid=self.profileid) + + @classmethod + async def fetch(cls, bot: LionBot, profile_id: int) -> Self: + profile_row = await bot.get_cog('ProfileCog').data.UserProfileRow.fetch(profile_id) + if profile_row is None: + raise ValueError("Provided profile_id does not exist.") + return cls(bot, profile_row) + + @classmethod + async def fetch_from_twitchid(cls, bot: LionBot, userid: int | str) -> Optional[Self]: + data = bot.get_cog('ProfileCog').data + rows = await data.TwitchProfileRow.fetch_where(userid=str(userid)) + if rows: + return await cls.fetch(bot, rows[0].profileid) + + @classmethod + async def fetch_from_discordid(cls, bot: LionBot, userid: int) -> Optional[Self]: + data = bot.get_cog('ProfileCog').data + rows = await data.DiscordProfileRow.fetch_where(userid=str(userid)) + if rows: + return await cls.fetch(bot, rows[0].profileid) + + @classmethod + async def create(cls, bot: LionBot, **kwargs) -> Self: + """ + Create a new empty profile with the given initial arguments. + + Profiles should usually be created using `create_from_discord` or `create_from_twitch` + to correctly setup initial profile preferences (e.g. name, avatar). + """ + # Create a new profile + data = bot.get_cog('ProfileCog').data + profile_row = await data.UserProfileRow.create(created_at=utc_now()) + profile = await cls.fetch(bot, profile_row.profileid) + return profile + + @classmethod + async def create_from_discord(cls, bot: LionBot, user: discord.Member | discord.User, **kwargs) -> Self: + """ + Create a new profile using the given Discord user as a base. + """ + profile = await cls.create(bot, **kwargs) + await profile.attach_discord(user) + return profile + + @classmethod + async def create_from_twitch(cls, bot: LionBot, user: twitchio.User, **kwargs) -> Self: + """ + Create a new profile using the given Twitch user as a base. + """ + profile = await cls.create(bot, **kwargs) + await profile.attach_twitch(user) + return profile diff --git a/src/twitch/userflow.py b/src/twitch/userflow.py index ce7d20dc..11c0fef9 100644 --- a/src/twitch/userflow.py +++ b/src/twitch/userflow.py @@ -47,7 +47,7 @@ class UserAuthFlow: self._setup_done.set() return await ws.receive_json() - async def run(self): + async def run(self) -> TwitchAuthData.UserAuthRow: if not self._setup_done.is_set(): raise ValueError("Cannot run UserAuthFlow before setup.") if self._comm_task is None: