keycloak-aplcache

diff --git a/audit/mongo/pom.xml b/audit/mongo/pom.xml
index 69a5a57..c62abf8 100755
--- a/audit/mongo/pom.xml
+++ b/audit/mongo/pom.xml
@@ -38,6 +38,11 @@
             <scope>test</scope>
         </dependency>
         <dependency>
+            <groupId>org.jboss.logging</groupId>
+            <artifactId>jboss-logging</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
             <groupId>org.mongodb</groupId>
             <artifactId>mongo-java-driver</artifactId>
             <scope>provided</scope>
@@ -85,9 +90,6 @@
                                 <keycloak.audit.mongo.db>${keycloak.audit.mongo.db}</keycloak.audit.mongo.db>
                                 <keycloak.audit.mongo.clearOnStartup>${keycloak.audit.mongo.clearOnStartup}</keycloak.audit.mongo.clearOnStartup>
                             </systemPropertyVariables>
-                            <dependenciesToScan>
-                                <dependency>org.keycloak:keycloak-model-tests</dependency>
-                            </dependenciesToScan>
                         </configuration>
                     </execution>
                     <execution>

diff --git a/audit/mongo/src/main/java/org/keycloak/audit/mongo/MongoAuditProviderFactory.java b/audit/mongo/src/main/java/org/keycloak/audit/mongo/MongoAuditProviderFactory.java
index b81afc5..5a3d95a 100644
--- a/audit/mongo/src/main/java/org/keycloak/audit/mongo/MongoAuditProviderFactory.java
+++ b/audit/mongo/src/main/java/org/keycloak/audit/mongo/MongoAuditProviderFactory.java
@@ -5,6 +5,7 @@ import com.mongodb.MongoClient;
 import com.mongodb.MongoCredential;
 import com.mongodb.ServerAddress;
 import com.mongodb.WriteConcern;
+import org.jboss.logging.Logger;
 import org.keycloak.Config;
 import org.keycloak.audit.AuditProvider;
 import org.keycloak.audit.AuditProviderFactory;
@@ -21,6 +22,8 @@ import java.util.Set;
  */
 public class MongoAuditProviderFactory implements AuditProviderFactory {
 
+    protected static final Logger logger = Logger.getLogger(MongoAuditProviderFactory.class);
+
     public static final String ID = "mongo";
     private MongoClient client;
     private DB db;
@@ -55,6 +58,8 @@ public class MongoAuditProviderFactory implements AuditProviderFactory {
             if (clearOnStartup) {
                 db.getCollection("audit").drop();
             }
+
+            logger.infof("Initialized mongo audit. host: %s, port: %d, db: %s, clearOnStartup: %b", host, port, dbName, clearOnStartup);
         } catch (UnknownHostException e) {
             throw new RuntimeException(e);
         }

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/impl/MongoStoreImpl.java b/model/mongo/src/main/java/org/keycloak/models/mongo/impl/MongoStoreImpl.java
index 79b0e95..a9c1d96 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/impl/MongoStoreImpl.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/impl/MongoStoreImpl.java
@@ -237,7 +237,7 @@ public class MongoStoreImpl implements MongoStore {
     public <T extends MongoIdentifiableEntity> T loadEntity(Class<T> type, String id, MongoStoreInvocationContext context) {
         // First look if we already read the object with this oid and type during this transaction. If yes, use it instead of DB lookup
         T cached = context.getLoadedEntity(type, id);
-        if (cached != null) return cached;
+        if (cached != null && type.isAssignableFrom(cached.getClass())) return cached;
 
         DBCollection dbCollection = getDBCollectionForType(type);
 

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 c30db32..b8dbd0a 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
@@ -22,7 +22,7 @@ import java.util.Set;
  */
 public class ApplicationAdapter extends ClientAdapter<MongoApplicationEntity> implements ApplicationModel {
 
-    public ApplicationAdapter(RealmModel realm, MongoApplicationEntity applicationEntity, MongoStoreInvocationContext invContext) {
+    public ApplicationAdapter(RealmAdapter realm, MongoApplicationEntity applicationEntity, MongoStoreInvocationContext invContext) {
         super(realm, applicationEntity, invContext);
     }
 

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 b7404fd..55c2bd9 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
@@ -9,12 +9,11 @@ import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.ClientModel;
 import org.keycloak.models.RealmModel;
-import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
 import org.keycloak.models.entities.ClientEntity;
 import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
-import org.keycloak.models.mongo.keycloak.entities.MongoClientUserSessionAssociationEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoUserSessionEntity;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -22,9 +21,9 @@ import org.keycloak.models.mongo.keycloak.entities.MongoClientUserSessionAssocia
 public class ClientAdapter<T extends MongoIdentifiableEntity> extends AbstractMongoAdapter<T> implements ClientModel {
 
     protected final T clientEntity;
-    private final RealmModel realm;
+    private final RealmAdapter realm;
 
-    public ClientAdapter(RealmModel realm, T clientEntity, MongoStoreInvocationContext invContext) {
+    public ClientAdapter(RealmAdapter realm, T clientEntity, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.clientEntity = clientEntity;
         this.realm = realm;
@@ -154,7 +153,7 @@ public class ClientAdapter<T extends MongoIdentifiableEntity> extends AbstractMo
     }
 
     @Override
-    public RealmModel getRealm() {
+    public RealmAdapter getRealm() {
         return realm;
     }
 
@@ -172,14 +171,13 @@ public class ClientAdapter<T extends MongoIdentifiableEntity> extends AbstractMo
     @Override
     public Set<UserSessionModel> getUserSessions() {
         DBObject query = new QueryBuilder()
-                .and("clientId").is(getId())
+                .and("associatedClientIds").is(getId())
                 .get();
-        List<MongoClientUserSessionAssociationEntity> associations = getMongoStore().loadEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
+        List<MongoUserSessionEntity> sessions = getMongoStore().loadEntities(MongoUserSessionEntity.class, query, invocationContext);
 
         Set<UserSessionModel> result = new HashSet<UserSessionModel>();
-        for (MongoClientUserSessionAssociationEntity association : associations) {
-            UserSessionModel session = realm.getUserSession(association.getSessionId());
-            result.add(session);
+        for (MongoUserSessionEntity session : sessions) {
+            result.add(new UserSessionAdapter(session, realm, invocationContext));
         }
         return result;
 

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSessionFactory.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSessionFactory.java
index 142cdbe..fa39098 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSessionFactory.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/MongoKeycloakSessionFactory.java
@@ -64,7 +64,7 @@ public class MongoKeycloakSessionFactory implements KeycloakSessionFactory {
         try {
             String host = config.get("host", ServerAddress.defaultHost());
             int port = config.getInt("port", ServerAddress.defaultPort());
-            String dbName = config.get("db", "keycloak-audit");
+            String dbName = config.get("db", "keycloak");
             boolean clearOnStartup = config.getBoolean("clearOnStartup", false);
 
             String user = config.get("user");
@@ -79,6 +79,8 @@ public class MongoKeycloakSessionFactory implements KeycloakSessionFactory {
             DB db = client.getDB(dbName);
 
             this.mongoStore = new MongoStoreImpl(db, clearOnStartup, MANAGED_ENTITY_TYPES);
+
+            logger.infof("Initialized mongo model. host: %s, port: %d, db: %s, clearOnStartup: %b", host, port, dbName, clearOnStartup);
         } catch (UnknownHostException e) {
             throw new RuntimeException(e);
         }

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 fbb2b3e..d5d6532 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
@@ -10,7 +10,7 @@ import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
  */
 public class OAuthClientAdapter extends ClientAdapter<MongoOAuthClientEntity> implements OAuthClientModel {
 
-    public OAuthClientAdapter(RealmModel realm, MongoOAuthClientEntity oauthClientEntity, MongoStoreInvocationContext invContext) {
+    public OAuthClientAdapter(RealmAdapter realm, MongoOAuthClientEntity oauthClientEntity, MongoStoreInvocationContext invContext) {
         super(realm, oauthClientEntity, invContext);
     }
 

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 da01022..da20338 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
@@ -299,6 +299,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
     @Override
     public void setSsoSessionIdleTimeout(int seconds) {
         realm.setSsoSessionIdleTimeout(seconds);
+        updateRealm();
     }
 
     @Override
@@ -309,6 +310,7 @@ public class RealmAdapter extends AbstractMongoAdapter<MongoRealmEntity> impleme
     @Override
     public void setSsoSessionMaxLifespan(int seconds) {
         realm.setSsoSessionMaxLifespan(seconds);
+        updateRealm();
     }
 
     @Override

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..41318bb 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,7 +7,6 @@ import java.util.Set;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
-import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleContainerModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
@@ -25,13 +24,13 @@ public class RoleAdapter extends AbstractMongoAdapter<MongoRoleEntity> implement
 
     private final MongoRoleEntity role;
     private RoleContainerModel roleContainer;
-    private RealmModel realm;
+    private RealmAdapter realm;
 
-    public RoleAdapter(RealmModel realm, MongoRoleEntity roleEntity, MongoStoreInvocationContext invContext) {
+    public RoleAdapter(RealmAdapter realm, MongoRoleEntity roleEntity, MongoStoreInvocationContext invContext) {
         this(realm, roleEntity, null, invContext);
     }
 
-    public RoleAdapter(RealmModel realm, MongoRoleEntity roleEntity, RoleContainerModel roleContainer, MongoStoreInvocationContext invContext) {
+    public RoleAdapter(RealmAdapter realm, MongoRoleEntity roleEntity, RoleContainerModel roleContainer, MongoStoreInvocationContext invContext) {
         super(invContext);
         this.role = roleEntity;
         this.roleContainer = roleContainer;

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserSessionAdapter.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserSessionAdapter.java
index 78dfcb8..bc23df2 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserSessionAdapter.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/adapters/UserSessionAdapter.java
@@ -2,27 +2,27 @@ package org.keycloak.models.mongo.keycloak.adapters;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.jboss.logging.Logger;
+import org.keycloak.models.ApplicationModel;
 import org.keycloak.models.ClientModel;
-import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserSessionModel;
+import org.keycloak.models.entities.ClientEntity;
 import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
-import org.keycloak.models.mongo.keycloak.entities.MongoClientUserSessionAssociationEntity;
-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.MongoApplicationEntity;
+import org.keycloak.models.mongo.keycloak.entities.MongoOAuthClientEntity;
 import org.keycloak.models.mongo.keycloak.entities.MongoUserSessionEntity;
 
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 
 /**
  * @author <a href="mailto:sthorger@redhat.com">Stian Thorgersen</a>
  */
 public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEntity> implements UserSessionModel {
 
+    private static final Logger logger = Logger.getLogger(RealmAdapter.class);
+
     private MongoUserSessionEntity entity;
     private RealmAdapter realm;
 
@@ -46,6 +46,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setId(String id) {
         entity.setId(id);
+        updateMongoEntity();
     }
 
     @Override
@@ -56,6 +57,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setUser(UserModel user) {
         entity.setUser(user.getId());
+        updateMongoEntity();
     }
 
     @Override
@@ -66,6 +68,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setIpAddress(String ipAddress) {
         entity.setIpAddress(ipAddress);
+        updateMongoEntity();
     }
 
     @Override
@@ -76,6 +79,7 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setStarted(int started) {
         entity.setStarted(started);
+        updateMongoEntity();
     }
 
     @Override
@@ -86,58 +90,39 @@ public class UserSessionAdapter extends AbstractMongoAdapter<MongoUserSessionEnt
     @Override
     public void setLastSessionRefresh(int seconds) {
         entity.setLastSessionRefresh(seconds);
+        updateMongoEntity();
     }
 
     @Override
     public void associateClient(ClientModel client) {
-        List<ClientModel> clients = getClientAssociations();
-        for (ClientModel ass : clients) {
-            if (ass.getId().equals(client.getId())) return;
-        }
-
-        MongoClientUserSessionAssociationEntity association = new MongoClientUserSessionAssociationEntity();
-        association.setClientId(client.getId());
-        association.setSessionId(getId());
-
-        getMongoStore().insertEntity(association, invocationContext);
+        getMongoStore().pushItemToList(entity, "associatedClientIds", client.getId(), true, invocationContext);
     }
 
     @Override
     public List<ClientModel> getClientAssociations() {
-        DBObject query = new QueryBuilder()
-                .and("sessionId").is(getId())
-                .get();
-        List<MongoClientUserSessionAssociationEntity> associations = getMongoStore().loadEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
-
-        List<ClientModel> result = new ArrayList<ClientModel>();
-        for (MongoClientUserSessionAssociationEntity association : associations) {
-            ClientModel client = realm.findClientById(association.getClientId());
-            result.add(client);
+        List<String> associatedClientIds = getMongoEntity().getAssociatedClientIds();
+
+        List<ClientModel> clients = new ArrayList<ClientModel>();
+        for (String clientId : associatedClientIds) {
+            // Try application first
+            ClientModel client = realm.getApplicationById(clientId);
+
+            // And then OAuthClient
+            if (client == null) {
+                client = realm.getOAuthClientById(clientId);
+            }
+
+            if (client != null) {
+                clients.add(client);
+            } else {
+                logger.warnf("Not found associated client with Id: %s", clientId);
+            }
         }
-        return result;
+        return clients;
     }
 
     @Override
     public void removeAssociatedClient(ClientModel client) {
-        DBObject query = new QueryBuilder()
-                .and("sessionId").is(getId())
-                .and("clientId").is(client.getId())
-                .get();
-        getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-
-        UserSessionAdapter that = (UserSessionAdapter) o;
-        return getId().equals(that.getId());
-    }
-
-    @Override
-    public int hashCode() {
-        return getId().hashCode();
+        getMongoStore().pullItemFromList(entity, "associatedClientIds", client.getId(), invocationContext);
     }
 }

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoApplicationEntity.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoApplicationEntity.java
index bb83e0e..c8dd6da 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoApplicationEntity.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoApplicationEntity.java
@@ -1,7 +1,10 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
+import java.util.List;
+
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.jboss.logging.Logger;
 import org.keycloak.models.entities.ApplicationEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
@@ -23,10 +26,13 @@ public class MongoApplicationEntity extends ApplicationEntity implements MongoId
                 .get();
         context.getMongoStore().removeEntities(MongoRoleEntity.class, query, context);
 
+        // Remove all session associations
         query = new QueryBuilder()
-                .and("clientId").is(getId())
+                .and("associatedClientIds").is(getId())
                 .get();
-        context.getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, context);
-
+        List<MongoUserSessionEntity> sessions = context.getMongoStore().loadEntities(MongoUserSessionEntity.class, query, context);
+        for (MongoUserSessionEntity session : sessions) {
+            context.getMongoStore().pullItemFromList(session, "associatedClientIds", getId(), context);
+        }
     }
 }

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoOAuthClientEntity.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoOAuthClientEntity.java
index a499e51..8950292 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoOAuthClientEntity.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoOAuthClientEntity.java
@@ -1,5 +1,7 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
+import java.util.List;
+
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.entities.OAuthClientEntity;
@@ -16,10 +18,14 @@ import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 public class MongoOAuthClientEntity extends OAuthClientEntity implements MongoIdentifiableEntity {
 
     @Override
-    public void afterRemove(MongoStoreInvocationContext invocationContext) {
+    public void afterRemove(MongoStoreInvocationContext context) {
+        // Remove all session associations
         DBObject query = new QueryBuilder()
-                .and("clientId").is(getId())
+                .and("associatedClientIds").is(getId())
                 .get();
-        invocationContext.getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, invocationContext);
+        List<MongoUserSessionEntity> sessions = context.getMongoStore().loadEntities(MongoUserSessionEntity.class, query, context);
+        for (MongoUserSessionEntity session : sessions) {
+            context.getMongoStore().pullItemFromList(session, "associatedClientIds", getId(), context);
+        }
     }
 }

diff --git a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoUserSessionEntity.java b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoUserSessionEntity.java
index 2389840..c728b27 100755
--- a/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoUserSessionEntity.java
+++ b/model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/MongoUserSessionEntity.java
@@ -1,7 +1,8 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
-import com.mongodb.DBObject;
-import com.mongodb.QueryBuilder;
+import java.util.ArrayList;
+import java.util.List;
+
 import org.keycloak.models.entities.AbstractIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
@@ -23,6 +24,8 @@ public class MongoUserSessionEntity extends AbstractIdentifiableEntity implement
 
     private int lastSessionRefresh;
 
+    private List<String> associatedClientIds = new ArrayList<String>();
+
     public String getRealmId() {
         return realmId;
     }
@@ -63,13 +66,16 @@ public class MongoUserSessionEntity extends AbstractIdentifiableEntity implement
         this.lastSessionRefresh = lastSessionRefresh;
     }
 
+    public List<String> getAssociatedClientIds() {
+        return associatedClientIds;
+    }
+
+    public void setAssociatedClientIds(List<String> associatedClientIds) {
+        this.associatedClientIds = associatedClientIds;
+    }
+
     @Override
     public void afterRemove(MongoStoreInvocationContext context) {
-        // Remove all roles, which belongs to this application
-        DBObject query = new QueryBuilder()
-                .and("sessionId").is(getId())
-                .get();
-        context.getMongoStore().removeEntities(MongoClientUserSessionAssociationEntity.class, query, context);
     }
 
 }

diff --git a/model/tests/src/test/java/org/keycloak/model/test/AdapterTest.java b/model/tests/src/test/java/org/keycloak/model/test/AdapterTest.java
index be26f44..84d4586 100755
--- a/model/tests/src/test/java/org/keycloak/model/test/AdapterTest.java
+++ b/model/tests/src/test/java/org/keycloak/model/test/AdapterTest.java
@@ -762,4 +762,49 @@ public class AdapterTest extends AbstractModelTest {
         assertNull(realmManager.getRealmByName("userSessions").getUserSession(userSession.getId()));
     }
 
+    @Test
+    public void userSessionAssociations() {
+        RealmModel realm = realmManager.createRealm("userSessions");
+        UserModel user = realm.addUser("userSessions1");
+        UserSessionModel userSession = realm.createUserSession(user, "127.0.0.1");
+
+        ApplicationModel app1 = realm.addApplication("app1");
+        ApplicationModel app2 = realm.addApplication("app2");
+        OAuthClientModel client1 = realm.addOAuthClient("client1");
+
+        Assert.assertEquals(0, userSession.getClientAssociations().size());
+
+        userSession.associateClient(app1);
+        userSession.associateClient(client1);
+
+        Assert.assertEquals(2, userSession.getClientAssociations().size());
+        Assert.assertTrue(app1.getUserSessions().contains(userSession));
+        Assert.assertFalse(app2.getUserSessions().contains(userSession));
+        Assert.assertTrue(client1.getUserSessions().contains(userSession));
+
+        commit();
+
+        // Refresh all
+        realm = realmManager.getRealm("userSessions");
+        userSession = realm.getUserSession(userSession.getId());
+        app1 = realm.getApplicationByName("app1");
+        client1 = realm.getOAuthClient("client1");
+
+        userSession.removeAssociatedClient(app1);
+        Assert.assertEquals(1, userSession.getClientAssociations().size());
+        Assert.assertEquals(client1, userSession.getClientAssociations().get(0));
+        Assert.assertFalse(app1.getUserSessions().contains(userSession));
+
+        commit();
+
+        // Refresh all
+        realm = realmManager.getRealm("userSessions");
+        userSession = realm.getUserSession(userSession.getId());
+        client1 = realm.getOAuthClient("client1");
+
+        userSession.removeAssociatedClient(client1);
+        Assert.assertEquals(0, userSession.getClientAssociations().size());
+        Assert.assertFalse(client1.getUserSessions().contains(userSession));
+    }
+
 }

diff --git a/testsuite/integration/pom.xml b/testsuite/integration/pom.xml
index 74bac64..513666f 100755
--- a/testsuite/integration/pom.xml
+++ b/testsuite/integration/pom.xml
@@ -439,7 +439,7 @@
             <id>mongo</id>
             <activation>
                 <property>
-                    <name>keycloak.model</name>
+                    <name>keycloak.model.provider</name>
                     <value>mongo</value>
                 </property>
             </activation>
@@ -471,13 +471,13 @@
                                         <keycloak.model.mongo.host>${keycloak.model.mongo.host}</keycloak.model.mongo.host>
                                         <keycloak.model.mongo.port>${keycloak.model.mongo.port}</keycloak.model.mongo.port>
                                         <keycloak.model.mongo.db>${keycloak.model.mongo.db}</keycloak.model.mongo.db>
+                                        <keycloak.model.mongo.clearOnStartup>${keycloak.model.mongo.clearOnStartup}</keycloak.model.mongo.clearOnStartup>
 
                                         <keycloak.audit.provider>mongo</keycloak.audit.provider>
                                         <keycloak.audit.mongo.host>${keycloak.model.mongo.host}</keycloak.audit.mongo.host>
                                         <keycloak.audit.mongo.port>${keycloak.model.mongo.port}</keycloak.audit.mongo.port>
                                         <keycloak.audit.mongo.db>${keycloak.model.mongo.db}</keycloak.audit.mongo.db>
-
-                                        <keycloak.model.mongo.clearOnStartup>${keycloak.model.mongo.clearOnStartup}</keycloak.model.mongo.clearOnStartup>
+                                        <keycloak.audit.mongo.clearOnStartup>${keycloak.model.mongo.clearOnStartup}</keycloak.audit.mongo.clearOnStartup>
                                     </systemPropertyVariables>
                                 </configuration>
                             </execution>

diff --git a/testsuite/integration/README.md b/testsuite/integration/README.md
index 003c215..a2033eb 100644
--- a/testsuite/integration/README.md
+++ b/testsuite/integration/README.md
@@ -58,6 +58,8 @@ By default it's using database `keycloak` on localhost/27017 and it uses already
 
     mvn exec:java -Pkeycloak-server -Dkeycloak.model.provider=mongo -Dkeycloak.model.mongo.host=localhost -Dkeycloak.model.mongo.port=27017 -Dkeycloak.model.mongo.db=keycloak -Dkeycloak.model.mongo.clearOnStartup=false
 
+Note that if you are using Mongo model, it would mean that Mongo will be used for audit as well. You may need to use audit related properties for configuration of Mongo if you want to override default ones (For example keycloak.audit.mongo.host, keycloak.audit.mongo.port etc)
+
 TOTP codes
 ----------
 

diff --git a/testsuite/integration/src/main/resources/META-INF/keycloak-server.json b/testsuite/integration/src/main/resources/META-INF/keycloak-server.json
index f0fd046..251dfa4 100644
--- a/testsuite/integration/src/main/resources/META-INF/keycloak-server.json
+++ b/testsuite/integration/src/main/resources/META-INF/keycloak-server.json
@@ -4,11 +4,23 @@
     },
 
     "audit": {
-        "provider": "${keycloak.audit.provider:jpa}"
+        "provider": "${keycloak.audit.provider,keycloak.model.provider:jpa}",
+        "mongo": {
+          "host": "${keycloak.audit.mongo.host:127.0.0.1}",
+          "port": "${keycloak.audit.mongo.port:27017}",
+          "db": "${keycloak.audit.mongo.db:keycloak-audit}",
+          "clearOnStartup": "${keycloak.audit.mongo.clearOnStartup:false}"
+        }
     },
 
     "model": {
-        "provider": "${keycloak.model.provider:jpa}"
+        "provider": "${keycloak.model.provider:jpa}",
+        "mongo": {
+          "host": "${keycloak.model.mongo.host:127.0.0.1}",
+          "port": "${keycloak.model.mongo.port:27017}",
+          "db": "${keycloak.model.mongo.db:keycloak}",
+          "clearOnStartup": "${keycloak.model.mongo.clearOnStartup:false}"
+        }
     },
 
     "timer": {