From 84fdd21a9d3b96fd6fd21b0b4fe307f3b34d81cb Mon Sep 17 00:00:00 2001 From: nex Date: Tue, 5 Dec 2023 14:37:59 +0000 Subject: [PATCH] Persistence! --- cogs/mcdonalds.py | 170 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 131 insertions(+), 39 deletions(-) diff --git a/cogs/mcdonalds.py b/cogs/mcdonalds.py index 1453290..9c1356f 100644 --- a/cogs/mcdonalds.py +++ b/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[^|]+)(#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: - if message.content.upper() != "MCDONALDS!": - await message.delete() - if (message.created_at.timestamp() - self.targets[author]) > 10: - await message.channel.send( - f"{message.author.mention} Please say `MCDONALDS!` to end commercial.", - delete_after=30 + async with McDataBase() as db: + if (last_info := await db.get_break(author.id)) is not None: + if message.content.upper() != "MCDONALDS!": + await message.delete() + if (message.created_at.timestamp() - last_info) > 10: + await message.channel.send( + 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(): - await ctx.respond(f"{member.mention} is already in a commercial break.") - return - 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" at the earliest." - ) + async with McDataBase() as db: + if await db.get_break(member.id) is not None: + await ctx.respond(f"{member.mention} is already in a commercial break.") return - else: - self.cooldown.pop(member, None) + elif (cooldown := await db.get_cooldown(member.id)) is not 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" at the earliest." + ) + return + else: + await db.remove_cooldown(member.id) - self.targets[member] = discord.utils.utcnow().timestamp() - await ctx.send( - f"{member.mention} Commercial break! Please say `MCDONALDS!` to end commercial.\n" - f"*This commercial break is sponsored by {ctx.user.mention}.*", - delete_after=300, - allowed_mentions=discord.AllowedMentions(users=True, roles=False, everyone=False) - ) - await ctx.respond("Commercial break started.", ephemeral=True) - await ctx.delete(delay=120) + await db.set_break(member.id, discord.utils.utcnow().timestamp()) + await ctx.send( + f"{member.mention} Commercial break! Please say `MCDONALDS!` to end commercial.\n" + f"*This commercial break is sponsored by {ctx.user.mention}.*", + delete_after=300, + allowed_mentions=discord.AllowedMentions(users=True, roles=False, everyone=False) + ) + await ctx.respond("Commercial break started.", ephemeral=True) + await ctx.delete(delay=120) def setup(bot):