diff --git a/cogs/loops/__init__.py b/cogs/loops/__init__.py index 5cd44f5..9ee014f 100644 --- a/cogs/loops/__init__.py +++ b/cogs/loops/__init__.py @@ -5,6 +5,7 @@ Initializing code for all loops. from nextcord.ext import commands from .bot_presence import BotPresence +from .guild_data import GuildData # This function will be called when this extension is loaded. @@ -12,3 +13,4 @@ from .bot_presence import BotPresence def setup(bot: commands.Bot): """Adds function to bot""" bot.add_cog(BotPresence(bot)) + bot.add_cog(GuildData(bot)) diff --git a/cogs/loops/guild_data.py b/cogs/loops/guild_data.py new file mode 100644 index 0000000..3ff9dd8 --- /dev/null +++ b/cogs/loops/guild_data.py @@ -0,0 +1,48 @@ +""" +Contains a loop that updates guild data in the database periodically. +""" + +from nextcord.ext import commands, tasks +from sqlalchemy import select +from sqlalchemy.orm import Session + +# pylint: disable=import-error +from utilities.sql_data_classes import Guild, engine + + +class GuildData(commands.Cog, name="Guild Data"): + """Class to contain loop.""" + + def __init__(self, bot: commands.bot) -> None: + self.bot = bot + + @commands.Cog.listener() + async def on_ready(self): + """Starts the loop when bot is ready.""" + await self.bot.wait_until_ready() + # pylint doesn't recognize that periodic_update_guild_data is subclassed from tasks.loop. + await self.periodic_update_guild_data.start() # pylint: disable=no-member + + @tasks.loop(hours=24) + async def periodic_update_guild_data(self): + """Loops every day, calls update_guild_data.""" + self.update_guild_data() + + def update_guild_data(self): + """ + Updates all known guild's data in the database. + """ + with Session(engine) as session: + for guild in self.bot.guilds: + # update given guild name for guild id + stmt = select(Guild).where(Guild.GuildId == guild.id) + g = session.scalar(stmt) + + # just in case checking if the guild doesn't exist. + if g is None: + new_guild = Guild(GuildId=guild.id, GuildName=guild.name) + session.add(new_guild) + else: + g.GuildName = guild.name + + session.commit()