✨ Add db system

Author: Paillat-dev

pdm.lock

@@ -21,6 +21,8 @@ dependencies = [
     "aiofile>=3.9.0",
     "sentry-sdk>=2.18.0",
     "aiocache[redis]>=0.12.3",
+    "tortoise-orm[asyncpg]>=0.23.0",
+    "aerich>=0.8.0",
 ]
 requires-python = "==3.12.*"
 readme = "README.md"
@@ -78,6 +80,7 @@ docstring-code-line-length = "dynamic"
 
 [tool.ruff.lint]
 select = ["ALL"]
+per-file-ignores = { "src/database/migrations/*"= ["INP001", "ARG001"] }
 extend-ignore = [
     "N999",
     "D104",
@@ -104,3 +107,8 @@ extend-ignore = [
     "PLR2004",
     "C901"
 ]
+
+[tool.aerich]
+tortoise_orm = "src.database.config.TORTOISE_ORM"
+location = "./src/database/migrations"
+src_folder = "./."

pyproject.toml

@@ -1,19 +1,23 @@
 # Copyright (c) NiceBots.xyz
 # SPDX-License-Identifier: MIT
 
+import contextlib
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, override
 
 import aiocache
 import discord
-from discord import Message
+from discord import Interaction, Message, WebhookMessage
 from discord.ext import bridge
 from discord.ext.bridge import (
     BridgeExtContext,
 )
 
 from src.i18n.classes import ExtensionTranslation, TranslationWrapper, apply_locale
 
+if TYPE_CHECKING:
+    from src.database.models import Guild, User
+
 logger = getLogger("bot")
 
 
@@ -22,6 +26,8 @@ def __init__(self, bot: "Bot", interaction: discord.Interaction) -> None:
         self.translations: TranslationWrapper = TranslationWrapper({}, "en-US")  # empty placeholder
         super().__init__(bot=bot, interaction=interaction)
         self.bot: Bot
+        self.user_obj: User | None = None
+        self.guild_obj: Guild | None = None
 
     @override
     def __setattr__(self, key: Any, value: Any) -> None:
@@ -33,22 +39,42 @@ def __setattr__(self, key: Any, value: Any) -> None:
         super().__setattr__(key, value)
 
 
+async def remove_reaction(user: discord.User, message: discord.Message, emoji: str) -> None:
+    await message.remove_reaction(emoji, user)
+
+
 class ExtContext(bridge.BridgeExtContext):
     def __init__(self, **kwargs: Any) -> None:
         self.translations: TranslationWrapper = TranslationWrapper({}, "en-US")  # empty placeholder
         super().__init__(**kwargs)
         self.bot: Bot
+        self.user_obj: User | None = None
+        self.guild_obj: Guild | None = None
 
     def load_translations(self) -> None:
         if hasattr(self.command, "translations") and self.command.translations:  # pyright: ignore[reportUnknownArgumentType,reportOptionalMemberAccess,reportAttributeAccessIssue]
             locale: str | None = None
             if guild := self.guild:
                 locale = guild.preferred_locale
             self.translations = apply_locale(
-                self.command.translations,
+                self.command.translations,  # pyright: ignore [reportAttributeAccessIssue, reportOptionalMemberAccess, reportUnknownArgumentType]
                 locale,
             )
 
+    @override
+    async def defer(self, *args: Any, **kwargs: Any) -> None:
+        await super().defer(*args, **kwargs)
+        with contextlib.suppress(Exception):
+            await self.message.add_reaction("🔄")
+
+    @override
+    async def respond(self, *args: Any, **kwargs: Any) -> "Interaction | WebhookMessage | Message":
+        r = await super().respond(*args, **kwargs)
+        with contextlib.suppress(Exception):
+            if self.me:
+                await remove_reaction(self.me, self.message, "🔄")
+        return r
+
 
 class Bot(bridge.Bot):
     def __init__(

src/custom/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) NiceBots

src/database/__init__.py

@@ -0,0 +1,40 @@
+# Copyright (c) NiceBots
+# SPDX-License-Identifier: MIT
+
+from logging import getLogger
+
+import aerich
+from tortoise import Tortoise
+
+from src.config import config
+
+logger = getLogger("bot").getChild("database")
+
+TORTOISE_ORM = {
+    "connections": {"default": config["db"]["url"]},
+    "apps": {
+        "models": {
+            "models": ["src.database.models", "aerich.models"],
+            "default_connection": "default",
+        }
+    },
+}
+
+
+async def init() -> None:
+    command = aerich.Command(
+        TORTOISE_ORM,
+        app="models",
+        location="./src/database/migrations/",
+    )
+    await command.init()
+    migrated = await command.upgrade(run_in_transaction=True)  # pyright: ignore[reportUnknownVariableType]
+    logger.success(f"Successfully migrated {migrated} migrations")  # pyright: ignore [reportAttributeAccessIssue]
+    await Tortoise.init(config=TORTOISE_ORM)
+
+
+async def shutdown() -> None:
+    await Tortoise.close_connections()
+
+
+__all__ = ["init", "shutdown"]