diff --git a/PluralKit.Core/DataFiles.cs b/PluralKit.Core/DataFiles.cs index 346ab224..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 @@ -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(), ModifiedNames = new List()}; + var hidMapping = new Dictionary(); // 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; @@ -134,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>>(); + foreach (var sw in data.Switches) + { + 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)); + var mapped = new Tuple>(timestamp, swMembers); + mappedSwitches.Add(mapped); + } + // Import switches + await _switches.RegisterSwitches(system, mappedSwitches); + + _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; diff --git a/PluralKit.Core/Stores.cs b/PluralKit.Core/Stores.cs index 9625a304..81566f82 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>> 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( + @"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> GetSwitches(PKSystem system, int count = 9999999) { // TODO: refactor the PKSwitch data structure to somehow include a hydrated member list