Add a count method to the OrganizationMembersResource

Closes #31388

Signed-off-by: Martin Kanis <mkanis@redhat.com>

integration/admin-client-jee/src/main/java/org/keycloak/admin/client/resource/OrganizationMembersResource.java

@@ -81,4 +81,9 @@ public interface OrganizationMembersResource {
     @Path("invite-existing-user")
     @Consumes(MediaType.APPLICATION_FORM_URLENCODED)
     Response inviteExistingUser(@FormParam("id") String id);
+
+    @Path("count")
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    Long count();
 }

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/idp/InfinispanIDPProvider.java

@@ -16,7 +16,6 @@
  */
 package org.keycloak.models.cache.infinispan.idp;
 
-import org.keycloak.models.cache.infinispan.CachedCount;
 import java.util.Map;
 import java.util.stream.Stream;
 import org.keycloak.models.IDPProvider;
@@ -24,6 +23,7 @@ import org.keycloak.models.IdentityProviderModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.cache.CacheRealmProvider;
+import org.keycloak.models.cache.infinispan.CachedCount;
 import org.keycloak.models.cache.infinispan.RealmCacheSession;
 
 public class InfinispanIDPProvider implements IDPProvider {
@@ -34,7 +34,7 @@ public class InfinispanIDPProvider implements IDPProvider {
     private final KeycloakSession session;
     private final IDPProvider idpDelegate;
     private final RealmCacheSession realmCache;
-    
+
     public InfinispanIDPProvider(KeycloakSession session) {
         this.session = session;
         this.idpDelegate = session.getProvider(IDPProvider.class, "jpa");

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProvider.java

@@ -33,6 +33,7 @@ import org.keycloak.organization.OrganizationProvider;
 public class InfinispanOrganizationProvider implements OrganizationProvider {
 
     private static final String ORG_COUNT_KEY_SUFFIX = ".org.count";
+    private static final String ORG_MEMBERS_COUNT_KEY_SUFFIX = ".members.count";
 
     private final KeycloakSession session;
     private final OrganizationProvider orgDelegate;
@@ -49,6 +50,10 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         return realm.getId() + ORG_COUNT_KEY_SUFFIX;
     }
 
+    public static String cacheKeyOrgMemberCount(RealmModel realm, OrganizationModel organization) {
+        return realm.getId() + ".org." + organization.getId() + ORG_MEMBERS_COUNT_KEY_SUFFIX;
+    }
+
     @Override
     public OrganizationModel create(String name, String alias) {
         registerCountInvalidation();
@@ -154,6 +159,24 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
         return orgDelegate.getMembersStream(organization, search, exact, first, max);
     }
 
+    @Override
+    public long getMembersCount(OrganizationModel organization) {
+        String cacheKey = cacheKeyOrgMemberCount(getRealm(), organization);
+        CachedCount cached = realmCache.getCache().get(cacheKey, CachedCount.class);
+
+        // cached and not invalidated
+        if (cached != null && !isInvalid(cacheKey)) {
+            return cached.getCount();
+        }
+
+        Long loaded = realmCache.getCache().getCurrentRevision(cacheKey);
+        long membersCount = orgDelegate.getMembersCount(organization);
+        cached = new CachedCount(loaded, getRealm(), cacheKey, membersCount);
+        realmCache.getCache().addRevisioned(cached, realmCache.getStartupRevision());
+
+        return membersCount;
+    }
+
     @Override
     public UserModel getMemberById(OrganizationModel organization, String id) {
         RealmModel realm = getRealm();
@@ -319,6 +342,7 @@ public class InfinispanOrganizationProvider implements OrganizationProvider {
     void registerMemberInvalidation(OrganizationModel organization, UserModel member) {
         realmCache.registerInvalidation(cacheKeyByMember(member));
         realmCache.registerInvalidation(cacheKeyMembership(getRealm(), organization, member));
+        realmCache.registerInvalidation(cacheKeyOrgMemberCount(getRealm(), organization));
     }
 
     private boolean isInvalid(String cacheKey) {

model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/organization/InfinispanOrganizationProviderFactory.java

@@ -49,7 +49,7 @@ public class InfinispanOrganizationProviderFactory implements OrganizationProvid
             if (e instanceof RealmModel.IdentityProviderRemovedEvent event) {
                 registerOrganizationInvalidation(event.getKeycloakSession(), event.getRemovedIdentityProvider());
             }
-            if (e instanceof UserModel.UserRemovedEvent event) {
+            if (e instanceof UserModel.UserPreRemovedEvent event) {
                 KeycloakSession session = event.getKeycloakSession();
                 InfinispanOrganizationProvider orgProvider = (InfinispanOrganizationProvider) session.getProvider(OrganizationProvider.class, getId());
                 orgProvider.getByMember(event.getUser()).forEach(organization -> orgProvider.registerMemberInvalidation(organization, event.getUser()));

model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java

@@ -70,7 +70,6 @@ import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -209,7 +208,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
     @Override
     public void addConsent(RealmModel realm, String userId, UserConsentModel consent) {
         String clientId = consent.getClient().getId();
-        
+
         long currentTime = Time.currentTimeMillis();
 
         UserConsentEntity consentEntity = new UserConsentEntity();
@@ -621,11 +620,11 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
         List<Predicate> predicates = new ArrayList<>();
 
         predicates.add(builder.equal(root.get("realmId"), realm.getId()));
-        
+
         for (String stringToSearch : search.trim().split("\\s+")) {
             predicates.add(builder.or(getSearchOptionPredicateArray(stringToSearch, builder, root)));
         }
-        
+
         queryBuilder.where(predicates.toArray(new Predicate[0]));
 
         return em.createQuery(queryBuilder).getSingleResult().intValue();
@@ -648,11 +647,11 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
         List<Predicate> predicates = new ArrayList<>();
 
         predicates.add(builder.equal(userJoin.get("realmId"), realm.getId()));
-        
+
         for (String stringToSearch : search.trim().split("\\s+")) {
             predicates.add(builder.or(getSearchOptionPredicateArray(stringToSearch, builder, userJoin)));
         }
-        
+
         predicates.add(groupMembership.get("groupId").in(groupIds));
 
         queryBuilder.where(predicates.toArray(new Predicate[0]));
@@ -693,7 +692,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
 
         List<Predicate> restrictions = predicates(params, root, Map.of());
         restrictions.add(cb.equal(root.get("realmId"), realm.getId()));
-        
+
         groupsWithPermissionsSubquery(countQuery, groupIds, root, restrictions);
 
         countQuery.where(restrictions.toArray(new Predicate[0]));
@@ -1055,7 +1054,7 @@ public class JpaUserProvider implements UserProvider, UserCredentialStore {
         CriteriaBuilder cb = em.getCriteriaBuilder();
 
         Subquery subquery = query.subquery(String.class);
-        
+
         Root<UserGroupMembershipEntity> from = subquery.from(UserGroupMembershipEntity.class);
 
         subquery.select(cb.literal(1));

model/jpa/src/main/java/org/keycloak/organization/jpa/JpaOrganizationProvider.java

@@ -26,6 +26,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Stream;
 
 import jakarta.persistence.EntityManager;
@@ -271,6 +272,14 @@ public class JpaOrganizationProvider implements OrganizationProvider {
         return userProvider.getGroupMembersStream(getRealm(), group, search, exact, first, max);
     }
 
+    @Override
+    public long getMembersCount(OrganizationModel organization) {
+        throwExceptionIfObjectIsNull(organization, "Organization");
+        String groupId = getOrganizationGroup(organization).getId();
+
+        return userProvider.getUsersCount(getRealm(), Set.of(groupId));
+    }
+
     @Override
     public UserModel getMemberById(OrganizationModel organization, String id) {
         throwExceptionIfObjectIsNull(organization, "Organization");

model/storage-private/src/main/java/org/keycloak/storage/UserStorageManager.java

@@ -360,6 +360,8 @@ public class UserStorageManager extends AbstractStorageManager<UserStorageProvid
             getFederatedStorage().preRemove(realm, user);
         }
 
+        publishUserPreRemovedEvent(realm, user);
+
         StorageId storageId = new StorageId(user.getId());
 
         if (storageId.getProviderId() == null) {
@@ -932,4 +934,23 @@ public class UserStorageManager extends AbstractStorageManager<UserStorageProvid
                 .anyMatch((org) -> (organizationProvider.isEnabled() && org.isManaged(delegate) && !org.isEnabled()) ||
                         (!organizationProvider.isEnabled() && org.isManaged(delegate)));
     }
+
+    private void publishUserPreRemovedEvent(RealmModel realm, UserModel user) {
+        session.getKeycloakSessionFactory().publish(new UserModel.UserPreRemovedEvent() {
+            @Override
+            public RealmModel getRealm() {
+                return realm;
+            }
+
+            @Override
+            public UserModel getUser() {
+                return user;
+            }
+
+            @Override
+            public KeycloakSession getKeycloakSession() {
+                return session;
+            }
+        });
+    }
 }

server-spi/src/main/java/org/keycloak/models/UserModel.java

@@ -56,6 +56,12 @@ public interface UserModel extends RoleMapperModel {
         KeycloakSession getKeycloakSession();
     }
 
+    interface UserPreRemovedEvent extends ProviderEvent {
+        RealmModel getRealm();
+        UserModel getUser();
+        KeycloakSession getKeycloakSession();
+    }
+
     String getId();
 
     // No default method here to allow Abstract subclasses where the username is provided in a different manner

server-spi/src/main/java/org/keycloak/organization/OrganizationProvider.java

@@ -130,6 +130,13 @@ public interface OrganizationProvider extends Provider {
      */
     Stream<UserModel> getMembersStream(OrganizationModel organization, String search, Boolean exact, Integer first, Integer max);
 
+    /**
+     * Returns number of members in the organization.
+     * @param organization the organization
+     * @return Number of members in the organization.
+     */
+    long getMembersCount(OrganizationModel organization);
+
     /**
      * Returns the member of the {@link OrganizationModel} by its {@code id}.
      *

services/src/main/java/org/keycloak/organization/admin/resource/OrganizationMemberResource.java

@@ -202,6 +202,16 @@ public class OrganizationMemberResource {
         });
     }
 
+    @Path("count")
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    @NoCache
+    @Tag(name = KeycloakOpenAPI.Admin.Tags.ORGANIZATIONS)
+    @Operation( summary = "Returns number of members in the organization.")
+    public Long count() {
+        return provider.getMembersCount(organization);
+    }
+
     private UserModel getMember(String id) {
         UserModel member = provider.getMemberById(organization, id);
 

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/cache/OrganizationCacheTest.java

@@ -20,6 +20,7 @@ package org.keycloak.testsuite.organization.cache;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.keycloak.models.cache.infinispan.organization.InfinispanOrganizationProvider.cacheKeyOrgMemberCount;
 
 import java.util.List;
 import java.util.Set;
@@ -33,6 +34,9 @@ import org.keycloak.models.OrganizationDomainModel;
 import org.keycloak.models.OrganizationModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.cache.CacheRealmProvider;
+import org.keycloak.models.cache.infinispan.RealmCacheSession;
+import org.keycloak.models.cache.infinispan.CachedCount;
 import org.keycloak.organization.OrganizationProvider;
 import org.keycloak.representations.idm.UserRepresentation;
 import org.keycloak.testsuite.arquillian.annotation.EnableFeature;
@@ -187,4 +191,92 @@ public class OrganizationCacheTest extends AbstractOrganizationTest {
             assertEquals(0, orgProvider.getByMember(member).count());
         });
     }
+
+    @Test
+    public void testMembersCount() {
+        getTestingClient().server(TEST_REALM_NAME).run((RunOnServer) session -> {
+            OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class);
+            OrganizationModel orgb = orgProvider.getByDomainName("orgb.org");
+            RealmModel realm = session.getContext().getRealm();
+            UserModel member = session.users().addUser(realm, "member");
+            member.setEnabled(true);
+            orgProvider.addMember(orgb, member);
+
+            String cachedKey = cacheKeyOrgMemberCount(realm, orgb);
+            RealmCacheSession realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
+            CachedCount cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+
+            // initially members count is not cached
+            assertNull(cached);
+
+            // members count is cached after first call of getMembersCount()
+            long membersCount = orgProvider.getMembersCount(orgb);
+            assertEquals(1, membersCount);
+            cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+            assertNotNull(cached);
+            assertEquals(1, cached.getCount());
+
+            UserModel user = session.users().addUser(realm, "another-member");
+            user.setEnabled(true);
+            orgProvider.addMember(orgb, user);
+        });
+
+        // addMember invalidates cached members count
+        getTestingClient().server(TEST_REALM_NAME).run((RunOnServer) session -> {
+            OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class);
+            OrganizationModel orgb = orgProvider.getByDomainName("orgb.org");
+            RealmModel realm = session.getContext().getRealm();
+
+            String cachedKey = cacheKeyOrgMemberCount(realm, orgb);
+            RealmCacheSession realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
+            CachedCount cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+
+            assertNull(cached);
+            assertEquals(2, orgProvider.getMembersCount(orgb));
+
+            cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+            assertNotNull(cached);
+            assertEquals(2, cached.getCount());
+
+            orgProvider.removeMember(orgb, session.users().getUserByUsername(realm, "another-member"));
+        });
+
+        // removeMember invalidates cached members count
+        getTestingClient().server(TEST_REALM_NAME).run((RunOnServer) session -> {
+            OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class);
+            OrganizationModel orgb = orgProvider.getByDomainName("orgb.org");
+            RealmModel realm = session.getContext().getRealm();
+
+            String cachedKey = cacheKeyOrgMemberCount(realm, orgb);
+            RealmCacheSession realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
+            CachedCount cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+
+            assertNull(cached);
+            assertEquals(1, orgProvider.getMembersCount(orgb));
+
+            cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+            assertNotNull(cached);
+            assertEquals(1, cached.getCount());
+
+            session.users().removeUser(realm, session.users().getUserByUsername(realm, "member"));
+        });
+
+        // remove user from the realm invalidates cached members count
+        getTestingClient().server(TEST_REALM_NAME).run((RunOnServer) session -> {
+            OrganizationProvider orgProvider = session.getProvider(OrganizationProvider.class);
+            OrganizationModel orgb = orgProvider.getByDomainName("orgb.org");
+            RealmModel realm = session.getContext().getRealm();
+
+            String cachedKey = cacheKeyOrgMemberCount(realm, orgb);
+            RealmCacheSession realmCache = (RealmCacheSession) session.getProvider(CacheRealmProvider.class);
+            CachedCount cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+
+            assertNull(cached);
+            assertEquals(0, orgProvider.getMembersCount(orgb));
+
+            cached = realmCache.getCache().get(cachedKey, CachedCount.class);
+            assertNotNull(cached);
+            assertEquals(0, cached.getCount());
+        });
+    }
 }

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/organization/member/OrganizationMemberTest.java

@@ -509,6 +509,17 @@ public class OrganizationMemberTest extends AbstractOrganizationTest {
         assertThat(orgb.members().getAll().size(), is(0));
     }
 
+    @Test
+    public void testMembersCount() {
+        OrganizationResource organization = testRealm().organizations().get(createOrganization().getId());
+
+        for (int i = 0; i < 10; i++) {
+            addMember(organization, "user" + i +"@neworg.org", "First" + i, "Last" + i);
+        }
+
+        assertEquals(10, (long) organization.members().count());
+    }
+
     private void loginViaNonOrgIdP(String idpAlias) {
         oauth.clientId("broker-app");
         loginPage.open(bc.consumerRealmName());