From ad2ce5fffd34b9ef6140bc2e3b014fb53d63eef8 Mon Sep 17 00:00:00 2001 From: SylteA Date: Sat, 16 Mar 2024 22:52:07 +0100 Subject: [PATCH] Add model and functionality to track command usage --- bot/core.py | 36 +++++++++++++++++-- bot/models/__init__.py | 2 ++ bot/models/command_usage.py | 11 ++++++ .../migrations/007_down__command_usage.sql | 1 + .../migrations/007_up__command_usage.sql | 10 ++++++ bot/models/model.py | 7 +++- 6 files changed, 63 insertions(+), 4 deletions(-) create mode 100644 bot/models/command_usage.py create mode 100644 bot/models/migrations/007_down__command_usage.sql create mode 100644 bot/models/migrations/007_up__command_usage.sql diff --git a/bot/core.py b/bot/core.py index cd733d32..d733058e 100644 --- a/bot/core.py +++ b/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] diff --git a/bot/models/__init__.py b/bot/models/__init__.py index 2f693454..cfd72391 100644 --- a/bot/models/__init__.py +++ b/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, diff --git a/bot/models/command_usage.py b/bot/models/command_usage.py new file mode 100644 index 00000000..63faab3a --- /dev/null +++ b/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 diff --git a/bot/models/migrations/007_down__command_usage.sql b/bot/models/migrations/007_down__command_usage.sql new file mode 100644 index 00000000..9b597edf --- /dev/null +++ b/bot/models/migrations/007_down__command_usage.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS command_usage; diff --git a/bot/models/migrations/007_up__command_usage.sql b/bot/models/migrations/007_up__command_usage.sql new file mode 100644 index 00000000..75cddca6 --- /dev/null +++ b/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 +) diff --git a/bot/models/model.py b/bot/models/model.py index 28a9dce6..32c29506 100644 --- a/bot/models/model.py +++ b/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")