diff --git a/PluralKit.Core/Database/Functions/functions.sql b/PluralKit.Core/Database/Functions/functions.sql index 8d060ce5..1be8aaa3 100644 --- a/PluralKit.Core/Database/Functions/functions.sql +++ b/PluralKit.Core/Database/Functions/functions.sql @@ -81,4 +81,31 @@ as $$ inner join members on members.system = systems.id left join member_guild on member_guild.member = members.id and member_guild.guild = guild_id where accounts.uid = account_id -$$ language sql stable rows 10; \ No newline at end of file +$$ language sql stable rows 10; + + +create function generate_hid() returns text as $$ + select string_agg(substr('abcdefghijklmnopqrstuvwxyz', ceil(random() * 26)::integer, 1), '') from generate_series(1, 5) +$$ language sql volatile; + + +create function find_free_system_hid() returns text as $$ +declare new_hid text; +begin + loop + new_hid := generate_hid(); + if not exists (select 1 from systems where hid = new_hid) then return new_hid; end if; + end loop; +end +$$ language plpgsql volatile; + + +create function find_free_member_hid() returns text as $$ +declare new_hid text; +begin + loop + new_hid := generate_hid(); + if not exists (select 1 from members where hid = new_hid) then return new_hid; end if; + end loop; +end +$$ language plpgsql volatile; \ No newline at end of file diff --git a/PluralKit.Core/Database/clean.sql b/PluralKit.Core/Database/clean.sql index d1a197a4..e2431f43 100644 --- a/PluralKit.Core/Database/clean.sql +++ b/PluralKit.Core/Database/clean.sql @@ -2,4 +2,7 @@ drop view if exists member_list; drop function if exists message_context; -drop function if exists proxy_members; \ No newline at end of file +drop function if exists proxy_members; +drop function if exists generate_hid; +drop function if exists find_free_system_hid; +drop function if exists find_free_member_hid; \ No newline at end of file diff --git a/PluralKit.Core/Services/PostgresDataStore.cs b/PluralKit.Core/Services/PostgresDataStore.cs index 0c1b976e..29fd660b 100644 --- a/PluralKit.Core/Services/PostgresDataStore.cs +++ b/PluralKit.Core/Services/PostgresDataStore.cs @@ -58,15 +58,9 @@ namespace PluralKit.Core { } public async Task CreateSystem(string systemName = null) { - string hid; - do - { - hid = StringUtils.GenerateHid(); - } while (await GetSystemByHid(hid) != null); - PKSystem system; using (var conn = await _conn.Obtain()) - system = await conn.QuerySingleAsync("insert into systems (hid, name) values (@Hid, @Name) returning *", new { Hid = hid, Name = systemName }); + system = await conn.QuerySingleAsync("insert into systems (hid, name) values (find_free_system_hid(), @Name) returning *", new { Name = systemName }); _logger.Information("Created system {System}", system.Id); // New system has no accounts, therefore nothing gets cached, therefore no need to invalidate caches right here @@ -147,16 +141,9 @@ namespace PluralKit.Core { } public async Task CreateMember(PKSystem system, string name) { - string hid; - do - { - hid = StringUtils.GenerateHid(); - } while (await GetMemberByHid(hid) != null); - PKMember member; using (var conn = await _conn.Obtain()) - member = await conn.QuerySingleAsync("insert into members (hid, system, name) values (@Hid, @SystemId, @Name) returning *", new { - Hid = hid, + member = await conn.QuerySingleAsync("insert into members (hid, system, name) values (find_free_member_hid(), @SystemId, @Name) returning *", new { SystemID = system.Id, Name = name }); diff --git a/PluralKit.Core/Utils/BulkImporter.cs b/PluralKit.Core/Utils/BulkImporter.cs index d5af0c74..99323097 100644 --- a/PluralKit.Core/Utils/BulkImporter.cs +++ b/PluralKit.Core/Utils/BulkImporter.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; using System.Collections.Immutable; +using System.Data; using System.Linq; using System.Threading.Tasks; @@ -69,7 +70,7 @@ namespace PluralKit.Core // if not, roll a new hid and we'll insert one with that // (we can't trust the hid given in the member, it might let us overwrite another system's members) var existingMember = FindExistingMemberInSystem(member.Hid, member.Name); - string newHid = existingMember?.Hid ?? await FindFreeHid(); + string newHid = existingMember?.Hid ?? await _conn.QuerySingleAsync("find_free_member_hid", commandType: CommandType.StoredProcedure); // Upsert member data and return the ID QueryBuilder qb = QueryBuilder.Upsert("members", "hid") @@ -112,19 +113,6 @@ namespace PluralKit.Core return null; } - private async Task FindFreeHid() - { - string hid; - do - { - hid = await _conn.QuerySingleOrDefaultAsync( - "select @Hid where not exists (select id from members where hid = @Hid)", - new {Hid = StringUtils.GenerateHid()}); - } while (hid == null); - - return hid; - } - /// /// Register switches in bulk. /// diff --git a/PluralKit.Core/Utils/StringUtils.cs b/PluralKit.Core/Utils/StringUtils.cs index dcdb70f6..d8f1fa82 100644 --- a/PluralKit.Core/Utils/StringUtils.cs +++ b/PluralKit.Core/Utils/StringUtils.cs @@ -6,18 +6,6 @@ namespace PluralKit.Core { public static class StringUtils { - public static string GenerateHid() - { - var rnd = new Random(); - var charset = "abcdefghijklmnopqrstuvwxyz"; - string hid = ""; - for (int i = 0; i < 5; i++) - { - hid += charset[rnd.Next(charset.Length)]; - } - return hid; - } - public static string GenerateToken() { var buf = new byte[48]; // Results in a 64-byte Base64 string (no padding)