Merge pull request #1 from nephanim/feature/import-switches

Bring performance branch in line with export/import changes

PluralKit.Core/DataFiles.cs

@@ -1,8 +1,10 @@
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
 using Newtonsoft.Json;
 using NodaTime;
+using NodaTime.Text;
 using Serilog;
 
 namespace PluralKit.Bot
@@ -49,6 +51,7 @@ namespace PluralKit.Bot
         {
             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,
@@ -72,6 +75,7 @@ namespace PluralKit.Bot
             // which probably means refactoring SystemStore.Save and friends etc
             
             var result = new ImportResult {AddedNames = new List<string>(), ModifiedNames = new List<string>()};
+            var hidMapping = new Dictionary<string, PKMember>();
 
             // If we don't already have a system to save to, create one
             if (system == null) system = await _systems.Create(data.Name);
@@ -115,8 +119,13 @@ namespace PluralKit.Bot
                     result.ModifiedNames.Add(dataMember.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);
+
                 // Apply member info
                 member.Name = dataMember.Name;
+                if (dataMember.DisplayName != null) member.DisplayName = dataMember.DisplayName;
                 if (dataMember.Description != null) member.Description = dataMember.Description;
                 if (dataMember.Color != null) member.Color = dataMember.Color;
                 if (dataMember.AvatarUrl != null) member.AvatarUrl = dataMember.AvatarUrl;
@@ -135,9 +144,21 @@ namespace PluralKit.Bot
                 await _members.Save(member);
             }
 
-            _logger.Information("Imported system {System}", system.Id);
+            // Re-map the switch members in the likely case IDs have changed
+            var mappedSwitches = new List<Tuple<Instant, ICollection<PKMember>>>();
+            foreach (var sw in data.Switches)
+            {
+                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));
+                var mapped = new Tuple<Instant, ICollection<PKMember>>(timestamp, swMembers);
+                mappedSwitches.Add(mapped);
+            }
+            // Import switches
+            await _switches.RegisterSwitches(system, mappedSwitches);
 
-            // TODO: import switches, too?
+            _logger.Information("Imported system {System}", system.Id);
 
             result.System = system;
             return result;
@@ -173,6 +194,7 @@ namespace PluralKit.Bot
     {
         [JsonProperty("id")] public string Id;
         [JsonProperty("name")] public string Name;
+        [JsonProperty("display_name")] public string DisplayName;
         [JsonProperty("description")] public string Description;
         [JsonProperty("birthday")] public string Birthday;
         [JsonProperty("pronouns")] public string Pronouns;

PluralKit.Core/Stores.cs

@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
@@ -294,6 +295,40 @@ namespace PluralKit {
             }
         }
 
+        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
+            using (var conn = await _conn.Obtain())
+            using (var tx = conn.BeginTransaction())
+            {
+                foreach (var s in switches)
+                {
+                    // First, we insert the switch itself
+                    var sw = await conn.QueryFirstOrDefaultAsync<PKSwitch>(
+                            @"insert into switches(system, timestamp)
+                            select @System, @Timestamp
+                            where not exists (
+                                select * from switches
+                                where system = @System and timestamp::timestamp(0) = @Timestamp
+                                limit 1
+                            )
+                            returning *",
+                        new { System = system.Id, Timestamp = s.Item1 });
+
+                    // If we inserted a switch, also insert each member in the switch in the switch_members table
+                    if (sw != null && s.Item2.Any())
+                        await conn.ExecuteAsync(
+                            "insert into switch_members(switch, member) select @Switch, * FROM unnest(@Members)",
+                            new { Switch = sw.Id, Members = s.Item2.Select(x => x.Id).ToArray() });
+                }
+
+                // Finally we commit the tx, since the using block will otherwise rollback it
+                tx.Commit();
+
+                _logger.Information("Registered {SwitchCount} switches in system {System}", switches.Count, system.Id);
+            }
+        }
+
         public async Task<IEnumerable<PKSwitch>> GetSwitches(PKSystem system, int count = 9999999)
         {
             // TODO: refactor the PKSwitch data structure to somehow include a hydrated member list