Merge pull request #299 from SylteA/feature/command_usage

Add model and functionality to track command usage

bot/core.py

@@ -11,7 +11,7 @@
 from discord.ext import commands, tasks
 
 from bot.config import settings
-from bot.models import GuildConfig
+from bot.models import GuildConfig, Model
 from bot.services import http, paste
 from bot.services.paste import Document
 from utils.errors import IgnorableException
@@ -51,7 +51,7 @@ def guild(self) -> discord.Guild | None:
 
     async def setup_hook(self) -> None:
         """Connect DB before bot is ready to assure that no calls are made before its ready"""
-        self.loop.create_task(self.when_online())
+        _ = self.loop.create_task(self.when_online())
         self.presence.start()
 
         update_health("running", True)
@@ -130,6 +130,36 @@ async def process_commands(self, message: discord.Message, /):
         log.info(f"{ctx.author} invoking command: {ctx.clean_prefix}{ctx.command.qualified_name}")
         await self.invoke(ctx)
 
+    async def on_app_command_completion(
+        self, interaction: discord.Interaction, command: app_commands.Command | app_commands.ContextMenu
+    ):
+        try:
+            if isinstance(command, app_commands.ContextMenu):
+                log.warning("ContextMenu finished but not handled by stats.")
+                return
+
+            log.info(f"{interaction.user.name} used command {command.qualified_name} -> {interaction.data}")
+
+            query = """
+                INSERT INTO command_usage
+                (user_id, guild_id, channel_id, interaction_id, command_id, command_name, options)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7)
+            """
+
+            res = await Model.pool.execute(
+                query,
+                interaction.user.id,
+                interaction.guild_id,
+                interaction.channel_id,
+                interaction.id,
+                int(interaction.data["id"]),  # noqa
+                command.qualified_name,
+                interaction.data["options"],
+            )
+            log.info(f"Command usage inserted {res}")
+        except Exception as e:
+            await self.on_error("on_app_command_completion", e)
+
     async def send_error(self, content: str, header: str, invoked_details_document: Document = None) -> None:
         def wrap(code: str) -> str:
             code = code.replace("`", "\u200b`")
@@ -181,7 +211,7 @@ async def on_app_command_error(self, interaction: "InteractionType", error: app_
     @tasks.loop(hours=24)
     async def presence(self):
         await self.wait_until_ready()
-        await self.change_presence(activity=discord.Game(name='use the prefix "tim."'))
+        await self.change_presence(activity=discord.Game(name="Now using slash commands!"))
 
 
 InteractionType = discord.Interaction[DiscordBot]

bot/models/__init__.py

@@ -1,3 +1,4 @@
+from .command_usage import CommandUsage
 from .custom_roles import CustomRole
 from .guild_configs import GuildConfig
 from .levelling_ignored_channels import IgnoredChannel
@@ -16,6 +17,7 @@
     Rep,
     Tag,
     User,
+    CommandUsage,
     LevellingUser,
     PersistedRole,
     IgnoredChannel,

bot/models/command_usage.py

@@ -0,0 +1,11 @@
+from .model import Model
+
+
+class CommandUsage(Model):
+    user_id: int
+    guild_id: int
+    channel_id: int | None
+    interaction_id: int
+    command_id: int
+    command_name: str
+    options: dict

bot/models/migrations/007_down__command_usage.sql

@@ -0,0 +1 @@
+DROP TABLE IF EXISTS command_usage;

bot/models/migrations/007_up__command_usage.sql

@@ -0,0 +1,10 @@
+CREATE TABLE IF NOT EXISTS command_usage
+(
+    user_id BIGINT NOT NULL,
+    guild_id BIGINT NOT NULL,
+    channel_id BIGINT,
+    interaction_id BIGINT PRIMARY KEY,
+    command_id BIGINT NOT NULL,
+    command_name TEXT NOT NULL,
+    options JSON NOT NULL
+)

bot/models/model.py

@@ -1,7 +1,9 @@
 import asyncio
+import json
 import logging
 from typing import ClassVar, List, Type, TypeVar, Union
 
+import asyncpg
 from asyncpg import Connection, Pool, Record, connect, create_pool
 from pydantic import BaseModel
 
@@ -32,8 +34,11 @@ async def create_pool(
         loop: asyncio.AbstractEventLoop = None,
         **kwargs,
     ) -> None:
+        async def init(con: asyncpg.Connection) -> None:
+            await con.set_type_codec("json", schema="pg_catalog", encoder=json.dumps, decoder=json.loads)
+
         cls.pool = await create_pool(
-            uri, min_size=min_con, max_size=max_con, loop=loop, record_class=CustomRecord, **kwargs
+            uri, min_size=min_con, max_size=max_con, loop=loop, record_class=CustomRecord, init=init, **kwargs
         )
         log.info(f"Established a pool with {min_con} - {max_con} connections\n")