refactor: sqlalchemy is setup, bot now boots properly again

Author: Snaacky

chiya/bot.py

@@ -8,10 +8,9 @@
 from discord.ext import commands
 from loguru import logger as log
 
-import database
+from chiya import database  # noqa
 from chiya.config import config
 
-
 bot = commands.Bot(
     activity=discord.Activity(type=discord.ActivityType.listening, name=config.bot.status),
     case_insensitive=True,
@@ -76,5 +75,4 @@ async def load_cogs():
     if sys.platform == "win32":
         asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
 
-    database.Database().setup()
     asyncio.run(main())

chiya/cogs/listeners/autoresponder.py

@@ -6,8 +6,8 @@
 from discord.ext import commands
 from loguru import logger as log
 
-from chiya import database
 from chiya.config import config
+from chiya.database import Highlight, Session
 from chiya.utils import embeds
 
 
@@ -17,12 +17,10 @@ def __init__(self, bot: commands.Bot) -> None:
         self.refresh_highlights()
 
     def refresh_highlights(self):
-        db = database.Database().get()
-        self.highlights = [
-            {"term": highlight["term"], "users": orjson.loads(highlight["users"])}
-            for highlight in db["highlights"].find()
-        ]
-        db.close()
+        with Session() as session:
+            self.highlights = [
+                {"term": row.term, "users": orjson.loads(row.users)} for row in session.query(Highlight).all()
+            ]
 
     async def active_members(self, channel: discord.TextChannel) -> set:
         """

chiya/cogs/listeners/highlight.py

@@ -1,10 +1,9 @@
-from datetime import datetime, timezone
+import arrow
 import discord
-
 from discord.ext import commands, tasks
 from loguru import logger as log
 
-from chiya import database
+from chiya.database import RemindMe, Session
 from chiya.utils import embeds
 
 
@@ -23,42 +22,39 @@ async def check_for_reminder(self) -> None:
         """
         await self.bot.wait_until_ready()
 
-        db = database.Database().get()
-        result = db["remind_me"].find(
-            sent=False, date_to_remind={"<": datetime.now(tz=timezone.utc).timestamp()}
-        )
+        with Session() as session:
+            result = session.query(RemindMe).filter(RemindMe.date_to_remind < arrow.utcnow().timestamp())
 
         if not result:
-            return db.close()
+            return
 
         for reminder in result:
             try:
-                user = await self.bot.fetch_user(reminder["author_id"])
+                user = await self.bot.fetch_user(reminder.author_id)
             except discord.errors.NotFound:
-                db["remind_me"].update(dict(id=reminder["id"], sent=True), ["id"])
-                log.warning(
-                    f"Reminder entry with ID {reminder['id']} has an invalid user ID: {reminder['author_id']}."
-                )
+                with Session() as session:
+                    result = session.query(RemindMe).filter_by(id=reminder.id).first()
+                    result.sent = True
+                    session.commit()
+                log.warning(f"Reminder entry with ID {reminder.id} has an invalid user ID: {reminder.author_id}.")
                 continue
 
             embed = embeds.make_embed(
                 title="Here is your reminder",
-                description=reminder["message"],
+                description=reminder.message,
                 color="blurple",
             )
 
             try:
                 channel = await user.create_dm()
                 await channel.send(embed=embed)
             except discord.Forbidden:
-                log.warning(
-                    f"Unable to post or DM {user}'s reminder {reminder['id']=}."
-                )
-
-            db["remind_me"].update(dict(id=reminder["id"], sent=True), ["id"])
+                log.warning(f"Unable to post or DM {user}'s reminder {reminder.id=}.")
 
-        db.commit()
-        db.close()
+            with Session() as session:
+                result = session.query(RemindMe).filter_by(id=reminder.id).first()
+                result.sent = True
+                session.commit()
 
 
 async def setup(bot: commands.Bot) -> None:

chiya/cogs/tasks/reminder.py

@@ -6,33 +6,12 @@
 
 
 Base = declarative_base()
-engine = create_engine(config.database.url, connect_args={"check_same_thread": False})
-session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
-
-
-class BaseModel(Base):
-    __abstract__ = True
-
-    def save(self):
-        session.add(self)
-        session.commit()
-        return self
-
-    def delete(self):
-        session.delete(self)
-        session.commit()
-        return self
-
-    def flush(self):
-        session.add(self)
-        session.flush()
-        return self
 
 
 class ModLog(Base):
     __tablename__ = "mod_logs"
 
-    id = Column(Integer, primar_key=True)
+    id = Column(Integer, primary_key=True)
     user_id = Column(BigInteger, nullable=False)
     mod_id = Column(BigInteger, nullable=False)
     timestamp = Column(BigInteger, nullable=False)
@@ -44,7 +23,7 @@ class ModLog(Base):
 class RemindMe(Base):
     __tablename__ = "remind_me"
 
-    id = Column(Integer, primar_key=True)
+    id = Column(Integer, primary_key=True)
     reminder_location = Column(BigInteger, nullable=False)
     author_id = Column(BigInteger, nullable=False)
     date_to_remind = Column(BigInteger, nullable=False)
@@ -80,3 +59,8 @@ class Highlight(Base):
     id = Column(Integer, primary_key=True)
     term = Column(Text, nullable=False)
     users = Column(Text, nullable=False)
+
+
+engine = create_engine(config.database.url, connect_args={"check_same_thread": False})
+Session = sessionmaker(bind=engine)
+Base.metadata.create_all(engine)