feat!: piccolo

Author: ooliver1

botbase/db/__init__.py

@@ -1,3 +1,2 @@
 from .blacklist import *
 from .commands import *
-from .metadata import *

botbase/db/blacklist.py

@@ -1,23 +1,15 @@
-from ormar import BigInteger, Model, String
-
-from .metadata import BaseMeta
+from piccolo.table import Table
+from piccolo.columns import BigInt, Text
 
 
 __all__ = ("BlacklistGuild", "BlacklistUser")
 
 
-class BlacklistGuild(Model):
-    class Meta(BaseMeta):
-        tablename = "blacklist_guilds"
-
-    # pyright: reportGeneralTypeIssues=false
-    id: int = BigInteger(primary_key=True, autoincrement=False)
-    reason: str = String(max_length=255, default="Unknown reason.")
-
+class BlacklistGuild(Table):
+    id = BigInt(primary_key=True)
+    reason = Text(default="Unknown reason.")
 
-class BlacklistUser(Model):
-    class Meta(BaseMeta):
-        tablename = "blacklist_users"
 
-    id: int = BigInteger(primary_key=True, autoincrement=False)
-    reason: str = String(max_length=255, default="Unknown reason.")
+class BlacklistUser(Table):
+    id = BigInt(primary_key=True)
+    reason = Text(default="Unknown reason.")

botbase/db/commands.py

@@ -1,26 +1,18 @@
-from ormar import BigInteger, Integer, Model, String
-from sqlalchemy import PrimaryKeyConstraint
+from logging import getLogger
 
-from .metadata import BaseMeta
+from piccolo.table import Table
+from piccolo.columns import BigInt, Text, Integer
 
+log = getLogger(__name__)
 __all__ = ("CommandLog",)
 
 
-class CommandLog(Model):
-    class Meta(BaseMeta):
-        tablename = "commands"
-        constraints = [PrimaryKeyConstraint("command", "guild", "channel", "member")]
+class CommandLog(Table):
+    command = Text(primary_key=True)
+    guild = BigInt()
+    channel = BigInt()
+    member = BigInt()
+    amount = Integer(default=1)
 
-    # pyright: reportGeneralTypeIssues=false
-    command: str = String(max_length=255, primary_key=True)
-    guild: int = BigInteger()
-    channel: int = BigInteger()
-    member: int = BigInteger()
-    amount: int = BigInteger(default=1)
 
-
-# ormar :(
-CommandLog.command.primary_key = True
-CommandLog.guild.primary_key = True
-CommandLog.channel.primary_key = True
-CommandLog.member.primary_key = True
+log.critical("Make sure to set command, guild, channel, member UNIQUE.")

botbase/db/metadata.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from logging import getLogger
-from typing import TYPE_CHECKING
 
 from nextcord import Embed, Guild, slash_command
 from nextcord.ext.application_checks import is_owner
@@ -11,10 +10,6 @@
 from ..models import CogBase
 from ..wraps import MyInter
 
-if TYPE_CHECKING:
-    from ..botbase import BotBase
-
-
 _log = getLogger(__name__)
 
 
@@ -35,24 +30,32 @@ async def add(
 
         if guild:
             self.guilds.add(obj)
-            await BlacklistGuild.objects.create(id=obj, reason=reason)
+            await BlacklistGuild(
+                {BlacklistGuild.id: obj, BlacklistGuild.reason: reason}
+            ).save()
         else:
             self.users.add(obj)
-            await BlacklistUser.objects.create(id=obj, reason=reason)
+            await BlacklistUser(
+                {BlacklistUser.id: obj, BlacklistUser.reason: reason}
+            ).save()
 
     async def remove(self, obj: int, guild: bool = False) -> None:
         assert isinstance(obj, int)
 
         if guild:
             self.guilds.discard(obj)
-            await BlacklistGuild.objects.delete(id=obj)
+            await BlacklistGuild.delete().where(BlacklistGuild.id == obj)
         else:
             self.users.discard(obj)
-            await BlacklistUser.objects.delete(id=obj)
+            await BlacklistUser.delete().where(BlacklistGuild.id == obj)
 
     async def load(self) -> None:
-        self.guilds = set(await BlacklistGuild.objects.values_list("id", flatten=True))
-        self.users = set(await BlacklistUser.objects.values_list("id", flatten=True))
+        self.guilds = set(
+            item["id"] for item in await BlacklistGuild.select(BlacklistGuild.id)
+        )
+        self.users = set(
+            item["id"] for item in await BlacklistUser.select(BlacklistUser.id)
+        )
 
 
 class BlacklistCog(CogBase[BotBase]):
@@ -62,7 +65,8 @@ def __init__(self, bot: BotBase) -> None:
         self.blacklist = Blacklist()
         bot.loop.create_task(self.load_blacklist())
         self.old_process_application_commands = bot.process_application_commands
-        bot.process_application_commands = self.check  # type: ignore  # MyInter + Interaction
+        # MyInter + Interaction
+        bot.process_application_commands = self.check  # type: ignore
 
         if bot.guild_ids:
             self.blacklist_.guild_ids_to_rollout.update(bot.guild_ids)

botbase/exts/blacklist.py

@@ -1,19 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-from ormar import NoMatch
-
 from ..botbase import BotBase
-from ..db import CommandLog, database
+from ..db import CommandLog
 from ..models import CogBase
 from ..wraps import MyInter
 
 
-if TYPE_CHECKING:
-    from ..botbase import BotBase
-
-
 class CommandLogging(CogBase[BotBase]):
     @CogBase.listener()
     async def on_application_command_completion(self, inter: MyInter):
@@ -26,20 +18,18 @@ async def on_application_command_completion(self, inter: MyInter):
         channel = inter.channel.id if inter.guild is not None else None
         member = inter.user.id
 
-        # Composite PK not in ormar ruins the whole SQL statement.
-        await database.execute(
+        # Incrementing amount without funky sub queries is not possible.
+        await CommandLog.raw(
             """
-            INSERT INTO commands (command, guild, channel, member, amount)
-            VALUES (:command, :guild, :channel, :member, 1)
+            INSERT INTO command_log (command, guild, channel, member, amount)
+            VALUES ({}, {}, {}, {}, 1)
             ON CONFLICT (command, guild, channel, member)
-            DO UPDATE SET amount = commands.amount + 1
+            DO UPDATE SET amount = command_log.amount + 1;
             """,
-            {
-                "command": cmd,
-                "guild": guild,
-                "channel": channel,
-                "member": member,
-            },
+            cmd,
+            guild,
+            channel,
+            member,
         )
 
 

botbase/exts/log_commands.py

@@ -10,15 +10,18 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.8"
-asyncpg = { version = ">=0.25,<0.28", optional = true }
 nextcord = "^2.0.0"
 nextcord-ext-menus = "^1.5.2"
 psutil = "^5.9.0"
-ormar = { version = ">=0.11.0,<0.13.0", optional = true, extras = ["postgres"] }
 python-dotenv = ">=0.16.0,<1.1.0"
+piccolo = { version = "^0.111.1", optional = true, extras = [
+    "orjson",
+    "uvloop",
+    "postgres",
+] }
 
 [tool.poetry.extras]
-db = ["asyncpg", "ormar"]
+db = ["piccolo"]
 
 [tool.poetry.group.dev.dependencies]
 black = ">=22.3,<24.0"