Add bulk default collection creation method (#6075)

src/Infrastructure.Dapper/AdminConsole/Helpers/BulkResourceCreationService.cs

@@ -0,0 +1,129 @@
+using System.Data;
+using Bit.Core.Entities;
+using Microsoft.Data.SqlClient;
+
+namespace Bit.Infrastructure.Dapper.AdminConsole.Helpers;
+
+public static class BulkResourceCreationService
+{
+    private const string _defaultErrorMessage = "Must have at least one record for bulk creation.";
+    public static async Task CreateCollectionsUsersAsync(SqlConnection connection, SqlTransaction transaction, IEnumerable<CollectionUser> collectionUsers, string errorMessage = _defaultErrorMessage)
+    {
+        using var bulkCopy = new SqlBulkCopy(connection, SqlBulkCopyOptions.KeepIdentity, transaction);
+        bulkCopy.DestinationTableName = "[dbo].[CollectionUser]";
+        var dataTable = BuildCollectionsUsersTable(bulkCopy, collectionUsers, errorMessage);
+        await bulkCopy.WriteToServerAsync(dataTable);
+    }
+
+    private static DataTable BuildCollectionsUsersTable(SqlBulkCopy bulkCopy, IEnumerable<CollectionUser> collectionUsers, string errorMessage)
+    {
+        var collectionUser = collectionUsers.FirstOrDefault();
+
+        if (collectionUser == null)
+        {
+            throw new ApplicationException(errorMessage);
+        }
+
+        var table = new DataTable("CollectionUserDataTable");
+
+        var collectionIdColumn = new DataColumn(nameof(collectionUser.CollectionId), collectionUser.CollectionId.GetType());
+        table.Columns.Add(collectionIdColumn);
+        var orgUserIdColumn = new DataColumn(nameof(collectionUser.OrganizationUserId), collectionUser.OrganizationUserId.GetType());
+        table.Columns.Add(orgUserIdColumn);
+        var readOnlyColumn = new DataColumn(nameof(collectionUser.ReadOnly), collectionUser.ReadOnly.GetType());
+        table.Columns.Add(readOnlyColumn);
+        var hidePasswordsColumn = new DataColumn(nameof(collectionUser.HidePasswords), collectionUser.HidePasswords.GetType());
+        table.Columns.Add(hidePasswordsColumn);
+        var manageColumn = new DataColumn(nameof(collectionUser.Manage), collectionUser.Manage.GetType());
+        table.Columns.Add(manageColumn);
+
+        foreach (DataColumn col in table.Columns)
+        {
+            bulkCopy.ColumnMappings.Add(col.ColumnName, col.ColumnName);
+        }
+
+        var keys = new DataColumn[2];
+        keys[0] = collectionIdColumn;
+        keys[1] = orgUserIdColumn;
+        table.PrimaryKey = keys;
+
+        foreach (var collectionUserRecord in collectionUsers)
+        {
+            var row = table.NewRow();
+
+            row[collectionIdColumn] = collectionUserRecord.CollectionId;
+            row[orgUserIdColumn] = collectionUserRecord.OrganizationUserId;
+            row[readOnlyColumn] = collectionUserRecord.ReadOnly;
+            row[hidePasswordsColumn] = collectionUserRecord.HidePasswords;
+            row[manageColumn] = collectionUserRecord.Manage;
+
+            table.Rows.Add(row);
+        }
+
+        return table;
+    }
+
+    public static async Task CreateCollectionsAsync(SqlConnection connection, SqlTransaction transaction, IEnumerable<Collection> collections, string errorMessage = _defaultErrorMessage)
+    {
+        using var bulkCopy = new SqlBulkCopy(connection, SqlBulkCopyOptions.KeepIdentity, transaction);
+        bulkCopy.DestinationTableName = "[dbo].[Collection]";
+        var dataTable = BuildCollectionsTable(bulkCopy, collections, errorMessage);
+        await bulkCopy.WriteToServerAsync(dataTable);
+    }
+
+    private static DataTable BuildCollectionsTable(SqlBulkCopy bulkCopy, IEnumerable<Collection> collections, string errorMessage)
+    {
+        var collection = collections.FirstOrDefault();
+
+        if (collection == null)
+        {
+            throw new ApplicationException(errorMessage);
+        }
+
+        var collectionsTable = new DataTable("CollectionDataTable");
+
+        var idColumn = new DataColumn(nameof(collection.Id), collection.Id.GetType());
+        collectionsTable.Columns.Add(idColumn);
+        var organizationIdColumn = new DataColumn(nameof(collection.OrganizationId), collection.OrganizationId.GetType());
+        collectionsTable.Columns.Add(organizationIdColumn);
+        var nameColumn = new DataColumn(nameof(collection.Name), collection.Name.GetType());
+        collectionsTable.Columns.Add(nameColumn);
+        var creationDateColumn = new DataColumn(nameof(collection.CreationDate), collection.CreationDate.GetType());
+        collectionsTable.Columns.Add(creationDateColumn);
+        var revisionDateColumn = new DataColumn(nameof(collection.RevisionDate), collection.RevisionDate.GetType());
+        collectionsTable.Columns.Add(revisionDateColumn);
+        var externalIdColumn = new DataColumn(nameof(collection.ExternalId), typeof(string));
+        collectionsTable.Columns.Add(externalIdColumn);
+        var typeColumn = new DataColumn(nameof(collection.Type), collection.Type.GetType());
+        collectionsTable.Columns.Add(typeColumn);
+        var defaultUserCollectionEmailColumn = new DataColumn(nameof(collection.DefaultUserCollectionEmail), typeof(string));
+        collectionsTable.Columns.Add(defaultUserCollectionEmailColumn);
+
+        foreach (DataColumn col in collectionsTable.Columns)
+        {
+            bulkCopy.ColumnMappings.Add(col.ColumnName, col.ColumnName);
+        }
+
+        var keys = new DataColumn[1];
+        keys[0] = idColumn;
+        collectionsTable.PrimaryKey = keys;
+
+        foreach (var collectionRecord in collections)
+        {
+            var row = collectionsTable.NewRow();
+
+            row[idColumn] = collectionRecord.Id;
+            row[organizationIdColumn] = collectionRecord.OrganizationId;
+            row[nameColumn] = collectionRecord.Name;
+            row[creationDateColumn] = collectionRecord.CreationDate;
+            row[revisionDateColumn] = collectionRecord.RevisionDate;
+            row[externalIdColumn] = collectionRecord.ExternalId;
+            row[typeColumn] = collectionRecord.Type;
+            row[defaultUserCollectionEmailColumn] = collectionRecord.DefaultUserCollectionEmail;
+
+            collectionsTable.Rows.Add(row);
+        }
+
+        return collectionsTable;
+    }
+}

src/Infrastructure.Dapper/Repositories/CollectionRepository.cs

@@ -2,9 +2,11 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json;
 using Bit.Core.Entities;
+using Bit.Core.Enums;
 using Bit.Core.Models.Data;
 using Bit.Core.Repositories;
 using Bit.Core.Settings;
+using Bit.Infrastructure.Dapper.AdminConsole.Helpers;
 using Dapper;
 using Microsoft.Data.SqlClient;
 
@@ -222,6 +224,8 @@ public class CollectionRepository : Repository<Collection, Guid>, ICollectionRep
     public async Task CreateAsync(Collection obj, IEnumerable<CollectionAccessSelection>? groups, IEnumerable<CollectionAccessSelection>? users)
     {
         obj.SetNewId();
+
+
         var objWithGroupsAndUsers = JsonSerializer.Deserialize<CollectionWithGroupsAndUsers>(JsonSerializer.Serialize(obj))!;
 
         objWithGroupsAndUsers.Groups = groups != null ? groups.ToArrayTVP() : Enumerable.Empty<CollectionAccessSelection>().ToArrayTVP();
@@ -322,6 +326,100 @@ public class CollectionRepository : Repository<Collection, Guid>, ICollectionRep
         }
     }
 
+    public async Task CreateDefaultCollectionsAsync(Guid organizationId, IEnumerable<Guid> affectedOrgUserIds, string defaultCollectionName)
+    {
+        if (!affectedOrgUserIds.Any())
+        {
+            return;
+        }
+
+        await using var connection = new SqlConnection(ConnectionString);
+        connection.Open();
+        await using var transaction = connection.BeginTransaction();
+        try
+        {
+            var orgUserIdWithDefaultCollection = await GetOrgUserIdsWithDefaultCollectionAsync(connection, transaction, organizationId);
+
+            var missingDefaultCollectionUserIds = affectedOrgUserIds.Except(orgUserIdWithDefaultCollection);
+
+            var (collectionUsers, collections) = BuildDefaultCollectionForUsers(organizationId, missingDefaultCollectionUserIds, defaultCollectionName);
+
+            if (!collectionUsers.Any() || !collections.Any())
+            {
+                return;
+            }
+
+            await BulkResourceCreationService.CreateCollectionsAsync(connection, transaction, collections);
+            await BulkResourceCreationService.CreateCollectionsUsersAsync(connection, transaction, collectionUsers);
+
+            transaction.Commit();
+        }
+        catch
+        {
+            transaction.Rollback();
+            throw;
+        }
+    }
+
+    private async Task<HashSet<Guid>> GetOrgUserIdsWithDefaultCollectionAsync(SqlConnection connection, SqlTransaction transaction, Guid organizationId)
+    {
+        const string sql = @"
+                    SELECT
+                        ou.Id AS OrganizationUserId
+                    FROM
+                        OrganizationUser ou
+                    INNER JOIN
+                        CollectionUser cu ON cu.OrganizationUserId = ou.Id
+                    INNER JOIN
+                        Collection c ON c.Id = cu.CollectionId
+                    WHERE
+                        ou.OrganizationId = @OrganizationId
+                        AND c.Type = @CollectionType;
+                ";
+
+        var organizationUserIds = await connection.QueryAsync<Guid>(
+            sql,
+            new { OrganizationId = organizationId, CollectionType = CollectionType.DefaultUserCollection },
+            transaction: transaction
+        );
+
+        return organizationUserIds.ToHashSet();
+    }
+
+    private (List<CollectionUser> collectionUser, List<Collection> collection) BuildDefaultCollectionForUsers(Guid organizationId, IEnumerable<Guid> missingDefaultCollectionUserIds, string defaultCollectionName)
+    {
+        var collectionUsers = new List<CollectionUser>();
+        var collections = new List<Collection>();
+
+        foreach (var orgUserId in missingDefaultCollectionUserIds)
+        {
+            var collectionId = Guid.NewGuid();
+
+            collections.Add(new Collection
+            {
+                Id = collectionId,
+                OrganizationId = organizationId,
+                Name = defaultCollectionName,
+                CreationDate = DateTime.UtcNow,
+                RevisionDate = DateTime.UtcNow,
+                Type = CollectionType.DefaultUserCollection,
+                DefaultUserCollectionEmail = null
+
+            });
+
+            collectionUsers.Add(new CollectionUser
+            {
+                CollectionId = collectionId,
+                OrganizationUserId = orgUserId,
+                ReadOnly = false,
+                HidePasswords = false,
+                Manage = true,
+            });
+        }
+
+        return (collectionUsers, collections);
+    }
+
     public class CollectionWithGroupsAndUsers : Collection
     {
         [DisallowNull]