From 8a689ac0f24f02a58d962c5a7ac732b3eefcde62 Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 18 Jan 2020 00:02:17 +0100 Subject: [PATCH] Upgrade various store methods to IAsyncEnumerable --- PluralKit.API/Controllers/SystemController.cs | 8 +++--- PluralKit.Bot/Commands/SwitchCommands.cs | 28 +++++++++---------- PluralKit.Bot/Commands/SystemCommands.cs | 9 ++++-- PluralKit.Bot/Services/EmbedService.cs | 13 +++++---- PluralKit.Core/PluralKit.Core.csproj | 1 + PluralKit.Core/Stores.cs | 28 +++++++++++-------- PluralKit.Core/Utils.cs | 13 +++++++++ 7 files changed, 62 insertions(+), 38 deletions(-) diff --git a/PluralKit.API/Controllers/SystemController.cs b/PluralKit.API/Controllers/SystemController.cs index 5c8d1a69..9e62bfe6 100644 --- a/PluralKit.API/Controllers/SystemController.cs +++ b/PluralKit.API/Controllers/SystemController.cs @@ -113,11 +113,11 @@ namespace PluralKit.API.Controllers var sw = await _data.GetLatestSwitch(system); if (sw == null) return NotFound("System has no registered switches."); - var members = await _data.GetSwitchMembers(sw); + var members = _data.GetSwitchMembers(sw); return Ok(new FrontersReturn { Timestamp = sw.Timestamp, - Members = members.Select(m => m.ToJson(_auth.ContextFor(system))) + Members = await members.Select(m => m.ToJson(_auth.ContextFor(system))).ToListAsync() }); } @@ -151,10 +151,10 @@ namespace PluralKit.API.Controllers var latestSwitch = await _data.GetLatestSwitch(_auth.CurrentSystem); if (latestSwitch != null) { - var latestSwitchMembers = await _data.GetSwitchMembers(latestSwitch); + var latestSwitchMembers = _data.GetSwitchMembers(latestSwitch); // Bail if this switch is identical to the latest one - if (latestSwitchMembers.Select(m => m.Hid).SequenceEqual(param.Members)) + if (await latestSwitchMembers.Select(m => m.Hid).SequenceEqualAsync(param.Members.ToAsyncEnumerable())) return BadRequest("New members identical to existing fronters."); } diff --git a/PluralKit.Bot/Commands/SwitchCommands.cs b/PluralKit.Bot/Commands/SwitchCommands.cs index 5bb91ea5..c0c586ac 100644 --- a/PluralKit.Bot/Commands/SwitchCommands.cs +++ b/PluralKit.Bot/Commands/SwitchCommands.cs @@ -58,9 +58,9 @@ namespace PluralKit.Bot.Commands var lastSwitch = await _data.GetLatestSwitch(ctx.System); if (lastSwitch != null) { - var lastSwitchMembers = await _data.GetSwitchMembers(lastSwitch); + var lastSwitchMembers = _data.GetSwitchMembers(lastSwitch); // Make sure the requested switch isn't identical to the last one - if (lastSwitchMembers.Select(m => m.Id).SequenceEqual(members.Select(m => m.Id))) + if (await lastSwitchMembers.Select(m => m.Id).SequenceEqualAsync(members.Select(m => m.Id).ToAsyncEnumerable())) throw Errors.SameSwitch(members); } @@ -86,13 +86,13 @@ namespace PluralKit.Bot.Commands if (time.ToInstant() > SystemClock.Instance.GetCurrentInstant()) throw Errors.SwitchTimeInFuture; // Fetch the last two switches for the system to do bounds checking on - var lastTwoSwitches = (await _data.GetSwitches(ctx.System, 2)).ToArray(); + var lastTwoSwitches = await _data.GetSwitches(ctx.System).Take(2).ToListAsync(); // If we don't have a switch to move, don't bother - if (lastTwoSwitches.Length == 0) throw Errors.NoRegisteredSwitches; + if (lastTwoSwitches.Count == 0) throw Errors.NoRegisteredSwitches; // If there's a switch *behind* the one we move, we check to make srue we're not moving the time further back than that - if (lastTwoSwitches.Length == 2) + if (lastTwoSwitches.Count == 2) { if (lastTwoSwitches[1].Timestamp > time.ToInstant()) throw Errors.SwitchMoveBeforeSecondLast(lastTwoSwitches[1].Timestamp.InZone(tz)); @@ -100,8 +100,8 @@ namespace PluralKit.Bot.Commands // Now we can actually do the move, yay! // But, we do a prompt to confirm. - var lastSwitchMembers = await _data.GetSwitchMembers(lastTwoSwitches[0]); - var lastSwitchMemberStr = string.Join(", ", lastSwitchMembers.Select(m => m.Name)); + var lastSwitchMembers = _data.GetSwitchMembers(lastTwoSwitches[0]); + var lastSwitchMemberStr = string.Join(", ", await lastSwitchMembers.Select(m => m.Name).ToListAsync()); var lastSwitchTimeStr = Formats.ZonedDateTimeFormat.Format(lastTwoSwitches[0].Timestamp.InZone(ctx.System.Zone)); var lastSwitchDeltaStr = Formats.DurationFormat.Format(SystemClock.Instance.GetCurrentInstant() - lastTwoSwitches[0].Timestamp); var newSwitchTimeStr = Formats.ZonedDateTimeFormat.Format(time); @@ -132,23 +132,23 @@ namespace PluralKit.Bot.Commands } // Fetch the last two switches for the system to do bounds checking on - var lastTwoSwitches = (await _data.GetSwitches(ctx.System, 2)).ToArray(); - if (lastTwoSwitches.Length == 0) throw Errors.NoRegisteredSwitches; + var lastTwoSwitches = await _data.GetSwitches(ctx.System).Take(2).ToListAsync(); + if (lastTwoSwitches.Count == 0) throw Errors.NoRegisteredSwitches; - var lastSwitchMembers = await _data.GetSwitchMembers(lastTwoSwitches[0]); - var lastSwitchMemberStr = string.Join(", ", lastSwitchMembers.Select(m => m.Name)); + var lastSwitchMembers = _data.GetSwitchMembers(lastTwoSwitches[0]); + var lastSwitchMemberStr = string.Join(", ", await lastSwitchMembers.Select(m => m.Name).ToListAsync()); var lastSwitchDeltaStr = Formats.DurationFormat.Format(SystemClock.Instance.GetCurrentInstant() - lastTwoSwitches[0].Timestamp); IUserMessage msg; - if (lastTwoSwitches.Length == 1) + if (lastTwoSwitches.Count == 1) { msg = await ctx.Reply( $"{Emojis.Warn} This will delete the latest switch ({lastSwitchMemberStr.SanitizeMentions()}, {lastSwitchDeltaStr} ago). You have no other switches logged. Is this okay?"); } else { - var secondSwitchMembers = await _data.GetSwitchMembers(lastTwoSwitches[1]); - var secondSwitchMemberStr = string.Join(", ", secondSwitchMembers.Select(m => m.Name)); + var secondSwitchMembers = _data.GetSwitchMembers(lastTwoSwitches[1]); + var secondSwitchMemberStr = string.Join(", ", await secondSwitchMembers.Select(m => m.Name).ToListAsync()); var secondSwitchDeltaStr = Formats.DurationFormat.Format(SystemClock.Instance.GetCurrentInstant() - lastTwoSwitches[1].Timestamp); msg = await ctx.Reply( $"{Emojis.Warn} This will delete the latest switch ({lastSwitchMemberStr.SanitizeMentions()}, {lastSwitchDeltaStr} ago). The next latest switch is {secondSwitchMemberStr.SanitizeMentions()} ({secondSwitchDeltaStr} ago). Is this okay?"); diff --git a/PluralKit.Bot/Commands/SystemCommands.cs b/PluralKit.Bot/Commands/SystemCommands.cs index f11d6e41..97bfc203 100644 --- a/PluralKit.Bot/Commands/SystemCommands.cs +++ b/PluralKit.Bot/Commands/SystemCommands.cs @@ -239,10 +239,13 @@ namespace PluralKit.Bot.Commands if (system == null) throw Errors.NoSystemError; ctx.CheckSystemPrivacy(system, system.FrontHistoryPrivacy); - var sws = (await _data.GetSwitches(system, 10)).ToList(); - if (sws.Count == 0) throw Errors.NoRegisteredSwitches; + var sws = _data.GetSwitches(system).Take(10); + var embed = await _embeds.CreateFrontHistoryEmbed(sws, system.Zone); - await ctx.Reply(embed: await _embeds.CreateFrontHistoryEmbed(sws, system.Zone)); + // Moving the count check to the CreateFrontHistoryEmbed function to avoid a double-iteration + // If embed == null, then there's no switches, so error + if (embed == null) throw Errors.NoRegisteredSwitches; + await ctx.Reply(embed: embed); } public async Task SystemFrontPercent(Context ctx, PKSystem system) diff --git a/PluralKit.Bot/Services/EmbedService.cs b/PluralKit.Bot/Services/EmbedService.cs index 528c430c..ec15546f 100644 --- a/PluralKit.Bot/Services/EmbedService.cs +++ b/PluralKit.Bot/Services/EmbedService.cs @@ -36,7 +36,7 @@ namespace PluralKit.Bot { var latestSwitch = await _data.GetLatestSwitch(system); if (latestSwitch != null && system.FrontPrivacy.CanAccess(ctx)) { - var switchMembers = (await _data.GetSwitchMembers(latestSwitch)).ToList(); + var switchMembers = await _data.GetSwitchMembers(latestSwitch).ToListAsync(); if (switchMembers.Count > 0) eb.AddField("Fronter".ToQuantity(switchMembers.Count(), ShowQuantityAs.None), string.Join(", ", switchMembers.Select(m => m.Name))); @@ -115,7 +115,7 @@ namespace PluralKit.Bot { public async Task CreateFronterEmbed(PKSwitch sw, DateTimeZone zone) { - var members = (await _data.GetSwitchMembers(sw)).ToList(); + var members = await _data.GetSwitchMembers(sw).ToListAsync(); var timeSinceSwitch = SystemClock.Instance.GetCurrentInstant() - sw.Timestamp; return new EmbedBuilder() .WithColor(members.FirstOrDefault()?.Color?.ToDiscordColor() ?? Color.Blue) @@ -124,15 +124,15 @@ namespace PluralKit.Bot { .Build(); } - public async Task CreateFrontHistoryEmbed(IEnumerable sws, DateTimeZone zone) + public async Task CreateFrontHistoryEmbed(IAsyncEnumerable sws, DateTimeZone zone) { var outputStr = ""; PKSwitch lastSw = null; - foreach (var sw in sws) + await foreach (var sw in sws) { // Fetch member list and format - var members = (await _data.GetSwitchMembers(sw)).ToList(); + var members = await _data.GetSwitchMembers(sw).ToListAsync(); var membersStr = members.Any() ? string.Join(", ", members.Select(m => m.Name)) : "no fronter"; var switchSince = SystemClock.Instance.GetCurrentInstant() - sw.Timestamp; @@ -156,6 +156,9 @@ namespace PluralKit.Bot { lastSw = sw; } + if (lastSw == null) + return null; + return new EmbedBuilder() .WithTitle("Past switches") .WithDescription(outputStr) diff --git a/PluralKit.Core/PluralKit.Core.csproj b/PluralKit.Core/PluralKit.Core.csproj index 708ab317..7352dd2f 100644 --- a/PluralKit.Core/PluralKit.Core.csproj +++ b/PluralKit.Core/PluralKit.Core.csproj @@ -25,6 +25,7 @@ + diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs index c4f37ed1..7b5d29ea 100644 --- a/PluralKit.Core/Stores.cs +++ b/PluralKit.Core/Stores.cs @@ -1,4 +1,5 @@ using System.Collections.Generic; +using System.Data.Common; using System.Linq; using System.Threading.Tasks; @@ -289,7 +290,7 @@ namespace PluralKit { /// Gets switches from a system. /// /// An enumerable of the *count* latest switches in the system, in latest-first order. May contain fewer elements than requested. - Task> GetSwitches(PKSystem system, int count); + IAsyncEnumerable GetSwitches(PKSystem system); /// /// Gets the latest (temporally; closest to now) switch of a given system. @@ -299,7 +300,7 @@ namespace PluralKit { /// /// Gets the members a given switch consists of. /// - Task> GetSwitchMembers(PKSwitch sw); + IAsyncEnumerable GetSwitchMembers(PKSwitch sw); /// /// Gets a list of fronters over a given period of time. @@ -787,7 +788,9 @@ namespace PluralKit { public async Task AddSwitchesBulk(PKSystem system, IEnumerable switches) { // Read existing switches to enforce unique timestamps - var priorSwitches = await GetSwitches(system); + var priorSwitches = new List(); + await foreach (var sw in GetSwitches(system)) priorSwitches.Add(sw); + var lastSwitchId = priorSwitches.Any() ? priorSwitches.Max(x => x.Id) : 0; @@ -855,12 +858,13 @@ namespace PluralKit { _logger.Information("Completed bulk import of switches for system {0}", system.Hid); } - public async Task> GetSwitches(PKSystem system, int count = 9999999) + public IAsyncEnumerable GetSwitches(PKSystem system) { // TODO: refactor the PKSwitch data structure to somehow include a hydrated member list // (maybe when we get caching in?) - using (var conn = await _conn.Obtain()) - return await conn.QueryAsync("select * from switches where system = @System order by timestamp desc limit @Count", new {System = system.Id, Count = count}); + return _conn.QueryStreamAsync( + "select * from switches where system = @System order by timestamp desc", + new {System = system.Id}); } public async Task> GetSwitchMembersList(PKSystem system, Instant start, Instant end) @@ -899,15 +903,15 @@ namespace PluralKit { } } - public async Task> GetSwitchMembers(PKSwitch sw) + public IAsyncEnumerable GetSwitchMembers(PKSwitch sw) { - using (var conn = await _conn.Obtain()) - return await conn.QueryAsync( - "select * from switch_members, members where switch_members.member = members.id and switch_members.switch = @Switch order by switch_members.id", - new {Switch = sw.Id}); + return _conn.QueryStreamAsync( + "select * from switch_members, members where switch_members.member = members.id and switch_members.switch = @Switch order by switch_members.id", + new {Switch = sw.Id}); } - public async Task GetLatestSwitch(PKSystem system) => (await GetSwitches(system, 1)).FirstOrDefault(); + public async Task GetLatestSwitch(PKSystem system) => + await GetSwitches(system).FirstOrDefaultAsync(); public async Task MoveSwitch(PKSwitch sw, Instant time) { diff --git a/PluralKit.Core/Utils.cs b/PluralKit.Core/Utils.cs index d19809c5..38a5b2f9 100644 --- a/PluralKit.Core/Utils.cs +++ b/PluralKit.Core/Utils.cs @@ -669,4 +669,17 @@ namespace PluralKit EventId = Guid.NewGuid(); } } + + public static class ConnectionUtils + { + public static async IAsyncEnumerable QueryStreamAsync(this DbConnectionFactory connFactory, string sql, object param) + { + using var conn = await connFactory.Obtain(); + + var reader = await conn.ExecuteReaderAsync(sql, param); + var parser = reader.GetRowParser(); + while (reader.Read()) + yield return parser(reader); + } + } }