liz/PluralKit
liz/PluralKit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(bot): store command message info in redis

Browse Source
This commit is contained in:
spiral 2022-06-19 20:28:55 -04:00
parent 33b77470ee
commit 5c055871e3
No known key found for this signature in database
GPG Key ID: 244A11E4B0BCF40E
6 changed files with 31 additions and 75 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
PluralKit.Bot/CommandSystem/Context/Context.cs
View File

@ -94,12 +94,12 @@ public class Context
AllowedMentions = mentions ?? new AllowedMentions()
});
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)
// This may need to be changed at some point but works well enough for now
// but since we can, we just store all sent messages for possible deletion
await _commandMessageService.RegisterMessage(msg.Id, msg.ChannelId, Author.Id);
}
// }
return msg;
}

10
PluralKit.Bot/Commands/Message.cs
View File

@ -2,6 +2,8 @@
using System.Text;
using System.Text.RegularExpressions;
using Autofac;
using Myriad.Builders;
using Myriad.Cache;
using Myriad.Extensions;
@ -304,14 +306,14 @@ public class ProxiedMessage
private async Task DeleteCommandMessage(Context ctx, ulong messageId)
{
var message = await ctx.Repository.GetCommandMessage(messageId);
if (message == null)
var (authorId, channelId) = await ctx.Services.Resolve<CommandMessageService>().GetCommandMessage(messageId);
if (authorId == null)
throw Errors.MessageNotFound(messageId);
if (message.AuthorId != ctx.Author.Id)
if (authorId != ctx.Author.Id)
throw new PKError("You can only delete command messages queried by this account.");
await ctx.Rest.DeleteMessage(message.ChannelId, message.MessageId);
await ctx.Rest.DeleteMessage(channelId!.Value, messageId);
if (ctx.Guild != null)
await ctx.Rest.DeleteMessage(ctx.Message);

10
PluralKit.Bot/Handlers/ReactionAdded.cs
View File

@ -73,10 +73,10 @@ public class ReactionAdded: IEventHandler<MessageReactionAddEvent>
return;
}
var commandMsg = await _commandMessageService.GetCommandMessage(evt.MessageId);
if (commandMsg != null)
var (authorId, _) = await _commandMessageService.GetCommandMessage(evt.MessageId);
if (authorId != null)
{
await HandleCommandDeleteReaction(evt, commandMsg);
await HandleCommandDeleteReaction(evt, authorId.Value);
return;
}
}
@ -141,11 +141,11 @@ public class ReactionAdded: IEventHandler<MessageReactionAddEvent>
await _repo.DeleteMessage(evt.MessageId);
}
private async ValueTask HandleCommandDeleteReaction(MessageReactionAddEvent evt, CommandMessage? msg)
private async ValueTask HandleCommandDeleteReaction(MessageReactionAddEvent evt, ulong? authorId)
{
// Can only delete your own message
// (except in DMs, where msg will be null)
if (msg != null && msg.AuthorId != evt.UserId)
if (authorId != null && authorId != evt.UserId)
return;
// todo: don't try to delete the user's own messages in DMs

26
PluralKit.Bot/Services/CommandMessageService.cs
View File

@ -8,16 +8,13 @@ namespace PluralKit.Bot;
public class CommandMessageService
{
private readonly IClock _clock;
private readonly IDatabase _db;
private readonly RedisService _redis;
private readonly ILogger _logger;
private readonly ModelRepository _repo;
private static readonly TimeSpan CommandMessageRetention = TimeSpan.FromHours(24);
public CommandMessageService(IDatabase db, ModelRepository repo, IClock clock, ILogger logger)
public CommandMessageService(RedisService redis, IClock clock, ILogger logger)
{
_db = db;
_repo = repo;
_clock = clock;
_redis = redis;
_logger = logger.ForContext<CommandMessageService>();
}
@ -27,9 +24,18 @@ public class CommandMessageService
"Registering command response {MessageId} from author {AuthorId} in {ChannelId}",
messageId, authorId, channelId
);
await _repo.SaveCommandMessage(messageId, channelId, authorId);
await _redis.Connection.GetDatabase().StringSetAsync(messageId.ToString(), $"{authorId}-{channelId}", expiry: CommandMessageRetention);
}
public async Task<CommandMessage?> GetCommandMessage(ulong messageId) =>
await _repo.GetCommandMessage(messageId);
public async Task<(ulong?, ulong?)> GetCommandMessage(ulong messageId)
{
var str = await _redis.Connection.GetDatabase().StringGetAsync(messageId.ToString());
if (str.HasValue)
{
var split = ((string)str).Split("-");
return (ulong.Parse(split[0]), ulong.Parse(split[1]));
}
return (null, null);
}
}

36
PluralKit.Core/Database/Repository/ModelRepository.CommandMessage.cs
View File

@ -1,36 +0,0 @@
using SqlKata;
namespace PluralKit.Core;
public partial class ModelRepository
{
public Task SaveCommandMessage(ulong messageId, ulong channelId, ulong authorId)
{
var query = new Query("command_messages").AsInsert(new
{
message_id = messageId,
channel_id = channelId,
author_id = authorId,
});
return _db.ExecuteQuery(query);
}
public Task<CommandMessage?> GetCommandMessage(ulong messageId)
{
var query = new Query("command_messages").Where("message_id", messageId);
return _db.QueryFirst<CommandMessage?>(query);
}
public Task<int> DeleteCommandMessagesBefore(ulong messageIdThreshold)
{
var query = new Query("command_messages").AsDelete().Where("message_id", "<", messageIdThreshold);
return _db.QueryFirst<int>(query);
}
}
public class CommandMessage
{
public ulong AuthorId { get; set; }
public ulong MessageId { get; set; }
public ulong ChannelId { get; set; }
}

16
PluralKit.ScheduledTasks/TaskHandler.cs
View File

@ -19,7 +19,6 @@ namespace PluralKit.ScheduledTasks;
public class TaskHandler
{
private static readonly Duration CommandMessageRetention = Duration.FromHours(24);
private readonly IDatabase _db;
private readonly RedisService _redis;
private readonly bool _useRedisMetrics;
@ -65,9 +64,6 @@ public class TaskHandler
if (_useRedisMetrics)
await CollectBotStats();
// Clean up message cache in postgres
await CleanupOldMessages();
stopwatch.Stop();
_logger.Information("Ran scheduled tasks in {Time}", stopwatch.ElapsedDuration());
}
@ -108,18 +104,6 @@ public class TaskHandler
_logger.Debug("Submitted metrics to backend");
}
private async Task CleanupOldMessages()
{
var deleteThresholdInstant = SystemClock.Instance.GetCurrentInstant() - CommandMessageRetention;
var deleteThresholdSnowflake = InstantToSnowflake(deleteThresholdInstant);
var deletedRows = await _repo.DeleteCommandMessagesBefore(deleteThresholdSnowflake);
_logger.Information(
"Pruned {DeletedRows} command messages older than retention {Retention} (older than {DeleteThresholdInstant} / {DeleteThresholdSnowflake})",
deletedRows, CommandMessageRetention, deleteThresholdInstant, deleteThresholdSnowflake);
}
// we don't have access to PluralKit.Bot here, so this needs to be vendored
public static ulong InstantToSnowflake(Instant time) =>
(ulong)(time - Instant.FromUtc(2015, 1, 1, 0, 0, 0)).TotalMilliseconds << 22;