Optimize database queries for proxy filtering
This commit is contained in:
		| @@ -18,16 +18,17 @@ namespace PluralKit.Bot { | ||||
|             _logger = logger.ForContext<LogChannelService>(); | ||||
|         } | ||||
|  | ||||
|         public async Task LogMessage(PKSystem system, PKMember member, ulong messageId, ulong originalMsgId, IGuildChannel originalChannel, IUser sender, string content) | ||||
|         public async Task LogMessage(PKSystem system, PKMember member, ulong messageId, ulong originalMsgId, IGuildChannel originalChannel, IUser sender, string content, GuildConfig? guildCfg = null) | ||||
|         { | ||||
|             var guildCfg = await _data.GetOrCreateGuildConfig(originalChannel.GuildId); | ||||
|              | ||||
|             if (guildCfg == null)  | ||||
|                 guildCfg = await _data.GetOrCreateGuildConfig(originalChannel.GuildId); | ||||
|  | ||||
|             // Bail if logging is disabled either globally or for this channel | ||||
|             if (guildCfg.LogChannel == null) return; | ||||
|             if (guildCfg.LogBlacklist.Contains(originalChannel.Id)) return; | ||||
|             if (guildCfg.Value.LogChannel == null) return; | ||||
|             if (guildCfg.Value.LogBlacklist.Contains(originalChannel.Id)) return; | ||||
|              | ||||
|             // Bail if we can't find the channel | ||||
|             if (!(await _client.GetChannelAsync(guildCfg.LogChannel.Value) is ITextChannel logChannel)) return; | ||||
|             if (!(await _client.GetChannelAsync(guildCfg.Value.LogChannel.Value) is ITextChannel logChannel)) return; | ||||
|  | ||||
|             var embed = _embed.CreateLoggedMessageEmbed(system, member, messageId, originalMsgId, sender, content, originalChannel); | ||||
|  | ||||
|   | ||||
| @@ -3,6 +3,8 @@ using System.Collections.Generic; | ||||
| using System.Linq; | ||||
| using System.Threading.Tasks; | ||||
|  | ||||
| using Dapper; | ||||
|  | ||||
| using Discord; | ||||
| using Discord.Net; | ||||
| using Discord.WebSocket; | ||||
| @@ -24,12 +26,13 @@ namespace PluralKit.Bot | ||||
|         private IDiscordClient _client; | ||||
|         private LogChannelService _logChannel; | ||||
|         private IDataStore _data; | ||||
|         private DbConnectionFactory _conn; | ||||
|         private EmbedService _embeds; | ||||
|         private ILogger _logger; | ||||
|         private WebhookExecutorService _webhookExecutor; | ||||
|         private ProxyCacheService _cache; | ||||
|          | ||||
|         public ProxyService(IDiscordClient client, LogChannelService logChannel, IDataStore data, EmbedService embeds, ILogger logger, ProxyCacheService cache, WebhookExecutorService webhookExecutor) | ||||
|         public ProxyService(IDiscordClient client, LogChannelService logChannel, IDataStore data, EmbedService embeds, ILogger logger, ProxyCacheService cache, WebhookExecutorService webhookExecutor, DbConnectionFactory conn) | ||||
|         { | ||||
|             _client = client; | ||||
|             _logChannel = logChannel; | ||||
| @@ -37,6 +40,7 @@ namespace PluralKit.Bot | ||||
|             _embeds = embeds; | ||||
|             _cache = cache; | ||||
|             _webhookExecutor = webhookExecutor; | ||||
|             _conn = conn; | ||||
|             _logger = logger.ForContext<ProxyService>(); | ||||
|         } | ||||
|  | ||||
| @@ -91,19 +95,16 @@ namespace PluralKit.Bot | ||||
|             var results = await _cache.GetResultsFor(message.Author.Id); | ||||
|             var match = GetProxyTagMatch(message.Content, results); | ||||
|             if (match == null) return; | ||||
|  | ||||
|             // Gather all "extra" data from DB at once | ||||
|             var aux = await _data.GetAuxillaryProxyInformation(channel.GuildId, match.System, match.Member); | ||||
|              | ||||
|             // And make sure the channel's not blacklisted from proxying. | ||||
|             var guildCfg = await _data.GetOrCreateGuildConfig(channel.GuildId); | ||||
|             if (guildCfg.Blacklist.Contains(channel.Id)) return; | ||||
|             if (aux.Guild.Blacklist.Contains(channel.Id)) return; | ||||
|              | ||||
|             // Make sure the system hasn't blacklisted the guild either | ||||
|             var systemGuildCfg = await _data.GetSystemGuildSettings(match.System, channel.GuildId); | ||||
|             if (!systemGuildCfg.ProxyEnabled) return; | ||||
|             if (!aux.SystemGuild.ProxyEnabled) return; | ||||
|              | ||||
|             // Also, check if the member has a guild nickname  | ||||
|             // TODO: roll this into the cached results as well as system/guild settings, maybe? Add a separate cache or something. | ||||
|             var memberGuildCfg = await _data.GetMemberGuildSettings(match.Member, channel.GuildId); | ||||
|  | ||||
|             // We know message.Channel can only be ITextChannel as PK doesn't work in DMs/groups | ||||
|             // Afterwards we ensure the bot has the right permissions, otherwise bail early | ||||
|             if (!await EnsureBotPermissions(channel)) return; | ||||
| @@ -113,7 +114,7 @@ namespace PluralKit.Bot | ||||
|                 return; | ||||
|              | ||||
|             // Get variables in order and all | ||||
|             var proxyName = match.Member.ProxyName(match.System.Tag, memberGuildCfg.DisplayName); | ||||
|             var proxyName = match.Member.ProxyName(match.System.Tag, aux.MemberGuild.DisplayName); | ||||
|             var avatarUrl = match.Member.AvatarUrl ?? match.System.AvatarUrl; | ||||
|              | ||||
|             // If the name's too long (or short), bail | ||||
| @@ -138,7 +139,7 @@ namespace PluralKit.Bot | ||||
|  | ||||
|             // Store the message in the database, and log it in the log channel (if applicable) | ||||
|             await _data.AddMessage(message.Author.Id, hookMessageId, message.Channel.Id, message.Id, match.Member); | ||||
|             await _logChannel.LogMessage(match.System, match.Member, hookMessageId, message.Id, message.Channel as IGuildChannel, message.Author, match.InnerText); | ||||
|             await _logChannel.LogMessage(match.System, match.Member, hookMessageId, message.Id, message.Channel as IGuildChannel, message.Author, match.InnerText, aux.Guild); | ||||
|  | ||||
|             // Wait a second or so before deleting the original message | ||||
|             await Task.Delay(1000); | ||||
|   | ||||
| @@ -74,6 +74,13 @@ namespace PluralKit { | ||||
|         public string DisplayName { get; set; } | ||||
|     } | ||||
|  | ||||
|     public class AuxillaryProxyInformation | ||||
|     { | ||||
|         public GuildConfig Guild { get; set; } | ||||
|         public SystemGuildSettings SystemGuild { get; set; } | ||||
|         public MemberGuildSettings MemberGuild { get; set; } | ||||
|     } | ||||
|  | ||||
|     public interface IDataStore | ||||
|     { | ||||
|         /// <summary> | ||||
| @@ -378,6 +385,8 @@ namespace PluralKit { | ||||
|         /// Saves the given guild configuration struct to the data store. | ||||
|         /// </summary> | ||||
|         Task SaveGuildConfig(GuildConfig cfg); | ||||
|  | ||||
|         Task<AuxillaryProxyInformation> GetAuxillaryProxyInformation(ulong guild, PKSystem system, PKMember member); | ||||
|     } | ||||
|      | ||||
|     public class PostgresDataStore: IDataStore { | ||||
| @@ -694,22 +703,24 @@ namespace PluralKit { | ||||
|             public ulong? LogChannel { get; set; } | ||||
|             public long[] LogBlacklist { get; set; } | ||||
|             public long[] Blacklist { get; set; } | ||||
|  | ||||
|             public GuildConfig Into() => | ||||
|                 new GuildConfig | ||||
|                 { | ||||
|                     Id = Id, | ||||
|                     LogChannel = LogChannel, | ||||
|                     LogBlacklist = new HashSet<ulong>(LogBlacklist?.Select(c => (ulong) c) ?? new ulong[] {}), | ||||
|                     Blacklist = new HashSet<ulong>(Blacklist?.Select(c => (ulong) c) ?? new ulong[]{}) | ||||
|                 }; | ||||
|         } | ||||
|  | ||||
|         public async Task<GuildConfig> GetOrCreateGuildConfig(ulong guild) | ||||
|         { | ||||
|             using (var conn = await _conn.Obtain()) | ||||
|             { | ||||
|                 var compat = await conn.QuerySingleOrDefaultAsync<DatabaseCompatibleGuildConfig>( | ||||
|                 return (await conn.QuerySingleOrDefaultAsync<DatabaseCompatibleGuildConfig>( | ||||
|                     "insert into servers (id) values (@Id) on conflict do nothing; select * from servers where id = @Id", | ||||
|                     new {Id = guild}); | ||||
|                 return new GuildConfig | ||||
|                 { | ||||
|                     Id = compat.Id, | ||||
|                     LogChannel = compat.LogChannel, | ||||
|                     LogBlacklist = new HashSet<ulong>(compat.LogBlacklist.Select(c => (ulong) c)), | ||||
|                     Blacklist = new HashSet<ulong>(compat.Blacklist.Select(c => (ulong) c)), | ||||
|                 }; | ||||
|                     new {Id = guild})).Into(); | ||||
|             } | ||||
|         } | ||||
|  | ||||
| @@ -725,7 +736,24 @@ namespace PluralKit { | ||||
|                 }); | ||||
|             _logger.Information("Updated guild configuration {@GuildCfg}", cfg); | ||||
|         } | ||||
|          | ||||
|  | ||||
|         public async Task<AuxillaryProxyInformation> GetAuxillaryProxyInformation(ulong guild, PKSystem system, PKMember member) | ||||
|         { | ||||
|             using var conn = await _conn.Obtain(); | ||||
|             var args = new {Guild = guild, System = system.Id, Member = member.Id}; | ||||
|              | ||||
|             var multi = await conn.QueryMultipleAsync(@" | ||||
|             select servers.* from servers where id = @Guild;  | ||||
|             select * from system_guild where guild = @Guild and system = @System; | ||||
|             select * from member_guild where guild = @Guild and member = @Member", args); | ||||
|             return new AuxillaryProxyInformation | ||||
|             { | ||||
|                 Guild = (await multi.ReadSingleOrDefaultAsync<DatabaseCompatibleGuildConfig>()).Into(), | ||||
|                 SystemGuild = await multi.ReadSingleOrDefaultAsync<SystemGuildSettings>() ?? new SystemGuildSettings(), | ||||
|                 MemberGuild = await multi.ReadSingleOrDefaultAsync<MemberGuildSettings>() ?? new MemberGuildSettings() | ||||
|             }; | ||||
|         } | ||||
|  | ||||
|         public async Task AddSwitch(PKSystem system, IEnumerable<PKMember> members) | ||||
|         { | ||||
|             // Use a transaction here since we're doing multiple executed commands in one | ||||
|   | ||||
		Reference in New Issue
	
	Block a user