[PM-32073] - Added Bulk Get Org Ability

src/Core/AdminConsole/AbilitiesCache/ExtendedOrganizationAbilityCacheService.cs

@@ -17,27 +17,30 @@ public class ExtendedOrganizationAbilityCacheService(
     IOrganizationRepository organizationRepository)
     : IOrganizationAbilityCacheService
 {
-
-    public async Task<OrganizationAbility?> GetOrganizationAbilityAsync(Guid orgId)
+    public async Task<OrganizationAbility?> GetOrganizationAbilityAsync(Guid orgId, CancellationToken cancellationToken = default)
     {
-        var cacheKey = BuildCacheKeyForOrganizationAbility(orgId);
         return await cache.GetOrSetAsync<OrganizationAbility?>(
-            cacheKey,
-            async _ => await organizationRepository.GetAbilityAsync(orgId));
+            orgId.ToString(),
+            async (_, _) => await organizationRepository.GetAbilityAsync(orgId),
+            token: cancellationToken);
     }
 
-    public async Task UpsertOrganizationAbilityAsync(Organization organization)
+    public async Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsync(IEnumerable<Guid> orgIds, CancellationToken cancellationToken = default)
     {
-        var cacheKey = BuildCacheKeyForOrganizationAbility(organization.Id);
-        await cache.SetAsync<OrganizationAbility?>(cacheKey, new OrganizationAbility(organization));
+        var tasks = orgIds.Distinct().Select(async orgId => (orgId, ability: await GetOrganizationAbilityAsync(orgId, cancellationToken)));
+        var results = await Task.WhenAll(tasks);
+        return results
+            .Where(r => r.ability != null)
+            .ToDictionary(r => r.orgId, r => r.ability!);
     }
 
-    public async Task DeleteOrganizationAbilityAsync(Guid organizationId)
+    public async Task UpsertOrganizationAbilityAsync(Organization organization, CancellationToken cancellationToken = default)
     {
-        var cacheKey = BuildCacheKeyForOrganizationAbility(organizationId);
-        await cache.RemoveAsync(cacheKey);
+        await cache.SetAsync<OrganizationAbility?>(organization.Id.ToString(), new OrganizationAbility(organization), token: cancellationToken);
     }
 
-    private static string BuildCacheKeyForOrganizationAbility(Guid organizationId)
-        => $"org-ability:{organizationId:N}";
+    public async Task DeleteOrganizationAbilityAsync(Guid organizationId, CancellationToken cancellationToken = default)
+    {
+        await cache.RemoveAsync(organizationId.ToString(), token: cancellationToken);
+    }
 }

src/Core/AdminConsole/AbilitiesCache/IOrganizationAbilityCacheService.cs

@@ -5,7 +5,8 @@ namespace Bit.Core.AdminConsole.AbilitiesCache;
 
 public interface IOrganizationAbilityCacheService
 {
-    Task<OrganizationAbility?> GetOrganizationAbilityAsync(Guid orgId);
-    Task UpsertOrganizationAbilityAsync(Organization organization);
-    Task DeleteOrganizationAbilityAsync(Guid organizationId);
+    Task<OrganizationAbility?> GetOrganizationAbilityAsync(Guid orgId, CancellationToken cancellationToken = default);
+    Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsync(IEnumerable<Guid> orgIds, CancellationToken cancellationToken = default);
+    Task UpsertOrganizationAbilityAsync(Organization organization, CancellationToken cancellationToken = default);
+    Task DeleteOrganizationAbilityAsync(Guid organizationId, CancellationToken cancellationToken = default);
 }

src/Core/AdminConsole/AbilitiesCache/OrganizationAbilityServiceCollectionsExtension.cs

@@ -8,5 +8,5 @@ public static class OrganizationAbilityServiceCollectionsExtension
     public static IServiceCollection AddOrganizationAbilityCache(this IServiceCollection serviceCollection,
         GlobalSettings globalSettings) =>
         serviceCollection.AddExtendedCache(ExtendedOrganizationAbilityCacheConstants.CacheName, globalSettings)
-            .AddScoped<IOrganizationAbilityCacheService, ExtendedOrganizationAbilityCacheService>();
+            .AddSingleton<IOrganizationAbilityCacheService, ExtendedOrganizationAbilityCacheService>();
 }

src/Core/Services/Implementations/FeatureRoutedCacheService.cs

@@ -8,7 +8,7 @@ namespace Bit.Core.Services.Implementations;
 
 public class FeatureRoutedCacheService(
     IVCurrentInMemoryApplicationCacheService inMemoryApplicationCacheService,
-    IOrganizationAbilityCacheService extendedCacheService,
+    IOrganizationAbilityCacheService extendedOrgAbilityCacheService,
     IProviderAbilityCacheService providerAbilityCacheService,
     IFeatureService featureService)
     : IApplicationCacheService
@@ -18,7 +18,7 @@ public Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsyn
 
     public Task<OrganizationAbility?> GetOrganizationAbilityAsync(Guid orgId) =>
         featureService.IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
-            ? extendedCacheService.GetOrganizationAbilityAsync(orgId)
+            ? extendedOrgAbilityCacheService.GetOrganizationAbilityAsync(orgId)
             : inMemoryApplicationCacheService.GetOrganizationAbilityAsync(orgId);
 
     public Task<IDictionary<Guid, ProviderAbility>> GetProviderAbilitiesAsync() =>
@@ -46,6 +46,11 @@ public async Task<IDictionary<Guid, ProviderAbility>> GetProviderAbilitiesAsync(
 
     public async Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbilitiesAsync(IEnumerable<Guid> orgIds)
     {
+        if (featureService.IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache))
+        {
+            return await extendedOrgAbilityCacheService.GetOrganizationAbilitiesAsync(orgIds);
+        }
+
         var allOrganizationAbilities = await inMemoryApplicationCacheService.GetOrganizationAbilitiesAsync();
         return orgIds
             .Distinct()
@@ -55,7 +60,7 @@ public async Task<IDictionary<Guid, OrganizationAbility>> GetOrganizationAbiliti
 
     public Task UpsertOrganizationAbilityAsync(Organization organization) =>
         featureService.IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
-            ? extendedCacheService.UpsertOrganizationAbilityAsync(organization)
+            ? extendedOrgAbilityCacheService.UpsertOrganizationAbilityAsync(organization)
             : inMemoryApplicationCacheService.UpsertOrganizationAbilityAsync(organization);
 
     public Task UpsertProviderAbilityAsync(Provider provider)
@@ -70,7 +75,7 @@ public Task UpsertProviderAbilityAsync(Provider provider)
 
     public Task DeleteOrganizationAbilityAsync(Guid organizationId) =>
         featureService.IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
-            ? extendedCacheService.DeleteOrganizationAbilityAsync(organizationId)
+            ? extendedOrgAbilityCacheService.DeleteOrganizationAbilityAsync(organizationId)
             : inMemoryApplicationCacheService.DeleteOrganizationAbilityAsync(organizationId);
 
     public Task DeleteProviderAbilityAsync(Guid providerId)
@@ -87,7 +92,7 @@ public async Task BaseUpsertOrganizationAbilityAsync(Organization organization)
     {
         if (featureService.IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache))
         {
-            await extendedCacheService.UpsertOrganizationAbilityAsync(organization);
+            await extendedOrgAbilityCacheService.UpsertOrganizationAbilityAsync(organization);
             return;
         }
 
@@ -105,7 +110,7 @@ public async Task BaseDeleteOrganizationAbilityAsync(Guid organizationId)
     {
         if (featureService.IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache))
         {
-            await extendedCacheService.DeleteOrganizationAbilityAsync(organizationId);
+            await extendedOrgAbilityCacheService.DeleteOrganizationAbilityAsync(organizationId);
             return;
         }
 

test/Api.IntegrationTest/AdminConsole/Controllers/OrganizationAbilityCacheTests.cs

@@ -1,8 +1,8 @@
 using System.Net;
-using Bit.Api.AdminConsole.Models.Request.Organizations;
 using Bit.Api.Auth.Models.Request.Accounts;
 using Bit.Api.IntegrationTest.Factories;
 using Bit.Api.IntegrationTest.Helpers;
+using Bit.Api.Models.Request.Organizations;
 using Bit.Core;
 using Bit.Core.AdminConsole.Entities;
 using Bit.Core.Billing.Enums;
@@ -76,57 +76,56 @@ public async Task Put_UpdatesOrganization_CacheReflectsUpdatedValues()
     {
         // Arrange - setup in InitializeAsync()
         await _loginHelper.LoginAsync(_ownerEmail);
-        var updateRequest = new OrganizationUpdateRequestModel
+
+        var cacheService = _factory.GetService<IApplicationCacheService>();
+        var abilityBefore = await cacheService.GetOrganizationAbilityAsync(_organization.Id);
+        Assert.NotNull(abilityBefore);
+        Assert.False(abilityBefore.LimitCollectionCreation);
+
+        var updateRequest = new OrganizationCollectionManagementUpdateRequestModel
         {
-            Name = "Updated Cache Test Org",
-            BillingEmail = "updated-cache@example.com"
+            LimitCollectionCreation = true,
+            LimitCollectionDeletion = false,
+            LimitItemDeletion = false,
+            AllowAdminAccessToAllCollectionItems = true
         };
 
-        // Act - update the organization via the HTTP endpoint
-        var response = await _client.PutAsJsonAsync($"/organizations/{_organization.Id}", updateRequest);
+        // Act - update collection management settings via the HTTP endpoint
+        var response = await _client.PutAsJsonAsync(
+            $"/organizations/{_organization.Id}/collection-management", updateRequest);
 
-        // Assert - endpoint succeeded and cache was updated
+        // Assert - endpoint succeeded and cache reflects the updated value
         Assert.Equal(HttpStatusCode.OK, response.StatusCode);
 
-        var cacheService = _factory.GetService<IApplicationCacheService>();
-        var ability = await cacheService.GetOrganizationAbilityAsync(_organization.Id);
-        Assert.NotNull(ability);
-        Assert.Equal(_organization.Id, ability.Id);
+        var abilityAfter = await cacheService.GetOrganizationAbilityAsync(_organization.Id);
+        Assert.NotNull(abilityAfter);
+        Assert.True(abilityAfter.LimitCollectionCreation);
     }
 
     [Fact]
     public async Task Delete_RemovesOrganization_CacheReturnsNull()
     {
-        // Arrange - create a separate org for deletion so we don't affect other tests
-        var deleteOwnerEmail = $"delete-test-{Guid.NewGuid()}@example.com";
-        await _factory.LoginWithNewAccount(deleteOwnerEmail);
-
-        var signUpResult = await OrganizationTestHelpers.SignUpAsync(
-            _factory,
-            plan: PlanType.EnterpriseAnnually,
-            ownerEmail: deleteOwnerEmail,
-            passwordManagerSeats: 5,
-            paymentMethod: PaymentMethodType.Card);
-        var orgToDelete = signUpResult.Item1;
+        // Arrange - setup in InitializeAsync()
+        await _loginHelper.LoginAsync(_ownerEmail);
 
         // Verify cache is populated before delete
         var cacheService = _factory.GetService<IApplicationCacheService>();
-        var abilityBeforeDelete = await cacheService.GetOrganizationAbilityAsync(orgToDelete.Id);
+        var abilityBeforeDelete = await cacheService.GetOrganizationAbilityAsync(_organization.Id);
         Assert.NotNull(abilityBeforeDelete);
 
         // Act - delete the organization via the HTTP endpoint
-        await _loginHelper.LoginAsync(deleteOwnerEmail);
+        await _loginHelper.LoginAsync(_ownerEmail);
         var deleteRequest = new SecretVerificationRequestModel
         {
             MasterPasswordHash = "master_password_hash"
         };
         var response = await _client.PostAsJsonAsync(
-            $"/organizations/{orgToDelete.Id}/delete", deleteRequest);
+            $"/organizations/{_organization.Id}/delete", deleteRequest);
 
         // Assert - endpoint succeeded and cache was cleared
         Assert.Equal(HttpStatusCode.OK, response.StatusCode);
 
-        var abilityAfterDelete = await cacheService.GetOrganizationAbilityAsync(orgToDelete.Id);
+        var abilityAfterDelete = await cacheService.GetOrganizationAbilityAsync(_organization.Id);
         Assert.Null(abilityAfterDelete);
     }
 }

test/Core.Test/Services/Implementations/FeatureRoutedCacheServiceTests.cs

@@ -315,12 +315,15 @@ public async Task GetProviderAbilitiesAsync_WhenDuplicateIdsProvided_DoesNotThro
     }
 
     [Theory, BitAutoData]
-    public async Task GetOrganizationAbilitiesAsync_ReturnsOnlyMatchingAbilities(
+    public async Task GetOrganizationAbilitiesAsync_WhenFlagOff_ReturnsFromInMemoryService(
         SutProvider<FeatureRoutedCacheService> sutProvider,
         OrganizationAbility matchedAbility,
         OrganizationAbility unmatchedAbility)
     {
         // Arrange
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
+            .Returns(false);
         var allAbilities = new Dictionary<Guid, OrganizationAbility>
         {
             [matchedAbility.Id] = matchedAbility,
@@ -336,14 +339,52 @@ public async Task GetOrganizationAbilitiesAsync_ReturnsOnlyMatchingAbilities(
         // Assert
         Assert.Single(result);
         Assert.Equal(matchedAbility, result[matchedAbility.Id]);
+        await sutProvider.GetDependency<IVCurrentInMemoryApplicationCacheService>()
+            .Received(1)
+            .GetOrganizationAbilitiesAsync();
+        await sutProvider.GetDependency<IOrganizationAbilityCacheService>()
+            .DidNotReceiveWithAnyArgs()
+            .GetOrganizationAbilitiesAsync(default);
+    }
+
+    [Theory, BitAutoData]
+    public async Task GetOrganizationAbilitiesAsync_WhenFlagOn_ReturnsFromExtendedCacheService(
+        SutProvider<FeatureRoutedCacheService> sutProvider,
+        OrganizationAbility ability)
+    {
+        // Arrange
+        var orgIds = new[] { ability.Id };
+        var expectedResult = new Dictionary<Guid, OrganizationAbility> { [ability.Id] = ability };
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
+            .Returns(true);
+        sutProvider.GetDependency<IOrganizationAbilityCacheService>()
+            .GetOrganizationAbilitiesAsync(orgIds)
+            .Returns(expectedResult);
+
+        // Act
+        var result = await sutProvider.Sut.GetOrganizationAbilitiesAsync(orgIds);
+
+        // Assert
+        Assert.Single(result);
+        Assert.Equal(ability, result[ability.Id]);
+        await sutProvider.GetDependency<IOrganizationAbilityCacheService>()
+            .Received(1)
+            .GetOrganizationAbilitiesAsync(orgIds);
+        await sutProvider.GetDependency<IVCurrentInMemoryApplicationCacheService>()
+            .DidNotReceiveWithAnyArgs()
+            .GetOrganizationAbilitiesAsync();
     }
 
     [Theory, BitAutoData]
-    public async Task GetOrganizationAbilitiesAsync_WhenDuplicateIdsProvided_DoesNotThrowAndReturnsSingleEntry(
+    public async Task GetOrganizationAbilitiesAsync_WhenFlagOff_WhenDuplicateIdsProvided_DoesNotThrowAndReturnsSingleEntry(
         SutProvider<FeatureRoutedCacheService> sutProvider,
         OrganizationAbility ability)
     {
         // Arrange
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
+            .Returns(false);
         var allAbilities = new Dictionary<Guid, OrganizationAbility> { [ability.Id] = ability };
         sutProvider.GetDependency<IVCurrentInMemoryApplicationCacheService>()
             .GetOrganizationAbilitiesAsync()
@@ -357,13 +398,15 @@ public async Task GetOrganizationAbilitiesAsync_WhenDuplicateIdsProvided_DoesNot
         Assert.Equal(ability, result[ability.Id]);
     }
 
-
     [Theory, BitAutoData]
-    public async Task GetOrganizationAbilitiesAsync_WhenNoIdsMatched_ReturnsEmptyDictionary(
+    public async Task GetOrganizationAbilitiesAsync_WhenFlagOff_WhenNoIdsMatched_ReturnsEmptyDictionary(
         SutProvider<FeatureRoutedCacheService> sutProvider,
         Guid missingOrgId)
     {
         // Arrange
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(FeatureFlagKeys.OrgAbilityExtendedCache)
+            .Returns(false);
         sutProvider.GetDependency<IVCurrentInMemoryApplicationCacheService>()
             .GetOrganizationAbilitiesAsync()
             .Returns(new Dictionary<Guid, OrganizationAbility>());