From 7406f18bf53c65f168c906b8d0871dda83a788ed Mon Sep 17 00:00:00 2001
From: Noko <nephanim@gmail.com>
Date: Sun, 29 Sep 2019 14:40:13 -0500
Subject: [PATCH] Import switches

- ImportSystem builds a mapping of data file HID to current system HID
- Switches in a data file are reconciled with system members' actual IDs using this mapping
- SwitchStore provides a RegisterSwitches method to register multiple switches
- RegisterSwitches only imports a switch if one does not exist with the same timestamp
- The number of switches created is logged
---
 PluralKit.Core/DataFiles.cs | 25 ++++++++++++++++++++++---
 PluralKit.Core/Stores.cs    | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs
index d5b7a24b..105caad1 100644
--- a/PluralKit.Core/DataFiles.cs
+++ b/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
@@ -73,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);
@@ -116,6 +119,10 @@ 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;
@@ -136,10 +143,22 @@ namespace PluralKit.Bot
 
                 await _members.Save(member);
             }
-            
-            _logger.Information("Imported system {System}", system.Id);
 
-            // TODO: import switches, too?
+            // 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);
+
+            _logger.Information("Imported system {System}", system.Id);
 
             result.System = system;
             return result;
diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs
index 293d6de4..b53e3872 100644
--- a/PluralKit.Core/Stores.cs
+++ b/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