diff --git a/litecord/blueprints/guilds.py b/litecord/blueprints/guilds.py index 0ad3a67..36c8f11 100644 --- a/litecord/blueprints/guilds.py +++ b/litecord/blueprints/guilds.py @@ -438,21 +438,13 @@ async def ack_guild(guild_id): return '', 204 -async def vanity_invite(guild_id: int) -> Optional[str]: - """Get the vanity invite for a guild.""" - return await app.db.fetchval(""" - SELECT code FROM vanity_invites - WHERE guild_id = $1 - """, guild_id) - - @bp.route('//vanity-url', methods=['GET']) async def get_vanity_url(guild_id: int): """Get the vanity url of a guild.""" user_id = await token_check() await guild_perm_check(user_id, guild_id, 'manage_guild') - inv_code = await vanity_invite(guild_id) + inv_code = await app.storage.vanity_invite(guild_id) if inv_code is None: return jsonify({'code': None}) @@ -478,7 +470,7 @@ async def change_vanity_url(guild_id: int): # store old vanity in a variable to delete it from # invites table - old_vanity = await vanity_invite(guild_id) + old_vanity = await app.storage.vanity_invite(guild_id) if old_vanity == inv_code: raise BadRequest('can not change to same invite') diff --git a/litecord/storage.py b/litecord/storage.py index faaed00..7442e31 100644 --- a/litecord/storage.py +++ b/litecord/storage.py @@ -153,6 +153,13 @@ class Storage: WHERE id = $1 """, guild_id) + async def vanity_invite(self, guild_id: int) -> Optional[str]: + """Get the vanity invite for a guild.""" + return await app.db.fetchval(""" + SELECT code FROM vanity_invites + WHERE guild_id = $1 + """, guild_id) + async def get_guild(self, guild_id: int, user_id=None) -> Optional[Dict]: """Get gulid payload.""" row = await self.db.fetchrow(""" @@ -176,6 +183,8 @@ class Storage: if user_id: drow['owner'] = drow['owner_id'] == str(user_id) + drow['vanity_url_code'] = await self.vanity_invite(guild_id) + return drow async def _member_basic(self, guild_id: int, member_id: int):