Persistence!

cogs/mcdonalds.py

@@ -1,5 +1,7 @@
 import asyncio
+import pathlib
 import re
+import typing
 
 import aiosqlite
 
@@ -7,12 +9,101 @@ import discord
 from discord.ext import commands
 
 
+class McDataBase:
+    def __init__(self):
+        self.db = pathlib.Path.home() / ".cache" / "McDataBase.db"
+        self._conn: typing.Optional[aiosqlite.Connection] = None
+
+    async def init_db(self):
+        if self._conn:
+            conn = self._conn
+            await conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS breaks (
+                    user_id INTEGER PRIMARY KEY,
+                    since FLOAT NOT NULL,
+                    author INTEGER NOT NULL
+                );
+                """
+            )
+            await conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS cooldowns (
+                    user_id INTEGER PRIMARY KEY,
+                    expires FLOAT NOT NULL
+                );
+                """
+            )
+
+    async def get_break(self, user_id: int) -> typing.Optional[float]:
+        async with self._conn.execute(
+            """
+            SELECT since FROM breaks WHERE user_id = ?;
+            """,
+            (user_id,)
+        ) as cursor:
+            return await cursor.fetchone()
+
+    async def set_break(self, user_id: int, since: float):
+        await self._conn.execute(
+            """
+            INSERT INTO breaks (user_id, since) VALUES (?, ?)
+            ON CONFLICT(user_id) DO UPDATE SET since = excluded.since
+            """,
+            (user_id, since)
+        )
+
+    async def remove_break(self, user_id: int):
+        now = discord.utils.utcnow().timestamp()
+        await self._conn.execute(
+            """
+            DELETE FROM breaks WHERE user_id = ?;
+            """,
+            (user_id,)
+        )
+        await self.set_cooldown(user_id, now)
+
+    async def get_cooldown(self, user_id: int) -> typing.Optional[float]:
+        async with self._conn.execute(
+            """
+            SELECT expires FROM cooldowns WHERE user_id = ?;
+            """,
+            (user_id,)
+        ) as cursor:
+            return await cursor.fetchone()
+
+    async def set_cooldown(self, user_id: int, expires: float):
+        await self._conn.execute(
+            """
+            INSERT INTO cooldowns (user_id, expires) VALUES (?, ?)
+            ON CONFLICT(user_id) DO UPDATE SET expires = excluded.expires;
+            """,
+            (user_id, expires)
+        )
+
+    async def remove_cooldown(self, user_id: int):
+        await self._conn.execute(
+            """
+            DELETE FROM cooldowns WHERE user_id = ?;
+            """,
+            (user_id,)
+        )
+
+    async def __aenter__(self) -> "McDataBase":
+        self._conn = await aiosqlite.connect(self.db)
+        await self.init_db()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._conn.commit()
+        await self._conn.close()
+        self._conn = None
+
+
 class McDonaldsCog(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
-        self.targets: dict[discord.Member, float] = {}
         self.lock = asyncio.Lock()
-        self.cooldown: dict[discord.Member, float] = {}
 
     @commands.Cog.listener()
     async def on_message(self, message: discord.Message):
@@ -22,29 +113,29 @@ class McDonaldsCog(commands.Cog):
             return
 
         async with self.lock:
-            NIGHTMARE_REGEX = re.compile(r"\|\| \|\|([^|]+)#0\|\| \|\|: bypassed.*")
+            NIGHTMARE_REGEX = re.compile(r"[|\s]*(?P<username>[^|]+)(#0)?[|\s]*:( bypassed.*)?")
             if m := NIGHTMARE_REGEX.match(message.content):
                 username = m.group(1)
                 member = discord.utils.get(message.guild.members, name=username)
                 if member:
                     author = member
 
-            if author in self.targets:
+            async with McDataBase() as db:
-                if message.content.upper() != "MCDONALDS!":
+                if (last_info := await db.get_break(author.id)) is not None:
-                    await message.delete()
+                    if message.content.upper() != "MCDONALDS!":
-                    if (message.created_at.timestamp() - self.targets[author]) > 10:
+                        await message.delete()
-                        await message.channel.send(
+                        if (message.created_at.timestamp() - last_info) > 10:
-                            f"{message.author.mention} Please say `MCDONALDS!` to end commercial.",
+                            await message.channel.send(
-                            delete_after=30
+                                f"{message.author.mention} Please say `MCDONALDS!` to end commercial.",
+                                delete_after=30
+                            )
+                            await db.set_break(author.id, message.created_at.timestamp())
+                    else:
+                        await db.remove_break(author.id)
+                        await message.reply(
+                            "Thank you. You may now resume your activity.",
+                            delete_after=120
                         )
-                        self.targets[author] = message.created_at.timestamp()
-                else:
-                    self.targets.pop(author, None)
-                    self.cooldown[author] = message.created_at.timestamp()
-                    await message.reply(
-                        "Thank you. You may now resume your activity.",
-                        delete_after=120
-                    )
 
     @commands.user_command(name="Commercial Break")
     async def commercial_break(self, ctx: discord.ApplicationContext, member: discord.Member):
@@ -56,29 +147,30 @@ class McDonaldsCog(commands.Cog):
         if member.bot or member == ctx.user:
             return await ctx.respond("No.", ephemeral=True)
 
-        if member in self.targets.keys():
+        async with McDataBase() as db:
-            await ctx.respond(f"{member.mention} is already in a commercial break.")
+            if await db.get_break(member.id) is not None:
-            return
+                await ctx.respond(f"{member.mention} is already in a commercial break.")
-        elif member in self.cooldown.keys():
-            expires = self.cooldown[member] + 300
-            if expires > discord.utils.utcnow().timestamp():
-                await ctx.respond(
-                    f"{member.mention} is not due another ad break yet. Their next commercial break will start "
-                    f"<t:{int(expires)}:R> at the earliest."
-                )
                 return
-            else:
+            elif (cooldown := await db.get_cooldown(member.id)) is not None:
-                self.cooldown.pop(member, None)
+                expires = cooldown + 300
+                if expires > discord.utils.utcnow().timestamp():
+                    await ctx.respond(
+                        f"{member.mention} is not due another ad break yet. Their next commercial break will start "
+                        f"<t:{int(expires)}:R> at the earliest."
+                    )
+                    return
+                else:
+                    await db.remove_cooldown(member.id)
 
-        self.targets[member] = discord.utils.utcnow().timestamp()
+            await db.set_break(member.id, discord.utils.utcnow().timestamp())
-        await ctx.send(
+            await ctx.send(
-            f"{member.mention} Commercial break! Please say `MCDONALDS!` to end commercial.\n"
+                f"{member.mention} Commercial break! Please say `MCDONALDS!` to end commercial.\n"
-            f"*This commercial break is sponsored by {ctx.user.mention}.*",
+                f"*This commercial break is sponsored by {ctx.user.mention}.*",
-            delete_after=300,
+                delete_after=300,
-            allowed_mentions=discord.AllowedMentions(users=True, roles=False, everyone=False)
+                allowed_mentions=discord.AllowedMentions(users=True, roles=False, everyone=False)
-        )
+            )
-        await ctx.respond("Commercial break started.", ephemeral=True)
+            await ctx.respond("Commercial break started.", ephemeral=True)
-        await ctx.delete(delay=120)
+            await ctx.delete(delay=120)
 
 
 def setup(bot):