diff --git a/PluralKit.API/Controllers/SystemController.cs b/PluralKit.API/Controllers/SystemController.cs index 9e62bfe6..bd6956f4 100644 --- a/PluralKit.API/Controllers/SystemController.cs +++ b/PluralKit.API/Controllers/SystemController.cs @@ -70,10 +70,11 @@ namespace PluralKit.API.Controllers if (!system.MemberListPrivacy.CanAccess(_auth.ContextFor(system))) return StatusCode(StatusCodes.Status403Forbidden, "Unauthorized to view member list."); - var members = await _data.GetSystemMembers(system); - return Ok(members + var members = _data.GetSystemMembers(system); + return Ok(await members .Where(m => m.MemberPrivacy.CanAccess(_auth.ContextFor(system))) - .Select(m => m.ToJson(_auth.ContextFor(system)))); + .Select(m => m.ToJson(_auth.ContextFor(system))) + .ToListAsync()); } [HttpGet("{hid}/switches")] diff --git a/PluralKit.Bot/Commands/MemberCommands.cs b/PluralKit.Bot/Commands/MemberCommands.cs index d4d57d6c..6946e6c6 100644 --- a/PluralKit.Bot/Commands/MemberCommands.cs +++ b/PluralKit.Bot/Commands/MemberCommands.cs @@ -65,7 +65,8 @@ namespace PluralKit.Bot.Commands var randGen = new System.Random(); //Maybe move this somewhere else in the file structure since it doesn't need to get created at every command - var members = (await _data.GetSystemMembers(ctx.System)).Where(m => m.MemberPrivacy == PrivacyLevel.Public).ToList(); + // TODO: don't buffer these, find something else to do ig + var members = await _data.GetSystemMembers(ctx.System).Where(m => m.MemberPrivacy == PrivacyLevel.Public).ToListAsync(); if (members == null || !members.Any()) throw Errors.NoMembersError; var randInt = randGen.Next(members.Count); diff --git a/PluralKit.Bot/Commands/SystemCommands.cs b/PluralKit.Bot/Commands/SystemCommands.cs index 97bfc203..1f4986c3 100644 --- a/PluralKit.Bot/Commands/SystemCommands.cs +++ b/PluralKit.Bot/Commands/SystemCommands.cs @@ -154,16 +154,22 @@ namespace PluralKit.Bot.Commands var authCtx = ctx.LookupContextFor(system); var shouldShowPrivate = authCtx == LookupContext.ByOwner && ctx.Match("all", "everyone", "private"); - var members = (await _data.GetSystemMembers(system)).ToList(); var embedTitle = system.Name != null ? $"Members of {system.Name.SanitizeMentions()} (`{system.Hid}`)" : $"Members of `{system.Hid}`"; - var membersToDisplay = members + var memberCountPublic = _data.GetSystemMemberCount(system, false); + var memberCountAll = _data.GetSystemMemberCount(system, true); + await Task.WhenAll(memberCountPublic, memberCountAll); + + var memberCountDisplayed = shouldShowPrivate ? memberCountAll.Result : memberCountPublic.Result; + + var members = _data.GetSystemMembers(system) .Where(m => m.MemberPrivacy == PrivacyLevel.Public || shouldShowPrivate) - .OrderBy(m => m.Name, StringComparer.InvariantCultureIgnoreCase).ToList(); - var anyMembersHidden = members.Any(m => m.MemberPrivacy == PrivacyLevel.Private && !shouldShowPrivate); + .OrderBy(m => m.Name, StringComparer.InvariantCultureIgnoreCase); + var anyMembersHidden = !shouldShowPrivate && memberCountPublic.Result != memberCountAll.Result; await ctx.Paginate( - membersToDisplay, + members, + memberCountDisplayed, 25, embedTitle, (eb, ms) => @@ -176,10 +182,12 @@ namespace PluralKit.Bot.Commands return $"[`{m.Hid}`] **{m.Name.SanitizeMentions()}**"; })); - var footer = $"{membersToDisplay.Count} total."; + var footer = $"{memberCountDisplayed} total."; if (anyMembersHidden && authCtx == LookupContext.ByOwner) footer += "Private members have been hidden. type \"pk;system list all\" to include them."; eb.WithFooter(footer); + + return Task.CompletedTask; }); } @@ -189,17 +197,23 @@ namespace PluralKit.Bot.Commands var authCtx = ctx.LookupContextFor(system); var shouldShowPrivate = authCtx == LookupContext.ByOwner && ctx.Match("all", "everyone", "private"); - - var members = (await _data.GetSystemMembers(system)).ToList(); - var embedTitle = system.Name != null ? $"Members of {system.Name} (`{system.Hid}`)" : $"Members of `{system.Hid}`"; - var membersToDisplay = members + var embedTitle = system.Name != null ? $"Members of {system.Name} (`{system.Hid}`)" : $"Members of `{system.Hid}`"; + + var memberCountPublic = _data.GetSystemMemberCount(system, false); + var memberCountAll = _data.GetSystemMemberCount(system, true); + await Task.WhenAll(memberCountPublic, memberCountAll); + + var memberCountDisplayed = shouldShowPrivate ? memberCountAll.Result : memberCountPublic.Result; + + var members = _data.GetSystemMembers(system) .Where(m => m.MemberPrivacy == PrivacyLevel.Public || shouldShowPrivate) - .OrderBy(m => m.Name, StringComparer.InvariantCultureIgnoreCase).ToList(); - var anyMembersHidden = members.Any(m => m.MemberPrivacy == PrivacyLevel.Private && !shouldShowPrivate); + .OrderBy(m => m.Name, StringComparer.InvariantCultureIgnoreCase); + var anyMembersHidden = !shouldShowPrivate && memberCountPublic.Result != memberCountAll.Result; await ctx.Paginate( - membersToDisplay, + members, + memberCountDisplayed, 5, embedTitle, (eb, ms) => { @@ -215,10 +229,11 @@ namespace PluralKit.Bot.Commands eb.AddField(m.Name, profile.Truncate(1024)); } - var footer = $"{membersToDisplay.Count} total."; + var footer = $"{memberCountDisplayed} total."; if (anyMembersHidden && authCtx == LookupContext.ByOwner) footer += " Private members have been hidden. type \"pk;system list full all\" to include them."; eb.WithFooter(footer); + return Task.CompletedTask; } ); } @@ -233,19 +248,72 @@ namespace PluralKit.Bot.Commands await ctx.Reply(embed: await _embeds.CreateFronterEmbed(sw, system.Zone)); } + + struct FrontHistoryEntry + { + public Instant? LastTime; + public PKSwitch ThisSwitch; + + public FrontHistoryEntry(Instant? lastTime, PKSwitch thisSwitch) + { + LastTime = lastTime; + ThisSwitch = thisSwitch; + } + } public async Task SystemFrontHistory(Context ctx, PKSystem system) { if (system == null) throw Errors.NoSystemError; ctx.CheckSystemPrivacy(system, system.FrontHistoryPrivacy); - var sws = _data.GetSwitches(system).Take(10); - var embed = await _embeds.CreateFrontHistoryEmbed(sws, system.Zone); + var sws = _data.GetSwitches(system) + .Scan(new FrontHistoryEntry(null, null), (lastEntry, newSwitch) => new FrontHistoryEntry(lastEntry.ThisSwitch?.Timestamp, newSwitch)); + var totalSwitches = await _data.GetSwitchCount(system); + if (totalSwitches == 0) throw Errors.NoRegisteredSwitches; - // 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); + var embedTitle = system.Name != null ? $"Front history of {system.Name} (`{system.Hid}`)" : $"Front history of `{system.Hid}`"; + + await ctx.Paginate( + sws, + totalSwitches, + 10, + embedTitle, + async (builder, switches) => + { + var outputStr = ""; + foreach (var entry in switches) + { + var lastSw = entry.LastTime; + + var sw = entry.ThisSwitch; + // Fetch member list and format + 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; + + // If this isn't the latest switch, we also show duration + string stringToAdd; + if (lastSw != null) + { + // Calculate the time between the last switch (that we iterated - ie. the next one on the timeline) and the current one + var switchDuration = lastSw.Value - sw.Timestamp; + stringToAdd = + $"**{membersStr}** ({Formats.ZonedDateTimeFormat.Format(sw.Timestamp.InZone(system.Zone))}, {Formats.DurationFormat.Format(switchSince)} ago, for {Formats.DurationFormat.Format(switchDuration)})\n"; + } + else + { + stringToAdd = + $"**{membersStr}** ({Formats.ZonedDateTimeFormat.Format(sw.Timestamp.InZone(system.Zone))}, {Formats.DurationFormat.Format(switchSince)} ago)\n"; + } + + if (outputStr.Length + stringToAdd.Length > EmbedBuilder.MaxDescriptionLength) break; + outputStr += stringToAdd; + } + + builder.Description = outputStr; + } + ); } public async Task SystemFrontPercent(Context ctx, PKSystem system) diff --git a/PluralKit.Bot/ContextUtils.cs b/PluralKit.Bot/ContextUtils.cs index 52564056..bcdcf611 100644 --- a/PluralKit.Bot/ContextUtils.cs +++ b/PluralKit.Bot/ContextUtils.cs @@ -4,7 +4,6 @@ using System.Linq; using System.Net; using System.Threading.Tasks; using Discord; -using Discord.Commands; using Discord.Net; using Discord.WebSocket; @@ -62,22 +61,30 @@ namespace PluralKit.Bot { return string.Equals(msg.Content, expectedReply, StringComparison.InvariantCultureIgnoreCase); } - public static async Task Paginate(this Context ctx, ICollection items, int itemsPerPage, string title, Action> renderer) { + public static async Task Paginate(this Context ctx, IAsyncEnumerable items, int totalCount, int itemsPerPage, string title, Func, Task> renderer) { // TODO: make this generic enough we can use it in Choose below + + var buffer = new List(); + await using var enumerator = items.GetAsyncEnumerator(); - var pageCount = (items.Count / itemsPerPage) + 1; - Embed MakeEmbedForPage(int page) { + var pageCount = (totalCount / itemsPerPage) + 1; + async Task MakeEmbedForPage(int page) + { + var bufferedItemsNeeded = (page + 1) * itemsPerPage; + while (buffer.Count < bufferedItemsNeeded && await enumerator.MoveNextAsync()) + buffer.Add(enumerator.Current); + var eb = new EmbedBuilder(); eb.Title = pageCount > 1 ? $"[{page+1}/{pageCount}] {title}" : title; - renderer(eb, items.Skip(page*itemsPerPage).Take(itemsPerPage)); + await renderer(eb, buffer.Skip(page*itemsPerPage).Take(itemsPerPage)); return eb.Build(); } try { - var msg = await ctx.Channel.SendMessageAsync(embed: MakeEmbedForPage(0)); + var msg = await ctx.Channel.SendMessageAsync(embed: await MakeEmbedForPage(0)); if (pageCount == 1) return; // If we only have one page, don't bother with the reaction/pagination logic, lol - var botEmojis = new[] { new Emoji("\u23EA"), new Emoji("\u2B05"), new Emoji("\u27A1"), new Emoji("\u23E9"), new Emoji(Emojis.Error) }; + IEmote[] botEmojis = { new Emoji("\u23EA"), new Emoji("\u2B05"), new Emoji("\u27A1"), new Emoji("\u23E9"), new Emoji(Emojis.Error) }; await msg.AddReactionsAsync(botEmojis); try { @@ -99,13 +106,13 @@ namespace PluralKit.Bot { if (await ctx.HasPermission(ChannelPermission.ManageMessages) && reaction.User.IsSpecified) await msg.RemoveReactionAsync(reaction.Emote, reaction.User.Value); // Edit the embed with the new page - await msg.ModifyAsync((mp) => mp.Embed = MakeEmbedForPage(currentPage)); + var embed = await MakeEmbedForPage(currentPage); + await msg.ModifyAsync((mp) => mp.Embed = embed); } } catch (TimeoutException) { // "escape hatch", clean up as if we hit X } - if (await ctx.HasPermission(ChannelPermission.ManageMessages)) await msg.RemoveAllReactionsAsync(); else await msg.RemoveReactionsAsync(ctx.Shard.CurrentUser, botEmojis); } diff --git a/PluralKit.Bot/Services/EmbedService.cs b/PluralKit.Bot/Services/EmbedService.cs index ec15546f..be914cf2 100644 --- a/PluralKit.Bot/Services/EmbedService.cs +++ b/PluralKit.Bot/Services/EmbedService.cs @@ -124,47 +124,6 @@ namespace PluralKit.Bot { .Build(); } - public async Task CreateFrontHistoryEmbed(IAsyncEnumerable sws, DateTimeZone zone) - { - var outputStr = ""; - - PKSwitch lastSw = null; - await foreach (var sw in sws) - { - // Fetch member list and format - 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; - - // If this isn't the latest switch, we also show duration - string stringToAdd; - if (lastSw != null) - { - // Calculate the time between the last switch (that we iterated - ie. the next one on the timeline) and the current one - var switchDuration = lastSw.Timestamp - sw.Timestamp; - stringToAdd = $"**{membersStr}** ({Formats.ZonedDateTimeFormat.Format(sw.Timestamp.InZone(zone))}, {Formats.DurationFormat.Format(switchSince)} ago, for {Formats.DurationFormat.Format(switchDuration)})\n"; - } - else - { - stringToAdd = $"**{membersStr}** ({Formats.ZonedDateTimeFormat.Format(sw.Timestamp.InZone(zone))}, {Formats.DurationFormat.Format(switchSince)} ago)\n"; - } - - if (outputStr.Length + stringToAdd.Length > EmbedBuilder.MaxDescriptionLength) break; - outputStr += stringToAdd; - - lastSw = sw; - } - - if (lastSw == null) - return null; - - return new EmbedBuilder() - .WithTitle("Past switches") - .WithDescription(outputStr) - .Build(); - } - public async Task CreateMessageInfoEmbed(FullMessage msg) { var channel = await _client.GetChannelAsync(msg.Message.Channel) as ITextChannel; diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs index f9530825..23ec3aaf 100644 --- a/PluralKit.Core/DataFiles.cs +++ b/PluralKit.Core/DataFiles.cs @@ -25,9 +25,10 @@ namespace PluralKit.Bot { // Export members var members = new List(); - var pkMembers = await _data.GetSystemMembers(system); // Read all members in the system + var pkMembers = _data.GetSystemMembers(system); // Read all members in the system var messageCounts = await _data.GetMemberMessageCountBulk(system); // Count messages proxied by all members in the system - members.AddRange(pkMembers.Select(m => new DataFileMember + + await foreach (var member in pkMembers.Select(m => new DataFileMember { Id = m.Hid, Name = m.Name, @@ -41,7 +42,7 @@ namespace PluralKit.Bot KeepProxy = m.KeepProxy, Created = Formats.TimestampExportFormat.Format(m.Created), MessageCount = messageCounts.Where(x => x.Member == m.Id).Select(x => x.MessageCount).FirstOrDefault() - })); + })) members.Add(member); // Export switches var switches = new List(); @@ -96,18 +97,25 @@ namespace PluralKit.Bot await _data.AddAccount(system, accountId); // Determine which members already exist and which ones need to be created - var existingMembers = await _data.GetSystemMembers(system); + var membersByHid = new Dictionary(); + var membersByName = new Dictionary(); + await foreach (var member in _data.GetSystemMembers(system)) + { + membersByHid[member.Hid] = member; + membersByName[member.Name] = member; + } + foreach (var d in data.Members) { - // Try to look up the member with the given ID - var match = existingMembers.FirstOrDefault(m => m.Hid.Equals(d.Id)); - if (match == null) - match = existingMembers.FirstOrDefault(m => m.Name.Equals(d.Name)); // Try with the name instead + PKMember match = null; + if (membersByHid.TryGetValue(d.Id, out var matchByHid)) match = matchByHid; // Try to look up the member with the given ID + else if (membersByName.TryGetValue(d.Id, out var matchByName)) match = matchByName; // Try with the name instead + if (match != null) { dataFileToMemberMapping.Add(d.Id, match); // Relate the data file ID to the PKMember for importing switches result.ModifiedNames.Add(d.Name); - } + } else { unmappedMembers.Add(d); // Track members that weren't found so we can create them all @@ -117,7 +125,7 @@ namespace PluralKit.Bot // If creating the unmatched members would put us over the member limit, abort before creating any members // new total: # in the system + (# in the file - # in the file that already exist) - if (data.Members.Count - dataFileToMemberMapping.Count + existingMembers.Count() > Limits.MaxMemberCount) + if (data.Members.Count - dataFileToMemberMapping.Count + membersByHid.Count > Limits.MaxMemberCount) { result.Success = false; result.Message = $"Import would exceed the maximum number of members ({Limits.MaxMemberCount})."; diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs index 7b5d29ea..bc398083 100644 --- a/PluralKit.Core/Stores.cs +++ b/PluralKit.Core/Stores.cs @@ -203,7 +203,7 @@ namespace PluralKit { /// Gets all members inside a given system. /// /// An enumerable of structs representing each member in the system, in no particular order. - Task> GetSystemMembers(PKSystem system); + IAsyncEnumerable GetSystemMembers(PKSystem system, bool orderByName = false); /// /// Gets the amount of messages proxied by a given member. /// @@ -292,6 +292,12 @@ namespace PluralKit { /// An enumerable of the *count* latest switches in the system, in latest-first order. May contain fewer elements than requested. IAsyncEnumerable GetSwitches(PKSystem system); + /// + /// Gets the total amount of switches in a given system. + /// + Task GetSwitchCount(PKSystem system); + + /// /// Gets the latest (temporally; closest to now) switch of a given system. /// @@ -388,7 +394,7 @@ namespace PluralKit { /// Task SaveGuildConfig(GuildConfig cfg); - Task GetAuxillaryProxyInformation(ulong guild, PKSystem system, PKMember member); + Task GetAuxillaryProxyInformation(ulong guild, PKSystem system, PKMember member); } public class PostgresDataStore: IDataStore { @@ -585,9 +591,11 @@ namespace PluralKit { return await conn.QueryFirstOrDefaultAsync("select * from members where lower(name) = lower(@Name) and system = @SystemID", new { Name = name, SystemID = system.Id }); } - public async Task> GetSystemMembers(PKSystem system) { - using (var conn = await _conn.Obtain()) - return await conn.QueryAsync("select * from members where system = @SystemID", new { SystemID = system.Id }); + public IAsyncEnumerable GetSystemMembers(PKSystem system, bool orderByName) + { + var sql = "select * from members where system = @SystemID"; + if (orderByName) sql += " order by lower(name) asc"; + return _conn.QueryStreamAsync(sql, new { SystemID = system.Id }); } public async Task SaveMember(PKMember member) { @@ -867,24 +875,30 @@ namespace PluralKit { new {System = system.Id}); } - public async Task> GetSwitchMembersList(PKSystem system, Instant start, Instant end) + public async Task GetSwitchCount(PKSystem system) + { + using var conn = await _conn.Obtain(); + return await conn.QuerySingleAsync("select count(*) from switches where system = @Id", system); + } + + public async IAsyncEnumerable GetSwitchMembersList(PKSystem system, Instant start, Instant end) { // Wrap multiple commands in a single transaction for performance - using (var conn = await _conn.Obtain()) - using (var tx = conn.BeginTransaction()) - { - // Find the time of the last switch outside the range as it overlaps the range - // If no prior switch exists, the lower bound of the range remains the start time - var lastSwitch = await conn.QuerySingleOrDefaultAsync( - @"SELECT COALESCE(MAX(timestamp), @Start) + using var conn = await _conn.Obtain(); + using var tx = conn.BeginTransaction(); + + // Find the time of the last switch outside the range as it overlaps the range + // If no prior switch exists, the lower bound of the range remains the start time + var lastSwitch = await conn.QuerySingleOrDefaultAsync( + @"SELECT COALESCE(MAX(timestamp), @Start) FROM switches WHERE switches.system = @System AND switches.timestamp < @Start", - new { System = system.Id, Start = start }); + new { System = system.Id, Start = start }); - // Then collect the time and members of all switches that overlap the range - var switchMembersEntries = await conn.QueryAsync( - @"SELECT switch_members.member, switches.timestamp + // Then collect the time and members of all switches that overlap the range + var switchMembersEntries = conn.QueryStreamAsync( + @"SELECT switch_members.member, switches.timestamp FROM switches LEFT JOIN switch_members ON switches.id = switch_members.switch @@ -895,12 +909,13 @@ namespace PluralKit { ) AND switches.timestamp < @End ORDER BY switches.timestamp DESC", - new { System = system.Id, Start = start, End = end, LastSwitch = lastSwitch }); + new { System = system.Id, Start = start, End = end, LastSwitch = lastSwitch }); - // Commit and return the list - tx.Commit(); - return switchMembersEntries; - } + // Yield each value here + await foreach (var entry in switchMembersEntries) + yield return entry; + + // Don't really need to worry about the transaction here, we're not doing any *writes* } public IAsyncEnumerable GetSwitchMembers(PKSwitch sw) @@ -938,8 +953,10 @@ namespace PluralKit { public async Task> GetPeriodFronters(PKSystem system, Instant periodStart, Instant periodEnd) { + // TODO: IAsyncEnumerable-ify this one + // Returns the timestamps and member IDs of switches overlapping the range, in chronological (newest first) order - var switchMembers = await GetSwitchMembersList(system, periodStart, periodEnd); + var switchMembers = await GetSwitchMembersList(system, periodStart, periodEnd).ToListAsync(); // query DB for all members involved in any of the switches above and collect into a dictionary for future use // this makes sure the return list has the same instances of PKMember throughout, which is important for the dictionary @@ -950,7 +967,7 @@ namespace PluralKit { memberObjects = ( await conn.QueryAsync( "select * from members where id = any(@Switches)", // lol postgres specific `= any()` syntax - new { Switches = switchMembers.Select(m => m.Member).Distinct().ToList() }) + new { Switches = switchMembers.Select(m => m.Member).Distinct() }) ).ToDictionary(m => m.Id); } diff --git a/PluralKit.Core/Utils.cs b/PluralKit.Core/Utils.cs index 38a5b2f9..c448ff0b 100644 --- a/PluralKit.Core/Utils.cs +++ b/PluralKit.Core/Utils.cs @@ -676,7 +676,15 @@ namespace PluralKit { using var conn = await connFactory.Obtain(); - var reader = await conn.ExecuteReaderAsync(sql, param); + await using var reader = (DbDataReader) await conn.ExecuteReaderAsync(sql, param); + var parser = reader.GetRowParser(); + while (reader.Read()) + yield return parser(reader); + } + + public static async IAsyncEnumerable QueryStreamAsync(this IDbConnection conn, string sql, object param) + { + await using var reader = (DbDataReader) await conn.ExecuteReaderAsync(sql, param); var parser = reader.GetRowParser(); while (reader.Read()) yield return parser(reader);