diff --git a/PluralKit.API/Controllers/MemberController.cs b/PluralKit.API/Controllers/MemberController.cs index 3bf30976..f420b7f9 100644 --- a/PluralKit.API/Controllers/MemberController.cs +++ b/PluralKit.API/Controllers/MemberController.cs @@ -38,6 +38,11 @@ namespace PluralKit.API.Controllers if (newMember.Name == null) return BadRequest("Member name cannot be null."); + // Enforce per-system member limit + var memberCount = await _members.MemberCount(system); + if (memberCount >= Limits.MaxMemberCount) + return BadRequest($"Member limit reached ({memberCount} / {Limits.MaxMemberCount})."); + // Explicit bounds checks if (newMember.Name != null && newMember.Name.Length > Limits.MaxMemberNameLength) return BadRequest($"Member name too long ({newMember.Name.Length} > {Limits.MaxMemberNameLength}."); diff --git a/PluralKit.Bot/Commands/ImportExportCommands.cs b/PluralKit.Bot/Commands/ImportExportCommands.cs index 1d138601..d07bb964 100644 --- a/PluralKit.Bot/Commands/ImportExportCommands.cs +++ b/PluralKit.Bot/Commands/ImportExportCommands.cs @@ -101,13 +101,16 @@ namespace PluralKit.Bot.Commands // If passed system is null, it'll create a new one // (and that's okay!) var result = await _dataFiles.ImportSystem(data, ctx.System, ctx.Author.Id); - - if (ctx.System != null) + if (!result.Success) + await ctx.Reply($"{Emojis.Error} The provided system profile could not be imported. {result.Message}"); + else if (ctx.System == null) { + // We didn't have a system prior to importing, so give them the new system's ID await ctx.Reply($"{Emojis.Success} PluralKit has created a system for you based on the given file. Your system ID is `{result.System.Hid}`. Type `pk;system` for more information."); } else { + // We already had a system, so show them what changed await ctx.Reply($"{Emojis.Success} Updated {result.ModifiedNames.Count} members, created {result.AddedNames.Count} members. Type `pk;system list` to check!"); } } diff --git a/PluralKit.Bot/Commands/MemberCommands.cs b/PluralKit.Bot/Commands/MemberCommands.cs index 8961b1f0..dbb80f59 100644 --- a/PluralKit.Bot/Commands/MemberCommands.cs +++ b/PluralKit.Bot/Commands/MemberCommands.cs @@ -45,13 +45,24 @@ namespace PluralKit.Bot.Commands if (!await ctx.PromptYesNo(msg)) throw new PKError("Member creation cancelled."); } + // Enforce per-system member limit + var memberCount = await _members.MemberCount(ctx.System); + if (memberCount >= Limits.MaxMemberCount) + throw Errors.MemberLimitReachedError; + // Create the member var member = await _members.Create(ctx.System, memberName); + memberCount++; // Send confirmation and space hint await ctx.Reply($"{Emojis.Success} Member \"{memberName.SanitizeMentions()}\" (`{member.Hid}`) registered! See the user guide for commands for editing this member: https://pluralkit.me/guide#member-management"); - if (memberName.Contains(" ")) await ctx.Reply($"{Emojis.Note} Note that this member's name contains spaces. You will need to surround it with \"double quotes\" when using commands referring to it, or just use the member's 5-character ID (which is `{member.Hid}`)."); - + if (memberName.Contains(" ")) + await ctx.Reply($"{Emojis.Note} Note that this member's name contains spaces. You will need to surround it with \"double quotes\" when using commands referring to it, or just use the member's 5-character ID (which is `{member.Hid}`)."); + if (memberCount >= Limits.MaxMemberCount) + await ctx.Reply($"{Emojis.Warn} You have reached the per-system member limit ({Limits.MaxMemberCount}). You will be unable to create additional members until existing members are deleted."); + else if (memberCount >= Limits.MaxMembersWarnThreshold) + await ctx.Reply($"{Emojis.Warn} You are approaching the per-system member limit ({memberCount} / {Limits.MaxMemberCount} members). Please review your member list for unused or duplicate members."); + await _proxyCache.InvalidateResultsForSystem(ctx.System); } diff --git a/PluralKit.Bot/Errors.cs b/PluralKit.Bot/Errors.cs index 2d5d79ed..74760569 100644 --- a/PluralKit.Bot/Errors.cs +++ b/PluralKit.Bot/Errors.cs @@ -22,7 +22,8 @@ namespace PluralKit.Bot { public static PKError DescriptionTooLongError(int length) => new PKError($"Description too long ({length}/{Limits.MaxDescriptionLength} characters)."); public static PKError MemberNameTooLongError(int length) => new PKError($"Member name too long ({length}/{Limits.MaxMemberNameLength} characters)."); public static PKError MemberPronounsTooLongError(int length) => new PKError($"Member pronouns too long ({length}/{Limits.MaxMemberNameLength} characters)."); - + public static PKError MemberLimitReachedError => new PKError($"System has reached the maximum number of members ({Limits.MaxMemberCount}). Please delete unused members first in order to create new ones."); + public static PKError InvalidColorError(string color) => new PKError($"\"{color.SanitizeMentions()}\" is not a valid color. Color must be in 6-digit RGB hex format (eg. #ff0000)."); public static PKError BirthdayParseError(string birthday) => new PKError($"\"{birthday.SanitizeMentions()}\" could not be parsed as a valid date. Try a format like \"2016-12-24\" or \"May 3 1996\"."); public static PKError ProxyMustHaveText => new PKSyntaxError("Example proxy message must contain the string 'text'."); diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs index 8c839269..8a9de1b3 100644 --- a/PluralKit.Core/DataFiles.cs +++ b/PluralKit.Core/DataFiles.cs @@ -5,6 +5,7 @@ using System.Threading.Tasks; using Newtonsoft.Json; using NodaTime; using NodaTime.Text; +using PluralKit.Core; using Serilog; namespace PluralKit.Bot @@ -70,38 +71,22 @@ namespace PluralKit.Bot }; } - private async Task ExportMember(PKMember member) => new DataFileMember - { - Id = member.Hid, - Name = member.Name, - DisplayName = member.DisplayName, - Description = member.Description, - Birthday = member.Birthday != null ? Formats.DateExportFormat.Format(member.Birthday.Value) : null, - Pronouns = member.Pronouns, - Color = member.Color, - AvatarUrl = member.AvatarUrl, - Prefix = member.Prefix, - Suffix = member.Suffix, - Created = Formats.TimestampExportFormat.Format(member.Created), - MessageCount = await _members.MessageCount(member) - }; - - private async Task ExportSwitch(PKSwitch sw) => new DataFileSwitch - { - Members = (await _switches.GetSwitchMembers(sw)).Select(m => m.Hid).ToList(), - Timestamp = Formats.TimestampExportFormat.Format(sw.Timestamp) - }; - public async Task ImportSystem(DataFileSystem data, PKSystem system, ulong accountId) { // TODO: make atomic, somehow - we'd need to obtain one IDbConnection and reuse it // which probably means refactoring SystemStore.Save and friends etc - - var result = new ImportResult {AddedNames = new List(), ModifiedNames = new List()}; - var hidMapping = new Dictionary(); + var result = new ImportResult { + AddedNames = new List(), + ModifiedNames = new List(), + Success = true // Assume success unless indicated otherwise + }; + var dataFileToMemberMapping = new Dictionary(); + var unmappedMembers = new List(); // If we don't already have a system to save to, create one - if (system == null) system = await _systems.Create(data.Name); + if (system == null) + system = await _systems.Create(data.Name); + result.System = system; // Apply system info system.Name = data.Name; @@ -110,41 +95,56 @@ namespace PluralKit.Bot if (data.AvatarUrl != null) system.AvatarUrl = data.AvatarUrl; if (data.TimeZone != null) system.UiTz = data.TimeZone ?? "UTC"; await _systems.Save(system); - + // Make sure to link the sender account, too await _systems.Link(system, accountId); - // Apply members - // TODO: parallelize? - foreach (var dataMember in data.Members) + // Determine which members already exist and which ones need to be created + var existingMembers = await _members.GetBySystem(system); + foreach (var d in data.Members) { - // If member's given an ID, we try to look up the member with the given ID - PKMember member = null; - if (dataMember.Id != null) + // 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 + if (match != null) { - member = await _members.GetByHid(dataMember.Id); - - // ...but if it's a different system's member, we just make a new one anyway - if (member != null && member.System != system.Id) member = null; - } - - // Try to look up by name, too - if (member == null) member = await _members.GetByName(system, dataMember.Name); - - // And if all else fails (eg. fresh import from Tupperbox, etc) we just make a member lol - if (member == null) - { - member = await _members.Create(system, dataMember.Name); - result.AddedNames.Add(dataMember.Name); + dataFileToMemberMapping.Add(d.Id, match); // Relate the data file ID to the PKMember for importing switches + result.ModifiedNames.Add(d.Name); } else { - result.ModifiedNames.Add(dataMember.Name); + unmappedMembers.Add(d); // Track members that weren't found so we can create them all + result.AddedNames.Add(d.Name); } + } - // Keep track of what the data file's member ID maps to for switch import - if (!hidMapping.ContainsKey(dataMember.Id)) - hidMapping.Add(dataMember.Id, member); + // 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) + { + result.Success = false; + result.Message = $"Import would exceed the maximum number of members ({Limits.MaxMemberCount})."; + result.AddedNames.Clear(); + result.ModifiedNames.Clear(); + return result; + } + + // Create all unmapped members in one transaction + // These consist of members from another PluralKit system or another framework (e.g. Tupperbox) + var membersToCreate = new Dictionary(); + unmappedMembers.ForEach(x => membersToCreate.Add(x.Id, x.Name)); + var newMembers = await _members.CreateMultiple(system, membersToCreate); + foreach (var member in newMembers) + dataFileToMemberMapping.Add(member.Key, member.Value); + + // Update members with data file properties + // TODO: parallelize? + foreach (var dataMember in data.Members) + { + dataFileToMemberMapping.TryGetValue(dataMember.Id, out PKMember member); + if (member == null) + continue; // Apply member info member.Name = dataMember.Name; @@ -161,7 +161,7 @@ namespace PluralKit.Bot if (dataMember.Birthday != null) { var birthdayParse = Formats.DateExportFormat.Parse(dataMember.Birthday); - member.Birthday = birthdayParse.Success ? (LocalDate?) birthdayParse.Value : null; + member.Birthday = birthdayParse.Success ? (LocalDate?)birthdayParse.Value : null; } await _members.Save(member); @@ -174,16 +174,15 @@ namespace PluralKit.Bot var timestamp = InstantPattern.ExtendedIso.Parse(sw.Timestamp).Value; var swMembers = new List(); swMembers.AddRange(sw.Members.Select(x => - hidMapping.FirstOrDefault(y => y.Key.Equals(x)).Value)); + dataFileToMemberMapping.FirstOrDefault(y => y.Key.Equals(x)).Value)); var mapped = new Tuple>(timestamp, swMembers); mappedSwitches.Add(mapped); } // Import switches - await _switches.RegisterSwitches(system, mappedSwitches); + if (mappedSwitches.Any()) + await _switches.BulkImportSwitches(system, mappedSwitches); - _logger.Information("Imported system {System}", system.Id); - - result.System = system; + _logger.Information("Imported system {System}", system.Hid); return result; } } @@ -193,6 +192,8 @@ namespace PluralKit.Bot public ICollection AddedNames; public ICollection ModifiedNames; public PKSystem System; + public bool Success; + public string Message; } public struct DataFileSystem diff --git a/PluralKit.Core/Limits.cs b/PluralKit.Core/Limits.cs index 8ced1346..5c1e5d8b 100644 --- a/PluralKit.Core/Limits.cs +++ b/PluralKit.Core/Limits.cs @@ -5,6 +5,8 @@ namespace PluralKit.Core { public static readonly int MaxSystemNameLength = 100; public static readonly int MaxSystemTagLength = MaxProxyNameLength - 1; + public static readonly int MaxMemberCount = 1000; + public static readonly int MaxMembersWarnThreshold = MaxMemberCount - 50; public static readonly int MaxDescriptionLength = 1000; public static readonly int MaxMemberNameLength = 100; // Fair bit larger than MaxProxyNameLength for bookkeeping public static readonly int MaxPronounsLength = 100; diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs index 44aeeff0..aaef353b 100644 --- a/PluralKit.Core/Stores.cs +++ b/PluralKit.Core/Stores.cs @@ -5,7 +5,7 @@ using System.Threading.Tasks; using App.Metrics.Logging; using Dapper; using NodaTime; - +using Npgsql; using PluralKit.Core; using Serilog; @@ -128,6 +128,37 @@ namespace PluralKit { return member; } + public async Task> CreateMultiple(PKSystem system, Dictionary names) + { + using (var conn = await _conn.Obtain()) + using (var tx = conn.BeginTransaction()) + { + var results = new Dictionary(); + foreach (var name in names) + { + string hid; + do + { + hid = await conn.QuerySingleOrDefaultAsync("SELECT @Hid WHERE NOT EXISTS (SELECT id FROM members WHERE hid = @Hid LIMIT 1)", new + { + Hid = Utils.GenerateHid() + }); + } while (hid == null); + var member = await conn.QuerySingleAsync("INSERT INTO members (hid, system, name) VALUES (@Hid, @SystemId, @Name) RETURNING *", new + { + Hid = hid, + SystemID = system.Id, + Name = name.Value + }); + results.Add(name.Key, member); + } + + tx.Commit(); + _logger.Information("Created {MemberCount} members for system {SystemID}", names.Count(), system.Hid); + return results; + } + } + public async Task GetByHid(string hid) { using (var conn = await _conn.Obtain()) return await conn.QuerySingleOrDefaultAsync("select * from members where hid = @Hid", new { Hid = hid.ToLower() }); @@ -314,6 +345,77 @@ namespace PluralKit { } } + public async Task BulkImportSwitches(PKSystem system, ICollection>> switches) + { + // Read existing switches to enforce unique timestamps + var priorSwitches = await GetSwitches(system); + var lastSwitchId = priorSwitches.Any() + ? priorSwitches.Max(x => x.Id) + : 0; + + using (var conn = (PerformanceTrackingConnection) await _conn.Obtain()) + { + using (var tx = conn.BeginTransaction()) + { + // Import switches in bulk + using (var importer = conn.BeginBinaryImport("COPY switches (system, timestamp) FROM STDIN (FORMAT BINARY)")) + { + foreach (var sw in switches) + { + // If there's already a switch at this time, move on + if (priorSwitches.Any(x => x.Timestamp.Equals(sw.Item1))) + continue; + + // Otherwise, add it to the importer + importer.StartRow(); + importer.Write(system.Id, NpgsqlTypes.NpgsqlDbType.Integer); + importer.Write(sw.Item1, NpgsqlTypes.NpgsqlDbType.Timestamp); + } + importer.Complete(); // Commits the copy operation so dispose won't roll it back + } + + // Get all switches that were created above and don't have members for ID lookup + var switchesWithoutMembers = + await conn.QueryAsync(@" + SELECT switches.* + FROM switches + LEFT JOIN switch_members + ON switch_members.switch = switches.id + WHERE switches.id > @LastSwitchId + AND switches.system = @System + AND switch_members.id IS NULL", new { LastSwitchId = lastSwitchId, System = system.Id }); + + // Import switch_members in bulk + using (var importer = conn.BeginBinaryImport("COPY switch_members (switch, member) FROM STDIN (FORMAT BINARY)")) + { + // Iterate over the switches we created above and set their members + foreach (var pkSwitch in switchesWithoutMembers) + { + // If this isn't in our import set, move on + var sw = switches.FirstOrDefault(x => x.Item1.Equals(pkSwitch.Timestamp)); + if (sw == null) + continue; + + // Loop through associated members to add each to the switch + foreach (var m in sw.Item2) + { + // Skip switch-outs - these don't have switch_members + if (m == null) + continue; + importer.StartRow(); + importer.Write(pkSwitch.Id, NpgsqlTypes.NpgsqlDbType.Integer); + importer.Write(m.Id, NpgsqlTypes.NpgsqlDbType.Integer); + } + } + importer.Complete(); // Commits the copy operation so dispose won't roll it back + } + tx.Commit(); + } + } + + _logger.Information("Completed bulk import of switches for system {0}", system.Hid); + } + public async Task RegisterSwitches(PKSystem system, ICollection>> switches) { // Use a transaction here since we're doing multiple executed commands in one diff --git a/PluralKit.Core/Utils.cs b/PluralKit.Core/Utils.cs index b71c5b6a..31cd3f1b 100644 --- a/PluralKit.Core/Utils.cs +++ b/PluralKit.Core/Utils.cs @@ -488,6 +488,11 @@ namespace PluralKit _impl.Open(); } + public NpgsqlBinaryImporter BeginBinaryImport(string copyFromCommand) + { + return _impl.BeginBinaryImport(copyFromCommand); + } + public string ConnectionString { get => _impl.ConnectionString;