diff --git a/payphone/cogs/admin/cog.py b/payphone/cogs/admin/cog.py index 794a6cd..4cbbf75 100644 --- a/payphone/cogs/admin/cog.py +++ b/payphone/cogs/admin/cog.py @@ -61,20 +61,25 @@ class Admin(base.Cog): for server_map in server_counts[:10]: if servers_map[server_map[0]].name is None: try: - servers_map[server_map[0]].name = ( - await self.bot.fetch_guild(server_map[0]) - ).name + guild = await self.bot.fetch_guild(server_map[0]) except errors.DiscordException: pass + servers_map[server_map[0]].name = guild.name + + vanity = await guild.vanity_invite() + + if vanity is not None: + servers_map[server_map[0]].vanity_code = vanity.code + for i, server_map in enumerate(server_counts): servers_map[server_map[0]].leaderboard_position = i + 1 servers_map[server_map[0]].leaderboard_calls_count = server_map[1] await models.Server.bulk_update( servers_map.values(), - ["name", "leaderboard_position", "leaderboard_calls_count"], + ["name", "vanity_code", "leaderboard_position", "leaderboard_calls_count"], 500, ) @@ -89,7 +94,15 @@ class Admin(base.Cog): after (discord.Guild): The server after it being updated. """ - await models.Server.update_or_create({"name": after.name}, id=after.id) + vanity = await after.vanity_invite() + vanity_code = None + + if vanity is not None: + vanity_code = vanity.code + + await models.Server.update_or_create( + {"name": after.name, "vanity_code": vanity_code}, id=after.id + ) @commands.group(invoke_without_command=True) @commands.guild_only() diff --git a/payphone/cogs/users/cog.py b/payphone/cogs/users/cog.py index a2e94cb..48019aa 100644 --- a/payphone/cogs/users/cog.py +++ b/payphone/cogs/users/cog.py @@ -512,7 +512,7 @@ class Users(base.Cog): .limit(10) ) server_rows = [ - f"**#{pos + 1} - {server.name or "Updating..."}**{" (Here!)" if ctx.guild is not None and server.id == ctx.guild.id else ""}\nParticipated in {server.leaderboard_calls_count} calls." + f"**#{pos + 1} - {server.name or "Updating..."}{f" .gg/{server.vanity_code}" if server.vanity_code else ""}**{" (Here!)" if ctx.guild is not None and server.id == ctx.guild.id else ""}\nParticipated in {server.leaderboard_calls_count} calls." for pos, server in enumerate(top_servers) ] diff --git a/payphone/db/migrations/models/11_20241019004248_update.py b/payphone/db/migrations/models/11_20241019004248_update.py new file mode 100644 index 0000000..03e9a27 --- /dev/null +++ b/payphone/db/migrations/models/11_20241019004248_update.py @@ -0,0 +1,11 @@ +from tortoise import BaseDBAsyncClient + + +async def upgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "server" ADD "vanity_code" VARCHAR(32);""" + + +async def downgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "server" DROP COLUMN "vanity_code";""" diff --git a/payphone/db/models.py b/payphone/db/models.py index fe51bab..05daf41 100644 --- a/payphone/db/models.py +++ b/payphone/db/models.py @@ -93,6 +93,7 @@ class Server(models.Model): prefix: str = fields.CharField(max_length=5, default="p.") leaderboard_position: int = fields.IntField(default=0, db_index=True) leaderboard_calls_count: int = fields.IntField(default=0) + vanity_code: str = fields.CharField(max_length=32, null=True, default=None) channels: fields.BackwardFKRelation["Channel"] profiles: fields.BackwardFKRelation["Profile"]