keycloak-aplcache

Added method UserProvider.getUsersCount(). Refactor export-import

7/17/2014 5:15:13 AM

Details

diff --git a/connections/mongo/src/main/java/org/keycloak/connections/mongo/impl/MongoStoreImpl.java b/connections/mongo/src/main/java/org/keycloak/connections/mongo/impl/MongoStoreImpl.java
index af4212b..09104ea 100755
--- a/connections/mongo/src/main/java/org/keycloak/connections/mongo/impl/MongoStoreImpl.java
+++ b/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);
-        return cursor.size();
+        Long count = dbCollection.count(query);
+
+        // For now, assume that int is sufficient
+        return count.intValue();
     }
 
     @Override
diff --git a/export-import/export-import-api/src/main/java/org/keycloak/exportimport/util/MultipleStepsExportProvider.java b/export-import/export-import-api/src/main/java/org/keycloak/exportimport/util/MultipleStepsExportProvider.java
index 392e073..bf07427 100755
--- a/export-import/export-import-api/src/main/java/org/keycloak/exportimport/util/MultipleStepsExportProvider.java
+++ b/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().getUsers(realm).size();
+                    usersHolder.totalCount = session.users().getUsersCount(realm);
                 }
             }
 
@@ -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
-                        List<UserModel> users = session.users().getUsers(realm);
-                        usersHolder.users = users.subList(usersHolder.currentPageStart, usersHolder.currentPageEnd);
+                        usersHolder.users = session.users().getUsers(realm, usersHolder.currentPageStart, usersHolder.currentPageEnd - usersHolder.currentPageStart);
 
                         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");
                     }
 
                 });
diff --git a/model/api/src/main/java/org/keycloak/models/UserProvider.java b/model/api/src/main/java/org/keycloak/models/UserProvider.java
index ffb0842..58436a2 100755
--- a/model/api/src/main/java/org/keycloak/models/UserProvider.java
+++ b/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);
diff --git a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/DefaultCacheUserProvider.java b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/DefaultCacheUserProvider.java
index 35521b4..5817452 100755
--- a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/DefaultCacheUserProvider.java
+++ b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/DefaultCacheUserProvider.java
@@ -185,6 +185,11 @@ public class DefaultCacheUserProvider implements CacheUserProvider {
     }
 
     @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);
     }
diff --git a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/NoCacheUserProvider.java b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/NoCacheUserProvider.java
index 687db36..7468f4d 100755
--- a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/NoCacheUserProvider.java
+++ b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/NoCacheUserProvider.java
@@ -67,6 +67,11 @@ public class NoCacheUserProvider implements CacheUserProvider {
     }
 
     @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);
     }
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
index 5ea1a82..e47394a 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaUserProvider.java
+++ b/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;
@@ -206,6 +207,17 @@ public class JpaUserProvider implements UserProvider {
     }
 
     @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);
         RealmEntity realmEntity = em.getReference(RealmEntity.class, realm.getId());
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoUserProvider.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoUserProvider.java
index de37f54..95914ae 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoUserProvider.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoUserProvider.java
@@ -117,6 +117,14 @@ public class MongoUserProvider implements UserProvider {
     }
 
     @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()
                 .and("realmId").is(realm.getId())