keycloak-aplcache

mongo refactor for caching

6/12/2014 5:51:19 PM

Details

diff --git a/model/api/src/main/java/org/keycloak/models/KeycloakSession.java b/model/api/src/main/java/org/keycloak/models/KeycloakSession.java
index 2fbb9bd..d0e4a34 100755
--- a/model/api/src/main/java/org/keycloak/models/KeycloakSession.java
+++ b/model/api/src/main/java/org/keycloak/models/KeycloakSession.java
@@ -40,7 +40,7 @@ public interface KeycloakSession extends Provider {
 
     UsernameLoginFailureModel getUserLoginFailure(String username, RealmModel realm);
     UsernameLoginFailureModel addUserLoginFailure(String username, RealmModel realm);
-    List<UsernameLoginFailureModel> getAllUserLoginFailures();
+    List<UsernameLoginFailureModel> getAllUserLoginFailures(RealmModel realm);
 
     UserSessionModel createUserSession(RealmModel realm, UserModel user, String ipAddress);
     UserSessionModel getUserSession(String id, RealmModel realm);
diff --git a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/CacheKeycloakSession.java b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/CacheKeycloakSession.java
index 7d909ec..9aec5d4 100755
--- a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/CacheKeycloakSession.java
+++ b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/CacheKeycloakSession.java
@@ -240,7 +240,7 @@ public class CacheKeycloakSession implements KeycloakSession {
     }
 
     @Override
-    public List<UsernameLoginFailureModel> getAllUserLoginFailures() {
+    public List<UsernameLoginFailureModel> getAllUserLoginFailures(RealmModel realm) {
         return null;  //To change body of implemented methods use File | Settings | File Templates.
     }
 
diff --git a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/RealmAdapter.java b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/RealmAdapter.java
index 9c1a4b6..0440f6d 100755
--- a/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/RealmAdapter.java
+++ b/model/invalidation-cache/model-adapters/src/main/java/org/keycloak/models/cache/RealmAdapter.java
@@ -894,7 +894,7 @@ public class RealmAdapter implements RealmModel {
     @Override
     public List<UsernameLoginFailureModel> getAllUserLoginFailures() {
         if (updated != null) return updated.getAllUserLoginFailures();
-        return cacheSession.getAllUserLoginFailures();
+        return cacheSession.getAllUserLoginFailures(this);
     }
 
     @Override
diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaKeycloakSession.java b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaKeycloakSession.java
index b7dad7d..b55b602 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/JpaKeycloakSession.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/JpaKeycloakSession.java
@@ -304,7 +304,7 @@ public class JpaKeycloakSession implements KeycloakSession {
     }
 
     @Override
-    public List<UsernameLoginFailureModel> getAllUserLoginFailures() {
+    public List<UsernameLoginFailureModel> getAllUserLoginFailures(RealmModel realm) {
         TypedQuery<UsernameLoginFailureEntity> query = em.createNamedQuery("getAllFailures", UsernameLoginFailureEntity.class);
         List<UsernameLoginFailureEntity> entities = query.getResultList();
         List<UsernameLoginFailureModel> models = new ArrayList<UsernameLoginFailureModel>();
diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java b/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java
index f972463..89ab556 100755
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/RealmAdapter.java
@@ -440,7 +440,7 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public List<UsernameLoginFailureModel> getAllUserLoginFailures() {
-        return session.getAllUserLoginFailures();
+        return session.getAllUserLoginFailures(this);
     }
 
     @Override
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ApplicationAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ApplicationAdapter.java
index 381fc77..9569b3d 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ApplicationAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ApplicationAdapter.java
@@ -4,6 +4,7 @@ import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.ApplicationModel;
 import org.keycloak.models.ClientModel;
+import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
@@ -22,8 +23,8 @@ import java.util.Set;
  */
 public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> implements ApplicationModel {
 
-    public ApplicationAdapter(RealmModel realm, MongoApplicationEntity applicationEntity, MongoStoreInvocationContext invContext) {
-        super(realm, applicationEntity, invContext);
+    public ApplicationAdapter(KeycloakSession session, RealmModel realm, MongoApplicationEntity applicationEntity, MongoStoreInvocationContext invContext) {
+        super(session, realm, applicationEntity, invContext);
     }
 
     @Override
@@ -118,7 +119,7 @@ public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> im
         if (role == null) {
             return null;
         } else {
-            return new RoleAdapter(getRealm(), role, invocationContext);
+            return new RoleAdapter(session, getRealm(), role, invocationContext);
         }
     }
 
@@ -136,7 +137,7 @@ public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> im
 
         getMongoStore().insertEntity(roleEntity, invocationContext);
 
-        return new RoleAdapter(getRealm(), roleEntity, this, invocationContext);
+        return new RoleAdapter(session, getRealm(), roleEntity, this, invocationContext);
     }
 
     @Override
@@ -153,7 +154,7 @@ public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> im
 
         Set<RoleModel> result = new HashSet<RoleModel>();
         for (MongoRoleEntity role : roles) {
-            result.add(new RoleAdapter(getRealm(), role, this, invocationContext));
+            result.add(new RoleAdapter(session, getRealm(), role, this, invocationContext));
         }
 
         return result;
@@ -166,7 +167,7 @@ public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> im
 
         for (MongoRoleEntity role : roles) {
             if (getId().equals(role.getApplicationId())) {
-                result.add(new RoleAdapter(getRealm(), role, this, invocationContext));
+                result.add(new RoleAdapter(session, getRealm(), role, this, invocationContext));
             }
         }
         return result;
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ClientAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ClientAdapter.java
index 96363e8..614204b 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ClientAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/ClientAdapter.java
@@ -8,6 +8,7 @@ import java.util.Set;
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.ClientModel;
+import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserSessionModel;
@@ -25,11 +26,13 @@ public abstract class ClientAdapter<T extends MongoIdentifiableEntity> extends A
 
     protected final T clientEntity;
     private final RealmModel realm;
+    protected  KeycloakSession session;
 
-    public ClientAdapter(RealmModel realm, T clientEntity, MongoStoreInvocationContext invContext) {
+    public ClientAdapter(KeycloakSession session, RealmModel realm, T clientEntity, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.clientEntity = clientEntity;
         this.realm = realm;
+        this.session = session;
     }
 
     @Override
@@ -173,23 +176,12 @@ public abstract class ClientAdapter<T extends MongoIdentifiableEntity> extends A
 
     @Override
     public Set<UserSessionModel> getUserSessions() {
-        DBObject query = new QueryBuilder()
-                .and("associatedClientIds").is(getId())
-                .get();
-        List<MongoUserSessionEntity> sessions = getMongoStore().loadEntities(MongoUserSessionEntity.class, query, invocationContext);
-
-        Set<UserSessionModel> result = new HashSet<UserSessionModel>();
-        for (MongoUserSessionEntity session : sessions) {
-            result.add(new UserSessionAdapter(session, realm, invocationContext));
-        }
-        return result;
-
+        return session.getUserSessions(realm, this);
     }
 
     @Override
     public int getActiveUserSessions() {
-        // todo, something more efficient like COUNT in JPAQL?
-        return getUserSessions().size();
+        return session.getActiveUserSessions(realm, this);
     }
 
     @Override
@@ -199,10 +191,10 @@ public abstract class ClientAdapter<T extends MongoIdentifiableEntity> extends A
 
         for (MongoRoleEntity role : roles) {
             if (realm.getId().equals(role.getRealmId())) {
-                result.add(new RoleAdapter(realm, role, realm, invocationContext));
+                result.add(new RoleAdapter(session, realm, role, realm, invocationContext));
             } else {
                 // Likely applicationRole, but we don't have this application yet
-                result.add(new RoleAdapter(realm, role, invocationContext));
+                result.add(new RoleAdapter(session, realm, role, invocationContext));
             }
         }
         return result;
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSession.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSession.java
index db58dae..018e4e6 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSession.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSession.java
@@ -15,17 +15,28 @@ import org.keycloak.models.SocialLinkModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.UsernameLoginFailureModel;
+import org.keycloak.models.entities.SocialLinkEntity;
 import org.keycloak.models.mongo.api.MongoStore;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.impl.context.TransactionMongoStoreInvocationContext;
+import org.keycloak.models.mongo.keycloak.entities.MongoApplicationEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
 import org.keycloak.models.mongo.keycloak.entities.MongoRealmEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoRoleEntity;
 import org.keycloak.models.mongo.keycloak.entities.MongoUserEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoUserSessionEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoUsernameLoginFailureEntity;
 import org.keycloak.models.utils.KeycloakModelUtils;
+import org.keycloak.util.Time;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.regex.Pattern;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -69,13 +80,13 @@ public class MongoKeycloakSession implements KeycloakSession {
 
         getMongoStore().insertEntity(newRealm, invocationContext);
 
-        return new RealmAdapter(newRealm, invocationContext);
+        return new RealmAdapter(this, newRealm, invocationContext);
     }
 
     @Override
     public RealmModel getRealm(String id) {
         MongoRealmEntity realmEntity = getMongoStore().loadEntity(MongoRealmEntity.class, id, invocationContext);
-        return realmEntity != null ? new RealmAdapter(realmEntity, invocationContext) : null;
+        return realmEntity != null ? new RealmAdapter(this, realmEntity, invocationContext) : null;
     }
 
     @Override
@@ -85,7 +96,7 @@ public class MongoKeycloakSession implements KeycloakSession {
 
         List<RealmModel> results = new ArrayList<RealmModel>();
         for (MongoRealmEntity realmEntity : realms) {
-            results.add(new RealmAdapter(realmEntity, invocationContext));
+            results.add(new RealmAdapter(this, realmEntity, invocationContext));
         }
         return results;
     }
@@ -98,7 +109,7 @@ public class MongoKeycloakSession implements KeycloakSession {
         MongoRealmEntity realm = getMongoStore().loadSingleEntity(MongoRealmEntity.class, query, invocationContext);
 
         if (realm == null) return null;
-        return new RealmAdapter(realm, invocationContext);
+        return new RealmAdapter(this, realm, invocationContext);
     }
 
     @Override
@@ -109,7 +120,7 @@ public class MongoKeycloakSession implements KeycloakSession {
         if (user == null || !realm.getId().equals(user.getRealmId())) {
             return null;
         } else {
-            return new UserAdapter(realm, user, invocationContext);
+            return new UserAdapter(this, realm, user, invocationContext);
         }
     }
 
@@ -124,7 +135,7 @@ public class MongoKeycloakSession implements KeycloakSession {
         if (user == null) {
             return null;
         } else {
-            return new UserAdapter(realm, user, invocationContext);
+            return new UserAdapter(this, realm, user, invocationContext);
         }
     }
 
@@ -139,7 +150,7 @@ public class MongoKeycloakSession implements KeycloakSession {
         if (user == null) {
             return null;
         } else {
-            return new UserAdapter(realm, user, invocationContext);
+            return new UserAdapter(this, realm, user, invocationContext);
         }
     }
 
@@ -154,106 +165,296 @@ public class MongoKeycloakSession implements KeycloakSession {
 
     @Override
     public UserModel getUserBySocialLink(SocialLinkModel socialLink, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        DBObject query = new QueryBuilder()
+                .and("socialLinks.socialProvider").is(socialLink.getSocialProvider())
+                .and("socialLinks.socialUserId").is(socialLink.getSocialUserId())
+                .and("realmId").is(realm.getId())
+                .get();
+        MongoUserEntity userEntity = getMongoStore().loadSingleEntity(MongoUserEntity.class, query, invocationContext);
+        return userEntity == null ? null : new UserAdapter(this, realm, userEntity, invocationContext);
+    }
+
+    protected List<UserModel> convertUserEntities(RealmModel realm, List<MongoUserEntity> userEntities) {
+        List<UserModel> userModels = new ArrayList<UserModel>();
+        for (MongoUserEntity user : userEntities) {
+            userModels.add(new UserAdapter(this, realm, user, invocationContext));
+        }
+        return userModels;
     }
 
+
     @Override
     public List<UserModel> getUsers(RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        DBObject query = new QueryBuilder()
+                .and("realmId").is(realm.getId())
+                .get();
+        List<MongoUserEntity> users = getMongoStore().loadEntities(MongoUserEntity.class, query, invocationContext);
+        return convertUserEntities(realm, users);
     }
 
     @Override
     public List<UserModel> searchForUser(String search, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
-    }
+        search = search.trim();
+        Pattern caseInsensitivePattern = Pattern.compile("(?i:" + search + ")");
+
+        QueryBuilder nameBuilder;
+        int spaceInd = search.lastIndexOf(" ");
+
+        // Case when we have search string like "ohn Bow". Then firstName must end with "ohn" AND lastName must start with "bow" (everything case-insensitive)
+        if (spaceInd != -1) {
+            String firstName = search.substring(0, spaceInd);
+            String lastName = search.substring(spaceInd + 1);
+            Pattern firstNamePattern = Pattern.compile("(?i:" + firstName + "$)");
+            Pattern lastNamePattern = Pattern.compile("(?i:^" + lastName + ")");
+            nameBuilder = new QueryBuilder().and(
+                    new QueryBuilder().put("firstName").regex(firstNamePattern).get(),
+                    new QueryBuilder().put("lastName").regex(lastNamePattern).get()
+            );
+        } else {
+            // Case when we have search without spaces like "foo". The firstName OR lastName could be "foo" (everything case-insensitive)
+            nameBuilder = new QueryBuilder().or(
+                    new QueryBuilder().put("firstName").regex(caseInsensitivePattern).get(),
+                    new QueryBuilder().put("lastName").regex(caseInsensitivePattern).get()
+            );
+        }
+
+        QueryBuilder builder = new QueryBuilder().and(
+                new QueryBuilder().and("realmId").is(realm.getId()).get(),
+                new QueryBuilder().or(
+                        new QueryBuilder().put("loginName").regex(caseInsensitivePattern).get(),
+                        new QueryBuilder().put("email").regex(caseInsensitivePattern).get(),
+                        nameBuilder.get()
+
+                ).get()
+        );
+
+        List<MongoUserEntity> users = getMongoStore().loadEntities(MongoUserEntity.class, builder.get(), invocationContext);
+        return convertUserEntities(realm, users);    }
 
     @Override
     public List<UserModel> searchForUserByAttributes(Map<String, String> attributes, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        QueryBuilder queryBuilder = new QueryBuilder()
+                .and("realmId").is(realm.getId());
+
+        for (Map.Entry<String, String> entry : attributes.entrySet()) {
+            if (entry.getKey().equals(UserModel.LOGIN_NAME)) {
+                queryBuilder.and("loginName").regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
+            } else if (entry.getKey().equalsIgnoreCase(UserModel.FIRST_NAME)) {
+                queryBuilder.and(UserModel.FIRST_NAME).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
+
+            } else if (entry.getKey().equalsIgnoreCase(UserModel.LAST_NAME)) {
+                queryBuilder.and(UserModel.LAST_NAME).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
+
+            } else if (entry.getKey().equalsIgnoreCase(UserModel.EMAIL)) {
+                queryBuilder.and(UserModel.EMAIL).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
+            }
+        }
+        List<MongoUserEntity> users = getMongoStore().loadEntities(MongoUserEntity.class, queryBuilder.get(), invocationContext);
+        return convertUserEntities(realm, users);
     }
 
     @Override
     public Set<SocialLinkModel> getSocialLinks(UserModel user, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        MongoUserEntity userEntity = ((UserAdapter) user).getUser();
+        List<SocialLinkEntity> linkEntities = userEntity.getSocialLinks();
+
+        if (linkEntities == null) {
+            return Collections.EMPTY_SET;
+        }
+
+        Set<SocialLinkModel> result = new HashSet<SocialLinkModel>();
+        for (SocialLinkEntity socialLinkEntity : linkEntities) {
+            SocialLinkModel model = new SocialLinkModel(socialLinkEntity.getSocialProvider(),
+                    socialLinkEntity.getSocialUserId(), socialLinkEntity.getSocialUsername());
+            result.add(model);
+        }
+        return result;
     }
 
+    private SocialLinkEntity findSocialLink(UserModel user, String socialProvider) {
+        MongoUserEntity userEntity = ((UserAdapter) user).getUser();
+        List<SocialLinkEntity> linkEntities = userEntity.getSocialLinks();
+        if (linkEntities == null) {
+            return null;
+        }
+
+        for (SocialLinkEntity socialLinkEntity : linkEntities) {
+            if (socialLinkEntity.getSocialProvider().equals(socialProvider)) {
+                return socialLinkEntity;
+            }
+        }
+        return null;
+    }
+
+
     @Override
     public SocialLinkModel getSocialLink(UserModel user, String socialProvider, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        SocialLinkEntity socialLinkEntity = findSocialLink(user, socialProvider);
+        return socialLinkEntity != null ? new SocialLinkModel(socialLinkEntity.getSocialProvider(), socialLinkEntity.getSocialUserId(), socialLinkEntity.getSocialUsername()) : null;
     }
 
     @Override
     public RoleModel getRoleById(String id, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        MongoRoleEntity role = getMongoStore().loadEntity(MongoRoleEntity.class, id, invocationContext);
+        if (role == null) return null;
+        if (role.getRealmId() != null && !role.getRealmId().equals(realm.getId())) return null;
+        if (role.getApplicationId() != null && realm.getApplicationById(role.getApplicationId()) == null) return null;
+        return new RoleAdapter(this, realm, role, null, invocationContext);
     }
 
     @Override
     public ApplicationModel getApplicationById(String id, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        MongoApplicationEntity appData = getMongoStore().loadEntity(MongoApplicationEntity.class, id, invocationContext);
+
+        // Check if application belongs to this realm
+        if (appData == null || !realm.getId().equals(appData.getRealmId())) {
+            return null;
+        }
+
+        return new ApplicationAdapter(this, realm, appData, invocationContext);
     }
 
     @Override
     public OAuthClientModel getOAuthClientById(String id, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        MongoOAuthClientEntity clientEntity = getMongoStore().loadEntity(MongoOAuthClientEntity.class, id, invocationContext);
+
+        // Check if client belongs to this realm
+        if (clientEntity == null || !realm.getId().equals(clientEntity.getRealmId())) return null;
+
+        return new OAuthClientAdapter(this, realm, clientEntity, invocationContext);
     }
 
     @Override
     public UsernameLoginFailureModel getUserLoginFailure(String username, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        DBObject query = new QueryBuilder()
+                .and("username").is(username)
+                .and("realmId").is(realm.getId())
+                .get();
+        MongoUsernameLoginFailureEntity user = getMongoStore().loadSingleEntity(MongoUsernameLoginFailureEntity.class, query, invocationContext);
+
+        if (user == null) {
+            return null;
+        } else {
+            return new UsernameLoginFailureAdapter(invocationContext, user);
+        }
     }
 
     @Override
     public UsernameLoginFailureModel addUserLoginFailure(String username, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        UsernameLoginFailureModel userLoginFailure = getUserLoginFailure(username, realm);
+        if (userLoginFailure != null) {
+            return userLoginFailure;
+        }
+
+        MongoUsernameLoginFailureEntity userEntity = new MongoUsernameLoginFailureEntity();
+        userEntity.setUsername(username);
+        userEntity.setRealmId(realm.getId());
+
+        getMongoStore().insertEntity(userEntity, invocationContext);
+        return new UsernameLoginFailureAdapter(invocationContext, userEntity);
     }
 
     @Override
-    public List<UsernameLoginFailureModel> getAllUserLoginFailures() {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+    public List<UsernameLoginFailureModel> getAllUserLoginFailures(RealmModel realm) {
+        DBObject query = new QueryBuilder()
+                .and("realmId").is(realm.getId())
+                .get();
+        List<MongoUsernameLoginFailureEntity> failures = getMongoStore().loadEntities(MongoUsernameLoginFailureEntity.class, query, invocationContext);
+
+        List<UsernameLoginFailureModel> result = new ArrayList<UsernameLoginFailureModel>();
+
+        if (failures == null) return result;
+        for (MongoUsernameLoginFailureEntity failure : failures) {
+            result.add(new UsernameLoginFailureAdapter(invocationContext, failure));
+        }
+
+        return result;
     }
 
     @Override
     public UserSessionModel createUserSession(RealmModel realm, UserModel user, String ipAddress) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        MongoUserSessionEntity entity = new MongoUserSessionEntity();
+        entity.setRealmId(realm.getId());
+        entity.setUser(user.getId());
+        entity.setIpAddress(ipAddress);
+
+        int currentTime = Time.currentTime();
+
+        entity.setStarted(currentTime);
+        entity.setLastSessionRefresh(currentTime);
+
+        getMongoStore().insertEntity(entity, invocationContext);
+        return new UserSessionAdapter(entity, realm, invocationContext);
     }
 
     @Override
     public UserSessionModel getUserSession(String id, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        MongoUserSessionEntity entity = getMongoStore().loadEntity(MongoUserSessionEntity.class, id, invocationContext);
+        if (entity == null) {
+            return null;
+        } else {
+            return new UserSessionAdapter(entity, realm, invocationContext);
+        }
     }
 
     @Override
     public List<UserSessionModel> getUserSessions(UserModel user, RealmModel realm) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        DBObject query = new BasicDBObject("user", user.getId());
+        List<UserSessionModel> sessions = new LinkedList<UserSessionModel>();
+        for (MongoUserSessionEntity e : getMongoStore().loadEntities(MongoUserSessionEntity.class, query, invocationContext)) {
+            sessions.add(new UserSessionAdapter(e, realm, invocationContext));
+        }
+        return sessions;
     }
 
     @Override
     public Set<UserSessionModel> getUserSessions(RealmModel realm, ClientModel client) {
-        return null;  //To change body of implemented methods use File | Settings | File Templates.
+        DBObject query = new QueryBuilder()
+                .and("associatedClientIds").is(client.getId())
+                .get();
+        List<MongoUserSessionEntity> sessions = getMongoStore().loadEntities(MongoUserSessionEntity.class, query, invocationContext);
+
+        Set<UserSessionModel> result = new HashSet<UserSessionModel>();
+        for (MongoUserSessionEntity session : sessions) {
+            result.add(new UserSessionAdapter(session, realm, invocationContext));
+        }
+        return result;
     }
 
     @Override
     public int getActiveUserSessions(RealmModel realm, ClientModel client) {
-        return 0;  //To change body of implemented methods use File | Settings | File Templates.
+        return getUserSessions(realm, client).size();
     }
 
     @Override
     public void removeUserSession(UserSessionModel session) {
-        //To change body of implemented methods use File | Settings | File Templates.
+        getMongoStore().removeEntity(((UserSessionAdapter) session).getMongoEntity(), invocationContext);
     }
 
     @Override
     public void removeUserSessions(RealmModel realm, UserModel user) {
-        //To change body of implemented methods use File | Settings | File Templates.
+        DBObject query = new BasicDBObject("user", user.getId());
+        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
     }
 
     @Override
-    public void removeExpiredUserSessions(RealmModel realm) {
-        //To change body of implemented methods use File | Settings | File Templates.
+    public void removeUserSessions(RealmModel realm) {
+        DBObject query = new BasicDBObject("realmId", realm.getId());
+        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
     }
 
     @Override
-    public void removeUserSessions(RealmModel realm) {
-        //To change body of implemented methods use File | Settings | File Templates.
+    public void removeExpiredUserSessions(RealmModel realm) {
+        int currentTime = Time.currentTime();
+        DBObject query = new QueryBuilder()
+                .and("started").lessThan(currentTime - realm.getSsoSessionMaxLifespan())
+                .get();
+
+        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
+        query = new QueryBuilder()
+                .and("lastSessionRefresh").lessThan(currentTime - realm.getSsoSessionIdleTimeout())
+                .get();
+
+        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
     }
 }
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/OAuthClientAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/OAuthClientAdapter.java
index 11d5670..c008aaa 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/OAuthClientAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/OAuthClientAdapter.java
@@ -1,5 +1,6 @@
 package org.keycloak.models.mongo.keycloak.adapters;
 
+import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.OAuthClientModel;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
@@ -10,8 +11,8 @@ import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
  */
 public class OAuthClientAdapter extends ClientAdapter<MongoOAuthClientEntity> implements OAuthClientModel {
 
-    public OAuthClientAdapter(RealmAdapter realm, MongoOAuthClientEntity oauthClientEntity, MongoStoreInvocationContext invContext) {
-        super(realm, oauthClientEntity, invContext);
+    public OAuthClientAdapter(KeycloakSession session, RealmModel realm, MongoOAuthClientEntity oauthClientEntity, MongoStoreInvocationContext invContext) {
+        super(session, realm, oauthClientEntity, invContext);
     }
 
     @Override
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java
index 7861fb6..5c187c2 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RealmAdapter.java
@@ -8,6 +8,7 @@ import org.keycloak.models.ApplicationModel;
 import org.keycloak.models.AuthenticationLinkModel;
 import org.keycloak.models.AuthenticationProviderModel;
 import org.keycloak.models.ClientModel;
+import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.OAuthClientModel;
 import org.keycloak.models.PasswordPolicy;
 import org.keycloak.models.RealmModel;
@@ -64,10 +65,12 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
     protected volatile transient PrivateKey privateKey;
 
     private volatile transient PasswordPolicy passwordPolicy;
+    private volatile transient KeycloakSession session;
 
-    public RealmAdapter(MongoRealmEntity realmEntity, MongoStoreInvocationContext invocationContext) {
+    public RealmAdapter(KeycloakSession session, MongoRealmEntity realmEntity, MongoStoreInvocationContext invocationContext) {
         super(invocationContext);
         this.realm = realmEntity;
+        this.session = session;
     }
 
     @Override
@@ -445,92 +448,34 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
     }
 
     @Override
-    public UserAdapter getUser(String name) {
-        DBObject query = new QueryBuilder()
-                .and("loginName").is(name)
-                .and("realmId").is(getId())
-                .get();
-        MongoUserEntity user = getMongoStore().loadSingleEntity(MongoUserEntity.class, query, invocationContext);
-
-        if (user == null) {
-            return null;
-        } else {
-            return new UserAdapter(this, user, invocationContext);
-        }
+    public UserModel getUser(String name) {
+        return session.getUserByUsername(name, this);
     }
 
     @Override
-    public UsernameLoginFailureAdapter getUserLoginFailure(String name) {
-        DBObject query = new QueryBuilder()
-                .and("username").is(name)
-                .and("realmId").is(getId())
-                .get();
-        MongoUsernameLoginFailureEntity user = getMongoStore().loadSingleEntity(MongoUsernameLoginFailureEntity.class, query, invocationContext);
-
-        if (user == null) {
-            return null;
-        } else {
-            return new UsernameLoginFailureAdapter(invocationContext, user);
-        }
+    public UsernameLoginFailureModel getUserLoginFailure(String name) {
+        return session.getUserLoginFailure(name, this);
     }
 
     @Override
-    public UsernameLoginFailureAdapter addUserLoginFailure(String username) {
-        UsernameLoginFailureAdapter userLoginFailure = getUserLoginFailure(username);
-        if (userLoginFailure != null) {
-            return userLoginFailure;
-        }
-
-        MongoUsernameLoginFailureEntity userEntity = new MongoUsernameLoginFailureEntity();
-        userEntity.setUsername(username);
-        userEntity.setRealmId(getId());
-
-        getMongoStore().insertEntity(userEntity, invocationContext);
-        return new UsernameLoginFailureAdapter(invocationContext, userEntity);
+    public UsernameLoginFailureModel addUserLoginFailure(String username) {
+        return session.addUserLoginFailure(username, this);
     }
 
+
     @Override
     public List<UsernameLoginFailureModel> getAllUserLoginFailures() {
-        DBObject query = new QueryBuilder()
-                .and("realmId").is(getId())
-                .get();
-        List<MongoUsernameLoginFailureEntity> failures = getMongoStore().loadEntities(MongoUsernameLoginFailureEntity.class, query, invocationContext);
-
-        List<UsernameLoginFailureModel> result = new ArrayList<UsernameLoginFailureModel>();
-
-        if (failures == null) return result;
-        for (MongoUsernameLoginFailureEntity failure : failures) {
-            result.add(new UsernameLoginFailureAdapter(invocationContext, failure));
-        }
-
-        return result;
+        return session.getAllUserLoginFailures(this);
     }
 
     @Override
     public UserModel getUserByEmail(String email) {
-        DBObject query = new QueryBuilder()
-                .and("email").is(email)
-                .and("realmId").is(getId())
-                .get();
-        MongoUserEntity user = getMongoStore().loadSingleEntity(MongoUserEntity.class, query, invocationContext);
-
-        if (user == null) {
-            return null;
-        } else {
-            return new UserAdapter(this, user, invocationContext);
-        }
+        return session.getUserByEmail(email, this);
     }
 
     @Override
     public UserModel getUserById(String id) {
-        MongoUserEntity user = getMongoStore().loadEntity(MongoUserEntity.class, id, invocationContext);
-
-        // Check that it's user from this realm
-        if (user == null || !getId().equals(user.getRealmId())) {
-            return null;
-        } else {
-            return new UserAdapter(this, user, invocationContext);
-        }
+        return session.getUserById(id, this);
     }
 
     @Override
@@ -564,7 +509,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         userEntity.setRealmId(getId());
 
         getMongoStore().insertEntity(userEntity, invocationContext);
-        return new UserAdapter(this, userEntity, invocationContext);
+        return new UserAdapter(session, this, userEntity, invocationContext);
     }
 
     @Override
@@ -586,7 +531,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         if (role == null) {
             return null;
         } else {
-            return new RoleAdapter(this, role, this, invocationContext);
+            return new RoleAdapter(session, this, role, this, invocationContext);
         }
     }
 
@@ -604,7 +549,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
         getMongoStore().insertEntity(roleEntity, invocationContext);
 
-        return new RoleAdapter(this, roleEntity, this, invocationContext);
+        return new RoleAdapter(session, this, roleEntity, this, invocationContext);
     }
 
     @Override
@@ -628,7 +573,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
         if (roles == null) return result;
         for (MongoRoleEntity role : roles) {
-            result.add(new RoleAdapter(this, role, this, invocationContext));
+            result.add(new RoleAdapter(session, this, role, this, invocationContext));
         }
 
         return result;
@@ -636,15 +581,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
     @Override
     public RoleModel getRoleById(String id) {
-        MongoRoleEntity role = getMongoStore().loadEntity(MongoRoleEntity.class, id, invocationContext);
-        if (role == null) return null;
-        if (role.getRealmId() != null) {
-            if (!role.getRealmId().equals(this.getId())) return null;
-        } else {
-            ApplicationModel app = getApplicationById(role.getApplicationId());
-            if (app == null) return null;
-        }
-        return new RoleAdapter(this, role, null, invocationContext);
+        return session.getRoleById(id, this);
     }
 
     @Override
@@ -696,14 +633,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
     @Override
     public ApplicationModel getApplicationById(String id) {
-        MongoApplicationEntity appData = getMongoStore().loadEntity(MongoApplicationEntity.class, id, invocationContext);
-
-        // Check if application belongs to this realm
-        if (appData == null || !getId().equals(appData.getRealmId())) {
-            return null;
-        }
-
-        return new ApplicationAdapter(this, appData, invocationContext);
+        return session.getApplicationById(id, this);
     }
 
     @Override
@@ -713,7 +643,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
                 .and("name").is(name)
                 .get();
         MongoApplicationEntity appEntity = getMongoStore().loadSingleEntity(MongoApplicationEntity.class, query, invocationContext);
-        return appEntity == null ? null : new ApplicationAdapter(this, appEntity, invocationContext);
+        return appEntity == null ? null : new ApplicationAdapter(session, this, appEntity, invocationContext);
     }
 
     @Override
@@ -734,7 +664,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
         List<ApplicationModel> result = new ArrayList<ApplicationModel>();
         for (MongoApplicationEntity appData : appDatas) {
-            result.add(new ApplicationAdapter(this, appData, invocationContext));
+            result.add(new ApplicationAdapter(session, this, appData, invocationContext));
         }
         return result;
     }
@@ -753,7 +683,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         appData.setEnabled(true);
         getMongoStore().insertEntity(appData, invocationContext);
 
-        return new ApplicationAdapter(this, appData, invocationContext);
+        return new ApplicationAdapter(session, this, appData, invocationContext);
     }
 
     @Override
@@ -774,7 +704,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         oauthClient.setName(name);
         getMongoStore().insertEntity(oauthClient, invocationContext);
 
-        return new OAuthClientAdapter(this, oauthClient, invocationContext);
+        return new OAuthClientAdapter(session, this, oauthClient, invocationContext);
     }
 
     @Override
@@ -789,17 +719,12 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
                 .and("name").is(name)
                 .get();
         MongoOAuthClientEntity oauthClient = getMongoStore().loadSingleEntity(MongoOAuthClientEntity.class, query, invocationContext);
-        return oauthClient == null ? null : new OAuthClientAdapter(this, oauthClient, invocationContext);
+        return oauthClient == null ? null : new OAuthClientAdapter(session, this, oauthClient, invocationContext);
     }
 
     @Override
     public OAuthClientModel getOAuthClientById(String id) {
-        MongoOAuthClientEntity clientEntity = getMongoStore().loadEntity(MongoOAuthClientEntity.class, id, invocationContext);
-
-        // Check if client belongs to this realm
-        if (clientEntity == null || !getId().equals(clientEntity.getRealmId())) return null;
-
-        return new OAuthClientAdapter(this, clientEntity, invocationContext);
+        return session.getOAuthClientById(id, this);
     }
 
     @Override
@@ -810,7 +735,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
         List<MongoOAuthClientEntity> results = getMongoStore().loadEntities(MongoOAuthClientEntity.class, query, invocationContext);
         List<OAuthClientModel> list = new ArrayList<OAuthClientModel>();
         for (MongoOAuthClientEntity data : results) {
-            list.add(new OAuthClientAdapter(this, data, invocationContext));
+            list.add(new OAuthClientAdapter(session, this, data, invocationContext));
         }
         return list;
     }
@@ -907,36 +832,17 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
     @Override
     public UserModel getUserBySocialLink(SocialLinkModel socialLink) {
-        DBObject query = new QueryBuilder()
-                .and("socialLinks.socialProvider").is(socialLink.getSocialProvider())
-                .and("socialLinks.socialUserId").is(socialLink.getSocialUserId())
-                .and("realmId").is(getId())
-                .get();
-        MongoUserEntity userEntity = getMongoStore().loadSingleEntity(MongoUserEntity.class, query, invocationContext);
-        return userEntity == null ? null : new UserAdapter(this, userEntity, invocationContext);
+        return session.getUserBySocialLink(socialLink, this);
     }
 
     @Override
     public Set<SocialLinkModel> getSocialLinks(UserModel user) {
-        MongoUserEntity userEntity = ((UserAdapter) user).getUser();
-        List<SocialLinkEntity> linkEntities = userEntity.getSocialLinks();
-
-        if (linkEntities == null) {
-            return Collections.EMPTY_SET;
-        }
-
-        Set<SocialLinkModel> result = new HashSet<SocialLinkModel>();
-        for (SocialLinkEntity socialLinkEntity : linkEntities) {
-            SocialLinkModel model = new SocialLinkModel(socialLinkEntity.getSocialProvider(), socialLinkEntity.getSocialUserId(), socialLinkEntity.getSocialUsername());
-            result.add(model);
-        }
-        return result;
+        return session.getSocialLinks(user, this);
     }
 
     @Override
     public SocialLinkModel getSocialLink(UserModel user, String socialProvider) {
-        SocialLinkEntity socialLinkEntity = findSocialLink(user, socialProvider);
-        return socialLinkEntity != null ? new SocialLinkModel(socialLinkEntity.getSocialProvider(), socialLinkEntity.getSocialUserId(), socialLinkEntity.getSocialUsername()) : null;
+        return session.getSocialLink(user, socialProvider, this);
     }
 
     @Override
@@ -990,82 +896,20 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
     @Override
     public List<UserModel> getUsers() {
-        DBObject query = new QueryBuilder()
-                .and("realmId").is(getId())
-                .get();
-        List<MongoUserEntity> users = getMongoStore().loadEntities(MongoUserEntity.class, query, invocationContext);
-        return convertUserEntities(users);
+        return session.getUsers(this);
     }
 
     @Override
     public List<UserModel> searchForUser(String search) {
-        search = search.trim();
-        Pattern caseInsensitivePattern = Pattern.compile("(?i:" + search + ")");
-
-        QueryBuilder nameBuilder;
-        int spaceInd = search.lastIndexOf(" ");
-
-        // Case when we have search string like "ohn Bow". Then firstName must end with "ohn" AND lastName must start with "bow" (everything case-insensitive)
-        if (spaceInd != -1) {
-            String firstName = search.substring(0, spaceInd);
-            String lastName = search.substring(spaceInd + 1);
-            Pattern firstNamePattern = Pattern.compile("(?i:" + firstName + "$)");
-            Pattern lastNamePattern = Pattern.compile("(?i:^" + lastName + ")");
-            nameBuilder = new QueryBuilder().and(
-                    new QueryBuilder().put("firstName").regex(firstNamePattern).get(),
-                    new QueryBuilder().put("lastName").regex(lastNamePattern).get()
-            );
-        } else {
-            // Case when we have search without spaces like "foo". The firstName OR lastName could be "foo" (everything case-insensitive)
-            nameBuilder = new QueryBuilder().or(
-                    new QueryBuilder().put("firstName").regex(caseInsensitivePattern).get(),
-                    new QueryBuilder().put("lastName").regex(caseInsensitivePattern).get()
-            );
-        }
-
-        QueryBuilder builder = new QueryBuilder().and(
-                new QueryBuilder().and("realmId").is(getId()).get(),
-                new QueryBuilder().or(
-                        new QueryBuilder().put("loginName").regex(caseInsensitivePattern).get(),
-                        new QueryBuilder().put("email").regex(caseInsensitivePattern).get(),
-                        nameBuilder.get()
-
-                ).get()
-        );
+        return session.searchForUser(search, this);
 
-        List<MongoUserEntity> users = getMongoStore().loadEntities(MongoUserEntity.class, builder.get(), invocationContext);
-        return convertUserEntities(users);
     }
 
     @Override
     public List<UserModel> searchForUserByAttributes(Map<String, String> attributes) {
-        QueryBuilder queryBuilder = new QueryBuilder()
-                .and("realmId").is(getId());
-
-        for (Map.Entry<String, String> entry : attributes.entrySet()) {
-            if (entry.getKey().equals(UserModel.LOGIN_NAME)) {
-                queryBuilder.and("loginName").regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
-            } else if (entry.getKey().equalsIgnoreCase(UserModel.FIRST_NAME)) {
-                queryBuilder.and(UserModel.FIRST_NAME).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
-
-            } else if (entry.getKey().equalsIgnoreCase(UserModel.LAST_NAME)) {
-                queryBuilder.and(UserModel.LAST_NAME).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
-
-            } else if (entry.getKey().equalsIgnoreCase(UserModel.EMAIL)) {
-                queryBuilder.and(UserModel.EMAIL).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
-            }
-        }
-        List<MongoUserEntity> users = getMongoStore().loadEntities(MongoUserEntity.class, queryBuilder.get(), invocationContext);
-        return convertUserEntities(users);
+        return session.searchForUserByAttributes(attributes, this);
     }
 
-    protected List<UserModel> convertUserEntities(List<MongoUserEntity> userEntities) {
-        List<UserModel> userModels = new ArrayList<UserModel>();
-        for (MongoUserEntity user : userEntities) {
-            userModels.add(new UserAdapter(this, user, invocationContext));
-        }
-        return userModels;
-    }
 
     @Override
     public Map<String, String> getSmtpConfig() {
@@ -1166,7 +1010,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
     @Override
     public ApplicationModel getMasterAdminApp() {
         MongoApplicationEntity appData = getMongoStore().loadEntity(MongoApplicationEntity.class, realm.getAdminAppId(), invocationContext);
-        return appData != null ? new ApplicationAdapter(this, appData, invocationContext) : null;
+        return appData != null ? new ApplicationAdapter(session, this, appData, invocationContext) : null;
     }
 
     @Override
@@ -1182,70 +1026,37 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
 
     @Override
     public UserSessionModel createUserSession(UserModel user, String ipAddress) {
-        MongoUserSessionEntity entity = new MongoUserSessionEntity();
-        entity.setRealmId(getId());
-        entity.setUser(user.getId());
-        entity.setIpAddress(ipAddress);
-
-        int currentTime = Time.currentTime();
-
-        entity.setStarted(currentTime);
-        entity.setLastSessionRefresh(currentTime);
-
-        getMongoStore().insertEntity(entity, invocationContext);
-        return new UserSessionAdapter(entity, this, invocationContext);
+        return session.createUserSession(this, user, ipAddress);
     }
 
     @Override
     public UserSessionModel getUserSession(String id) {
-        MongoUserSessionEntity entity = getMongoStore().loadEntity(MongoUserSessionEntity.class, id, invocationContext);
-        if (entity == null) {
-            return null;
-        } else {
-            return new UserSessionAdapter(entity, this, invocationContext);
-        }
+        return session.getUserSession(id, this);
     }
 
     @Override
     public List<UserSessionModel> getUserSessions(UserModel user) {
-        DBObject query = new BasicDBObject("user", user.getId());
-        List<UserSessionModel> sessions = new LinkedList<UserSessionModel>();
-        for (MongoUserSessionEntity e : getMongoStore().loadEntities(MongoUserSessionEntity.class, query, invocationContext)) {
-            sessions.add(new UserSessionAdapter(e, this, invocationContext));
-        }
-        return sessions;
+        return session.getUserSessions(user, this);
     }
 
     @Override
     public void removeUserSession(UserSessionModel session) {
-        getMongoStore().removeEntity(((UserSessionAdapter) session).getMongoEntity(), invocationContext);
+        this.session.removeUserSession(session);
     }
 
     @Override
     public void removeUserSessions(UserModel user) {
-        DBObject query = new BasicDBObject("user", user.getId());
-        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
+        this.session.removeUserSessions(this, user);
     }
 
     @Override
     public void removeUserSessions() {
-        DBObject query = new BasicDBObject("realmId", getId());
-        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
+        this.session.removeUserSessions(this);
     }
 
     @Override
     public void removeExpiredUserSessions() {
-        int currentTime = Time.currentTime();
-        DBObject query = new QueryBuilder()
-                .and("started").lessThan(currentTime - realm.getSsoSessionMaxLifespan())
-                .get();
-
-        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
-        query = new QueryBuilder()
-                .and("lastSessionRefresh").lessThan(currentTime - realm.getSsoSessionIdleTimeout())
-                .get();
-
-        getMongoStore().removeEntities(MongoUserSessionEntity.class, query, invocationContext);
+        this.session.removeExpiredUserSessions(this);
     }
 
 }
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RoleAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RoleAdapter.java
index 3d8ef1a..6053df3 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RoleAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/RoleAdapter.java
@@ -7,6 +7,7 @@ import java.util.Set;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleContainerModel;
 import org.keycloak.models.RoleModel;
@@ -26,16 +27,18 @@ public class RoleAdapter extends AbstractMongoAdapter<MongoRoleEntity> implement
     private final MongoRoleEntity role;
     private RoleContainerModel roleContainer;
     private RealmModel realm;
+    private KeycloakSession session;
 
-    public RoleAdapter(RealmModel realm, MongoRoleEntity roleEntity, MongoStoreInvocationContext invContext) {
-        this(realm, roleEntity, null, invContext);
+    public RoleAdapter(KeycloakSession session, RealmModel realm, MongoRoleEntity roleEntity, MongoStoreInvocationContext invContext) {
+        this(session, realm, roleEntity, null, invContext);
     }
 
-    public RoleAdapter(RealmModel realm, MongoRoleEntity roleEntity, RoleContainerModel roleContainer, MongoStoreInvocationContext invContext) {
+    public RoleAdapter(KeycloakSession session, RealmModel realm, MongoRoleEntity roleEntity, RoleContainerModel roleContainer, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.role = roleEntity;
         this.roleContainer = roleContainer;
         this.realm = realm;
+        this.session = session;
     }
 
     @Override
@@ -97,7 +100,7 @@ public class RoleAdapter extends AbstractMongoAdapter<MongoRoleEntity> implement
 
         Set<RoleModel> set = new HashSet<RoleModel>();
         for (MongoRoleEntity childRole : childRoles) {
-            set.add(new RoleAdapter(realm, childRole, invocationContext));
+            set.add(new RoleAdapter(session, realm, childRole, invocationContext));
         }
         return set;
     }
@@ -111,13 +114,13 @@ public class RoleAdapter extends AbstractMongoAdapter<MongoRoleEntity> implement
                 if (realm == null) {
                     throw new IllegalStateException("Realm with id: " + role.getRealmId() + " doesn't exists");
                 }
-                roleContainer = new RealmAdapter(realm, invocationContext);
+                roleContainer = new RealmAdapter(session, realm, invocationContext);
             } else if (role.getApplicationId() != null) {
                 MongoApplicationEntity appEntity = getMongoStore().loadEntity(MongoApplicationEntity.class, role.getApplicationId(), invocationContext);
                 if (appEntity == null) {
                     throw new IllegalStateException("Application with id: " + role.getApplicationId() + " doesn't exists");
                 }
-                roleContainer = new ApplicationAdapter(realm, appEntity, invocationContext);
+                roleContainer = new ApplicationAdapter(session, realm, appEntity, invocationContext);
             } else {
                 throw new IllegalStateException("Both realmId and applicationId are null for role: " + this);
             }
diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java
index 0416cc0..b26c50a 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserAdapter.java
@@ -3,6 +3,7 @@ package org.keycloak.models.mongo.keycloak.adapters;
 import org.keycloak.models.ApplicationModel;
 import org.keycloak.models.AuthenticationLinkModel;
 import org.keycloak.models.ClientModel;
+import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserCredentialModel;
@@ -33,11 +34,13 @@ public class UserAdapter extends AbstractMongoAdapter<MongoUserEntity> implement
 
     private final MongoUserEntity user;
     private final RealmModel realm;
+    private final KeycloakSession session;
 
-    public UserAdapter(RealmModel realm, MongoUserEntity userEntity, MongoStoreInvocationContext invContext) {
+    public UserAdapter(KeycloakSession session, RealmModel realm, MongoUserEntity userEntity, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.user = userEntity;
         this.realm = realm;
+        this.session = session;
     }
 
     @Override
@@ -282,10 +285,10 @@ public class UserAdapter extends AbstractMongoAdapter<MongoUserEntity> implement
 
         for (MongoRoleEntity role : roles) {
             if (realm.getId().equals(role.getRealmId())) {
-                result.add(new RoleAdapter(realm, role, realm, invocationContext));
+                result.add(new RoleAdapter(session, realm, role, realm, invocationContext));
             } else {
                 // Likely applicationRole, but we don't have this application yet
-                result.add(new RoleAdapter(realm, role, invocationContext));
+                result.add(new RoleAdapter(session, realm, role, invocationContext));
             }
         }
         return result;
@@ -321,7 +324,7 @@ public class UserAdapter extends AbstractMongoAdapter<MongoUserEntity> implement
 
         for (MongoRoleEntity role : roles) {
             if (app.getId().equals(role.getApplicationId())) {
-                result.add(new RoleAdapter(realm, role, app, invocationContext));
+                result.add(new RoleAdapter(session, realm, role, app, invocationContext));
             }
         }
         return result;