[KEYCLOAK-5660] - Adding UserQueryProvider.getUsersCount(realm, inclu…

…deServiceAccount) method
filipeferraz · Nov 30, 2017 · 17748d5 · 17748d5
1 parent 4caa40f
commit 17748d5
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 7 deletions.
diff --git a/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/UserCacheSession.java b/model/infinispan/src/main/java/org/keycloak/models/cache/infinispan/UserCacheSession.java
@@ -597,9 +597,14 @@ public List<UserModel> getUsers(RealmModel realm, boolean includeServiceAccounts
         return getDelegate().getUsers(realm, includeServiceAccounts);
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm, boolean includeServiceAccount) {
+        return getDelegate().getUsersCount(realm, includeServiceAccount);
+    }
+
     @Override
     public int getUsersCount(RealmModel realm) {
-        return getDelegate().getUsersCount(realm);
+        return getUsersCount(realm, false);
     }
 
     @Override

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java
@@ -596,13 +596,24 @@ public List<UserModel> getUsers(RealmModel realm, boolean includeServiceAccounts
     }
 
     @Override
-    public int getUsersCount(RealmModel realm) {
-        Object count = em.createNamedQuery("getRealmUserCount")
+    public int getUsersCount(RealmModel realm, boolean includeServiceAccount) {
+        String namedQuery = "getRealmUserCountExcludeServiceAccount";
+
+        if (includeServiceAccount) {
+            namedQuery = "getRealmUserCount";
+        }
+
+        Object count = em.createNamedQuery(namedQuery)
                 .setParameter("realmId", realm.getId())
                 .getSingleResult();
         return ((Number)count).intValue();
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm) {
+        return getUsersCount(realm, false);
+    }
+
     @Override
     public List<UserModel> getUsers(RealmModel realm) {
         return getUsers(realm, false);

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/entities/UserEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/entities/UserEntity.java
@@ -48,7 +48,8 @@
         @NamedQuery(name="getRealmUserByLastName", query="select u from UserEntity u where u.lastName = :lastName and u.realmId = :realmId"),
         @NamedQuery(name="getRealmUserByFirstLastName", query="select u from UserEntity u where u.firstName = :first and u.lastName = :last and u.realmId = :realmId"),
         @NamedQuery(name="getRealmUserByServiceAccount", query="select u from UserEntity u where u.serviceAccountClientLink = :clientInternalId and u.realmId = :realmId"),
-        @NamedQuery(name="getRealmUserCount", query="select count(u) from UserEntity u where u.realmId = :realmId and (u.serviceAccountClientLink is null)"),
+        @NamedQuery(name="getRealmUserCount", query="select count(u) from UserEntity u where u.realmId = :realmId"),
+        @NamedQuery(name="getRealmUserCountExcludeServiceAccount", query="select count(u) from UserEntity u where u.realmId = :realmId and (u.serviceAccountClientLink is null)"),
         @NamedQuery(name="getRealmUsersByAttributeNameAndValue", query="select u from UserEntity u join u.attributes attr " +
                 "where u.realmId = :realmId and attr.name = :name and attr.value = :value"),
         @NamedQuery(name="deleteUsersByRealm", query="delete from UserEntity u where u.realmId = :realmId"),

diff --git a/server-spi/src/main/java/org/keycloak/storage/user/UserQueryProvider.java b/server-spi/src/main/java/org/keycloak/storage/user/UserQueryProvider.java
@@ -35,8 +35,25 @@
  */
 public interface UserQueryProvider {
 
+    /**
+     * Returns the number of users, without consider any service account.
+     *
+     * @param realm the realm
+     * @return the number of users
+     */
     int getUsersCount(RealmModel realm);
 
+    /**
+     * Returns the number of users.
+     *
+     * @param realm the realm
+     * @param includeServiceAccount if true, the number of users will also include service accounts. Otherwise, only the number of users.
+     * @return the number of users
+     */
+    default int getUsersCount(RealmModel realm, boolean includeServiceAccount) {
+        throw new RuntimeException("Not implemented");
+    }
+
     List<UserModel> getUsers(RealmModel realm);
     List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResults);
 

diff --git a/services/src/main/java/org/keycloak/exportimport/util/MultipleStepsExportProvider.java b/services/src/main/java/org/keycloak/exportimport/util/MultipleStepsExportProvider.java
@@ -81,7 +81,7 @@ protected void runExportImportTask(KeycloakSession session) throws IOException {
 
                 // Count total number of users
                 if (!exportUsersIntoRealmFile) {
-                    usersHolder.totalCount = session.users().getUsersCount(realm);
+                    usersHolder.totalCount = session.users().getUsersCount(realm, true);
                     federatedUsersHolder.totalCount = session.userFederatedStorage().getStoredUsersCount(realm);
                 }
             }

diff --git a/services/src/main/java/org/keycloak/storage/UserStorageManager.java b/services/src/main/java/org/keycloak/storage/UserStorageManager.java
@@ -399,14 +399,19 @@ public List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResult
     }
 
     @Override
-    public int getUsersCount(RealmModel realm) {
-        int size = localStorage().getUsersCount(realm);
+    public int getUsersCount(RealmModel realm, boolean includeServiceAccount) {
+        int size = localStorage().getUsersCount(realm, includeServiceAccount);
         for (UserQueryProvider provider : getStorageProviders(session, realm, UserQueryProvider.class)) {
             size += provider.getUsersCount(realm);
         }
         return size;
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm) {
+        return getUsersCount(realm, false);
+    }
+
     @FunctionalInterface
     interface PaginatedQuery {
         List<UserModel> query(Object provider, int first, int max);

diff --git a/...uite/integration-deprecated/src/test/java/org/keycloak/testsuite/model/UserModelTest.java b/...uite/integration-deprecated/src/test/java/org/keycloak/testsuite/model/UserModelTest.java
@@ -357,6 +357,9 @@ public void testServiceAccountLink() throws Exception {
         Assert.assertTrue(users.contains(user1));
         Assert.assertTrue(users.contains(user2));
 
+        Assert.assertEquals(2, session.users().getUsersCount(realm, true));
+        Assert.assertEquals(1, session.users().getUsersCount(realm, false));
+
         // Remove client
         new ClientManager(realmManager).removeClient(realm, client);
         commit();