Refactor command message deletion

This commit is contained in:
Ske 2020-10-23 12:18:28 +02:00
parent 134829b263
commit 81cd5496d5
7 changed files with 117 additions and 46 deletions

View File

@ -36,13 +36,13 @@ namespace PluralKit.Bot
private readonly PeriodicStatCollector _collector; private readonly PeriodicStatCollector _collector;
private readonly IMetrics _metrics; private readonly IMetrics _metrics;
private readonly ErrorMessageService _errorMessageService; private readonly ErrorMessageService _errorMessageService;
private readonly IDatabase _db; private readonly CommandMessageService _commandMessageService;
private bool _hasReceivedReady = false; private bool _hasReceivedReady = false;
private Timer _periodicTask; // Never read, just kept here for GC reasons private Timer _periodicTask; // Never read, just kept here for GC reasons
public Bot(DiscordShardedClient client, ILifetimeScope services, ILogger logger, PeriodicStatCollector collector, IMetrics metrics, public Bot(DiscordShardedClient client, ILifetimeScope services, ILogger logger, PeriodicStatCollector collector, IMetrics metrics,
ErrorMessageService errorMessageService, IDatabase db) ErrorMessageService errorMessageService, CommandMessageService commandMessageService)
{ {
_client = client; _client = client;
_logger = logger.ForContext<Bot>(); _logger = logger.ForContext<Bot>();
@ -50,7 +50,7 @@ namespace PluralKit.Bot
_collector = collector; _collector = collector;
_metrics = metrics; _metrics = metrics;
_errorMessageService = errorMessageService; _errorMessageService = errorMessageService;
_db = db; _commandMessageService = commandMessageService;
} }
public void Init() public void Init()
@ -183,7 +183,7 @@ namespace PluralKit.Bot
await UpdateBotStatus(); await UpdateBotStatus();
// Clean up message cache in postgres // Clean up message cache in postgres
await _db.Execute(conn => conn.QueryAsync("select from cleanup_command_message()")); await _commandMessageService.CleanupOldMessages();
// Collect some stats, submit them to the metrics backend // Collect some stats, submit them to the metrics backend
await _collector.CollectStats(); await _collector.CollectStats();

View File

@ -28,6 +28,7 @@ namespace PluralKit.Bot
private readonly ModelRepository _repo; private readonly ModelRepository _repo;
private readonly PKSystem _senderSystem; private readonly PKSystem _senderSystem;
private readonly IMetrics _metrics; private readonly IMetrics _metrics;
private readonly CommandMessageService _commandMessageService;
private Command _currentCommand; private Command _currentCommand;
@ -44,6 +45,7 @@ namespace PluralKit.Bot
_repo = provider.Resolve<ModelRepository>(); _repo = provider.Resolve<ModelRepository>();
_metrics = provider.Resolve<IMetrics>(); _metrics = provider.Resolve<IMetrics>();
_provider = provider; _provider = provider;
_commandMessageService = provider.Resolve<CommandMessageService>();
_parameters = new Parameters(message.Content.Substring(commandParseOffset)); _parameters = new Parameters(message.Content.Substring(commandParseOffset));
} }
@ -73,10 +75,14 @@ namespace PluralKit.Bot
if (embed != null && !this.BotHasAllPermissions(Permissions.EmbedLinks)) if (embed != null && !this.BotHasAllPermissions(Permissions.EmbedLinks))
throw new PKError("PluralKit does not have permission to send embeds in this channel. Please ensure I have the **Embed Links** permission enabled."); throw new PKError("PluralKit does not have permission to send embeds in this channel. Please ensure I have the **Embed Links** permission enabled.");
var msg = await Channel.SendMessageFixedAsync(text, embed: embed, mentions: mentions); var msg = await Channel.SendMessageFixedAsync(text, embed: embed, mentions: mentions);
if (embed != null) if (embed != null)
{
// Sensitive information that might want to be deleted by :x: reaction is typically in an embed format (member cards, for example) // Sensitive information that might want to be deleted by :x: reaction is typically in an embed format (member cards, for example)
// This may need to be changed at some point but works well enough for now // This may need to be changed at some point but works well enough for now
await _db.Execute(conn => _repo.SaveCommandMessage(conn, msg.Id, Author.Id)); await _commandMessageService.RegisterMessage(msg.Id, Author.Id);
}
return msg; return msg;
} }

View File

@ -15,14 +15,16 @@ namespace PluralKit.Bot
{ {
private readonly IDatabase _db; private readonly IDatabase _db;
private readonly ModelRepository _repo; private readonly ModelRepository _repo;
private readonly CommandMessageService _commandMessageService;
private readonly EmbedService _embeds; private readonly EmbedService _embeds;
private readonly ILogger _logger; private readonly ILogger _logger;
public ReactionAdded(EmbedService embeds, ILogger logger, IDatabase db, ModelRepository repo) public ReactionAdded(EmbedService embeds, ILogger logger, IDatabase db, ModelRepository repo, CommandMessageService commandMessageService)
{ {
_embeds = embeds; _embeds = embeds;
_db = db; _db = db;
_repo = repo; _repo = repo;
_commandMessageService = commandMessageService;
_logger = logger.ForContext<ReactionAdded>(); _logger = logger.ForContext<ReactionAdded>();
} }
@ -44,39 +46,53 @@ namespace PluralKit.Bot
// Ignore reactions from bots (we can't DM them anyway) // Ignore reactions from bots (we can't DM them anyway)
if (evt.User.IsBot) return; if (evt.User.IsBot) return;
Task<FullMessage> GetMessage() =>
_db.Execute(c => _repo.GetMessage(c, evt.Message.Id));
FullMessage msg;
CommandMessage cmdmsg;
switch (evt.Emoji.Name) switch (evt.Emoji.Name)
{ {
// Message deletion // Message deletion
case "\u274C": // Red X case "\u274C": // Red X
if ((msg = await GetMessage()) != null) {
await HandleDeleteReaction(evt, msg); await using var conn = await _db.Obtain();
else if ((cmdmsg = await _db.Execute(conn => _repo.GetCommandMessage(conn, evt.Message.Id))) != null) var msg = await _repo.GetMessage(conn, evt.Message.Id);
await HandleCommandDeleteReaction(evt, cmdmsg); if (msg != null)
{
await HandleProxyDeleteReaction(evt, msg);
break;
}
var commandMsg = await _commandMessageService.GetCommandMessage(conn, evt.Message.Id);
if (commandMsg != null)
await HandleCommandDeleteReaction(evt, commandMsg);
break; break;
}
case "\u2753": // Red question mark case "\u2753": // Red question mark
case "\u2754": // White question mark case "\u2754": // White question mark
if ((msg = await GetMessage()) != null) {
await using var conn = await _db.Obtain();
var msg = await _repo.GetMessage(conn, evt.Message.Id);
if (msg != null)
await HandleQueryReaction(evt, msg); await HandleQueryReaction(evt, msg);
break; break;
}
case "\U0001F514": // Bell case "\U0001F514": // Bell
case "\U0001F6CE": // Bellhop bell case "\U0001F6CE": // Bellhop bell
case "\U0001F3D3": // Ping pong paddle (lol) case "\U0001F3D3": // Ping pong paddle (lol)
case "\u23F0": // Alarm clock case "\u23F0": // Alarm clock
case "\u2757": // Exclamation mark case "\u2757": // Exclamation mark
if ((msg = await GetMessage()) != null) {
await using var conn = await _db.Obtain();
var msg = await _repo.GetMessage(conn, evt.Message.Id);
if (msg != null)
await HandlePingReaction(evt, msg); await HandlePingReaction(evt, msg);
break; break;
}
} }
} }
private async ValueTask HandleDeleteReaction(MessageReactionAddEventArgs evt, FullMessage msg) private async ValueTask HandleProxyDeleteReaction(MessageReactionAddEventArgs evt, FullMessage msg)
{ {
if (!evt.Channel.BotHasAllPermissions(Permissions.ManageMessages)) return; if (!evt.Channel.BotHasAllPermissions(Permissions.ManageMessages)) return;
@ -97,10 +113,12 @@ namespace PluralKit.Bot
private async ValueTask HandleCommandDeleteReaction(MessageReactionAddEventArgs evt, CommandMessage msg) private async ValueTask HandleCommandDeleteReaction(MessageReactionAddEventArgs evt, CommandMessage msg)
{ {
if (!evt.Channel.BotHasAllPermissions(Permissions.ManageMessages)) return; if (!evt.Channel.BotHasAllPermissions(Permissions.ManageMessages))
return;
// Can only delete your own message // Can only delete your own message
if (msg.author_id != evt.User.Id) return; if (msg.AuthorId != evt.User.Id)
return;
try try
{ {

View File

@ -74,6 +74,7 @@ namespace PluralKit.Bot
builder.RegisterType<LastMessageCacheService>().AsSelf().SingleInstance(); builder.RegisterType<LastMessageCacheService>().AsSelf().SingleInstance();
builder.RegisterType<LoggerCleanService>().AsSelf().SingleInstance(); builder.RegisterType<LoggerCleanService>().AsSelf().SingleInstance();
builder.RegisterType<ErrorMessageService>().AsSelf().SingleInstance(); builder.RegisterType<ErrorMessageService>().AsSelf().SingleInstance();
builder.RegisterType<CommandMessageService>().AsSelf().SingleInstance();
// Sentry stuff // Sentry stuff
builder.Register(_ => new Scope(null)).AsSelf().InstancePerLifetimeScope(); builder.Register(_ => new Scope(null)).AsSelf().InstancePerLifetimeScope();

View File

@ -0,0 +1,50 @@
using System.Threading.Tasks;
using NodaTime;
using PluralKit.Core;
using Serilog;
namespace PluralKit.Bot
{
public class CommandMessageService
{
private static readonly Duration CommandMessageRetention = Duration.FromSeconds(2);
private readonly IDatabase _db;
private readonly ModelRepository _repo;
private readonly IClock _clock;
private readonly ILogger _logger;
public CommandMessageService(IDatabase db, ModelRepository repo, IClock clock, ILogger logger)
{
_db = db;
_repo = repo;
_clock = clock;
_logger = logger;
}
public async Task RegisterMessage(ulong messageId, ulong authorId)
{
_logger.Debug("Registering command response {MessageId} from author {AuthorId}", messageId, authorId);
await _db.Execute(conn => _repo.SaveCommandMessage(conn, messageId, authorId));
}
public async Task<CommandMessage> GetCommandMessage(IPKConnection conn, ulong messageId)
{
return await _repo.GetCommandMessage(conn, messageId);
}
public async Task CleanupOldMessages()
{
var deleteThresholdInstant = _clock.GetCurrentInstant() - CommandMessageRetention;
var deleteThresholdSnowflake = DiscordUtils.InstantToSnowflake(deleteThresholdInstant);
var deletedRows = await _db.Execute(conn => _repo.DeleteCommandMessagesBefore(conn, deleteThresholdSnowflake));
_logger.Information("Pruned {DeletedRows} command messages older than retention {Retention} (older than {DeleteThresholdInstant} / {DeleteThresholdSnowflake})",
deletedRows, CommandMessageRetention, deleteThresholdInstant, deleteThresholdSnowflake);
}
}
}

View File

@ -1,17 +1,10 @@
-- SCHEMA VERSION 11: (insert date) -- -- SCHEMA VERSION 11: 2020-10-23 --
-- Create command message table -- -- Create command message table --
create table command_message create table command_messages
( (
message_id bigint primary key, message_id bigint primary key not null,
author_id bigint not null, author_id bigint not null
timestamp timestamp not null default now()
); );
create function cleanup_command_message() returns void as $$
begin
delete from command_message where timestamp < now() - interval '2 hours';
end;
$$ language plpgsql;
update info set schema_version = 11; update info set schema_version = 11;

View File

@ -1,5 +1,3 @@
using System.Collections.Generic;
using System.Data;
using System.Threading.Tasks; using System.Threading.Tasks;
using Dapper; using Dapper;
@ -8,17 +6,22 @@ namespace PluralKit.Core
{ {
public partial class ModelRepository public partial class ModelRepository
{ {
public Task SaveCommandMessage(IPKConnection conn, ulong message_id, ulong author_id) => public Task SaveCommandMessage(IPKConnection conn, ulong messageId, ulong authorId) =>
conn.QueryAsync("insert into command_message (message_id, author_id) values (@Message, @Author)", conn.QueryAsync("insert into command_messages (message_id, author_id) values (@Message, @Author)",
new {Message = message_id, Author = author_id }); new {Message = messageId, Author = authorId });
public Task<CommandMessage> GetCommandMessage(IPKConnection conn, ulong message_id) => public Task<CommandMessage> GetCommandMessage(IPKConnection conn, ulong messageId) =>
conn.QuerySingleOrDefaultAsync<CommandMessage>("select message_id, author_id from command_message where message_id = @Message", conn.QuerySingleOrDefaultAsync<CommandMessage>("select * from command_messages where message_id = @Message",
new {Message = message_id}); new {Message = messageId});
}
public Task<int> DeleteCommandMessagesBefore(IPKConnection conn, ulong messageIdThreshold) =>
conn.ExecuteAsync("delete from command_messages where message_id < @Threshold",
new {Threshold = messageIdThreshold});
}
public class CommandMessage public class CommandMessage
{ {
public ulong author_id { get; set; } public ulong AuthorId { get; set; }
public ulong MessageId { get; set; }
} }
} }