Mongo: Refactoring. All unit tests and testsuite are passing with Mongo.

model/mongo/src/main/java/org/keycloak/models/mongo/api/AbstractMongoEntity.java

@@ -1,12 +0,0 @@
-package org.keycloak.models.mongo.api;
-
-/**
- * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
- */
-public abstract class AbstractMongoEntity implements MongoEntity {
-
-    @Override
-    public void afterRemove(MongoStore mongoStore) {
-        // Empty by default
-    }
-}

model/mongo/src/main/java/org/keycloak/models/mongo/api/AbstractMongoIdentifiableEntity.java

@@ -0,0 +1,50 @@
+package org.keycloak.models.mongo.api;
+
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class AbstractMongoIdentifiableEntity implements MongoIdentifiableEntity {
+
+    private String id;
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(String id) {
+        this.id = id;
+    }
+
+    @Override
+    public void afterRemove(MongoStore mongoStore, MongoStoreInvocationContext invocationContext) {
+        // Empty by default
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == this) return true;
+
+        if (this.id == null) return false;
+
+        if (o == null || getClass() != o.getClass()) return false;
+
+        AbstractMongoIdentifiableEntity that = (AbstractMongoIdentifiableEntity) o;
+
+        if (!getId().equals(that.getId())) return false;
+
+        return true;
+
+    }
+
+    @Override
+    public int hashCode() {
+        return id!=null ? id.hashCode() : super.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return String.format("%s [ id=%s ]", getClass().getSimpleName(), getId());
+    }
+}

model/mongo/src/main/java/org/keycloak/models/mongo/api/MongoEntity.java

@@ -1,16 +1,11 @@
 package org.keycloak.models.mongo.api;
 
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+
 /**
  * Base interface for object, which is persisted in Mongo
  *
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 public interface MongoEntity {
-
-    /**
-     * Lifecycle callback, which is called after removal of this object from Mongo.
-     * It may be useful for triggering removal of wired objects.
-     */
-    void afterRemove(MongoStore mongoStore);
-
 }

model/mongo/src/main/java/org/keycloak/models/mongo/api/MongoId.java

@@ -1,18 +0,0 @@
-package org.keycloak.models.mongo.api;
-
-import java.lang.annotation.Documented;
-import java.lang.annotation.Retention;
-import java.lang.annotation.Target;
-
-import static java.lang.annotation.ElementType.FIELD;
-import static java.lang.annotation.ElementType.METHOD;
-import static java.lang.annotation.RetentionPolicy.RUNTIME;
-
-/**
- * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
- */
-@Target({METHOD, FIELD})
-@Documented
-@Retention(RUNTIME)
-public @interface MongoId {
-}

model/mongo/src/main/java/org/keycloak/models/mongo/api/MongoIdentifiableEntity.java

@@ -0,0 +1,21 @@
+package org.keycloak.models.mongo.api;
+
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+
+/**
+ * Entity with Id
+ *
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public interface MongoIdentifiableEntity extends MongoEntity {
+
+    public String getId();
+
+    public void setId(String id);
+
+    /**
+     * Lifecycle callback, which is called after removal of this object from Mongo.
+     * It may be useful for triggering removal of wired objects.
+     */
+    void afterRemove(MongoStore mongoStore, MongoStoreInvocationContext invocationContext);
+}

model/mongo/src/main/java/org/keycloak/models/mongo/api/MongoStore.java

@@ -1,6 +1,7 @@
 package org.keycloak.models.mongo.api;
 
 import com.mongodb.DBObject;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 
 import java.util.List;
 
@@ -14,30 +15,29 @@ public interface MongoStore {
      *
      * @param object to update
      */
-    void insertObject(MongoEntity object);
+    void insertObject(MongoIdentifiableEntity object, MongoStoreInvocationContext context);
 
     /**
      * Update existing object
      *
      * @param object to update
      */
-    void updateObject(MongoEntity object);
+    void updateObject(MongoIdentifiableEntity object, MongoStoreInvocationContext context);
 
 
-    <T extends MongoEntity> T loadObject(Class<T> type, String oid);
+    <T extends MongoIdentifiableEntity> T loadObject(Class<T> type, String oid, MongoStoreInvocationContext context);
 
-    <T extends MongoEntity> T loadSingleObject(Class<T> type, DBObject query);
+    <T extends MongoIdentifiableEntity> T loadSingleObject(Class<T> type, DBObject query, MongoStoreInvocationContext context);
 
-    <T extends MongoEntity> List<T> loadObjects(Class<T> type, DBObject query);
+    <T extends MongoIdentifiableEntity> List<T> loadObjects(Class<T> type, DBObject query, MongoStoreInvocationContext context);
 
-    // Object must have filled oid
+    boolean removeObject(MongoIdentifiableEntity object, MongoStoreInvocationContext context);
-    boolean removeObject(MongoEntity object);
 
-    boolean removeObject(Class<? extends MongoEntity> type, String oid);
+    boolean removeObject(Class<? extends MongoIdentifiableEntity> type, String id, MongoStoreInvocationContext context);
 
-    boolean removeObjects(Class<? extends MongoEntity> type, DBObject query);
+    boolean removeObjects(Class<? extends MongoIdentifiableEntity> type, DBObject query, MongoStoreInvocationContext context);
 
-    <S> boolean pushItemToList(MongoEntity object, String listPropertyName, S itemToPush, boolean skipIfAlreadyPresent);
+    <S> boolean pushItemToList(MongoIdentifiableEntity object, String listPropertyName, S itemToPush, boolean skipIfAlreadyPresent, MongoStoreInvocationContext context);
 
-    <S> void pullItemFromList(MongoEntity object, String listPropertyName, S itemToPull);
+    <S> boolean pullItemFromList(MongoIdentifiableEntity object, String listPropertyName, S itemToPull, MongoStoreInvocationContext context);
 }

model/mongo/src/main/java/org/keycloak/models/mongo/api/context/MongoStoreInvocationContext.java

@@ -0,0 +1,25 @@
+package org.keycloak.models.mongo.api.context;
+
+import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public interface MongoStoreInvocationContext {
+
+    void addLoadedObject(MongoIdentifiableEntity entity);
+
+    <T extends MongoIdentifiableEntity> T getLoadedObject(Class<T> type, String id);
+
+    void addUpdateTask(MongoIdentifiableEntity entityToUpdate, MongoTask task);
+
+    void addRemovedObject(MongoIdentifiableEntity entityToRemove);
+
+    void beforeDBSearch(Class<? extends MongoIdentifiableEntity> entityType);
+
+    void begin();
+
+    void commit();
+
+    void rollback();
+}

model/mongo/src/main/java/org/keycloak/models/mongo/api/context/MongoTask.java

@@ -0,0 +1,13 @@
+package org.keycloak.models.mongo.api.context;
+
+import org.keycloak.models.mongo.api.MongoStore;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public interface MongoTask {
+
+    void execute();
+
+    boolean isFullUpdate();
+}

model/mongo/src/main/java/org/keycloak/models/mongo/impl/MongoStoreImpl.java

@@ -11,8 +11,10 @@ import org.jboss.logging.Logger;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
+import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+import org.keycloak.models.mongo.api.context.MongoTask;
 import org.keycloak.models.mongo.api.types.Converter;
 import org.keycloak.models.mongo.api.types.ConverterContext;
 import org.keycloak.models.mongo.api.types.TypeConverter;
@@ -105,7 +107,7 @@ public class MongoStoreImpl implements MongoStore {
     }
 
     @Override
-    public void insertObject(MongoEntity object) {
+    public void insertObject(MongoIdentifiableEntity object, MongoStoreInvocationContext context) {
         Class<? extends MongoEntity> clazz = object.getClass();
 
         // Find annotations for ID, for all the properties and for the name of the collection.
@@ -116,8 +118,7 @@ public class MongoStoreImpl implements MongoStore {
 
         DBCollection dbCollection = database.getCollection(objectInfo.getDbCollectionName());
 
-        Property<String> oidProperty = objectInfo.getOidProperty();
+        String currentId = object.getId();
-        String currentId = oidProperty == null ? null : oidProperty.getValue(object);
 
         // Inserting object, which already has oid property set. So we need to set "_id"
         if (currentId != null) {
@@ -126,48 +127,73 @@ public class MongoStoreImpl implements MongoStore {
 
         dbCollection.insert(dbObject);
 
-        // Add oid to value of given object
+        // Add id to value of given object
-        if (currentId == null && oidProperty != null) {
-            oidProperty.setValue(object, dbObject.getString("_id"));
-        }
-    }
-
-    @Override
-    public void updateObject(MongoEntity object) {
-        Class<? extends MongoEntity> clazz = object.getClass();
-        ObjectInfo objectInfo = getObjectInfo(clazz);
-        BasicDBObject dbObject = typeConverter.convertApplicationObjectToDBObject(object, BasicDBObject.class);
-        DBCollection dbCollection = database.getCollection(objectInfo.getDbCollectionName());
-
-        Property<String> oidProperty = objectInfo.getOidProperty();
-        String currentId = oidProperty == null ? null : oidProperty.getValue(object);
-
         if (currentId == null) {
-            throw new IllegalStateException("Can't update object without id: " + object);
+            object.setId(dbObject.getString("_id"));
-        } else {
-            BasicDBObject query = new BasicDBObject("_id", getObjectId(currentId));
-            dbCollection.update(query, dbObject);
         }
+
+        // Treat object as if it is read (It is already submited to transaction)
+        context.addLoadedObject(object);
+    }
+
+    @Override
+    public void updateObject(final MongoIdentifiableEntity object, MongoStoreInvocationContext context) {
+        MongoTask fullUpdateTask = new MongoTask() {
+
+            @Override
+            public void execute() {
+                Class<? extends MongoEntity> clazz = object.getClass();
+                ObjectInfo objectInfo = getObjectInfo(clazz);
+                BasicDBObject dbObject = typeConverter.convertApplicationObjectToDBObject(object, BasicDBObject.class);
+                DBCollection dbCollection = database.getCollection(objectInfo.getDbCollectionName());
+
+                String currentId = object.getId();
+
+                if (currentId == null) {
+                    throw new IllegalStateException("Can't update object without id: " + object);
+                } else {
+                    BasicDBObject query = new BasicDBObject("_id", getObjectId(currentId));
+                    dbCollection.update(query, dbObject);
+                }
+            }
+
+            @Override
+            public boolean isFullUpdate() {
+                return true;
+            }
+        };
+
+        // update is just added to context and postponed
+        context.addUpdateTask(object, fullUpdateTask);
     }
 
 
     @Override
-    public <T extends MongoEntity> T loadObject(Class<T> type, String oid) {
+    public <T extends MongoIdentifiableEntity> T loadObject(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.getLoadedObject(type, id);
+        if (cached != null) return cached;
+
         DBCollection dbCollection = getDBCollectionForType(type);
 
-        BasicDBObject idQuery = new BasicDBObject("_id", getObjectId(oid));
+        BasicDBObject idQuery = new BasicDBObject("_id", getObjectId(id));
         DBObject dbObject = dbCollection.findOne(idQuery);
 
         if (dbObject == null) return null;
 
         ConverterContext<Object> converterContext = new ConverterContext<Object>(dbObject, type, null);
-        return (T)typeConverter.convertDBObjectToApplicationObject(converterContext);
+        T converted = (T)typeConverter.convertDBObjectToApplicationObject(converterContext);
+
+        // Now add it to loaded objects
+        context.addLoadedObject(converted);
+
+        return converted;
     }
 
 
     @Override
-    public <T extends MongoEntity> T loadSingleObject(Class<T> type, DBObject query) {
+    public <T extends MongoIdentifiableEntity> T loadSingleObject(Class<T> type, DBObject query, MongoStoreInvocationContext context) {
-        List<T> result = loadObjects(type, query);
+        List<T> result = loadObjects(type, query, context);
         if (result.size() > 1) {
             throw new IllegalStateException("There are " + result.size() + " results for type=" + type + ", query=" + query + ". We expect just one");
         } else if (result.size() == 1) {
@@ -180,47 +206,43 @@ public class MongoStoreImpl implements MongoStore {
 
 
     @Override
-    public <T extends MongoEntity> List<T> loadObjects(Class<T> type, DBObject query) {
+    public <T extends MongoIdentifiableEntity> List<T> loadObjects(Class<T> type, DBObject query, MongoStoreInvocationContext context) {
-        DBCollection dbCollection = getDBCollectionForType(type);
+        // First we should execute all pending tasks before searching DB
+        context.beforeDBSearch(type);
 
+        DBCollection dbCollection = getDBCollectionForType(type);
         DBCursor cursor = dbCollection.find(query);
 
-        return convertCursor(type, cursor);
+        return convertCursor(type, cursor, context);
     }
 
 
     @Override
-    public boolean removeObject(MongoEntity object) {
+    public boolean removeObject(MongoIdentifiableEntity object, MongoStoreInvocationContext context) {
-        Class<? extends MongoEntity> type = object.getClass();
+        return removeObject(object.getClass(), object.getId(), context);
-        ObjectInfo objectInfo = getObjectInfo(type);
-
-        Property<String> idProperty = objectInfo.getOidProperty();
-        String oid = idProperty.getValue(object);
-
-        return removeObject(type, oid);
     }
 
 
     @Override
-    public boolean removeObject(Class<? extends MongoEntity> type, String oid) {
+    public boolean removeObject(Class<? extends MongoIdentifiableEntity> type, String id, MongoStoreInvocationContext context) {
-        MongoEntity found = loadObject(type, oid);
+        MongoIdentifiableEntity found = loadObject(type, id, context);
         if (found == null) {
             return false;
         } else {
             DBCollection dbCollection = getDBCollectionForType(type);
-            BasicDBObject dbQuery = new BasicDBObject("_id", getObjectId(oid));
+            BasicDBObject dbQuery = new BasicDBObject("_id", getObjectId(id));
             dbCollection.remove(dbQuery);
-            logger.info("Object of type: " + type + ", oid: " + oid + " removed from MongoDB.");
+            logger.info("Object of type: " + type + ", id: " + id + " removed from MongoDB.");
 
-            found.afterRemove(this);
+            context.addRemovedObject(found);
             return true;
         }
     }
 
 
     @Override
-    public boolean removeObjects(Class<? extends MongoEntity> type, DBObject query) {
+    public boolean removeObjects(Class<? extends MongoIdentifiableEntity> type, DBObject query, MongoStoreInvocationContext context) {
-        List<? extends MongoEntity> foundObjects = loadObjects(type, query);
+        List<? extends MongoIdentifiableEntity> foundObjects = loadObjects(type, query, context);
         if (foundObjects.size() == 0) {
             return false;
         } else {
@@ -228,23 +250,18 @@ public class MongoStoreImpl implements MongoStore {
             dbCollection.remove(query);
             logger.info("Removed " + foundObjects.size() + " objects of type: " + type + ", query: " + query);
 
-            for (MongoEntity found : foundObjects) {
+            for (MongoIdentifiableEntity found : foundObjects) {
-                found.afterRemove(this);
+                context.addRemovedObject(found);;
             }
             return true;
         }
     }
 
     @Override
-    public <S> boolean pushItemToList(MongoEntity object, String listPropertyName, S itemToPush, boolean skipIfAlreadyPresent) {
+    public <S> boolean pushItemToList(final MongoIdentifiableEntity object, final String listPropertyName, S itemToPush, boolean skipIfAlreadyPresent, MongoStoreInvocationContext context) {
-        Class<? extends MongoEntity> type = object.getClass();
+        final Class<? extends MongoEntity> type = object.getClass();
         ObjectInfo objectInfo = getObjectInfo(type);
 
-        Property<String> oidProperty = getObjectInfo(type).getOidProperty();
-        if (oidProperty == null) {
-            throw new IllegalArgumentException("List pushes not supported for properties without oid");
-        }
-
         // Add item to list directly in this object
         Property<Object> listProperty = objectInfo.getPropertyByName(listPropertyName);
         if (listProperty == null) {
@@ -257,34 +274,44 @@ public class MongoStoreImpl implements MongoStore {
             listProperty.setValue(object, list);
         }
 
-        // Return if item is already in list
+        // Skip if item is already in list
         if (skipIfAlreadyPresent && list.contains(itemToPush)) {
             return false;
         }
 
+        // Update java object
         list.add(itemToPush);
 
-        // Push item to DB. We always convert whole list, so it's not so optimal...TODO: use $push if possible
+        // Add update of list to pending tasks
-        BasicDBList dbList = typeConverter.convertApplicationObjectToDBObject(list, BasicDBList.class);
+        final List<S> listt = list;
+        context.addUpdateTask(object, new MongoTask() {
+
+            @Override
+            public void execute() {
+                // Now DB update of new list with usage of $set
+                BasicDBList dbList = typeConverter.convertApplicationObjectToDBObject(listt, BasicDBList.class);
+
+                BasicDBObject query = new BasicDBObject("_id", getObjectId(object.getId()));
+                BasicDBObject listObject = new BasicDBObject(listPropertyName, dbList);
+                BasicDBObject setCommand = new BasicDBObject("$set", listObject);
+                getDBCollectionForType(type).update(query, setCommand);
+            }
+
+            @Override
+            public boolean isFullUpdate() {
+                return false;
+            }
+        });
 
-        BasicDBObject query = new BasicDBObject("_id", getObjectId(oidProperty.getValue(object)));
-        BasicDBObject listObject = new BasicDBObject(listPropertyName, dbList);
-        BasicDBObject setCommand = new BasicDBObject("$set", listObject);
-        getDBCollectionForType(type).update(query, setCommand);
         return true;
     }
 
 
     @Override
-    public <S> void pullItemFromList(MongoEntity object, String listPropertyName, S itemToPull) {
+    public <S> boolean pullItemFromList(final MongoIdentifiableEntity object, final String listPropertyName, final S itemToPull, MongoStoreInvocationContext context) {
-        Class<? extends MongoEntity> type = object.getClass();
+        final Class<? extends MongoEntity> type = object.getClass();
         ObjectInfo objectInfo = getObjectInfo(type);
 
-        Property<String> oidProperty = getObjectInfo(type).getOidProperty();
-        if (oidProperty == null) {
-            throw new IllegalArgumentException("List pulls not supported for properties without oid");
-        }
-
         // Remove item from list directly in this object
         Property<Object> listProperty = objectInfo.getPropertyByName(listPropertyName);
         if (listProperty == null) {
@@ -293,15 +320,33 @@ public class MongoStoreImpl implements MongoStore {
         List<S> list = (List<S>)listProperty.getValue(object);
 
         // If list is null, we skip both object and DB update
-        if (list != null) {
+        if (list == null || !list.contains(itemToPull)) {
+            return false;
+        } else {
+
+            // Update java object
             list.remove(itemToPull);
 
-            // Pull item from DB
+            // Add update of list to pending tasks
-            Object dbItemToPull = typeConverter.convertApplicationObjectToDBObject(itemToPull, Object.class);
+            context.addUpdateTask(object, new MongoTask() {
-            BasicDBObject query = new BasicDBObject("_id", getObjectId(oidProperty.getValue(object)));
+
-            BasicDBObject pullObject = new BasicDBObject(listPropertyName, dbItemToPull);
+                @Override
-            BasicDBObject pullCommand = new BasicDBObject("$pull", pullObject);
+                public void execute() {
-            getDBCollectionForType(type).update(query, pullCommand);
+                    // Pull item from DB
+                    Object dbItemToPull = typeConverter.convertApplicationObjectToDBObject(itemToPull, Object.class);
+                    BasicDBObject query = new BasicDBObject("_id", getObjectId(object.getId()));
+                    BasicDBObject pullObject = new BasicDBObject(listPropertyName, dbItemToPull);
+                    BasicDBObject pullCommand = new BasicDBObject("$pull", pullObject);
+                    getDBCollectionForType(type).update(query, pullCommand);
+                }
+
+                @Override
+                public boolean isFullUpdate() {
+                    return false;
+                }
+            });
+
+            return true;
         }
     }
 
@@ -317,14 +362,12 @@ public class MongoStoreImpl implements MongoStore {
     public ObjectInfo getObjectInfo(Class<? extends MongoEntity> objectClass) {
         ObjectInfo objectInfo = objectInfoCache.get(objectClass);
         if (objectInfo == null) {
-            Property<String> idProperty = PropertyQueries.<String>createQuery(objectClass).addCriteria(new AnnotatedPropertyCriteria(MongoId.class)).getFirstResult();
-
             List<Property<Object>> properties = PropertyQueries.createQuery(objectClass).addCriteria(new AnnotatedPropertyCriteria(MongoField.class)).getResultList();
 
             MongoCollection classAnnotation = objectClass.getAnnotation(MongoCollection.class);
 
             String dbCollectionName = classAnnotation==null ? null : classAnnotation.collectionName();
-            objectInfo = new ObjectInfo(objectClass, dbCollectionName, idProperty, properties);
+            objectInfo = new ObjectInfo(objectClass, dbCollectionName, properties);
 
             ObjectInfo existing = objectInfoCache.putIfAbsent(objectClass, objectInfo);
             if (existing != null) {
@@ -335,14 +378,23 @@ public class MongoStoreImpl implements MongoStore {
         return objectInfo;
     }
 
-    private <T extends MongoEntity> List<T> convertCursor(Class<T> type, DBCursor cursor) {
+    protected <T extends MongoIdentifiableEntity> List<T> convertCursor(Class<T> type, DBCursor cursor, MongoStoreInvocationContext context) {
         List<T> result = new ArrayList<T>();
 
         try {
             for (DBObject dbObject : cursor) {
-                ConverterContext<Object> converterContext = new ConverterContext<Object>(dbObject, type, null);
+                // First look if we already have loaded object cached. If yes, we will use cached instance
-                T converted = (T)typeConverter.convertDBObjectToApplicationObject(converterContext);
+                String id = dbObject.get("_id").toString();
-                result.add(converted);
+                T object = context.getLoadedObject(type, id);
+
+                if (object == null) {
+                    // So convert and use fresh instance from DB
+                    ConverterContext<Object> converterContext = new ConverterContext<Object>(dbObject, type, null);
+                    object = (T)typeConverter.convertDBObjectToApplicationObject(converterContext);
+                    context.addLoadedObject(object);
+                }
+
+                result.add(object);
             }
         } finally {
             cursor.close();
@@ -351,14 +403,14 @@ public class MongoStoreImpl implements MongoStore {
         return result;
     }
 
-    private DBCollection getDBCollectionForType(Class<? extends MongoEntity> type) {
+    protected DBCollection getDBCollectionForType(Class<? extends MongoEntity> type) {
         ObjectInfo objectInfo = getObjectInfo(type);
         String dbCollectionName = objectInfo.getDbCollectionName();
         return dbCollectionName==null ? null : database.getCollection(objectInfo.getDbCollectionName());
     }
 
     // We allow ObjectId to be both "ObjectId" or "String".
-    private Object getObjectId(String idAsString) {
+    protected Object getObjectId(String idAsString) {
         if (ObjectId.isValid(idAsString)) {
             return new ObjectId(idAsString);
         } else {

model/mongo/src/main/java/org/keycloak/models/mongo/impl/ObjectInfo.java

@@ -18,14 +18,11 @@ public class ObjectInfo {
 
     private final String dbCollectionName;
 
-    private final Property<String> oidProperty;
-
     private final Map<String, Property<Object>> properties;
 
-    public ObjectInfo(Class<? extends MongoEntity> objectClass, String dbCollectionName, Property<String> oidProperty, List<Property<Object>> properties) {
+    public ObjectInfo(Class<? extends MongoEntity> objectClass, String dbCollectionName, List<Property<Object>> properties) {
         this.objectClass = objectClass;
         this.dbCollectionName = dbCollectionName;
-        this.oidProperty = oidProperty;
 
         Map<String, Property<Object>> props= new HashMap<String, Property<Object>>();
         for (Property<Object> property : properties) {
@@ -42,10 +39,6 @@ public class ObjectInfo {
         return dbCollectionName;
     }
 
-    public Property<String> getOidProperty() {
-        return oidProperty;
-    }
-
     public Collection<Property<Object>> getProperties() {
         return properties.values();
     }

model/mongo/src/main/java/org/keycloak/models/mongo/impl/context/SimpleMongoStoreInvocationContext.java

@@ -0,0 +1,56 @@
+package org.keycloak.models.mongo.impl.context;
+
+import org.keycloak.models.mongo.api.MongoEntity;
+import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
+import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+import org.keycloak.models.mongo.api.context.MongoTask;
+
+/**
+ * Context, which is not doing any postponing of tasks and does not cache anything
+ *
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class SimpleMongoStoreInvocationContext implements MongoStoreInvocationContext {
+
+    private final MongoStore store;
+
+    public SimpleMongoStoreInvocationContext(MongoStore store) {
+        this.store = store;
+    }
+
+    @Override
+    public void addLoadedObject(MongoIdentifiableEntity entity) {
+    }
+
+    @Override
+    public <T extends MongoIdentifiableEntity> T getLoadedObject(Class<T> type, String id) {
+        return null;
+    }
+
+    @Override
+    public void addUpdateTask(MongoIdentifiableEntity entityToUpdate, MongoTask task) {
+        task.execute();
+    }
+
+    @Override
+    public void addRemovedObject(MongoIdentifiableEntity entityToRemove) {
+        entityToRemove.afterRemove(store, this);
+    }
+
+    @Override
+    public void beforeDBSearch(Class<? extends MongoIdentifiableEntity> entityType) {
+    }
+
+    @Override
+    public void begin() {
+    }
+
+    @Override
+    public void commit() {
+    }
+
+    @Override
+    public void rollback() {
+    }
+}

model/mongo/src/main/java/org/keycloak/models/mongo/impl/context/TransactionMongoStoreInvocationContext.java

@@ -0,0 +1,129 @@
+package org.keycloak.models.mongo.impl.context;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
+import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+import org.keycloak.models.mongo.api.context.MongoTask;
+
+/**
+ * Invocation context, which has some very basic support for transactions, and is able to cache loaded objects.
+ * It always execute all pending update tasks before start searching for other objects
+ *
+ * It's per-request object (not thread safe)
+ *
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class TransactionMongoStoreInvocationContext implements MongoStoreInvocationContext {
+
+    // Assumption is that all objects has unique ID (unique across all the types)
+    private Map<String, MongoIdentifiableEntity> loadedObjects = new HashMap<String, MongoIdentifiableEntity>();
+
+    private Map<MongoIdentifiableEntity, Set<MongoTask>> pendingUpdateTasks = new HashMap<MongoIdentifiableEntity, Set<MongoTask>>();
+
+    private final MongoStore mongoStore;
+
+    public TransactionMongoStoreInvocationContext(MongoStore mongoStore) {
+        this.mongoStore = mongoStore;
+    }
+
+    @Override
+    public void addLoadedObject(MongoIdentifiableEntity entity) {
+        loadedObjects.put(entity.getId(), entity);
+    }
+
+    @Override
+    public <T extends MongoIdentifiableEntity> T getLoadedObject(Class<T> type, String id) {
+        return (T)loadedObjects.get(id);
+    }
+
+    @Override
+    public void addUpdateTask(MongoIdentifiableEntity entityToUpdate, MongoTask task) {
+        if (!loadedObjects.containsValue(entityToUpdate)) {
+            throw new IllegalStateException("Entity " + entityToUpdate + " not found in loaded objects");
+        }
+
+        Set<MongoTask> currentObjectTasks = pendingUpdateTasks.get(entityToUpdate);
+        if (currentObjectTasks == null) {
+            currentObjectTasks = new HashSet<MongoTask>();
+            pendingUpdateTasks.put(entityToUpdate, currentObjectTasks);
+        } else {
+            // if task is full update, then remove all other tasks as we need to do full update of object anyway
+            if (task.isFullUpdate()) {
+                currentObjectTasks.clear();
+            } else {
+                // If it already contains task for fullUpdate, then we don't need to add ours as we need to do full update of object anyway
+                for (MongoTask current : currentObjectTasks) {
+                     if (current.isFullUpdate()) {
+                         return;
+                     }
+                }
+            }
+        }
+
+        currentObjectTasks.add(task);
+    }
+
+    @Override
+    public void addRemovedObject(MongoIdentifiableEntity entityToRemove) {
+        // Remove all pending tasks and object from cache
+        pendingUpdateTasks.remove(entityToRemove);
+        loadedObjects.remove(entityToRemove.getId());
+
+        entityToRemove.afterRemove(mongoStore, this);
+    }
+
+    @Override
+    public void beforeDBSearch(Class<? extends MongoIdentifiableEntity> entityType) {
+        // Now execute pending update tasks of type, which will be searched
+        Set<MongoIdentifiableEntity> toRemove = new HashSet<MongoIdentifiableEntity>();
+
+        for (MongoIdentifiableEntity currentEntity : pendingUpdateTasks.keySet()) {
+            if (currentEntity.getClass().equals(entityType)) {
+                Set<MongoTask> mongoTasks = pendingUpdateTasks.get(currentEntity);
+                for (MongoTask currentTask : mongoTasks) {
+                    currentTask.execute();
+                }
+
+                toRemove.add(currentEntity);
+            }
+        }
+
+        // Now remove all done tasks
+        for (MongoIdentifiableEntity entity : toRemove) {
+            pendingUpdateTasks.remove(entity);
+        }
+    }
+
+    @Override
+    public void begin() {
+        loadedObjects.clear();
+        pendingUpdateTasks.clear();
+    }
+
+    @Override
+    public void commit() {
+        loadedObjects.clear();
+
+        // Now execute all pending update tasks
+        for (Set<MongoTask> mongoTasks : pendingUpdateTasks.values()) {
+            for (MongoTask currentTask : mongoTasks) {
+                currentTask.execute();
+            }
+        }
+
+        // And clear it
+        pendingUpdateTasks.clear();
+    }
+
+    @Override
+    public void rollback() {
+        // Just clear the map without executions of tasks
+        loadedObjects.clear();
+        pendingUpdateTasks.clear();
+    }
+}

model/mongo/src/main/java/org/keycloak/models/mongo/impl/types/BasicDBObjectConverter.java

@@ -8,6 +8,7 @@ import java.util.List;
 import com.mongodb.BasicDBObject;
 import org.jboss.logging.Logger;
 import org.keycloak.models.mongo.api.MongoEntity;
+import org.keycloak.models.mongo.api.MongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.types.Converter;
 import org.keycloak.models.mongo.api.types.ConverterContext;
 import org.keycloak.models.mongo.api.types.TypeConverter;
@@ -55,9 +56,8 @@ public class BasicDBObjectConverter<S extends MongoEntity> implements Converter<
 
             if ("_id".equals(key)) {
                 // Current property is "id"
-                Property<String> idProperty = objectInfo.getOidProperty();
+                if (object instanceof MongoIdentifiableEntity) {
-                if (idProperty != null) {
+                    ((MongoIdentifiableEntity)object).setId(value.toString());
-                    idProperty.setValue(object, value.toString());
                 }
 
             } else if ((property = objectInfo.getPropertyByName(key)) != null) {

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

@@ -0,0 +1,38 @@
+package org.keycloak.models.mongo.keycloak.adapters;
+
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
+import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public abstract class AbstractAdapter {
+
+    protected MongoStore mongoStore;
+    protected MongoStoreInvocationContext invocationContext;
+
+    public AbstractAdapter(MongoStore mongoStore, MongoStoreInvocationContext invocationContext) {
+        this.mongoStore = mongoStore;
+        this.invocationContext = invocationContext;
+    }
+
+    public abstract AbstractMongoIdentifiableEntity getMongoEntity();
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == this) return true;
+
+        if (o == null || getClass() != o.getClass()) return false;
+
+        AbstractAdapter that = (AbstractAdapter) o;
+
+        if (getMongoEntity() == null && that.getMongoEntity() == null) return true;
+        return getMongoEntity().equals(that.getMongoEntity());
+    }
+
+    @Override
+    public int hashCode() {
+        return getMongoEntity()!=null ? getMongoEntity().hashCode() : super.hashCode();
+    }
+}

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

@@ -5,7 +5,9 @@ import com.mongodb.QueryBuilder;
 import org.keycloak.models.ApplicationModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.entities.ApplicationEntity;
 import org.keycloak.models.mongo.keycloak.entities.RoleEntity;
 import org.keycloak.models.mongo.keycloak.entities.UserEntity;
@@ -19,37 +21,35 @@ import java.util.Set;
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class ApplicationAdapter implements ApplicationModel {
+public class ApplicationAdapter extends AbstractAdapter implements ApplicationModel {
 
     private final ApplicationEntity application;
-    private final MongoStore mongoStore;
-
     private UserAdapter resourceUser;
 
-    public ApplicationAdapter(ApplicationEntity applicationEntity, MongoStore mongoStore) {
+    public ApplicationAdapter(ApplicationEntity applicationEntity, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
-        this(applicationEntity, null, mongoStore);
+        this(applicationEntity, null, mongoStore, invContext);
     }
 
-    public ApplicationAdapter(ApplicationEntity applicationEntity, UserAdapter resourceUser, MongoStore mongoStore) {
+    public ApplicationAdapter(ApplicationEntity applicationEntity, UserAdapter resourceUser, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
+        super(mongoStore, invContext);
         this.application = applicationEntity;
         this.resourceUser = resourceUser;
-        this.mongoStore = mongoStore;
     }
 
     @Override
     public void updateApplication() {
-        mongoStore.updateObject(application);
+        mongoStore.updateObject(application, invocationContext);
     }
 
     @Override
     public UserAdapter getApplicationUser() {
         // This is not thread-safe. Assumption is that ApplicationAdapter instance is per-client object
         if (resourceUser == null) {
-            UserEntity userEntity = mongoStore.loadObject(UserEntity.class, application.getResourceUserId());
+            UserEntity userEntity = mongoStore.loadObject(UserEntity.class, application.getResourceUserId(), invocationContext);
             if (userEntity == null) {
                 throw new IllegalStateException("User " + application.getResourceUserId() + " not found");
             }
-            resourceUser = new UserAdapter(userEntity, mongoStore);
+            resourceUser = new UserAdapter(userEntity, mongoStore, invocationContext);
         }
 
         return resourceUser;
@@ -116,21 +116,21 @@ public class ApplicationAdapter implements ApplicationModel {
                 .and("name").is(name)
                 .and("applicationId").is(getId())
                 .get();
-        RoleEntity role = mongoStore.loadSingleObject(RoleEntity.class, query);
+        RoleEntity role = mongoStore.loadSingleObject(RoleEntity.class, query, invocationContext);
         if (role == null) {
             return null;
         } else {
-            return new RoleAdapter(role, this, mongoStore);
+            return new RoleAdapter(role, this, mongoStore, invocationContext);
         }
     }
 
     @Override
     public RoleModel getRoleById(String id) {
-        RoleEntity role = mongoStore.loadObject(RoleEntity.class, id);
+        RoleEntity role = mongoStore.loadObject(RoleEntity.class, id, invocationContext);
         if (role == null) {
             return null;
         } else {
-            return new RoleAdapter(role, this, mongoStore);
+            return new RoleAdapter(role, this, mongoStore, invocationContext);
         }
     }
 
@@ -145,13 +145,13 @@ public class ApplicationAdapter implements ApplicationModel {
         roleEntity.setName(name);
         roleEntity.setApplicationId(getId());
 
-        mongoStore.insertObject(roleEntity);
+        mongoStore.insertObject(roleEntity, invocationContext);
-        return new RoleAdapter(roleEntity, this, mongoStore);
+        return new RoleAdapter(roleEntity, this, mongoStore, invocationContext);
     }
 
     @Override
     public boolean removeRoleById(String id) {
-        return mongoStore.removeObject(RoleEntity.class ,id);
+        return mongoStore.removeObject(RoleEntity.class ,id, invocationContext);
     }
 
     @Override
@@ -159,11 +159,11 @@ public class ApplicationAdapter implements ApplicationModel {
         DBObject query = new QueryBuilder()
                 .and("applicationId").is(getId())
                 .get();
-        List<RoleEntity> roles = mongoStore.loadObjects(RoleEntity.class, query);
+        List<RoleEntity> roles = mongoStore.loadObjects(RoleEntity.class, query, invocationContext);
 
         Set<RoleModel> result = new HashSet<RoleModel>();
         for (RoleEntity role : roles) {
-            result.add(new RoleAdapter(role, this, mongoStore));
+            result.add(new RoleAdapter(role, this, mongoStore, invocationContext));
         }
 
         return result;
@@ -172,11 +172,11 @@ public class ApplicationAdapter implements ApplicationModel {
     @Override
     public Set<RoleModel> getApplicationRoleMappings(UserModel user) {
         Set<RoleModel> result = new HashSet<RoleModel>();
-        List<RoleEntity> roles = MongoModelUtils.getAllRolesOfUser(user, mongoStore);
+        List<RoleEntity> roles = MongoModelUtils.getAllRolesOfUser(user, mongoStore, invocationContext);
 
         for (RoleEntity role : roles) {
             if (getId().equals(role.getApplicationId())) {
-                result.add(new RoleAdapter(role, this, mongoStore));
+                result.add(new RoleAdapter(role, this, mongoStore, invocationContext));
             }
         }
         return result;
@@ -185,17 +185,17 @@ public class ApplicationAdapter implements ApplicationModel {
     @Override
     public void addScope(RoleModel role) {
         UserAdapter appUser = getApplicationUser();
-        mongoStore.pushItemToList(appUser.getUser(), "scopeIds", role.getId(), true);
+        mongoStore.pushItemToList(appUser.getUser(), "scopeIds", role.getId(), true, invocationContext);
     }
 
     @Override
     public Set<RoleModel> getApplicationScopeMappings(UserModel user) {
         Set<RoleModel> result = new HashSet<RoleModel>();
-        List<RoleEntity> roles = MongoModelUtils.getAllScopesOfUser(user, mongoStore);
+        List<RoleEntity> roles = MongoModelUtils.getAllScopesOfUser(user, mongoStore, invocationContext);
 
         for (RoleEntity role : roles) {
             if (getId().equals(role.getApplicationId())) {
-                result.add(new RoleAdapter(role, this, mongoStore));
+                result.add(new RoleAdapter(role, this, mongoStore, invocationContext));
             }
         }
         return result;
@@ -213,7 +213,7 @@ public class ApplicationAdapter implements ApplicationModel {
             addRole(name);
         }
 
-        mongoStore.pushItemToList(application, "defaultRoles", name, true);
+        mongoStore.pushItemToList(application, "defaultRoles", name, true, invocationContext);
     }
 
     @Override
@@ -232,16 +232,7 @@ public class ApplicationAdapter implements ApplicationModel {
     }
 
     @Override
-    public boolean equals(Object o) {
+    public AbstractMongoIdentifiableEntity getMongoEntity() {
-        if (o == null) return false;
+        return application;
-        if (o == this) return true;
-        if (!(o instanceof ApplicationAdapter)) return false;
-        ApplicationAdapter app = (ApplicationAdapter)o;
-        return app.getId().equals(getId());
-    }
-
-    @Override
-    public int hashCode() {
-        return getId().hashCode();
     }
 }

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

@@ -8,6 +8,9 @@ import org.keycloak.models.KeycloakTransaction;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
+import org.keycloak.models.mongo.impl.context.SimpleMongoStoreInvocationContext;
+import org.keycloak.models.mongo.impl.context.TransactionMongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.entities.RealmEntity;
 import org.keycloak.models.utils.KeycloakModelUtils;
 
@@ -19,20 +22,25 @@ import java.util.List;
  */
 public class MongoKeycloakSession implements KeycloakSession {
 
-    private static final MongoKeycloakTransaction PLACEHOLDER = new MongoKeycloakTransaction();
+    private final MongoStoreInvocationContext invocationContext;
+    private final MongoKeycloakTransaction transaction;
     private final MongoStore mongoStore;
 
     public MongoKeycloakSession(MongoStore mongoStore) {
         this.mongoStore = mongoStore;
+        // this.invocationContext = new SimpleMongoStoreInvocationContext(mongoStore);
+        this.invocationContext = new TransactionMongoStoreInvocationContext(mongoStore);
+        this.transaction = new MongoKeycloakTransaction(invocationContext);
     }
 
     @Override
     public KeycloakTransaction getTransaction() {
-        return PLACEHOLDER;
+        return transaction;
     }
 
     @Override
     public void close() {
+        // TODO
     }
 
     @Override
@@ -50,26 +58,25 @@ public class MongoKeycloakSession implements KeycloakSession {
         newRealm.setId(id);
         newRealm.setName(name);
 
-        mongoStore.insertObject(newRealm);
+        mongoStore.insertObject(newRealm, invocationContext);
 
-        RealmAdapter realm = new RealmAdapter(newRealm, mongoStore);
+        return new RealmAdapter(newRealm, mongoStore, invocationContext);
-        return realm;
     }
 
     @Override
     public RealmModel getRealm(String id) {
-        RealmEntity realmEntity = mongoStore.loadObject(RealmEntity.class, id);
+        RealmEntity realmEntity = mongoStore.loadObject(RealmEntity.class, id, invocationContext);
-        return realmEntity != null ? new RealmAdapter(realmEntity, mongoStore) : null;
+        return realmEntity != null ? new RealmAdapter(realmEntity, mongoStore, invocationContext) : null;
     }
 
     @Override
     public List<RealmModel> getRealms(UserModel admin) {
         DBObject query = new BasicDBObject();
-        List<RealmEntity> realms = mongoStore.loadObjects(RealmEntity.class, query);
+        List<RealmEntity> realms = mongoStore.loadObjects(RealmEntity.class, query, invocationContext);
 
         List<RealmModel> results = new ArrayList<RealmModel>();
         for (RealmEntity realmEntity : realms) {
-            results.add(new RealmAdapter(realmEntity, mongoStore));
+            results.add(new RealmAdapter(realmEntity, mongoStore, invocationContext));
         }
         return results;
     }
@@ -79,14 +86,14 @@ public class MongoKeycloakSession implements KeycloakSession {
         DBObject query = new QueryBuilder()
                 .and("name").is(name)
                 .get();
-        RealmEntity realm = mongoStore.loadSingleObject(RealmEntity.class, query);
+        RealmEntity realm = mongoStore.loadSingleObject(RealmEntity.class, query, invocationContext);
 
         if (realm == null) return null;
-        return new RealmAdapter(realm, mongoStore);
+        return new RealmAdapter(realm, mongoStore, invocationContext);
     }
 
     @Override
     public boolean removeRealm(String id) {
-        return mongoStore.removeObject(RealmEntity.class, id);
+        return mongoStore.removeObject(RealmEntity.class, id, invocationContext);
     }
 }

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

@@ -1,39 +1,60 @@
 package org.keycloak.models.mongo.keycloak.adapters;
 
 import org.keycloak.models.KeycloakTransaction;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 public class MongoKeycloakTransaction implements KeycloakTransaction {
 
+    private final MongoStoreInvocationContext invocationContext;
+
+    private boolean started = false;
+    private boolean rollbackOnly = false;
+
+    public MongoKeycloakTransaction(MongoStoreInvocationContext invocationContext) {
+        this.invocationContext = invocationContext;
+    }
+
     @Override
     public void begin() {
-        //To change body of implemented methods use File | Settings | File Templates.
+        if (started) {
+            throw new IllegalStateException("Transaction already started");
+        }
+        started = true;
+        invocationContext.begin();
     }
 
     @Override
     public void commit() {
-        //To change body of implemented methods use File | Settings | File Templates.
+        if (!started) {
+            throw new IllegalStateException("Transaction not yet started");
+        }
+        if (rollbackOnly) {
+            throw new IllegalStateException("Can't commit as transaction marked for rollback");
+        }
+
+        invocationContext.commit();
     }
 
     @Override
     public void rollback() {
-        //To change body of implemented methods use File | Settings | File Templates.
+        invocationContext.rollback();
     }
 
     @Override
     public void setRollbackOnly() {
-        //To change body of implemented methods use File | Settings | File Templates.
+        this.rollbackOnly = true;
     }
 
     @Override
     public boolean getRollbackOnly() {
-        return false;  //To change body of implemented methods use File | Settings | File Templates.
+        return rollbackOnly;
     }
 
     @Override
     public boolean isActive() {
-        return true;
+        return started;
     }
 }

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

@@ -2,27 +2,28 @@ package org.keycloak.models.mongo.keycloak.adapters;
 
 import org.keycloak.models.OAuthClientModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.entities.OAuthClientEntity;
 import org.keycloak.models.mongo.keycloak.entities.UserEntity;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class OAuthClientAdapter implements OAuthClientModel {
+public class OAuthClientAdapter extends AbstractAdapter implements OAuthClientModel {
 
     private final OAuthClientEntity delegate;
     private UserAdapter oauthAgent;
-    private final MongoStore mongoStore;
 
-    public OAuthClientAdapter(OAuthClientEntity oauthClientEntity, UserAdapter oauthAgent, MongoStore mongoStore) {
+    public OAuthClientAdapter(OAuthClientEntity oauthClientEntity, UserAdapter oauthAgent, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
+        super(mongoStore, invContext);
         this.delegate = oauthClientEntity;
         this.oauthAgent = oauthAgent;
-        this.mongoStore = mongoStore;
     }
 
-    public OAuthClientAdapter(OAuthClientEntity oauthClientEntity, MongoStore mongoStore) {
+    public OAuthClientAdapter(OAuthClientEntity oauthClientEntity, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
-        this(oauthClientEntity, null, mongoStore);
+        this(oauthClientEntity, null, mongoStore, invContext);
     }
 
     @Override
@@ -34,10 +35,14 @@ public class OAuthClientAdapter implements OAuthClientModel {
     public UserModel getOAuthAgent() {
         // This is not thread-safe. Assumption is that OAuthClientAdapter instance is per-client object
         if (oauthAgent == null) {
-            UserEntity user = mongoStore.loadObject(UserEntity.class, delegate.getOauthAgentId());
+            UserEntity user = mongoStore.loadObject(UserEntity.class, delegate.getOauthAgentId(), invocationContext);
-            oauthAgent = user!=null ? new UserAdapter(user, mongoStore) : null;
+            oauthAgent = user!=null ? new UserAdapter(user, mongoStore, invocationContext) : null;
         }
         return oauthAgent;
     }
 
+    @Override
+    public AbstractMongoIdentifiableEntity getMongoEntity() {
+        return delegate;
+    }
 }

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

@@ -12,7 +12,9 @@ import org.keycloak.models.RoleModel;
 import org.keycloak.models.SocialLinkModel;
 import org.keycloak.models.UserCredentialModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.entities.ApplicationEntity;
 import org.keycloak.models.mongo.keycloak.entities.CredentialEntity;
 import org.keycloak.models.mongo.keycloak.entities.OAuthClientEntity;
@@ -30,6 +32,7 @@ import java.security.PrivateKey;
 import java.security.PublicKey;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -40,21 +43,20 @@ import java.util.regex.Pattern;
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class RealmAdapter implements RealmModel {
+public class RealmAdapter extends AbstractAdapter implements RealmModel {
 
     private static final Logger logger = Logger.getLogger(RealmAdapter.class);
 
     private final RealmEntity realm;
-    private final MongoStore mongoStore;
 
     protected volatile transient PublicKey publicKey;
     protected volatile transient PrivateKey privateKey;
 
     private volatile transient PasswordPolicy passwordPolicy;
 
-    public RealmAdapter(RealmEntity realmEntity, MongoStore mongoStore) {
+    public RealmAdapter(RealmEntity realmEntity, MongoStore mongoStore, MongoStoreInvocationContext invocationContext) {
+        super(mongoStore, invocationContext);
         this.realm = realmEntity;
-        this.mongoStore = mongoStore;
     }
 
     @Override
@@ -250,18 +252,40 @@ public class RealmAdapter implements RealmModel {
         setPrivateKeyPem(privateKeyPem);
     }
 
+    @Override
+    public String getLoginTheme() {
+        return realm.getLoginTheme();
+    }
+
+    @Override
+    public void setLoginTheme(String name) {
+        realm.setLoginTheme(name);
+        updateRealm();
+    }
+
+    @Override
+    public String getAccountTheme() {
+        return realm.getAccountTheme();
+    }
+
+    @Override
+    public void setAccountTheme(String name) {
+        realm.setAccountTheme(name);
+        updateRealm();
+    }
+
     @Override
     public UserAdapter getUser(String name) {
         DBObject query = new QueryBuilder()
                 .and("loginName").is(name)
                 .and("realmId").is(getId())
                 .get();
-        UserEntity user = mongoStore.loadSingleObject(UserEntity.class, query);
+        UserEntity user = mongoStore.loadSingleObject(UserEntity.class, query, invocationContext);
 
         if (user == null) {
             return null;
         } else {
-            return new UserAdapter(user, mongoStore);
+            return new UserAdapter(user, mongoStore, invocationContext);
         }
     }
 
@@ -271,12 +295,12 @@ public class RealmAdapter implements RealmModel {
                 .and("email").is(email)
                 .and("realmId").is(getId())
                 .get();
-        UserEntity user = mongoStore.loadSingleObject(UserEntity.class, query);
+        UserEntity user = mongoStore.loadSingleObject(UserEntity.class, query, invocationContext);
 
         if (user == null) {
             return null;
         } else {
-            return new UserAdapter(user, mongoStore);
+            return new UserAdapter(user, mongoStore, invocationContext);
         }
     }
 
@@ -308,8 +332,8 @@ public class RealmAdapter implements RealmModel {
         userEntity.setEnabled(true);
         userEntity.setRealmId(getId());
 
-        mongoStore.insertObject(userEntity);
+        mongoStore.insertObject(userEntity, invocationContext);
-        return new UserAdapter(userEntity, mongoStore);
+        return new UserAdapter(userEntity, mongoStore, invocationContext);
     }
 
     @Override
@@ -318,7 +342,7 @@ public class RealmAdapter implements RealmModel {
                 .and("loginName").is(name)
                 .and("realmId").is(getId())
                 .get();
-        return mongoStore.removeObjects(UserEntity.class, query);
+        return mongoStore.removeObjects(UserEntity.class, query, invocationContext);
     }
 
     @Override
@@ -327,11 +351,11 @@ public class RealmAdapter implements RealmModel {
                 .and("name").is(name)
                 .and("realmId").is(getId())
                 .get();
-        RoleEntity role = mongoStore.loadSingleObject(RoleEntity.class, query);
+        RoleEntity role = mongoStore.loadSingleObject(RoleEntity.class, query, invocationContext);
         if (role == null) {
             return null;
         } else {
-            return new RoleAdapter(role, this, mongoStore);
+            return new RoleAdapter(role, this, mongoStore, invocationContext);
         }
     }
 
@@ -348,13 +372,13 @@ public class RealmAdapter implements RealmModel {
         roleEntity.setName(name);
         roleEntity.setRealmId(getId());
 
-        mongoStore.insertObject(roleEntity);
+        mongoStore.insertObject(roleEntity, invocationContext);
-        return new RoleAdapter(roleEntity, this, mongoStore);
+        return new RoleAdapter(roleEntity, this, mongoStore, invocationContext);
     }
 
     @Override
     public boolean removeRoleById(String id) {
-        return mongoStore.removeObject(RoleEntity.class ,id);
+        return mongoStore.removeObject(RoleEntity.class ,id, invocationContext);
     }
 
     @Override
@@ -362,13 +386,13 @@ public class RealmAdapter implements RealmModel {
         DBObject query = new QueryBuilder()
                 .and("realmId").is(getId())
                 .get();
-        List<RoleEntity> roles = mongoStore.loadObjects(RoleEntity.class, query);
+        List<RoleEntity> roles = mongoStore.loadObjects(RoleEntity.class, query, invocationContext);
 
         Set<RoleModel> result = new HashSet<RoleModel>();
 
         if (roles == null) return result;
         for (RoleEntity role : roles) {
-            result.add(new RoleAdapter(role, this, mongoStore));
+            result.add(new RoleAdapter(role, this, mongoStore, invocationContext));
         }
 
         return result;
@@ -376,11 +400,11 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public RoleModel getRoleById(String id) {
-        RoleEntity role = mongoStore.loadObject(RoleEntity.class, id);
+        RoleEntity role = mongoStore.loadObject(RoleEntity.class, id, invocationContext);
         if (role == null) {
             return null;
         } else {
-            return new RoleAdapter(role, this, mongoStore);
+            return new RoleAdapter(role, this, mongoStore, invocationContext);
         }
     }
 
@@ -396,7 +420,7 @@ public class RealmAdapter implements RealmModel {
             addRole(name);
         }
 
-        mongoStore.pushItemToList(realm, "defaultRoles", name, true);
+        mongoStore.pushItemToList(realm, "defaultRoles", name, true, invocationContext);
     }
 
     @Override
@@ -417,15 +441,14 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public ApplicationModel getApplicationById(String id) {
-        ApplicationEntity appData = mongoStore.loadObject(ApplicationEntity.class, id);
+        ApplicationEntity appData = mongoStore.loadObject(ApplicationEntity.class, id, invocationContext);
 
         // Check if application belongs to this realm
         if (appData == null || !getId().equals(appData.getRealmId())) {
             return null;
         }
 
-        ApplicationModel model = new ApplicationAdapter(appData, mongoStore);
+        return new ApplicationAdapter(appData, mongoStore, invocationContext);
-        return model;
     }
 
     @Override
@@ -434,8 +457,8 @@ public class RealmAdapter implements RealmModel {
                 .and("realmId").is(getId())
                 .and("name").is(name)
                 .get();
-        ApplicationEntity appEntity = mongoStore.loadSingleObject(ApplicationEntity.class, query);
+        ApplicationEntity appEntity = mongoStore.loadSingleObject(ApplicationEntity.class, query, invocationContext);
-        return appEntity==null ? null : new ApplicationAdapter(appEntity, mongoStore);
+        return appEntity==null ? null : new ApplicationAdapter(appEntity, mongoStore, invocationContext);
     }
 
     @Override
@@ -452,11 +475,11 @@ public class RealmAdapter implements RealmModel {
         DBObject query = new QueryBuilder()
                 .and("realmId").is(getId())
                 .get();
-        List<ApplicationEntity> appDatas = mongoStore.loadObjects(ApplicationEntity.class, query);
+        List<ApplicationEntity> appDatas = mongoStore.loadObjects(ApplicationEntity.class, query, invocationContext);
 
         List<ApplicationModel> result = new ArrayList<ApplicationModel>();
         for (ApplicationEntity appData : appDatas) {
-            result.add(new ApplicationAdapter(appData, mongoStore));
+            result.add(new ApplicationAdapter(appData, mongoStore, invocationContext));
         }
         return result;
     }
@@ -470,15 +493,14 @@ public class RealmAdapter implements RealmModel {
         appData.setRealmId(getId());
         appData.setEnabled(true);
         appData.setResourceUserId(resourceUser.getUser().getId());
-        mongoStore.insertObject(appData);
+        mongoStore.insertObject(appData, invocationContext);
 
-        ApplicationModel resource = new ApplicationAdapter(appData, resourceUser, mongoStore);
+        return new ApplicationAdapter(appData, resourceUser, mongoStore, invocationContext);
-        return resource;
     }
 
     @Override
     public boolean removeApplication(String id) {
-        return mongoStore.removeObject(ApplicationEntity.class, id);
+        return mongoStore.removeObject(ApplicationEntity.class, id, invocationContext);
     }
 
     @Override
@@ -495,20 +517,20 @@ public class RealmAdapter implements RealmModel {
     @Override
     public void grantRole(UserModel user, RoleModel role) {
         UserEntity userEntity = ((UserAdapter)user).getUser();
-        mongoStore.pushItemToList(userEntity, "roleIds", role.getId(), true);
+        mongoStore.pushItemToList(userEntity, "roleIds", role.getId(), true, invocationContext);
     }
 
     @Override
     public Set<RoleModel> getRoleMappings(UserModel user) {
         Set<RoleModel> result = new HashSet<RoleModel>();
-        List<RoleEntity> roles = MongoModelUtils.getAllRolesOfUser(user, mongoStore);
+        List<RoleEntity> roles = MongoModelUtils.getAllRolesOfUser(user, mongoStore, invocationContext);
 
         for (RoleEntity role : roles) {
             if (getId().equals(role.getRealmId())) {
-                result.add(new RoleAdapter(role, this, mongoStore));
+                result.add(new RoleAdapter(role, this, mongoStore, invocationContext));
             } else {
                 // Likely applicationRole, but we don't have this application yet
-                result.add(new RoleAdapter(role, mongoStore));
+                result.add(new RoleAdapter(role, mongoStore, invocationContext));
             }
         }
         return result;
@@ -533,20 +555,20 @@ public class RealmAdapter implements RealmModel {
     @Override
     public void deleteRoleMapping(UserModel user, RoleModel role) {
         UserEntity userEntity = ((UserAdapter)user).getUser();
-        mongoStore.pullItemFromList(userEntity, "roleIds", role.getId());
+        mongoStore.pullItemFromList(userEntity, "roleIds", role.getId(), invocationContext);
     }
 
     @Override
     public Set<RoleModel> getScopeMappings(UserModel user) {
         Set<RoleModel> result = new HashSet<RoleModel>();
-        List<RoleEntity> roles = MongoModelUtils.getAllScopesOfUser(user, mongoStore);
+        List<RoleEntity> roles = MongoModelUtils.getAllScopesOfUser(user, mongoStore, invocationContext);
 
         for (RoleEntity role : roles) {
             if (getId().equals(role.getRealmId())) {
-                result.add(new RoleAdapter(role, this, mongoStore));
+                result.add(new RoleAdapter(role, this, mongoStore, invocationContext));
             } else {
                 // Likely applicationRole, but we don't have this application yet
-                result.add(new RoleAdapter(role, mongoStore));
+                result.add(new RoleAdapter(role, mongoStore, invocationContext));
             }
         }
         return result;
@@ -571,13 +593,13 @@ public class RealmAdapter implements RealmModel {
     @Override
     public void addScopeMapping(UserModel agent, RoleModel role) {
         UserEntity userEntity = ((UserAdapter)agent).getUser();
-        mongoStore.pushItemToList(userEntity, "scopeIds", role.getId(), true);
+        mongoStore.pushItemToList(userEntity, "scopeIds", role.getId(), true, invocationContext);
     }
 
     @Override
     public void deleteScopeMapping(UserModel user, RoleModel role) {
         UserEntity userEntity = ((UserAdapter)user).getUser();
-        mongoStore.pullItemFromList(userEntity, "scopeIds", role.getId());
+        mongoStore.pullItemFromList(userEntity, "scopeIds", role.getId(), invocationContext);
     }
 
     @Override
@@ -588,14 +610,14 @@ public class RealmAdapter implements RealmModel {
         oauthClient.setOauthAgentId(oauthAgent.getUser().getId());
         oauthClient.setRealmId(getId());
         oauthClient.setName(name);
-        mongoStore.insertObject(oauthClient);
+        mongoStore.insertObject(oauthClient, invocationContext);
 
-        return new OAuthClientAdapter(oauthClient, oauthAgent, mongoStore);
+        return new OAuthClientAdapter(oauthClient, oauthAgent, mongoStore, invocationContext);
     }
 
     @Override
     public boolean removeOAuthClient(String id) {
-        return mongoStore.removeObject(OAuthClientEntity.class, id);
+        return mongoStore.removeObject(OAuthClientEntity.class, id, invocationContext);
     }
 
     @Override
@@ -606,15 +628,15 @@ public class RealmAdapter implements RealmModel {
                 .and("realmId").is(getId())
                 .and("oauthAgentId").is(user.getUser().getId())
                 .get();
-        OAuthClientEntity oauthClient = mongoStore.loadSingleObject(OAuthClientEntity.class, query);
+        OAuthClientEntity oauthClient = mongoStore.loadSingleObject(OAuthClientEntity.class, query, invocationContext);
-        return oauthClient == null ? null : new OAuthClientAdapter(oauthClient, user, mongoStore);
+        return oauthClient == null ? null : new OAuthClientAdapter(oauthClient, user, mongoStore, invocationContext);
     }
 
     @Override
     public OAuthClientModel getOAuthClientById(String id) {
-        OAuthClientEntity clientEntity = mongoStore.loadObject(OAuthClientEntity.class, id);
+        OAuthClientEntity clientEntity = mongoStore.loadObject(OAuthClientEntity.class, id, invocationContext);
         if (clientEntity == null) return null;
-        return new OAuthClientAdapter(clientEntity, mongoStore);
+        return new OAuthClientAdapter(clientEntity, mongoStore, invocationContext);
     }
 
     @Override
@@ -622,10 +644,10 @@ public class RealmAdapter implements RealmModel {
         DBObject query = new QueryBuilder()
                 .and("realmId").is(getId())
                 .get();
-        List<OAuthClientEntity> results = mongoStore.loadObjects(OAuthClientEntity.class, query);
+        List<OAuthClientEntity> results = mongoStore.loadObjects(OAuthClientEntity.class, query, invocationContext);
         List<OAuthClientModel> list = new ArrayList<OAuthClientModel>();
         for (OAuthClientEntity data : results) {
-            list.add(new OAuthClientAdapter(data, mongoStore));
+            list.add(new OAuthClientAdapter(data, mongoStore, invocationContext));
         }
         return list;
     }
@@ -686,7 +708,7 @@ public class RealmAdapter implements RealmModel {
             }
         }
         for (RequiredCredentialEntity entity : toRemove) {
-            creds.remove(entity);
+            credsEntities.remove(entity);
         }
         for (String cred : creds) {
             logger.info("updating cred: " + cred);
@@ -773,39 +795,31 @@ public class RealmAdapter implements RealmModel {
         }
         credentialEntity.setDevice(cred.getDevice());
 
-        mongoStore.updateObject(userEntity);
+        mongoStore.updateObject(userEntity, invocationContext);
     }
 
     @Override
     public UserModel getUserBySocialLink(SocialLinkModel socialLink) {
         DBObject query = new QueryBuilder()
-                .and("socialProvider").is(socialLink.getSocialProvider())
+                .and("socialLinks.socialProvider").is(socialLink.getSocialProvider())
-                .and("socialUsername").is(socialLink.getSocialUsername())
+                .and("socialLinks.socialUsername").is(socialLink.getSocialUsername())
                 .and("realmId").is(getId())
                 .get();
-        SocialLinkEntity socialLinkEntity = mongoStore.loadSingleObject(SocialLinkEntity.class, query);
+        UserEntity userEntity = mongoStore.loadSingleObject(UserEntity.class, query, invocationContext);
-
+        return userEntity==null ? null : new UserAdapter(userEntity, mongoStore, invocationContext);
-        if (socialLinkEntity == null) {
-            return null;
-        } else {
-            UserEntity userEntity = mongoStore.loadObject(UserEntity.class, socialLinkEntity.getUserId());
-            // TODO: Programmatically remove binding if userEntity doesn't exists? (There are more similar places where this should be handled)
-            return userEntity==null ? null : new UserAdapter(userEntity, mongoStore);
-        }
     }
 
     @Override
     public Set<SocialLinkModel> getSocialLinks(UserModel user) {
         UserEntity userEntity = ((UserAdapter)user).getUser();
-        String userId = userEntity.getId();
+        List<SocialLinkEntity> linkEntities = userEntity.getSocialLinks();
 
-        DBObject query = new QueryBuilder()
+        if (linkEntities == null) {
-                .and("userId").is(userId)
+            return Collections.EMPTY_SET;
-                .get();
+        }
-        List<SocialLinkEntity> dbSocialLinks = mongoStore.loadObjects(SocialLinkEntity.class, query);
 
         Set<SocialLinkModel> result = new HashSet<SocialLinkModel>();
-        for (SocialLinkEntity socialLinkEntity : dbSocialLinks) {
+        for (SocialLinkEntity socialLinkEntity : linkEntities) {
             SocialLinkModel model = new SocialLinkModel(socialLinkEntity.getSocialProvider(), socialLinkEntity.getSocialUsername());
             result.add(model);
         }
@@ -818,26 +832,22 @@ public class RealmAdapter implements RealmModel {
         SocialLinkEntity socialLinkEntity = new SocialLinkEntity();
         socialLinkEntity.setSocialProvider(socialLink.getSocialProvider());
         socialLinkEntity.setSocialUsername(socialLink.getSocialUsername());
-        socialLinkEntity.setUserId(userEntity.getId());
-        socialLinkEntity.setRealmId(getId());
 
-        mongoStore.insertObject(socialLinkEntity);
+        mongoStore.pushItemToList(userEntity, "socialLinks", socialLinkEntity, true, invocationContext);
     }
 
     @Override
     public void removeSocialLink(UserModel user, SocialLinkModel socialLink) {
+        SocialLinkEntity socialLinkEntity = new SocialLinkEntity();
+        socialLinkEntity.setSocialProvider(socialLink.getSocialProvider());
+        socialLinkEntity.setSocialUsername(socialLink.getSocialUsername());
+
         UserEntity userEntity = ((UserAdapter)user).getUser();
-        String userId = userEntity.getId();
+        mongoStore.pullItemFromList(userEntity, "socialLinks", socialLinkEntity, invocationContext);
-        DBObject query = new QueryBuilder()
-                .and("socialProvider").is(socialLink.getSocialProvider())
-                .and("socialUsername").is(socialLink.getSocialUsername())
-                .and("userId").is(userId)
-                .get();
-        mongoStore.removeObjects(SocialLinkEntity.class, query);
     }
 
     protected void updateRealm() {
-        mongoStore.updateObject(realm);
+        mongoStore.updateObject(realm, invocationContext);
     }
 
     protected RequiredCredentialModel initRequiredCredentialModel(String type) {
@@ -853,7 +863,7 @@ public class RealmAdapter implements RealmModel {
         DBObject query = new QueryBuilder()
                 .and("realmId").is(getId())
                 .get();
-        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, query);
+        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, query, invocationContext);
         return convertUserEntities(users);
     }
 
@@ -893,7 +903,7 @@ public class RealmAdapter implements RealmModel {
                 ).get()
         );
 
-        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, builder.get());
+        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, builder.get(), invocationContext);
         return convertUserEntities(users);
     }
 
@@ -915,14 +925,14 @@ public class RealmAdapter implements RealmModel {
                 queryBuilder.and(UserModel.EMAIL).regex(Pattern.compile("(?i:" + entry.getValue() + "$)"));
             }
         }
-        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, queryBuilder.get());
+        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, queryBuilder.get(), invocationContext);
         return convertUserEntities(users);
     }
 
     protected List<UserModel> convertUserEntities(List<UserEntity> userEntities) {
         List<UserModel> userModels = new ArrayList<UserModel>();
         for (UserEntity user : userEntities) {
-            userModels.add(new UserAdapter(user, mongoStore));
+            userModels.add(new UserAdapter(user, mongoStore, invocationContext));
         }
         return userModels;
     }
@@ -950,15 +960,7 @@ public class RealmAdapter implements RealmModel {
     }
 
     @Override
-    public boolean equals(Object o) {
+    public AbstractMongoIdentifiableEntity getMongoEntity() {
-        if (o == null) return false;
+        return realm;
-        if (!(o instanceof RealmAdapter)) return false;
-        RealmAdapter r = (RealmAdapter)o;
-        return r.getId().equals(getId());
-    }
-
-    @Override
-    public int hashCode() {
-        return getId().hashCode();
     }
 }

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

@@ -9,7 +9,9 @@ import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.keycloak.models.RoleContainerModel;
 import org.keycloak.models.RoleModel;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.entities.ApplicationEntity;
 import org.keycloak.models.mongo.keycloak.entities.RealmEntity;
 import org.keycloak.models.mongo.keycloak.entities.RoleEntity;
@@ -21,17 +23,17 @@ import org.keycloak.models.utils.KeycloakModelUtils;
  *
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class RoleAdapter implements RoleModel {
+public class RoleAdapter extends AbstractAdapter implements RoleModel {
 
     private final RoleEntity role;
     private RoleContainerModel roleContainer;
-    private final MongoStore mongoStore;
 
-    public RoleAdapter(RoleEntity roleEntity, MongoStore mongoStore) {
+    public RoleAdapter(RoleEntity roleEntity, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
-        this(roleEntity, null, mongoStore);
+        this(roleEntity, null, mongoStore, invContext);
     }
 
-    public RoleAdapter(RoleEntity roleEntity, RoleContainerModel roleContainer, MongoStore mongoStore) {
+    public RoleAdapter(RoleEntity roleEntity, RoleContainerModel roleContainer, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
+        super(mongoStore, invContext);
         this.role = roleEntity;
         this.roleContainer = roleContainer;
         this.mongoStore = mongoStore;
@@ -66,27 +68,21 @@ public class RoleAdapter implements RoleModel {
 
     @Override
     public boolean isComposite() {
-        return role.isComposite();
+        return role.getCompositeRoleIds() != null && role.getCompositeRoleIds().size() > 0;
-    }
-
-    @Override
-    public void setComposite(boolean flag) {
-        role.setComposite(flag);
-        updateRole();
     }
 
     protected void updateRole() {
-        mongoStore.updateObject(role);
+        mongoStore.updateObject(role, invocationContext);
     }
 
     @Override
     public void addCompositeRole(RoleModel childRole) {
-        mongoStore.pushItemToList(role, "compositeRoleIds", childRole.getId(), true);
+        mongoStore.pushItemToList(role, "compositeRoleIds", childRole.getId(), true, invocationContext);
     }
 
     @Override
     public void removeCompositeRole(RoleModel childRole) {
-        mongoStore.pullItemFromList(role, "compositeRoleIds", childRole.getId());
+        mongoStore.pullItemFromList(role, "compositeRoleIds", childRole.getId(), invocationContext);
     }
 
     @Override
@@ -98,11 +94,11 @@ public class RoleAdapter implements RoleModel {
         DBObject query = new QueryBuilder()
                 .and("_id").in(MongoModelUtils.convertStringsToObjectIds(role.getCompositeRoleIds()))
                 .get();
-        List<RoleEntity> childRoles = mongoStore.loadObjects(RoleEntity.class, query);
+        List<RoleEntity> childRoles = mongoStore.loadObjects(RoleEntity.class, query, invocationContext);
 
         Set<RoleModel> set = new HashSet<RoleModel>();
         for (RoleEntity childRole : childRoles) {
-            set.add(new RoleAdapter(childRole, roleContainer, mongoStore));
+            set.add(new RoleAdapter(childRole, mongoStore, invocationContext));
         }
         return set;
     }
@@ -112,17 +108,17 @@ public class RoleAdapter implements RoleModel {
         if (roleContainer == null) {
             // Compute it
             if (role.getRealmId() != null) {
-                RealmEntity realm = mongoStore.loadObject(RealmEntity.class, role.getRealmId());
+                RealmEntity realm = mongoStore.loadObject(RealmEntity.class, role.getRealmId(), invocationContext);
                 if (realm == null) {
                     throw new IllegalStateException("Realm with id: " + role.getRealmId() + " doesn't exists");
                 }
-                roleContainer = new RealmAdapter(realm, mongoStore);
+                roleContainer = new RealmAdapter(realm, mongoStore, invocationContext);
             } else if (role.getApplicationId() != null) {
-                ApplicationEntity appEntity = mongoStore.loadObject(ApplicationEntity.class, role.getApplicationId());
+                ApplicationEntity appEntity = mongoStore.loadObject(ApplicationEntity.class, role.getApplicationId(), invocationContext);
                 if (appEntity == null) {
                     throw new IllegalStateException("Application with id: " + role.getApplicationId() + " doesn't exists");
                 }
-                roleContainer = new ApplicationAdapter(appEntity, mongoStore);
+                roleContainer = new ApplicationAdapter(appEntity, mongoStore, invocationContext);
             } else {
                 throw new IllegalStateException("Both realmId and applicationId are null for role: " + this);
             }
@@ -144,19 +140,7 @@ public class RoleAdapter implements RoleModel {
     }
 
     @Override
-    public boolean equals(Object o) {
+    public AbstractMongoIdentifiableEntity getMongoEntity() {
-        if (this == o) return true;
+        return role;
-        if (o == null || getClass() != o.getClass()) return false;
-
-        RoleAdapter that = (RoleAdapter) o;
-
-        if (!getId().equals(that.getId())) return false;
-
-        return true;
-    }
-
-    @Override
-    public int hashCode() {
-        return getId().hashCode();
     }
 }

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

@@ -1,8 +1,9 @@
 package org.keycloak.models.mongo.keycloak.adapters;
 
 import org.keycloak.models.UserModel;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoStore;
-import org.keycloak.models.mongo.keycloak.entities.CredentialEntity;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.entities.UserEntity;
 
 import java.util.ArrayList;
@@ -18,14 +19,13 @@ import java.util.Set;
  *
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class UserAdapter implements UserModel {
+public class UserAdapter extends AbstractAdapter implements UserModel {
 
     private final UserEntity user;
-    private final MongoStore mongoStore;
 
-    public UserAdapter(UserEntity userEntity, MongoStore mongoStore) {
+    public UserAdapter(UserEntity userEntity, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
+        super(mongoStore, invContext);
         this.user = userEntity;
-        this.mongoStore = mongoStore;
     }
 
     @Override
@@ -139,12 +139,12 @@ public class UserAdapter implements UserModel {
 
     @Override
     public void addWebOrigin(String webOrigin) {
-        mongoStore.pushItemToList(user, "webOrigins", webOrigin, true);
+        mongoStore.pushItemToList(user, "webOrigins", webOrigin, true, invocationContext);
     }
 
     @Override
     public void removeWebOrigin(String webOrigin) {
-        mongoStore.pullItemFromList(user, "webOrigins", webOrigin);
+        mongoStore.pullItemFromList(user, "webOrigins", webOrigin, invocationContext);
     }
 
     @Override
@@ -166,12 +166,12 @@ public class UserAdapter implements UserModel {
 
     @Override
     public void addRedirectUri(String redirectUri) {
-        mongoStore.pushItemToList(user, "redirectUris", redirectUri, true);
+        mongoStore.pushItemToList(user, "redirectUris", redirectUri, true, invocationContext);
     }
 
     @Override
     public void removeRedirectUri(String redirectUri) {
-        mongoStore.pullItemFromList(user, "redirectUris", redirectUri);
+        mongoStore.pullItemFromList(user, "redirectUris", redirectUri, invocationContext);
     }
 
     @Override
@@ -185,12 +185,12 @@ public class UserAdapter implements UserModel {
 
     @Override
     public void addRequiredAction(RequiredAction action) {
-        mongoStore.pushItemToList(user, "requiredActions", action, true);
+        mongoStore.pushItemToList(user, "requiredActions", action, true, invocationContext);
     }
 
     @Override
     public void removeRequiredAction(RequiredAction action) {
-        mongoStore.pullItemFromList(user, "requiredActions", action);
+        mongoStore.pullItemFromList(user, "requiredActions", action, invocationContext);
     }
 
     @Override
@@ -205,6 +205,11 @@ public class UserAdapter implements UserModel {
     }
 
     protected void updateUser() {
-        mongoStore.updateObject(user);
+        mongoStore.updateObject(user, invocationContext);
+    }
+
+    @Override
+    public AbstractMongoIdentifiableEntity getMongoEntity() {
+        return user;
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/ApplicationEntity.java

@@ -5,19 +5,19 @@ import java.util.List;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @MongoCollection(collectionName = "applications")
-public class ApplicationEntity implements MongoEntity {
+public class ApplicationEntity extends AbstractMongoIdentifiableEntity implements MongoEntity {
 
-    private String id;
     private String name;
     private boolean enabled;
     private boolean surrogateAuthRequired;
@@ -30,15 +30,6 @@ public class ApplicationEntity implements MongoEntity {
     // We are using names of defaultRoles (not ids)
     private List<String> defaultRoles = new ArrayList<String>();
 
-    @MongoId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @MongoField
     public String getName() {
         return name;
@@ -112,14 +103,14 @@ public class ApplicationEntity implements MongoEntity {
     }
 
     @Override
-    public void afterRemove(MongoStore mongoStore) {
+    public void afterRemove(MongoStore mongoStore, MongoStoreInvocationContext invContext) {
         // Remove resourceUser of this application
-        mongoStore.removeObject(UserEntity.class, resourceUserId);
+        mongoStore.removeObject(UserEntity.class, resourceUserId, invContext);
 
         // Remove all roles, which belongs to this application
         DBObject query = new QueryBuilder()
-                .and("applicationId").is(id)
+                .and("applicationId").is(getId())
                 .get();
-        mongoStore.removeObjects(RoleEntity.class, query);
+        mongoStore.removeObjects(RoleEntity.class, query, invContext);
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/CredentialEntity.java

@@ -1,12 +1,12 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
-import org.keycloak.models.mongo.api.AbstractMongoEntity;
+import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class CredentialEntity extends AbstractMongoEntity {
+public class CredentialEntity implements MongoEntity {
 
     private String type;
     private String value;

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/OAuthClientEntity.java

@@ -1,32 +1,23 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @MongoCollection(collectionName = "oauthClients")
-public class OAuthClientEntity implements MongoEntity {
+public class OAuthClientEntity extends AbstractMongoIdentifiableEntity implements MongoEntity {
 
-    private String id;
     private String name;
 
     private String oauthAgentId;
     private String realmId;
 
-    @MongoId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @MongoField
     public String getName() {
         return name;
@@ -55,8 +46,8 @@ public class OAuthClientEntity implements MongoEntity {
     }
 
     @Override
-    public void afterRemove(MongoStore mongoStore) {
+    public void afterRemove(MongoStore mongoStore, MongoStoreInvocationContext invContext) {
         // Remove user of this oauthClient
-        mongoStore.removeObject(UserEntity.class, oauthAgentId);
+        mongoStore.removeObject(UserEntity.class, oauthAgentId, invContext);
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/RealmEntity.java

@@ -2,11 +2,12 @@ package org.keycloak.models.mongo.keycloak.entities;
 
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -17,9 +18,7 @@ import java.util.Map;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @MongoCollection(collectionName = "realms")
-public class RealmEntity implements MongoEntity {
+public class RealmEntity extends AbstractMongoIdentifiableEntity implements MongoEntity {
-
-    private String id;
 
     private String name;
     private boolean enabled;
@@ -38,6 +37,9 @@ public class RealmEntity implements MongoEntity {
     private String publicKeyPem;
     private String privateKeyPem;
 
+    private String loginTheme;
+    private String accountTheme;
+
     // We are using names of defaultRoles (not ids)
     private List<String> defaultRoles = new ArrayList<String>();
 
@@ -48,15 +50,6 @@ public class RealmEntity implements MongoEntity {
     private Map<String, String> smtpConfig = new HashMap<String, String>();
     private Map<String, String> socialConfig = new HashMap<String, String>();
 
-    @MongoId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @MongoField
     public String getName() {
         return name;
@@ -183,6 +176,24 @@ public class RealmEntity implements MongoEntity {
         this.privateKeyPem = privateKeyPem;
     }
 
+    @MongoField
+    public String getLoginTheme() {
+        return loginTheme;
+    }
+
+    public void setLoginTheme(String loginTheme) {
+        this.loginTheme = loginTheme;
+    }
+
+    @MongoField
+    public String getAccountTheme() {
+        return accountTheme;
+    }
+
+    public void setAccountTheme(String accountTheme) {
+        this.accountTheme = accountTheme;
+    }
+
     @MongoField
     public List<String> getDefaultRoles() {
         return defaultRoles;
@@ -238,18 +249,18 @@ public class RealmEntity implements MongoEntity {
     }
 
     @Override
-    public void afterRemove(MongoStore mongoStore) {
+    public void afterRemove(MongoStore mongoStore, MongoStoreInvocationContext invContext) {
         DBObject query = new QueryBuilder()
-                .and("realmId").is(id)
+                .and("realmId").is(getId())
                 .get();
 
         // Remove all users of this realm
-        mongoStore.removeObjects(UserEntity.class, query);
+        mongoStore.removeObjects(UserEntity.class, query, invContext);
 
         // Remove all roles of this realm
-        mongoStore.removeObjects(RoleEntity.class, query);
+        mongoStore.removeObjects(RoleEntity.class, query, invContext);
 
         // Remove all applications of this realm
-        mongoStore.removeObjects(ApplicationEntity.class, query);
+        mongoStore.removeObjects(ApplicationEntity.class, query, invContext);
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/RequiredCredentialEntity.java

@@ -1,12 +1,12 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
-import org.keycloak.models.mongo.api.AbstractMongoEntity;
+import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class RequiredCredentialEntity extends AbstractMongoEntity {
+public class RequiredCredentialEntity implements MongoEntity {
 
     private String type;
     private boolean input;

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/RoleEntity.java

@@ -3,11 +3,12 @@ package org.keycloak.models.mongo.keycloak.entities;
 import com.mongodb.DBObject;
 import com.mongodb.QueryBuilder;
 import org.jboss.logging.Logger;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 
 import java.util.List;
 
@@ -15,29 +16,18 @@ import java.util.List;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @MongoCollection(collectionName = "roles")
-public class RoleEntity implements MongoEntity {
+public class RoleEntity extends AbstractMongoIdentifiableEntity implements MongoEntity {
 
     private static final Logger logger = Logger.getLogger(RoleEntity.class);
 
-    private String id;
     private String name;
     private String description;
-    private boolean composite;
 
     private List<String> compositeRoleIds;
 
     private String realmId;
     private String applicationId;
 
-    @MongoId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @MongoField
     public String getName() {
         return name;
@@ -56,15 +46,6 @@ public class RoleEntity implements MongoEntity {
         this.description = description;
     }
 
-    @MongoField
-    public boolean isComposite() {
-        return composite;
-    }
-
-    public void setComposite(boolean composite) {
-        this.composite = composite;
-    }
-
     @MongoField
     public List<String> getCompositeRoleIds() {
         return compositeRoleIds;
@@ -93,46 +74,46 @@ public class RoleEntity implements MongoEntity {
     }
 
     @Override
-    public void afterRemove(MongoStore mongoStore) {
+    public void afterRemove(MongoStore mongoStore, MongoStoreInvocationContext invContext) {
         // Remove this role from all users, which has it
         DBObject query = new QueryBuilder()
-                .and("roleIds").is(id)
+                .and("roleIds").is(getId())
                 .get();
 
-        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, query);
+        List<UserEntity> users = mongoStore.loadObjects(UserEntity.class, query, invContext);
         for (UserEntity user : users) {
             logger.info("Removing role " + getName() + " from user " + user.getLoginName());
-            mongoStore.pullItemFromList(user, "roleIds", getId());
+            mongoStore.pullItemFromList(user, "roleIds", getId(), invContext);
         }
 
         // Remove this scope from all users, which has it
         query = new QueryBuilder()
-                .and("scopeIds").is(id)
+                .and("scopeIds").is(getId())
                 .get();
 
-        users = mongoStore.loadObjects(UserEntity.class, query);
+        users = mongoStore.loadObjects(UserEntity.class, query, invContext);
         for (UserEntity user : users) {
             logger.info("Removing scope " + getName() + " from user " + user.getLoginName());
-            mongoStore.pullItemFromList(user, "scopeIds", getId());
+            mongoStore.pullItemFromList(user, "scopeIds", getId(), invContext);
         }
 
         // Remove defaultRoles from realm
         if (realmId != null) {
-            RealmEntity realmEntity = mongoStore.loadObject(RealmEntity.class, realmId);
+            RealmEntity realmEntity = mongoStore.loadObject(RealmEntity.class, realmId, invContext);
 
             // Realm might be already removed at this point
             if (realmEntity != null) {
-                mongoStore.pullItemFromList(realmEntity, "defaultRoles", getId());
+                mongoStore.pullItemFromList(realmEntity, "defaultRoles", getId(), invContext);
             }
         }
 
         // Remove defaultRoles from application
         if (applicationId != null) {
-            ApplicationEntity appEntity = mongoStore.loadObject(ApplicationEntity.class, applicationId);
+            ApplicationEntity appEntity = mongoStore.loadObject(ApplicationEntity.class, applicationId, invContext);
 
             // Application might be already removed at this point
             if (appEntity != null) {
-                mongoStore.pullItemFromList(appEntity, "defaultRoles", getId());
+                mongoStore.pullItemFromList(appEntity, "defaultRoles", getId(), invContext);
             }
         }
 
@@ -140,9 +121,9 @@ public class RoleEntity implements MongoEntity {
         query = new QueryBuilder()
                 .and("compositeRoleIds").is(getId())
                 .get();
-        List<RoleEntity> parentRoles = mongoStore.loadObjects(RoleEntity.class, query);
+        List<RoleEntity> parentRoles = mongoStore.loadObjects(RoleEntity.class, query, invContext);
         for (RoleEntity role : parentRoles) {
-            mongoStore.pullItemFromList(role, "compositeRoleIds", getId());
+            mongoStore.pullItemFromList(role, "compositeRoleIds", getId(), invContext);
         }
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/SocialLinkEntity.java

@@ -1,21 +1,15 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
-import org.keycloak.models.mongo.api.AbstractMongoEntity;
+import org.keycloak.models.mongo.api.MongoEntity;
-import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoField;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-@MongoCollection(collectionName = "socialLinks")
+public class SocialLinkEntity implements MongoEntity {
-public class SocialLinkEntity extends AbstractMongoEntity {
 
     private String socialUsername;
     private String socialProvider;
-    private String userId;
-    // realmId is needed to allow searching as combination socialUsername+socialProvider may not be unique
-    // (Same user could have mapped same facebook account to username "foo" in "realm1" and to username "bar" in "realm2")
-    private String realmId;
 
     @MongoField
     public String getSocialUsername() {
@@ -35,21 +29,30 @@ public class SocialLinkEntity extends AbstractMongoEntity {
         this.socialProvider = socialProvider;
     }
 
-    @MongoField
+    @Override
-    public String getUserId() {
+    public boolean equals(Object o) {
-        return userId;
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        SocialLinkEntity that = (SocialLinkEntity) o;
+
+        if (socialProvider != null && (that.socialProvider == null || !socialProvider.equals(that.socialProvider))) return false;
+        if (socialUsername != null && (that.socialUsername == null || !socialUsername.equals(that.socialUsername))) return false;
+        if (socialProvider == null && that.socialProvider != null)return false;
+        if (socialUsername == null && that.socialUsername != null) return false;
+
+        return true;
     }
 
-    public void setUserId(String userId) {
+    @Override
-        this.userId = userId;
+    public int hashCode() {
-    }
+        int code = 1;
-
+        if (socialUsername != null) {
-    @MongoField
+            code = code * 13;
-    public String getRealmId() {
+        }
-        return realmId;
+        if (socialProvider != null) {
-    }
+            code = code * 17;
-
+        }
-    public void setRealmId(String realmId) {
+        return code;
-        this.realmId = realmId;
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/keycloak/entities/UserEntity.java

@@ -1,16 +1,12 @@
 package org.keycloak.models.mongo.keycloak.entities;
 
-import com.mongodb.DBObject;
-import com.mongodb.QueryBuilder;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
-import org.keycloak.models.mongo.api.MongoStore;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -18,9 +14,8 @@ import java.util.Map;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @MongoCollection(collectionName = "users")
-public class UserEntity implements MongoEntity {
+public class UserEntity extends AbstractMongoIdentifiableEntity implements MongoEntity {
 
-    private String id;
     private String loginName;
     private String firstName;
     private String lastName;
@@ -39,15 +34,7 @@ public class UserEntity implements MongoEntity {
     private List<String> redirectUris;
     private List<UserModel.RequiredAction> requiredActions;
     private List<CredentialEntity> credentials = new ArrayList<CredentialEntity>();
-
+    private List<SocialLinkEntity> socialLinks;
-    @MongoId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
 
     @MongoField
     public String getLoginName() {
@@ -184,13 +171,12 @@ public class UserEntity implements MongoEntity {
         this.credentials = credentials;
     }
 
-    @Override
+    @MongoField
-    public void afterRemove(MongoStore mongoStore) {
+    public List<SocialLinkEntity> getSocialLinks() {
-        DBObject query = new QueryBuilder()
+        return socialLinks;
-                .and("userId").is(id)
+    }
-                .get();
 
-        // Remove social links of this user
+    public void setSocialLinks(List<SocialLinkEntity> socialLinks) {
-        mongoStore.removeObjects(SocialLinkEntity.class, query);
+        this.socialLinks = socialLinks;
     }
 }

model/mongo/src/main/java/org/keycloak/models/mongo/utils/MongoModelUtils.java

@@ -10,6 +10,7 @@ import com.mongodb.QueryBuilder;
 import org.bson.types.ObjectId;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.keycloak.adapters.UserAdapter;
 import org.keycloak.models.mongo.keycloak.entities.RoleEntity;
 import org.keycloak.models.mongo.keycloak.entities.UserEntity;
@@ -28,7 +29,7 @@ public class MongoModelUtils {
     }
 
     // Get everything including both application and realm roles
-    public static List<RoleEntity> getAllRolesOfUser(UserModel user, MongoStore mongoStore) {
+    public static List<RoleEntity> getAllRolesOfUser(UserModel user, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
         UserEntity userEntity = ((UserAdapter)user).getUser();
         List<String> roleIds = userEntity.getRoleIds();
 
@@ -39,11 +40,11 @@ public class MongoModelUtils {
         DBObject query = new QueryBuilder()
                 .and("_id").in(convertStringsToObjectIds(roleIds))
                 .get();
-        return mongoStore.loadObjects(RoleEntity.class, query);
+        return mongoStore.loadObjects(RoleEntity.class, query, invContext);
     }
 
     // Get everything including both application and realm scopes
-    public static List<RoleEntity> getAllScopesOfUser(UserModel user, MongoStore mongoStore) {
+    public static List<RoleEntity> getAllScopesOfUser(UserModel user, MongoStore mongoStore, MongoStoreInvocationContext invContext) {
         UserEntity userEntity = ((UserAdapter)user).getUser();
         List<String> scopeIds = userEntity.getScopeIds();
 
@@ -54,6 +55,6 @@ public class MongoModelUtils {
         DBObject query = new QueryBuilder()
                 .and("_id").in(convertStringsToObjectIds(scopeIds))
                 .get();
-        return mongoStore.loadObjects(RoleEntity.class, query);
+        return mongoStore.loadObjects(RoleEntity.class, query, invContext);
     }
 }

model/mongo/src/test/java/org/keycloak/models/mongo/test/Address.java

@@ -1,6 +1,6 @@
 package org.keycloak.models.mongo.test;
 
-import org.keycloak.models.mongo.api.AbstractMongoEntity;
+import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoField;
 
 import java.util.List;
@@ -8,7 +8,7 @@ import java.util.List;
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public class Address extends AbstractMongoEntity {
+public class Address implements MongoEntity {
 
     private String street;
     private int number;

model/mongo/src/test/java/org/keycloak/models/mongo/test/MongoDBModelTest.java

@@ -10,7 +10,10 @@ import org.junit.Before;
 import org.junit.Test;
 import org.keycloak.models.mongo.api.MongoEntity;
 import org.keycloak.models.mongo.api.MongoStore;
+import org.keycloak.models.mongo.api.context.MongoStoreInvocationContext;
 import org.keycloak.models.mongo.impl.MongoStoreImpl;
+import org.keycloak.models.mongo.impl.context.SimpleMongoStoreInvocationContext;
+import org.keycloak.models.mongo.impl.context.TransactionMongoStoreInvocationContext;
 
 import java.net.UnknownHostException;
 import java.util.ArrayList;
@@ -28,7 +31,7 @@ public class MongoDBModelTest {
     };
 
     private MongoClient mongoClient;
-    private MongoStore mongoDB;
+    private MongoStore mongoStore;
 
     @Before
     public void before() throws Exception {
@@ -37,7 +40,7 @@ public class MongoDBModelTest {
             mongoClient = new MongoClient("localhost", 27017);
 
             DB db = mongoClient.getDB("keycloakTest");
-            mongoDB = new MongoStoreImpl(db, true, MANAGED_DATA_TYPES);
+            mongoStore = new MongoStoreImpl(db, true, MANAGED_DATA_TYPES);
 
         } catch (UnknownHostException e) {
             throw new RuntimeException(e);
@@ -51,23 +54,25 @@ public class MongoDBModelTest {
 
     @Test
     public void mongoModelTest() throws Exception {
+        MongoStoreInvocationContext context = new TransactionMongoStoreInvocationContext(mongoStore);
+
         // Add some user
         Person john = new Person();
         john.setFirstName("john");
         john.setAge(25);
         john.setGender(Person.Gender.MALE);
 
-        mongoDB.insertObject(john);
+        mongoStore.insertObject(john, context);
 
         // Add another user
         Person mary = new Person();
         mary.setFirstName("mary");
-        mary.setKids(Arrays.asList("Peter", "Paul", "Wendy"));
+        mary.setKids(asList("Peter", "Paul", "Wendy"));
 
         Address addr1 = new Address();
         addr1.setStreet("Elm");
         addr1.setNumber(5);
-        addr1.setFlatNumbers(Arrays.asList("flat1", "flat2"));
+        addr1.setFlatNumbers(asList("flat1", "flat2"));
         Address addr2 = new Address();
         List<Address> addresses = new ArrayList<Address>();
         addresses.add(addr1);
@@ -76,14 +81,14 @@ public class MongoDBModelTest {
         mary.setAddresses(addresses);
         mary.setMainAddress(addr1);
         mary.setGender(Person.Gender.FEMALE);
-        mary.setGenders(Arrays.asList(Person.Gender.FEMALE));
+        mary.setGenders(asList(Person.Gender.FEMALE));
 
-        mongoDB.insertObject(mary);
+        mongoStore.insertObject(mary, context);
 
-        Assert.assertEquals(2, mongoDB.loadObjects(Person.class, new QueryBuilder().get()).size());
+        Assert.assertEquals(2, mongoStore.loadObjects(Person.class, new QueryBuilder().get(), context).size());
 
         DBObject query = new QueryBuilder().and("addresses.flatNumbers").is("flat1").get();
-        List<Person> persons = mongoDB.loadObjects(Person.class, query);
+        List<Person> persons = mongoStore.loadObjects(Person.class, query, context);
         Assert.assertEquals(1, persons.size());
         mary = persons.get(0);
         Assert.assertEquals(mary.getFirstName(), "mary");
@@ -92,15 +97,15 @@ public class MongoDBModelTest {
         Assert.assertEquals(Address.class, mary.getAddresses().get(0).getClass());
 
         // Test push/pull
-        mongoDB.pushItemToList(mary, "kids", "Pauline", true);
+        mongoStore.pushItemToList(mary, "kids", "Pauline", true, context);
-        mongoDB.pullItemFromList(mary, "kids", "Paul");
+        mongoStore.pullItemFromList(mary, "kids", "Paul", context);
 
         Address addr3 = new Address();
         addr3.setNumber(6);
         addr3.setStreet("Broadway");
-        mongoDB.pushItemToList(mary, "addresses", addr3, true);
+        mongoStore.pushItemToList(mary, "addresses", addr3, true, context);
 
-        mary = mongoDB.loadObject(Person.class, mary.getId());
+        mary = mongoStore.loadObject(Person.class, mary.getId(), context);
         Assert.assertEquals(3, mary.getKids().size());
         Assert.assertTrue(mary.getKids().contains("Pauline"));
         Assert.assertFalse(mary.getKids().contains("Paul"));
@@ -116,18 +121,26 @@ public class MongoDBModelTest {
         mary.addAttribute("attr1", "value1");
         mary.addAttribute("attr2", "value2");
         mary.addAttribute("attr.some3", "value3");
-        mongoDB.updateObject(mary);
+        mongoStore.updateObject(mary, context);
 
-        mary = mongoDB.loadObject(Person.class, mary.getId());
+        mary = mongoStore.loadObject(Person.class, mary.getId(), context);
         Assert.assertEquals(3, mary.getAttributes().size());
 
         mary.removeAttribute("attr2");
         mary.removeAttribute("nonExisting");
-        mongoDB.updateObject(mary);
+        mongoStore.updateObject(mary, context);
 
-        mary = mongoDB.loadObject(Person.class, mary.getId());
+        mary = mongoStore.loadObject(Person.class, mary.getId(), context);
         Assert.assertEquals(2, mary.getAttributes().size());
         Assert.assertEquals("value1", mary.getAttributes().get("attr1"));
         Assert.assertEquals("value3", mary.getAttributes().get("attr.some3"));
+
+        context.commit();
+    }
+
+    private <T> List<T> asList(T... objects) {
+        List<T> list = new ArrayList<T>();
+        list.addAll(Arrays.asList(objects));
+        return list;
     }
 }

model/mongo/src/test/java/org/keycloak/models/mongo/test/Person.java

@@ -1,9 +1,8 @@
 package org.keycloak.models.mongo.test;
 
-import org.keycloak.models.mongo.api.AbstractMongoEntity;
+import org.keycloak.models.mongo.api.AbstractMongoIdentifiableEntity;
 import org.keycloak.models.mongo.api.MongoCollection;
 import org.keycloak.models.mongo.api.MongoField;
-import org.keycloak.models.mongo.api.MongoId;
 
 import java.util.HashMap;
 import java.util.List;
@@ -13,9 +12,8 @@ import java.util.Map;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @MongoCollection(collectionName = "persons")
-public class Person extends AbstractMongoEntity {
+public class Person extends AbstractMongoIdentifiableEntity {
 
-    private String id;
     private String firstName;
     private int age;
     private List<String> kids;
@@ -25,16 +23,6 @@ public class Person extends AbstractMongoEntity {
     private List<Gender> genders;
     private Map<String, String> attributes = new HashMap<String, String>();
 
-
-    @MongoId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @MongoField
     public String getFirstName() {
         return firstName;

services/src/main/java/org/keycloak/services/managers/AuthenticationManager.java

@@ -11,6 +11,7 @@ import org.keycloak.models.RealmModel;
 import org.keycloak.models.RequiredCredentialModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.representations.SkeletonKeyToken;
 import org.keycloak.representations.idm.CredentialRepresentation;
 import org.keycloak.services.resources.AccountService;
@@ -41,7 +42,7 @@ public class AuthenticationManager {
 
     public SkeletonKeyToken createIdentityToken(RealmModel realm, String username) {
         SkeletonKeyToken token = new SkeletonKeyToken();
-        token.id(RealmManager.generateId());
+        token.id(KeycloakModelUtils.generateId());
         token.issuedNow();
         token.principal(username);
         token.audience(realm.getName());

services/src/main/java/org/keycloak/services/managers/RealmManager.java

@@ -13,6 +13,7 @@ import org.keycloak.models.SocialLinkModel;
 import org.keycloak.models.UserCredentialModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserModel.RequiredAction;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.representations.idm.ApplicationRepresentation;
 import org.keycloak.representations.idm.CredentialRepresentation;
 import org.keycloak.representations.idm.OAuthClientRepresentation;
@@ -44,11 +45,6 @@ import java.util.concurrent.atomic.AtomicLong;
  */
 public class RealmManager {
     protected static final Logger logger = Logger.getLogger(RealmManager.class);
-    private static AtomicLong counter = new AtomicLong(1);
-
-    public static String generateId() {
-        return counter.getAndIncrement() + "-" + System.currentTimeMillis();
-    }
 
     protected KeycloakSession identitySession;
 
@@ -73,7 +69,7 @@ public class RealmManager {
     }
 
     public RealmModel createRealm(String id, String name) {
-        if (id == null) id = generateId();
+        if (id == null) id = KeycloakModelUtils.generateId();
         RealmModel realm = identitySession.createRealm(id, name);
         realm.setName(name);
         realm.addRole(Constants.APPLICATION_ROLE);
@@ -166,7 +162,7 @@ public class RealmManager {
     public RealmModel importRealm(RealmRepresentation rep, UserModel realmCreator) {
         String id = rep.getId();
         if (id == null) {
-            id = generateId();
+            id = KeycloakModelUtils.generateId();
         }
         RealmModel realm = createRealm(id, rep.getRealm());
         importRealm(rep, realm);

services/src/main/java/org/keycloak/services/managers/TokenManager.java

@@ -3,9 +3,11 @@ package org.keycloak.services.managers;
 import org.jboss.resteasy.logging.Logger;
 import org.keycloak.jose.jws.JWSBuilder;
 import org.keycloak.models.ApplicationModel;
+import org.keycloak.models.Constants;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
+import org.keycloak.models.utils.KeycloakModelUtils;
 import org.keycloak.representations.SkeletonKeyScope;
 import org.keycloak.representations.SkeletonKeyToken;
 import org.keycloak.util.Base64Url;
@@ -14,6 +16,7 @@ import org.keycloak.util.JsonSerialization;
 import javax.ws.rs.core.MultivaluedMap;
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -132,7 +135,7 @@ public class TokenManager {
 
     protected SkeletonKeyToken initToken(RealmModel realm, UserModel client, UserModel user) {
         SkeletonKeyToken token = new SkeletonKeyToken();
-        token.id(RealmManager.generateId());
+        token.id(KeycloakModelUtils.generateId());
         token.principal(user.getLoginName());
         token.audience(realm.getName());
         token.issuedNow();
@@ -219,7 +222,7 @@ public class TokenManager {
 
     public SkeletonKeyToken createAccessToken(RealmModel realm, UserModel user) {
         SkeletonKeyToken token = new SkeletonKeyToken();
-        token.id(RealmManager.generateId());
+        token.id(KeycloakModelUtils.generateId());
         token.issuedNow();
         token.principal(user.getLoginName());
         token.audience(realm.getName());

services/src/test/java/org/keycloak/services/managers/AuthenticationManagerTest.java

@@ -10,7 +10,7 @@ import org.keycloak.models.UserModel.RequiredAction;
 import org.keycloak.models.utils.TimeBasedOTP;
 import org.keycloak.representations.idm.CredentialRepresentation;
 import org.keycloak.services.managers.AuthenticationManager.AuthenticationStatus;
-import org.keycloak.test.common.AbstractKeycloakTest;
+import org.keycloak.test.AbstractKeycloakTest;
 
 import javax.ws.rs.core.MultivaluedHashMap;
 import javax.ws.rs.core.MultivaluedMap;

services/src/test/java/org/keycloak/test/AbstractKeycloakTest.java

@@ -1,10 +1,8 @@
-package org.keycloak.test.common;
+package org.keycloak.test;
 
 import org.jboss.resteasy.logging.Logger;
 import org.junit.After;
-import org.junit.AfterClass;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.keycloak.models.KeycloakSession;
@@ -17,7 +15,6 @@ import org.keycloak.services.utils.ModelProviderUtils;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.ServiceLoader;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>

services/src/test/java/org/keycloak/test/AdapterTest.java

@@ -17,7 +17,6 @@ import org.keycloak.representations.idm.CredentialRepresentation;
 import org.keycloak.services.managers.ApplianceBootstrap;
 import org.keycloak.services.managers.OAuthClientManager;
 import org.keycloak.services.managers.RealmManager;
-import org.keycloak.test.common.AbstractKeycloakTest;
 
 import java.util.ArrayList;
 import java.util.Arrays;

services/src/test/java/org/keycloak/test/ApplicationModelTest.java

@@ -9,7 +9,6 @@ import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.representations.idm.ApplicationRepresentation;
 import org.keycloak.services.managers.ApplicationManager;
-import org.keycloak.test.common.AbstractKeycloakTest;
 
 import java.util.Iterator;
 import java.util.List;

services/src/test/java/org/keycloak/test/CompositeRolesModelTest.java

@@ -0,0 +1,105 @@
+package org.keycloak.test;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.keycloak.models.ApplicationModel;
+import org.keycloak.models.RealmModel;
+import org.keycloak.models.RoleModel;
+import org.keycloak.models.UserModel;
+import org.keycloak.representations.idm.RealmRepresentation;
+import org.keycloak.services.managers.RealmManager;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class CompositeRolesModelTest extends AbstractKeycloakTest {
+
+    public CompositeRolesModelTest(String providerId) {
+        super(providerId);
+    }
+
+    @Before
+    public void before() throws Exception {
+        super.before();
+        RealmManager manager = realmManager;
+        RealmRepresentation rep = AbstractKeycloakServerTest.loadJson("testcomposites.json");
+        RealmModel realm = manager.createRealm("Test", rep.getRealm());
+        manager.importRealm(rep, realm);
+    }
+
+    @Test
+    public void testAppComposites() {
+        Set<RoleModel> requestedRoles = getRequestedRoles("APP_COMPOSITE_APPLICATION", "APP_COMPOSITE_USER");
+        Assert.assertEquals(2, requestedRoles.size());
+
+        RoleModel expectedRole1 = getRole("APP_ROLE_APPLICATION", "APP_ROLE_1");
+        RoleModel expectedRole2 = getRole("realm", "REALM_ROLE_1");
+
+        assertContains(requestedRoles, expectedRole1);
+        assertContains(requestedRoles, expectedRole2);
+    }
+
+    // TODO: more tests...
+
+    // Same algorithm as in TokenManager.createAccessCode
+    private Set<RoleModel> getRequestedRoles(String applicationName, String username) {
+        Set<RoleModel> requestedRoles = new HashSet<RoleModel>();
+
+        RealmModel realm = realmManager.getRealm("Test");
+        UserModel user = realm.getUser(username);
+        ApplicationModel application = realm.getApplicationByName(applicationName);
+
+        Set<RoleModel> roleMappings = realm.getRoleMappings(user);
+        Set<RoleModel> scopeMappings = realm.getScopeMappings(application.getApplicationUser());
+        Set<RoleModel> appRoles = application.getRoles();
+        if (appRoles != null) scopeMappings.addAll(appRoles);
+
+        for (RoleModel role : roleMappings) {
+            if (role.getContainer().equals(application)) requestedRoles.add(role);
+            for (RoleModel desiredRole : scopeMappings) {
+                Set<RoleModel> visited = new HashSet<RoleModel>();
+                applyScope(role, desiredRole, visited, requestedRoles);
+            }
+        }
+
+        return requestedRoles;
+    }
+
+    private static void applyScope(RoleModel role, RoleModel scope, Set<RoleModel> visited, Set<RoleModel> requested) {
+        if (visited.contains(scope)) return;
+        visited.add(scope);
+        if (role.hasRole(scope)) {
+            requested.add(scope);
+            return;
+        }
+        if (!scope.isComposite()) return;
+
+        for (RoleModel contained : scope.getComposites()) {
+            applyScope(role, contained, visited, requested);
+        }
+    }
+
+    private RoleModel getRole(String appName, String roleName) {
+        RealmModel realm = realmManager.getRealm("Test");
+        if ("realm".equals(appName)) {
+            return realm.getRole(roleName);
+        }  else {
+            return realm.getApplicationByName(appName).getRole(roleName);
+        }
+    }
+
+    private void assertContains(Set<RoleModel> requestedRoles, RoleModel expectedRole) {
+        Assert.assertTrue(requestedRoles.contains(expectedRole));
+
+        // Check if requestedRole has correct role container
+        for (RoleModel role : requestedRoles) {
+            if (role.equals(expectedRole)) {
+                Assert.assertEquals(role.getContainer(), expectedRole.getContainer());
+            }
+        }
+    }
+}

services/src/test/java/org/keycloak/test/ImportTest.java

@@ -5,6 +5,7 @@ import org.junit.FixMethodOrder;
 import org.junit.Test;
 import org.junit.runners.MethodSorters;
 import org.keycloak.models.ApplicationModel;
+import org.keycloak.models.Constants;
 import org.keycloak.models.RealmModel;
 import org.keycloak.models.RequiredCredentialModel;
 import org.keycloak.models.RoleModel;
@@ -12,9 +13,9 @@ import org.keycloak.models.SocialLinkModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.representations.idm.RealmRepresentation;
 import org.keycloak.services.managers.RealmManager;
-import org.keycloak.test.common.AbstractKeycloakTest;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -56,14 +57,68 @@ public class ImportTest extends AbstractKeycloakTest {
         List<ApplicationModel> resources = realm.getApplications();
         Assert.assertEquals(3, resources.size());
 
-        // Test scope relationship
+        // Test applications imported
-        ApplicationModel application = realm.getApplicationNameMap().get("Application");
+        ApplicationModel application = realm.getApplicationByName("Application");
-        UserModel oauthClient = realm.getUser("oauthclient");
+        ApplicationModel otherApp = realm.getApplicationByName("OtherApp");
+        ApplicationModel accountApp = realm.getApplicationByName(Constants.ACCOUNT_APPLICATION);
+        ApplicationModel nonExisting = realm.getApplicationByName("NonExisting");
         Assert.assertNotNull(application);
+        Assert.assertNotNull(otherApp);
+        Assert.assertNull(nonExisting);
+        Map<String, ApplicationModel> apps = realm.getApplicationNameMap();
+        Assert.assertEquals(3, apps.size());
+        Assert.assertTrue(apps.values().contains(application));
+        Assert.assertTrue(apps.values().contains(otherApp));
+        Assert.assertTrue(apps.values().contains(accountApp));
+        realm.getApplications().containsAll(apps.values());
+
+        // Test finding applications by ID
+        Assert.assertNull(realm.getApplicationById("982734"));
+        Assert.assertEquals(application, realm.getApplicationById(application.getId()));
+
+
+        // Test role mappings
+        UserModel admin = realm.getUser("admin");
+        Set<RoleModel> allRoles = realm.getRoleMappings(admin);
+        Assert.assertEquals(5, allRoles.size());
+        Assert.assertTrue(allRoles.contains(realm.getRole("admin")));
+        Assert.assertTrue(allRoles.contains(application.getRole("app-admin")));
+        Assert.assertTrue(allRoles.contains(otherApp.getRole("otherapp-admin")));
+        Assert.assertTrue(allRoles.contains(accountApp.getRole(Constants.ACCOUNT_PROFILE_ROLE)));
+        Assert.assertTrue(allRoles.contains(accountApp.getRole(Constants.ACCOUNT_MANAGE_ROLE)));
+
+        UserModel wburke = realm.getUser("wburke");
+        allRoles = realm.getRoleMappings(wburke);
+        Assert.assertEquals(4, allRoles.size());
+        Assert.assertFalse(allRoles.contains(realm.getRole("admin")));
+        Assert.assertTrue(allRoles.contains(application.getRole("app-user")));
+        Assert.assertTrue(allRoles.contains(otherApp.getRole("otherapp-user")));
+
+        Assert.assertEquals(0, realm.getRealmRoleMappings(wburke).size());
+
+        Set<RoleModel> realmRoles = realm.getRealmRoleMappings(admin);
+        Assert.assertEquals(1, realmRoles.size());
+        Assert.assertEquals("admin", realmRoles.iterator().next().getName());
+
+        Set<RoleModel> appRoles = application.getApplicationRoleMappings(admin);
+        Assert.assertEquals(1, appRoles.size());
+        Assert.assertEquals("app-admin", appRoles.iterator().next().getName());
+
+
+        // Test scope relationship
+        UserModel oauthClient = realm.getUser("oauthclient");
         Assert.assertNotNull(oauthClient);
+        Set<RoleModel> allScopes = realm.getScopeMappings(oauthClient);
+        Assert.assertEquals(2, allScopes.size());
+        Assert.assertTrue(allScopes.contains(realm.getRole("admin")));
+        Assert.assertTrue(allScopes.contains(application.getRole("app-user")));
+
+        Set<RoleModel> realmScopes = realm.getRealmScopeMappings(oauthClient);
+        Assert.assertTrue(realmScopes.contains(realm.getRole("admin")));
+
         Set<RoleModel> appScopes = application.getApplicationScopeMappings(oauthClient);
-        RoleModel appUserRole = application.getRole("user");
+        Assert.assertTrue(appScopes.contains(application.getRole("app-user")));
-        Assert.assertTrue(appScopes.contains(appUserRole));
+
 
         // Test social linking
         UserModel socialUser = realm.getUser("mySocialUser");
@@ -86,6 +141,8 @@ public class ImportTest extends AbstractKeycloakTest {
         Assert.assertEquals(foundSocialUser.getLoginName(), socialUser.getLoginName());
         Assert.assertNull(realm.getUserBySocialLink(new SocialLinkModel("facebook", "not-existing")));
 
+
+
     }
 
     @Test

services/src/test/java/org/keycloak/test/ModelTest.java

@@ -7,7 +7,6 @@ import org.keycloak.models.RealmModel;
 import org.keycloak.models.RoleModel;
 import org.keycloak.representations.idm.RealmRepresentation;
 import org.keycloak.services.managers.ModelToRepresentation;
-import org.keycloak.test.common.AbstractKeycloakTest;
 
 import java.util.HashMap;
 import java.util.Iterator;

services/src/test/java/org/keycloak/test/UserModelTest.java

@@ -7,7 +7,6 @@ import org.keycloak.models.RoleModel;
 import org.keycloak.models.UserModel;
 import org.keycloak.models.UserModel.RequiredAction;
 import org.keycloak.services.managers.RealmManager;
-import org.keycloak.test.common.AbstractKeycloakTest;
 
 import java.util.Iterator;
 import java.util.List;

services/src/test/resources/testcomposites.json

@@ -0,0 +1,231 @@
+{
+    "id": "Test",
+    "realm": "Test",
+    "enabled": true,
+    "tokenLifespan": 600,
+    "accessCodeLifespan": 600,
+    "accessCodeLifespanUserAction": 600,
+    "sslNotRequired": true,
+    "registrationAllowed": true,
+    "resetPasswordAllowed": true,
+    "requiredCredentials": [ "password" ],
+    "requiredApplicationCredentials": [ "password" ],
+    "requiredOAuthClientCredentials": [ "password" ],
+    "smtpServer": {
+        "from": "auto@keycloak.org",
+        "host": "localhost",
+        "port":"3025"
+    },
+    "users" : [
+        {
+            "username" : "REALM_COMPOSITE_1_USER",
+            "enabled": true,
+            "email" : "test-user@localhost",
+            "credentials" : [
+                { "type" : "password",
+                    "value" : "password" }
+            ]
+        },
+        {
+            "username" : "REALM_ROLE_1_USER",
+            "enabled": true,
+            "email" : "test-user@localhost",
+            "credentials" : [
+                { "type" : "password",
+                    "value" : "password" }
+            ]
+        },
+        {
+            "username" : "REALM_APP_COMPOSITE_USER",
+            "enabled": true,
+            "email" : "test-user@localhost",
+            "credentials" : [
+                { "type" : "password",
+                    "value" : "password" }
+            ]
+        },
+        {
+            "username" : "REALM_APP_ROLE_USER",
+            "enabled": true,
+            "email" : "test-user@localhost",
+            "credentials" : [
+                { "type" : "password",
+                    "value" : "password" }
+            ]
+        },
+        {
+            "username" : "APP_COMPOSITE_USER",
+            "enabled": true,
+            "email" : "test-user@localhost",
+            "credentials" : [
+                { "type" : "password",
+                    "value" : "password" }
+            ]
+        }
+    ],
+    "oauthClients" : [
+        {
+            "name" : "third-party",
+            "enabled": true,
+            "credentials" : [
+                { "type" : "password",
+                    "value" : "password" }
+            ]
+        }
+    ],
+    "roleMappings": [
+        {
+            "username": "REALM_COMPOSITE_1_USER",
+            "roles": ["REALM_COMPOSITE_1"]
+        },
+        {
+            "username": "REALM_ROLE_1_USER",
+            "roles": ["REALM_ROLE_1"]
+        },
+        {
+            "username": "REALM_APP_COMPOSITE_USER",
+            "roles": ["REALM_APP_COMPOSITE_ROLE"]
+        },
+        {
+            "username": "APP_COMPOSITE_USER",
+            "roles": ["REALM_APP_COMPOSITE_ROLE", "REALM_COMPOSITE_1"]
+        }
+    ],
+    "scopeMappings": [
+        {
+            "username": "REALM_COMPOSITE_1_APPLICATION",
+            "roles": ["REALM_COMPOSITE_1"]
+        },
+        {
+            "username": "REALM_ROLE_1_APPLICATION",
+            "roles": ["REALM_ROLE_1"]
+        }
+    ],
+    "applications": [
+        {
+            "name": "REALM_COMPOSITE_1_APPLICATION",
+            "enabled": true,
+            "baseUrl": "http://localhost:8081/app",
+            "adminUrl": "http://localhost:8081/app/logout",
+            "credentials": [
+                {
+                    "type": "password",
+                    "value": "password"
+                }
+            ]
+         },
+        {
+            "name": "REALM_ROLE_1_APPLICATION",
+            "enabled": true,
+            "baseUrl": "http://localhost:8081/app",
+            "adminUrl": "http://localhost:8081/app/logout",
+            "credentials": [
+                {
+                    "type": "password",
+                    "value": "password"
+                }
+            ]
+        },
+        {
+            "name": "APP_ROLE_APPLICATION",
+            "enabled": true,
+            "baseUrl": "http://localhost:8081/app",
+            "adminUrl": "http://localhost:8081/app/logout",
+            "credentials": [
+                {
+                    "type": "password",
+                    "value": "password"
+                }
+            ]
+        },
+        {
+            "name": "APP_COMPOSITE_APPLICATION",
+            "enabled": true,
+            "baseUrl": "http://localhost:8081/app",
+            "adminUrl": "http://localhost:8081/app/logout",
+            "credentials": [
+                {
+                    "type": "password",
+                    "value": "password"
+                }
+            ]
+        }
+    ],
+    "roles" : {
+        "realm" : [
+            {
+                "name": "REALM_ROLE_1"
+            },
+            {
+                "name": "REALM_ROLE_2"
+            },
+            {
+                "name": "REALM_ROLE_3"
+            },
+            {
+                "name": "REALM_COMPOSITE_1",
+                "composites": {
+                    "realm": ["REALM_ROLE_1"]
+                }
+            },
+            {
+                "name": "REALM_APP_COMPOSITE_ROLE",
+                "composites": {
+                    "application": {
+                        "APP_ROLE_APPLICATION" :[
+                            "APP_ROLE_1"
+                        ]
+                    }
+                }
+            }
+        ],
+        "application" : {
+            "APP_ROLE_APPLICATION" : [
+                {
+                    "name": "APP_ROLE_1"
+                },
+                {
+                    "name": "APP_ROLE_2"
+                }
+            ],
+            "APP_COMPOSITE_APPLICATION" : [
+                {
+                    "name": "APP_COMPOSITE_ROLE",
+                    "composites": {
+                        "realm" : [
+                            "REALM_ROLE_1",
+                            "REALM_ROLE_2",
+                            "REALM_ROLE_3"
+                        ],
+                        "application": {
+                            "APP_ROLE_APPLICATION" :[
+                                "APP_ROLE_1"
+                            ]
+                        }
+                    }
+                },
+                {
+                    "name": "APP_ROLE_2"
+                }
+            ]
+        }
+
+    },
+
+    "applicationRoleMappings": {
+        "APP_ROLE_APPLICATION": [
+            {
+                "username": "REALM_APP_ROLE_USER",
+                "roles": ["APP_ROLE_2"]
+            }
+        ]
+    },
+    "applicationScopeMappings": {
+        "APP_ROLE_APPLICATION": [
+            {
+                "username": "APP_COMPOSITE_APPLICATION",
+                "roles": ["APP_ROLE_2"]
+            }
+        ]
+    }
+}

services/src/test/resources/testrealm.json

@@ -43,16 +43,6 @@
                 }
             ]
         },
-        {
-            "username": "oauthclient",
-            "enabled": true,
-            "credentials": [
-                {
-                    "type": "password",
-                    "value": "clientpassword"
-                }
-            ]
-        },
         {
             "username": "mySocialUser",
             "enabled": true
@@ -88,6 +78,16 @@
         }
 
     ],
+    "oauthClients" : [
+        {
+            "name" : "oauthclient",
+            "enabled": true,
+            "credentials" : [
+                { "type" : "password",
+                  "value" : "clientpassword" }
+            ]
+        }
+    ],
     "roles" : {
         "realm" : [
             {
@@ -97,18 +97,18 @@
         "application" : {
             "Application" : [
                 {
-                    "name": "admin"
+                    "name": "app-admin"
                 },
                 {
-                    "name": "user"
+                    "name": "app-user"
                 }
             ],
             "OtherApp" : [
                 {
-                    "name": "admin"
+                    "name": "otherapp-admin"
                 },
                 {
-                    "name": "user"
+                    "name": "otherapp-user"
                 }
             ]
         }
@@ -119,25 +119,31 @@
             "roles": ["admin"]
         }
     ],
+    "scopeMappings": [
+        {
+            "username": "oauthclient",
+            "roles": ["admin"]
+        }
+    ],
     "applicationRoleMappings": {
         "Application": [
             {
                 "username": "wburke",
-                "roles": ["user"]
+                "roles": ["app-user"]
             },
             {
                 "username": "admin",
-                "roles": ["admin"]
+                "roles": ["app-admin"]
             }
         ],
         "OtherApp": [
             {
                 "username": "wburke",
-                "roles": ["user"]
+                "roles": ["otherapp-user"]
             },
             {
                 "username": "admin",
-                "roles": ["admin"]
+                "roles": ["otherapp-admin"]
             }
         ]
     },
@@ -145,7 +151,7 @@
         "Application": [
             {
                 "username": "oauthclient",
-                "roles": ["user"]
+                "roles": ["app-user"]
             }
         ]
 

testsuite/integration/pom.xml

@@ -251,10 +251,22 @@
             <groupId>org.seleniumhq.selenium</groupId>
             <artifactId>selenium-chrome-driver</artifactId>
         </dependency>
+
+        <!-- Mongo dependencies specified here and not in mongo profile, just to allow running tests from IDE -->
+        <dependency>
+            <groupId>org.keycloak</groupId>
+            <artifactId>keycloak-model-mongo</artifactId>
+            <version>${project.version}</version>
+        </dependency>
         <dependency>
             <groupId>org.mongodb</groupId>
             <artifactId>mongo-java-driver</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.picketlink</groupId>
+            <artifactId>picketlink-common</artifactId>
+        </dependency>
+
     </dependencies>
     <build>
         <plugins>
@@ -326,5 +338,23 @@
                 </plugins>
             </build>
         </profile>
+
+        <profile>
+            <id>mongo</id>
+            <activation>
+                <property>
+                    <name>keycloak.model</name>
+                    <value>mongo</value>
+                </property>
+            </activation>
+
+            <properties>
+                <keycloak.mongo.host>localhost</keycloak.mongo.host>
+                <keycloak.mongo.port>27017</keycloak.mongo.port>
+                <keycloak.mongo.db>keycloak</keycloak.mongo.db>
+                <keycloak.mongo.clearCollectionsOnStartup>true</keycloak.mongo.clearCollectionsOnStartup>
+            </properties>
+
+        </profile>
     </profiles>
 </project>

testsuite/integration/src/main/java/org/keycloak/testutils/KeycloakServer.java

@@ -273,7 +273,7 @@ public class KeycloakServer {
 
         server.deploy(di);
 
-        factory = KeycloakApplication.createSessionFactory();
+        factory = ((KeycloakApplication)deployment.getApplication()).getFactory();
 
         setupDefaultRealm();