Merge pull request #134 from nephanim/feature/member-limit

Add members per system cap & improve import performance

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}.");

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 (!result.Success)
-                    if (ctx.System != null)
+                        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!");
                     }
                 }

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);
         }
         

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'.");

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<DataFileMember> 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<DataFileSwitch> ExportSwitch(PKSwitch sw) => new DataFileSwitch
-        {
-            Members = (await _switches.GetSwitchMembers(sw)).Select(m => m.Hid).ToList(),
-            Timestamp = Formats.TimestampExportFormat.Format(sw.Timestamp)
-        };
-
         public async Task<ImportResult> 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 {
-            var result = new ImportResult {AddedNames = new List<string>(), ModifiedNames = new List<string>()};
+                AddedNames = new List<string>(),
-            var hidMapping = new Dictionary<string, PKMember>();
+                ModifiedNames = new List<string>(),
+                Success = true // Assume success unless indicated otherwise
+            };
+            var dataFileToMemberMapping = new Dictionary<string, PKMember>();
+            var unmappedMembers = new List<DataFileMember>();
 
             // 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
+            // Determine which members already exist and which ones need to be created
-            // TODO: parallelize?
+            var existingMembers = await _members.GetBySystem(system);
-            foreach (var dataMember in data.Members)
+            foreach (var d in data.Members)
             {
-                // If member's given an ID, we try to look up the member with the given ID
+                // Try to look up the member with the given ID
-                PKMember member = null;
+                var match = existingMembers.FirstOrDefault(m => m.Hid.Equals(d.Id));
-                if (dataMember.Id != null)
+                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);
+                    dataFileToMemberMapping.Add(d.Id, match); // Relate the data file ID to the PKMember for importing switches
-
+                    result.ModifiedNames.Add(d.Name);
-                    // ...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);
                 }
                 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 creating the unmatched members would put us over the member limit, abort before creating any members
-                if (!hidMapping.ContainsKey(dataMember.Id))
+            // new total: # in the system + (# in the file - # in the file that already exist)
-                    hidMapping.Add(dataMember.Id, member);
+            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<string, string>();
+            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<PKMember>();
                 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<Instant, ICollection<PKMember>>(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);
+            _logger.Information("Imported system {System}", system.Hid);
-
-            result.System = system;
             return result;
         }
     }
@@ -193,6 +192,8 @@ namespace PluralKit.Bot
         public ICollection<string> AddedNames;
         public ICollection<string> ModifiedNames;
         public PKSystem System;
+        public bool Success;
+        public string Message;
     }
 
     public struct DataFileSystem

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;

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<Dictionary<string,PKMember>> CreateMultiple(PKSystem system, Dictionary<string,string> names)
+        {
+            using (var conn = await _conn.Obtain())
+            using (var tx = conn.BeginTransaction())
+            {
+                var results = new Dictionary<string, PKMember>();
+                foreach (var name in names)
+                {
+                    string hid;
+                    do
+                    {
+                        hid = await conn.QuerySingleOrDefaultAsync<string>("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<PKMember>("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<PKMember> GetByHid(string hid) {
             using (var conn = await _conn.Obtain())
                 return await conn.QuerySingleOrDefaultAsync<PKMember>("select * from members where hid = @Hid", new { Hid = hid.ToLower() });
@@ -314,6 +345,77 @@ namespace PluralKit {
             }
         }
 
+        public async Task BulkImportSwitches(PKSystem system, ICollection<Tuple<Instant, ICollection<PKMember>>> 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<PKSwitch>(@"
+                        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<Tuple<Instant, ICollection<PKMember>>> switches)
         {
             // Use a transaction here since we're doing multiple executed commands in one

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;