Added method UserProvider.getUsersCount(). Refactor export-import with users pagination available

connections/mongo/src/main/java/org/keycloak/connections/mongo/impl/MongoStoreImpl.java

@@ -304,8 +304,10 @@ public class MongoStoreImpl implements MongoStore {
         context.beforeDBSearch(type);
 
         DBCollection dbCollection = getDBCollectionForType(type);
-        DBCursor cursor = dbCollection.find(query);
+        Long count = dbCollection.count(query);
-        return cursor.size();
+
+        // For now, assume that int is sufficient
+        return count.intValue();
     }
 
     @Override

export-import/export-import-api/src/main/java/org/keycloak/exportimport/util/MultipleStepsExportProvider.java

@@ -56,8 +56,7 @@ public abstract class MultipleStepsExportProvider implements ExportProvider {
 
                 // Count total number of users
                 if (!exportUsersIntoSameFile) {
-                    // TODO: getUsersCount method on model
+                    usersHolder.totalCount = session.users().getUsersCount(realm);
-                    usersHolder.totalCount = session.users().getUsers(realm).size();
                 }
             }
 
@@ -82,13 +81,11 @@ public abstract class MultipleStepsExportProvider implements ExportProvider {
                     @Override
                     public void run(KeycloakSession session) throws IOException {
                         RealmModel realm = session.realms().getRealmByName(realmName);
-                        // TODO: pagination
+                        usersHolder.users = session.users().getUsers(realm, usersHolder.currentPageStart, usersHolder.currentPageEnd - usersHolder.currentPageStart);
-                        List<UserModel> users = session.users().getUsers(realm);
-                        usersHolder.users = users.subList(usersHolder.currentPageStart, usersHolder.currentPageEnd);
 
                         writeUsers(realmName + "-users-" + (usersHolder.currentPageStart / countPerPage) + ".json", session, realm, usersHolder.users);
 
-                        logger.info("Users " + usersHolder.currentPageStart + "-" + usersHolder.currentPageEnd + " exported");
+                        logger.info("Users " + usersHolder.currentPageStart + "-" + (usersHolder.currentPageEnd -1) + " exported");
                     }
 
                 });

model/api/src/main/java/org/keycloak/models/UserProvider.java

@@ -25,6 +25,7 @@ public interface UserProvider extends Provider {
     UserModel getUserByEmail(String email, RealmModel realm);
     UserModel getUserBySocialLink(SocialLinkModel socialLink, RealmModel realm);
     List<UserModel> getUsers(RealmModel realm);
+    int getUsersCount(RealmModel realm);
     List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResults);
     List<UserModel> searchForUser(String search, RealmModel realm);
     List<UserModel> searchForUser(String search, RealmModel realm, int firstResult, int maxResults);

model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/DefaultCacheUserProvider.java

@@ -184,6 +184,11 @@ public class DefaultCacheUserProvider implements CacheUserProvider {
         return getDelegate().getUsers(realm);
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm) {
+        return getDelegate().getUsersCount(realm);
+    }
+
     @Override
     public List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResults) {
         return getDelegate().getUsers(realm, firstResult, maxResults);

model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/NoCacheUserProvider.java

@@ -66,6 +66,11 @@ public class NoCacheUserProvider implements CacheUserProvider {
         return getDelegate().getUsers(realm);
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm) {
+        return getDelegate().getUsersCount(realm);
+    }
+
     @Override
     public List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResults) {
         return getDelegate().getUsers(realm, firstResult, maxResults);

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

@@ -18,6 +18,7 @@ import org.keycloak.models.utils.CredentialValidation;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
 import javax.persistence.EntityManager;
+import javax.persistence.Query;
 import javax.persistence.TypedQuery;
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -205,6 +206,17 @@ public class JpaUserProvider implements UserProvider {
         return getUsers(realm, -1, -1);
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm) {
+        RealmEntity realmEntity = em.getReference(RealmEntity.class, realm.getId());
+
+        // TODO: named query?
+        Object count = em.createQuery("select count(u) from UserEntity u where u.realm = :realm")
+                .setParameter("realm", realmEntity)
+                .getSingleResult();
+        return ((Number)count).intValue();
+    }
+
     @Override
     public List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResults) {
         TypedQuery<UserEntity> query = em.createQuery("select u from UserEntity u where u.realm = :realm order by u.username", UserEntity.class);

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoUserProvider.java

@@ -116,6 +116,14 @@ public class MongoUserProvider implements UserProvider {
         return getUsers(realm, -1, -1);
     }
 
+    @Override
+    public int getUsersCount(RealmModel realm) {
+        DBObject query = new QueryBuilder()
+                .and("realmId").is(realm.getId())
+                .get();
+        return getMongoStore().countEntities(MongoUserEntity.class, query, invocationContext);
+    }
+
     @Override
     public List<UserModel> getUsers(RealmModel realm, int firstResult, int maxResults) {
         DBObject query = new QueryBuilder()