keycloak-memoizeit

Added system properties to support switch between picketlink

9/6/2013 6:36:30 AM

Changes

Details

diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractAttributedNoSQLObject.java b/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractAttributedNoSQLObject.java
index efdddee..a3a6ddb 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractAttributedNoSQLObject.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractAttributedNoSQLObject.java
@@ -7,7 +7,7 @@ import java.util.Map;
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-public abstract class AbstractAttributedNoSQLObject implements AttributedNoSQLObject {
+public abstract class AbstractAttributedNoSQLObject extends AbstractNoSQLObject implements AttributedNoSQLObject {
 
     // Simple hashMap for now (no thread-safe)
     private Map<String, String> attributes = new HashMap<String, String>();
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractNoSQLObject.java b/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractNoSQLObject.java
new file mode 100644
index 0000000..692839b
--- /dev/null
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/AbstractNoSQLObject.java
@@ -0,0 +1,12 @@
+package org.keycloak.services.models.nosql.api;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public abstract class AbstractNoSQLObject implements NoSQLObject {
+
+    @Override
+    public void afterRemove(NoSQL noSQL) {
+        // Empty by default
+    }
+}
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQL.java b/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQL.java
index c3081b0..1b28fd1 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQL.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQL.java
@@ -3,6 +3,7 @@ package org.keycloak.services.models.nosql.api;
 import java.util.List;
 
 import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
+import org.keycloak.services.models.nosql.api.query.NoSQLQueryBuilder;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -26,4 +27,6 @@ public interface NoSQL {
     void removeObject(Class<? extends NoSQLObject> type, String oid);
 
     void removeObjects(Class<? extends NoSQLObject> type, NoSQLQuery query);
+
+    NoSQLQueryBuilder createQueryBuilder();
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLField.java b/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLField.java
index c3e0586..8c5d2a2 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLField.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLField.java
@@ -16,7 +16,5 @@ import static java.lang.annotation.RetentionPolicy.RUNTIME;
 @Retention(RUNTIME)
 public @interface NoSQLField {
 
-    String fieldName() default "";
-
-    // TODO: add lazy loading?
+    // TODO: fieldName add lazy loading?
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLObject.java b/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLObject.java
index 1f430b6..7d31b68 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLObject.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/NoSQLObject.java
@@ -1,9 +1,16 @@
 package org.keycloak.services.models.nosql.api;
 
 /**
- * Just marker interface
+ * Base interface for object, which is persisted in NoSQL database
  *
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 public interface NoSQLObject {
+
+    /**
+     * Lifecycle callback, which is called after removal of this object from NoSQL database.
+     * It may be useful for triggering removal of wired objects.
+     */
+    void afterRemove(NoSQL noSQL);
+
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/query/NoSQLQueryBuilder.java b/services/src/main/java/org/keycloak/services/models/nosql/api/query/NoSQLQueryBuilder.java
index aa4eb48..341a115 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/query/NoSQLQueryBuilder.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/query/NoSQLQueryBuilder.java
@@ -12,14 +12,6 @@ public abstract class NoSQLQueryBuilder {
 
     protected NoSQLQueryBuilder() {};
 
-    public static NoSQLQueryBuilder create(Class<? extends NoSQLQueryBuilder> builderClass) {
-        try {
-            return builderClass.newInstance();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     public NoSQLQuery build() {
         return new NoSQLQuery(queryAttributes);
     }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/types/ConverterKey.java b/services/src/main/java/org/keycloak/services/models/nosql/api/types/ConverterKey.java
index f4b6fe0..68fc3cb 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/types/ConverterKey.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/types/ConverterKey.java
@@ -27,4 +27,5 @@ class ConverterKey {
         ConverterKey tc = (ConverterKey)obj;
         return tc.applicationObjectType.equals(this.applicationObjectType) && tc.dbObjectType.equals(this.dbObjectType);
     }
+
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/api/types/TypeConverter.java b/services/src/main/java/org/keycloak/services/models/nosql/api/types/TypeConverter.java
index 7bdb384..3368e4f 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/api/types/TypeConverter.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/api/types/TypeConverter.java
@@ -11,6 +11,7 @@ import java.util.Map;
  */
 public class TypeConverter {
 
+    // TODO: Thread-safety support (maybe...)
     private Map<ConverterKey, Converter<?, ?>> converterRegistry = new HashMap<ConverterKey, Converter<?, ?>>();
 
     public <T, S> void addConverter(Converter<T, S> converter) {
@@ -18,14 +19,19 @@ public class TypeConverter {
         converterRegistry.put(converterKey, converter);
     }
 
-    public <T, S> T convertDBObjectToApplicationObject(S dbObject, Class<T> expectedApplicationObjectType, Class<S> expectedDBObjectType) {
+    public <T, S> T convertDBObjectToApplicationObject(S dbObject, Class<T> expectedApplicationObjectType) {
+        // TODO: Not type safe as it expects that S type of converter must exactly match type of dbObject. Converter lookup should be more flexible
+        Class<S> expectedDBObjectType = (Class<S>)dbObject.getClass();
         Converter<T, S> converter = getConverter(expectedApplicationObjectType, expectedDBObjectType);
         return converter.convertDBObjectToApplicationObject(dbObject);
     }
 
-    public <T, S> S convertApplicationObjectToDBObject(T applicationobject, Class<T> expectedApplicationObjectType, Class<S> expectedDBObjectType) {
+    public <T, S> S convertApplicationObjectToDBObject(T applicationObject, Class<S> expectedDBObjectType) {
+        // TODO: Not type safe as it expects that T type of converter must exactly match type of applicationObject. Converter lookup should be more flexible
+        Class<T> expectedApplicationObjectType = (Class<T>)applicationObject.getClass();
         Converter<T, S> converter = getConverter(expectedApplicationObjectType, expectedDBObjectType);
-        return converter.convertApplicationObjectToDBObject(applicationobject);
+
+        return converter.convertApplicationObjectToDBObject(applicationObject);
     }
 
     private <T, S> Converter<T, S> getConverter( Class<T> expectedApplicationObjectType, Class<S> expectedDBObjectType) {
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBImpl.java b/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBImpl.java
index fd6efa7..0385aee 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBImpl.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBImpl.java
@@ -12,20 +12,21 @@ import com.mongodb.DBCollection;
 import com.mongodb.DBCursor;
 import com.mongodb.DBObject;
 import org.bson.types.ObjectId;
-import org.keycloak.services.models.nosql.api.AttributedNoSQLObject;
+import org.jboss.resteasy.logging.Logger;
 import org.keycloak.services.models.nosql.api.NoSQL;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
 import org.keycloak.services.models.nosql.api.NoSQLObject;
 import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
+import org.keycloak.services.models.nosql.api.query.NoSQLQueryBuilder;
 import org.keycloak.services.models.nosql.api.types.Converter;
 import org.keycloak.services.models.nosql.api.types.TypeConverter;
 import org.keycloak.services.models.nosql.impl.types.BasicDBListToStringArrayConverter;
+import org.keycloak.services.models.nosql.impl.types.NoSQLObjectConverter;
 import org.picketlink.common.properties.Property;
 import org.picketlink.common.properties.query.AnnotatedPropertyCriteria;
 import org.picketlink.common.properties.query.PropertyQueries;
-import org.picketlink.common.reflection.Types;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -33,59 +34,61 @@ import org.picketlink.common.reflection.Types;
 public class MongoDBImpl implements NoSQL {
 
     private final DB database;
-    // private static final Logger logger = Logger.getLogger(MongoDBImpl.class);
+    private static final Logger logger = Logger.getLogger(MongoDBImpl.class);
 
     private final TypeConverter typeConverter;
+    private ConcurrentMap<Class<? extends NoSQLObject>, ObjectInfo> objectInfoCache =
+            new ConcurrentHashMap<Class<? extends NoSQLObject>, ObjectInfo>();
 
-    public MongoDBImpl(DB database) {
+    public MongoDBImpl(DB database, boolean removeAllObjectsAtStartup, Class<? extends NoSQLObject>[] managedDataTypes) {
         this.database = database;
 
         typeConverter = new TypeConverter();
         typeConverter.addConverter(new BasicDBListToStringArrayConverter());
-    }
-
-    private ConcurrentMap<Class<? extends NoSQLObject>, ObjectInfo<? extends NoSQLObject>> objectInfoCache =
-            new ConcurrentHashMap<Class<? extends NoSQLObject>, ObjectInfo<? extends NoSQLObject>>();
-
+        for (Class<? extends NoSQLObject> type : managedDataTypes) {
+            typeConverter.addConverter(new NoSQLObjectConverter(this, typeConverter, type));
+            getObjectInfo(type);
+        }
 
+        if (removeAllObjectsAtStartup) {
+            for (Class<? extends NoSQLObject> type : managedDataTypes) {
+                ObjectInfo objectInfo = getObjectInfo(type);
+                String collectionName = objectInfo.getDbCollectionName();
+                if (collectionName != null) {
+                    logger.debug("Removing all objects of type " + type);
+
+                    DBCollection dbCollection = this.database.getCollection(collectionName);
+                    dbCollection.remove(new BasicDBObject());
+                }  else {
+                    logger.debug("Skip removing objects of type " + type + " as it doesn't have it's own collection");
+                }
+            }
+            logger.info("All objects successfully removed from MongoDB");
+        }
+    }
 
     @Override
     public void saveObject(NoSQLObject object) {
-        Class<?> clazz = object.getClass();
+        Class<? extends NoSQLObject> clazz = object.getClass();
 
         // Find annotations for ID, for all the properties and for the name of the collection.
         ObjectInfo objectInfo = getObjectInfo(clazz);
 
         // Create instance of BasicDBObject and add all declared properties to it (properties with null value probably should be skipped)
-        BasicDBObject dbObject = new BasicDBObject();
-        List<Property<Object>> props = objectInfo.getProperties();
-        for (Property<Object> property : props) {
-            String propName = property.getName();
-            Object propValue = property.getValue(object);
-
-
-            dbObject.append(propName, propValue);
-        }
-
-        // Adding attributes
-        if (object instanceof AttributedNoSQLObject) {
-            AttributedNoSQLObject attributedObject = (AttributedNoSQLObject)object;
-            Map<String, String> attributes = attributedObject.getAttributes();
-            for (Map.Entry<String, String> attribute : attributes.entrySet()) {
-                dbObject.append(attribute.getKey(), attribute.getValue());
-            }
-        }
+        BasicDBObject dbObject = typeConverter.convertApplicationObjectToDBObject(object, BasicDBObject.class);
 
         DBCollection dbCollection = database.getCollection(objectInfo.getDbCollectionName());
 
         // Decide if we should insert or update (based on presence of oid property in original object)
         Property<String> oidProperty = objectInfo.getOidProperty();
-        String currentId = oidProperty.getValue(object);
+        String currentId = oidProperty == null ? null : oidProperty.getValue(object);
         if (currentId == null) {
             dbCollection.insert(dbObject);
 
             // Add oid to value of given object
-            oidProperty.setValue(object, dbObject.getString("_id"));
+            if (oidProperty != null) {
+                oidProperty.setValue(object, dbObject.getString("_id"));
+            }
         } else {
             BasicDBObject setCommand = new BasicDBObject("$set", dbObject);
             BasicDBObject query = new BasicDBObject("_id", new ObjectId(currentId));
@@ -100,7 +103,7 @@ public class MongoDBImpl implements NoSQL {
         BasicDBObject idQuery = new BasicDBObject("_id", new ObjectId(oid));
         DBObject dbObject = dbCollection.findOne(idQuery);
 
-        return convertObject(type, dbObject);
+        return typeConverter.convertDBObjectToApplicationObject(dbObject, type);
     }
 
     @Override
@@ -129,7 +132,7 @@ public class MongoDBImpl implements NoSQL {
     @Override
     public void removeObject(NoSQLObject object) {
         Class<? extends NoSQLObject> type = object.getClass();
-        ObjectInfo<?> objectInfo = getObjectInfo(type);
+        ObjectInfo objectInfo = getObjectInfo(type);
 
         Property<String> idProperty = objectInfo.getOidProperty();
         String oid = idProperty.getValue(object);
@@ -139,18 +142,39 @@ public class MongoDBImpl implements NoSQL {
 
     @Override
     public void removeObject(Class<? extends NoSQLObject> type, String oid) {
-        DBCollection dbCollection = getDBCollectionForType(type);
+        NoSQLObject found = loadObject(type, oid);
+        if (found == null) {
+            logger.warn("Object of type: " + type + ", oid: " + oid + " doesn't exist in MongoDB. Skip removal");
+        } else {
+            DBCollection dbCollection = getDBCollectionForType(type);
+            BasicDBObject dbQuery = new BasicDBObject("_id", new ObjectId(oid));
+            dbCollection.remove(dbQuery);
+            logger.info("Object of type: " + type + ", oid: " + oid + " removed from MongoDB.");
 
-        BasicDBObject dbQuery = new BasicDBObject("_id", new ObjectId(oid));
-        dbCollection.remove(dbQuery);
+            found.afterRemove(this);
+        }
     }
 
     @Override
     public void removeObjects(Class<? extends NoSQLObject> type, NoSQLQuery query) {
-        DBCollection dbCollection = getDBCollectionForType(type);
-        BasicDBObject dbQuery = getDBQueryFromQuery(query);
+        List<? extends NoSQLObject> foundObjects = loadObjects(type, query);
+        if (foundObjects.size() == 0) {
+            logger.info("Not found any objects of type: " + type + ", query: " + query);
+        } else {
+            DBCollection dbCollection = getDBCollectionForType(type);
+            BasicDBObject dbQuery = getDBQueryFromQuery(query);
+            dbCollection.remove(dbQuery);
+            logger.info("Removed " + foundObjects.size() + " objects of type: " + type + ", query: " + query);
+
+            for (NoSQLObject found : foundObjects) {
+                found.afterRemove(this);
+            }
+        }
+    }
 
-        dbCollection.remove(dbQuery);
+    @Override
+    public NoSQLQueryBuilder createQueryBuilder() {
+        return new MongoDBQueryBuilder();
     }
 
     // Possibility to add user-defined converters
@@ -158,26 +182,19 @@ public class MongoDBImpl implements NoSQL {
         typeConverter.addConverter(converter);
     }
 
-    private <T extends NoSQLObject> ObjectInfo<T> getObjectInfo(Class<?> objectClass) {
-        ObjectInfo<T> objectInfo = (ObjectInfo<T>)objectInfoCache.get(objectClass);
+    public ObjectInfo getObjectInfo(Class<? extends NoSQLObject> objectClass) {
+        ObjectInfo objectInfo = objectInfoCache.get(objectClass);
         if (objectInfo == null) {
             Property<String> idProperty = PropertyQueries.<String>createQuery(objectClass).addCriteria(new AnnotatedPropertyCriteria(NoSQLId.class)).getFirstResult();
-            if (idProperty == null) {
-                // TODO: should be allowed to have NoSQLObject classes without declared NoSQLId annotation?
-                throw new IllegalStateException("Class " + objectClass + " doesn't have property with declared annotation " + NoSQLId.class);
-            }
 
             List<Property<Object>> properties = PropertyQueries.createQuery(objectClass).addCriteria(new AnnotatedPropertyCriteria(NoSQLField.class)).getResultList();
 
             NoSQLCollection classAnnotation = objectClass.getAnnotation(NoSQLCollection.class);
-            if (classAnnotation == null) {
-                throw new IllegalStateException("Class " + objectClass + " doesn't have annotation " + NoSQLCollection.class);
-            }
 
-            String dbCollectionName = classAnnotation.collectionName();
-            objectInfo = new ObjectInfo<T>((Class<T>)objectClass, dbCollectionName, idProperty, properties);
+            String dbCollectionName = classAnnotation==null ? null : classAnnotation.collectionName();
+            objectInfo = new ObjectInfo(objectClass, dbCollectionName, idProperty, properties);
 
-            ObjectInfo existing = objectInfoCache.putIfAbsent((Class<T>)objectClass, objectInfo);
+            ObjectInfo existing = objectInfoCache.putIfAbsent(objectClass, objectInfo);
             if (existing != null) {
                 objectInfo = existing;
             }
@@ -186,75 +203,23 @@ public class MongoDBImpl implements NoSQL {
         return objectInfo;
     }
 
-
-    private <T extends NoSQLObject> T convertObject(Class<T> type, DBObject dbObject) {
-        if (dbObject == null) {
-            return null;
-        }
-
-        ObjectInfo<T> objectInfo = getObjectInfo(type);
-
-        T object;
-        try {
-            object = type.newInstance();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-
-        for (String key : dbObject.keySet()) {
-            Object value = dbObject.get(key);
-            Property<Object> property;
-
-            if ("_id".equals(key)) {
-                // Current property is "id"
-                Property<String> idProperty = objectInfo.getOidProperty();
-                idProperty.setValue(object, value.toString());
-
-            } else if ((property = objectInfo.getPropertyByName(key)) != null) {
-                // It's declared property with @DBField annotation
-                Class<?> expectedType = property.getJavaClass();
-                Class actualType = value != null ? value.getClass() : expectedType;
-
-                // handle primitives
-                expectedType = Types.boxedClass(expectedType);
-                actualType = Types.boxedClass(actualType);
-
-                if (actualType.isAssignableFrom(expectedType)) {
-                    property.setValue(object, value);
-                } else {
-                    // we need to convert
-                    Object convertedValue = typeConverter.convertDBObjectToApplicationObject(value, expectedType, actualType);
-                    property.setValue(object, convertedValue);
-                }
-
-            } else if (object instanceof AttributedNoSQLObject) {
-                // It's attributed object and property is not declared, so we will call setAttribute
-                ((AttributedNoSQLObject)object).setAttribute(key, value.toString());
-
-            } else {
-                // Show warning if it's unknown
-                // TODO: logging
-                // logger.warn("Property with key " + key + " not known for type " + type);
-                System.err.println("Property with key " + key + " not known for type " + type);
-            }
-        }
-
-        return object;
-    }
-
     private <T extends NoSQLObject> List<T> convertCursor(Class<T> type, DBCursor cursor) {
         List<T> result = new ArrayList<T>();
 
-        for (DBObject dbObject : cursor) {
-            T converted = convertObject(type, dbObject);
-            result.add(converted);
+        try {
+            for (DBObject dbObject : cursor) {
+                T converted = typeConverter.convertDBObjectToApplicationObject(dbObject, type);
+                result.add(converted);
+            }
+        } finally {
+            cursor.close();
         }
 
         return result;
     }
 
     private DBCollection getDBCollectionForType(Class<? extends NoSQLObject> type) {
-        ObjectInfo<?> objectInfo = getObjectInfo(type);
+        ObjectInfo objectInfo = getObjectInfo(type);
         return database.getCollection(objectInfo.getDbCollectionName());
     }
 
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBQueryBuilder.java b/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBQueryBuilder.java
index b3af732..94b6196 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBQueryBuilder.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/impl/MongoDBQueryBuilder.java
@@ -9,6 +9,8 @@ import org.keycloak.services.models.nosql.api.query.NoSQLQueryBuilder;
  */
 public class MongoDBQueryBuilder extends NoSQLQueryBuilder {
 
+    protected MongoDBQueryBuilder() {};
+
     @Override
     public NoSQLQueryBuilder inCondition(String name, Object[] values) {
         if (values == null) {
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/impl/ObjectInfo.java b/services/src/main/java/org/keycloak/services/models/nosql/impl/ObjectInfo.java
index 867ac12..27f782c 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/impl/ObjectInfo.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/impl/ObjectInfo.java
@@ -8,9 +8,9 @@ import org.picketlink.common.properties.Property;
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
-class ObjectInfo<T extends NoSQLObject> {
+public class ObjectInfo {
 
-    private final Class<T> objectClass;
+    private final Class<? extends NoSQLObject> objectClass;
 
     private final String dbCollectionName;
 
@@ -18,14 +18,14 @@ class ObjectInfo<T extends NoSQLObject> {
 
     private final List<Property<Object>> properties;
 
-    public ObjectInfo(Class<T> objectClass, String dbCollectionName, Property<String> oidProperty, List<Property<Object>> properties) {
+    public ObjectInfo(Class<? extends NoSQLObject> objectClass, String dbCollectionName, Property<String> oidProperty, List<Property<Object>> properties) {
         this.objectClass = objectClass;
         this.dbCollectionName = dbCollectionName;
         this.oidProperty = oidProperty;
         this.properties = properties;
     }
 
-    public Class<T> getObjectClass() {
+    public Class<? extends NoSQLObject> getObjectClass() {
         return objectClass;
     }
 
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/impl/types/BasicDBListToStringArrayConverter.java b/services/src/main/java/org/keycloak/services/models/nosql/impl/types/BasicDBListToStringArrayConverter.java
index 0f9d5d5..0c3f176 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/impl/types/BasicDBListToStringArrayConverter.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/impl/types/BasicDBListToStringArrayConverter.java
@@ -13,7 +13,7 @@ public class BasicDBListToStringArrayConverter implements Converter<Object, Basi
     private static final String[] PLACEHOLDER = new String[] {};
 
     @Override
-    public Object convertDBObjectToApplicationObject(BasicDBList dbObject) {
+    public String[] convertDBObjectToApplicationObject(BasicDBList dbObject) {
         return dbObject.toArray(PLACEHOLDER);
     }
 
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/impl/types/NoSQLObjectConverter.java b/services/src/main/java/org/keycloak/services/models/nosql/impl/types/NoSQLObjectConverter.java
new file mode 100644
index 0000000..a3974b3
--- /dev/null
+++ b/services/src/main/java/org/keycloak/services/models/nosql/impl/types/NoSQLObjectConverter.java
@@ -0,0 +1,133 @@
+package org.keycloak.services.models.nosql.impl.types;
+
+import java.util.List;
+import java.util.Map;
+
+import com.mongodb.BasicDBObject;
+import org.keycloak.services.models.nosql.api.AttributedNoSQLObject;
+import org.keycloak.services.models.nosql.api.NoSQLObject;
+import org.keycloak.services.models.nosql.api.types.Converter;
+import org.keycloak.services.models.nosql.api.types.TypeConverter;
+import org.keycloak.services.models.nosql.impl.MongoDBImpl;
+import org.keycloak.services.models.nosql.impl.ObjectInfo;
+import org.picketlink.common.properties.Property;
+import org.picketlink.common.reflection.Types;
+
+/**
+ * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
+ */
+public class NoSQLObjectConverter<T extends NoSQLObject> implements Converter<T, BasicDBObject> {
+
+    private final MongoDBImpl mongoDBImpl;
+    private final TypeConverter typeConverter;
+    private final Class<T> expectedNoSQLObjectType;
+
+    public NoSQLObjectConverter(MongoDBImpl mongoDBImpl, TypeConverter typeConverter, Class<T> expectedNoSQLObjectType) {
+        this.mongoDBImpl = mongoDBImpl;
+        this.typeConverter = typeConverter;
+        this.expectedNoSQLObjectType = expectedNoSQLObjectType;
+    }
+
+    @Override
+    public T convertDBObjectToApplicationObject(BasicDBObject dbObject) {
+        if (dbObject == null) {
+            return null;
+        }
+
+        ObjectInfo objectInfo = mongoDBImpl.getObjectInfo(expectedNoSQLObjectType);
+
+        T object;
+        try {
+            object = expectedNoSQLObjectType.newInstance();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        for (String key : dbObject.keySet()) {
+            Object value = dbObject.get(key);
+            Property<Object> property;
+
+            if ("_id".equals(key)) {
+                // Current property is "id"
+                Property<String> idProperty = objectInfo.getOidProperty();
+                if (idProperty != null) {
+                    idProperty.setValue(object, value.toString());
+                }
+
+            } else if ((property = objectInfo.getPropertyByName(key)) != null) {
+                // It's declared property with @DBField annotation
+                setPropertyValue(object, value, property);
+
+            } else if (object instanceof AttributedNoSQLObject) {
+                // It's attributed object and property is not declared, so we will call setAttribute
+                ((AttributedNoSQLObject)object).setAttribute(key, value.toString());
+
+            } else {
+                // Show warning if it's unknown
+                // TODO: logging
+                // logger.warn("Property with key " + key + " not known for type " + type);
+                System.err.println("Property with key " + key + " not known for type " + expectedNoSQLObjectType);
+            }
+        }
+
+        return object;
+    }
+
+    private void setPropertyValue(NoSQLObject object, Object valueFromDB, Property property) {
+        Class<?> expectedType = property.getJavaClass();
+        Class actualType = valueFromDB != null ? valueFromDB.getClass() : expectedType;
+
+        // handle primitives
+        expectedType = Types.boxedClass(expectedType);
+        actualType = Types.boxedClass(actualType);
+
+        if (actualType.isAssignableFrom(expectedType)) {
+            property.setValue(object, valueFromDB);
+        } else {
+            // we need to convert
+            Object convertedValue = typeConverter.convertDBObjectToApplicationObject(valueFromDB, expectedType);
+            property.setValue(object, convertedValue);
+        }
+    }
+
+    @Override
+    public BasicDBObject convertApplicationObjectToDBObject(T applicationObject) {
+        ObjectInfo objectInfo = mongoDBImpl.getObjectInfo(applicationObject.getClass());
+
+        // Create instance of BasicDBObject and add all declared properties to it (properties with null value probably should be skipped)
+        BasicDBObject dbObject = new BasicDBObject();
+        List<Property<Object>> props = objectInfo.getProperties();
+        for (Property<Object> property : props) {
+            String propName = property.getName();
+            Object propValue = property.getValue(applicationObject);
+
+            // Check if we have noSQLObject, which is indication that we need to convert recursively
+            if (propValue instanceof NoSQLObject) {
+                propValue = typeConverter.convertApplicationObjectToDBObject(propValue, BasicDBObject.class);
+            }
+
+            dbObject.append(propName, propValue);
+        }
+
+        // Adding attributes
+        if (applicationObject instanceof AttributedNoSQLObject) {
+            AttributedNoSQLObject attributedObject = (AttributedNoSQLObject)applicationObject;
+            Map<String, String> attributes = attributedObject.getAttributes();
+            for (Map.Entry<String, String> attribute : attributes.entrySet()) {
+                dbObject.append(attribute.getKey(), attribute.getValue());
+            }
+        }
+
+        return dbObject;
+    }
+
+    @Override
+    public Class<T> getApplicationObjectType() {
+        return expectedNoSQLObjectType;
+    }
+
+    @Override
+    public Class<BasicDBObject> getDBObjectType() {
+        return BasicDBObject.class;
+    }
+}
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/impl/Utils.java b/services/src/main/java/org/keycloak/services/models/nosql/impl/Utils.java
index 1398e27..26ab30a 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/impl/Utils.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/impl/Utils.java
@@ -53,4 +53,23 @@ public class Utils {
         return false;
     }
 
+    public static <T> T[] removeItemFromArray(T[] inputArray, T item) {
+        if (item == null) {
+            throw new IllegalArgumentException("item must be non-null");
+        }
+
+        if (inputArray == null) {
+            return inputArray;
+        } else {
+            T[] outputArray = (T[])Array.newInstance(item.getClass(), inputArray.length - 1);
+            int counter = 0;
+            for (T object : inputArray) {
+                if (!item.equals(object)) {
+                    outputArray[counter++] = object;
+                }
+            }
+            return outputArray;
+        }
+    }
+
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/ApplicationAdapter.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/ApplicationAdapter.java
index fc8a6e1..dd10400 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/ApplicationAdapter.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/ApplicationAdapter.java
@@ -94,7 +94,7 @@ public class ApplicationAdapter implements ApplicationModel {
 
     @Override
     public RoleAdapter getRole(String name) {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("name", name)
                 .andCondition("applicationId", getId())
                 .build();
@@ -122,7 +122,7 @@ public class ApplicationAdapter implements ApplicationModel {
 
     @Override
     public List<RoleModel> getRoles() {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("applicationId", getId())
                 .build();
         List<RoleData> roles = noSQL.loadObjects(RoleData.class, query);
@@ -142,7 +142,7 @@ public class ApplicationAdapter implements ApplicationModel {
 
         Set<String> result = new HashSet<String>();
 
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .inCondition("_id", roleIds)
                 .build();
         List<RoleData> roles = noSQL.loadObjects(RoleData.class, query);
@@ -184,7 +184,7 @@ public class ApplicationAdapter implements ApplicationModel {
 
         Set<String> result = new HashSet<String>();
 
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .inCondition("_id", scopeIds)
                 .build();
         List<RoleData> roles = noSQL.loadObjects(RoleData.class, query);
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java
index a0ff435..114f7b7 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/MongoDBSessionFactory.java
@@ -4,6 +4,8 @@ import java.net.UnknownHostException;
 
 import com.mongodb.DB;
 import com.mongodb.MongoClient;
+import org.jboss.resteasy.logging.Logger;
+import org.keycloak.services.managers.RealmManager;
 import org.keycloak.services.models.KeycloakSession;
 import org.keycloak.services.models.KeycloakSessionFactory;
 import org.keycloak.services.models.nosql.api.NoSQL;
@@ -26,8 +28,9 @@ import org.keycloak.services.models.nosql.keycloak.data.credentials.PasswordData
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 public class MongoDBSessionFactory implements KeycloakSessionFactory {
+    protected static final Logger logger = Logger.getLogger(MongoDBSessionFactory.class);
 
-    private static final Class<?>[] MANAGED_DATA_TYPES = {
+    private static final Class<? extends NoSQLObject>[] MANAGED_DATA_TYPES = (Class<? extends NoSQLObject>[])new Class<?>[] {
             RealmData.class,
             UserData.class,
             RoleData.class,
@@ -41,25 +44,17 @@ public class MongoDBSessionFactory implements KeycloakSessionFactory {
     private final NoSQL mongoDB;
 
     public MongoDBSessionFactory(String host, int port, String dbName, boolean removeAllObjectsAtStartup) {
+        logger.info(String.format("Going to use MongoDB database. host: %s, port: %d, databaseName: %s, removeAllObjectsAtStartup: %b", host, port, dbName, removeAllObjectsAtStartup));
         try {
             // TODO: authentication support
             mongoClient = new MongoClient(host, port);
 
             DB db = mongoClient.getDB(dbName);
-            mongoDB = new MongoDBImpl(db);
+            mongoDB = new MongoDBImpl(db, removeAllObjectsAtStartup, MANAGED_DATA_TYPES);
 
         } catch (UnknownHostException e) {
             throw new RuntimeException(e);
         }
-
-        if (removeAllObjectsAtStartup) {
-            NoSQLQuery emptyQuery = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class).build();
-            for (Class<?> type : MANAGED_DATA_TYPES) {
-                mongoDB.removeObjects((Class<? extends NoSQLObject>)type, emptyQuery);
-            }
-            // TODO: logging
-            System.out.println("All objects successfully removed from DB");
-        }
     }
 
     @Override
@@ -69,6 +64,7 @@ public class MongoDBSessionFactory implements KeycloakSessionFactory {
 
     @Override
     public void close() {
+        logger.info("Closing MongoDB client");
         mongoClient.close();
     }
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/NoSQLSession.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/NoSQLSession.java
index 9e2619e..838a6a4 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/NoSQLSession.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/NoSQLSession.java
@@ -55,7 +55,7 @@ public class NoSQLSession implements KeycloakSession {
 
     @Override
     public RealmModel getRealm(String id) {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("id", id)
                 .build();
         RealmData realmData = noSQL.loadSingleObject(RealmData.class, query);
@@ -65,7 +65,7 @@ public class NoSQLSession implements KeycloakSession {
     @Override
     public List<RealmModel> getRealms(UserModel admin) {
         String userId = ((UserAdapter)admin).getUser().getId();
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("realmAdmins", userId)
                 .build();
         List<RealmData> realms = noSQL.loadObjects(RealmData.class, query);
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java
index 5181c42..f7e3e83 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/adapters/RealmAdapter.java
@@ -250,7 +250,7 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public UserAdapter getUser(String name) {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("loginName", name)
                 .andCondition("realmId", getOid())
                 .build();
@@ -280,7 +280,7 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public RoleAdapter getRole(String name) {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("name", name)
                 .andCondition("realmId", getOid())
                 .build();
@@ -308,7 +308,7 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public List<RoleModel> getRoles() {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("realmId", getOid())
                 .build();
         List<RoleData> roles = noSQL.loadObjects(RoleData.class, query);
@@ -325,7 +325,7 @@ public class RealmAdapter implements RealmModel {
     public List<RoleModel> getDefaultRoles() {
         String[] defaultRoles = realm.getDefaultRoles();
 
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .inCondition("_id", defaultRoles)
                 .build();
         List<RoleData> defaultRolesData = noSQL.loadObjects(RoleData.class, query);
@@ -393,7 +393,7 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public List<ApplicationModel> getApplications() {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("realmId", getOid())
                 .build();
         List<ApplicationData> appDatas = noSQL.loadObjects(ApplicationData.class, query);
@@ -456,7 +456,7 @@ public class RealmAdapter implements RealmModel {
 
         Set<String> result = new HashSet<String>();
 
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .inCondition("_id", roleIds)
                 .build();
         List<RoleData> roles = noSQL.loadObjects(RoleData.class, query);
@@ -492,7 +492,7 @@ public class RealmAdapter implements RealmModel {
 
         Set<String> result = new HashSet<String>();
 
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .inCondition("_id", scopeIds)
                 .build();
         List<RoleData> roles = noSQL.loadObjects(RoleData.class, query);
@@ -639,7 +639,7 @@ public class RealmAdapter implements RealmModel {
     }
 
     protected List<RequiredCredentialData> getRequiredCredentialsData(int credentialType) {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("realmId", getOid())
                 .andCondition("clientType", credentialType)
                 .build();
@@ -681,7 +681,7 @@ public class RealmAdapter implements RealmModel {
 
     @Override
     public UserModel getUserBySocialLink(SocialLinkModel socialLink) {
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("socialProvider", socialLink.getSocialProvider())
                 .andCondition("socialUsername", socialLink.getSocialUsername())
                 .build();
@@ -701,7 +701,7 @@ public class RealmAdapter implements RealmModel {
         UserData userData = ((UserAdapter)user).getUser();
         String userId = userData.getId();
 
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("userId", userId)
                 .build();
         List<SocialLinkData> dbSocialLinks = noSQL.loadObjects(SocialLinkData.class, query);
@@ -729,7 +729,7 @@ public class RealmAdapter implements RealmModel {
     public void removeSocialLink(UserModel user, SocialLinkModel socialLink) {
         UserData userData = ((UserAdapter)user).getUser();
         String userId = userData.getId();
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("socialProvider", socialLink.getSocialProvider())
                 .andCondition("socialUsername", socialLink.getSocialUsername())
                 .andCondition("userId", userId)
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java
index 32a84d7..7709987 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/credentials/PasswordCredentialHandler.java
@@ -57,7 +57,7 @@ public class PasswordCredentialHandler {
         // If the user for the provided username cannot be found we fail validation
         if (user != null) {
             if (user.isEnabled()) {
-                NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+                NoSQLQuery query = noSQL.createQueryBuilder()
                         .andCondition("userId", user.getId())
                         .build();
                 PasswordData passwordData = noSQL.loadSingleObject(PasswordData.class, query);
@@ -87,7 +87,7 @@ public class PasswordCredentialHandler {
                        Date effectiveDate, Date expiryDate) {
 
         // Try to look if user already has password
-        NoSQLQuery query = NoSQLQueryBuilder.create(MongoDBQueryBuilder.class)
+        NoSQLQuery query = noSQL.createQueryBuilder()
                 .andCondition("userId", user.getId())
                 .build();
 
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/ApplicationData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/ApplicationData.java
index 7035f9a..108f1d0 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/ApplicationData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/ApplicationData.java
@@ -1,9 +1,11 @@
 package org.keycloak.services.models.nosql.keycloak.data;
 
+import org.keycloak.services.models.nosql.api.NoSQL;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
 import org.keycloak.services.models.nosql.api.NoSQLObject;
+import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -82,4 +84,16 @@ public class ApplicationData implements NoSQLObject {
     public void setRealmId(String realmId) {
         this.realmId = realmId;
     }
+
+    @Override
+    public void afterRemove(NoSQL noSQL) {
+        // Remove resourceUser of this application
+        noSQL.removeObject(UserData.class, resourceUserId);
+
+        // Remove all roles, which belongs to this application
+        NoSQLQuery query = noSQL.createQueryBuilder()
+                .andCondition("applicationId", id)
+                .build();
+        noSQL.removeObjects(RoleData.class, query);
+    }
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/PasswordData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/PasswordData.java
index a5ffad4..4834316 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/PasswordData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/credentials/PasswordData.java
@@ -2,6 +2,7 @@ package org.keycloak.services.models.nosql.keycloak.data.credentials;
 
 import java.util.Date;
 
+import org.keycloak.services.models.nosql.api.AbstractNoSQLObject;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
@@ -11,9 +12,8 @@ import org.keycloak.services.models.nosql.api.NoSQLObject;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @NoSQLCollection(collectionName = "passwordCredentials")
-public class PasswordData implements NoSQLObject {
+public class PasswordData extends AbstractNoSQLObject {
 
-    private String id;
     private Date effectiveDate = new Date();
     private Date expiryDate;
     private String encodedHash;
@@ -21,15 +21,6 @@ public class PasswordData implements NoSQLObject {
 
     private String userId;
 
-    @NoSQLId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @NoSQLField
     public Date getEffectiveDate() {
         return effectiveDate;
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RealmData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RealmData.java
index a550c6a..42f8ab2 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RealmData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RealmData.java
@@ -4,10 +4,12 @@ import java.security.SecureRandom;
 import java.util.Random;
 import java.util.UUID;
 
+import org.keycloak.services.models.nosql.api.NoSQL;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
 import org.keycloak.services.models.nosql.api.NoSQLObject;
+import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -169,4 +171,22 @@ public class RealmData implements NoSQLObject {
         this.realmAdmins = realmAdmins;
     }
 
+    @Override
+    public void afterRemove(NoSQL noSQL) {
+        NoSQLQuery query = noSQL.createQueryBuilder()
+                .andCondition("realmId", oid)
+                .build();
+
+        // Remove all users of this realm
+        noSQL.removeObjects(UserData.class, query);
+
+        // Remove all requiredCredentials of this realm
+        noSQL.removeObjects(RequiredCredentialData.class, query);
+
+        // Remove all roles of this realm
+        noSQL.removeObjects(RoleData.class, query);
+
+        // Remove all applications of this realm
+        noSQL.removeObjects(ApplicationData.class, query);
+    }
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RequiredCredentialData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RequiredCredentialData.java
index d20b595..124e507 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RequiredCredentialData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RequiredCredentialData.java
@@ -1,5 +1,6 @@
 package org.keycloak.services.models.nosql.keycloak.data;
 
+import org.keycloak.services.models.nosql.api.AbstractNoSQLObject;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
@@ -9,7 +10,7 @@ import org.keycloak.services.models.nosql.api.NoSQLObject;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @NoSQLCollection(collectionName = "requiredCredentials")
-public class RequiredCredentialData implements NoSQLObject {
+public class RequiredCredentialData extends AbstractNoSQLObject {
 
     public static final int CLIENT_TYPE_USER = 1;
     public static final int CLIENT_TYPE_RESOURCE = 2;
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RoleData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RoleData.java
index 9e00ef1..89138e6 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RoleData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/RoleData.java
@@ -1,9 +1,15 @@
 package org.keycloak.services.models.nosql.keycloak.data;
 
+import java.util.List;
+
+import org.jboss.resteasy.logging.Logger;
+import org.keycloak.services.models.nosql.api.NoSQL;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
 import org.keycloak.services.models.nosql.api.NoSQLObject;
+import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
+import org.keycloak.services.models.nosql.impl.Utils;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -11,6 +17,8 @@ import org.keycloak.services.models.nosql.api.NoSQLObject;
 @NoSQLCollection(collectionName = "roles")
 public class RoleData implements NoSQLObject {
 
+    private static final Logger logger = Logger.getLogger(RoleData.class);
+
     private String id;
     private String name;
     private String description;
@@ -62,4 +70,36 @@ public class RoleData implements NoSQLObject {
     public void setApplicationId(String applicationId) {
         this.applicationId = applicationId;
     }
+
+    @Override
+    public void afterRemove(NoSQL noSQL) {
+        // Remove this role from all users, which has it
+        NoSQLQuery query = noSQL.createQueryBuilder()
+                .andCondition("roleIds", id)
+                .build();
+
+        List<UserData> users = noSQL.loadObjects(UserData.class, query);
+        for (UserData user : users) {
+            logger.info("Removing role " + getName() + " from user " + user.getLoginName());
+            String[] roleIds = user.getRoleIds();
+            String[] newRoleIds = Utils.removeItemFromArray(roleIds, getId());
+            user.setRoleIds(newRoleIds);
+            noSQL.saveObject(user);
+        }
+
+        // Remove this scope from all users, which has it
+        query = noSQL.createQueryBuilder()
+                .andCondition("scopeIds", id)
+                .build();
+
+        users = noSQL.loadObjects(UserData.class, query);
+        for (UserData user : users) {
+            logger.info("Removing scope " + getName() + " from user " + user.getLoginName());
+            String[] scopeIds = user.getScopeIds();
+            String[] newScopeIds = Utils.removeItemFromArray(scopeIds, getId());
+            user.setScopeIds(newScopeIds);
+            noSQL.saveObject(user);
+        }
+
+    }
 }
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/SocialLinkData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/SocialLinkData.java
index 0ed96a4..7cfe6f5 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/SocialLinkData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/SocialLinkData.java
@@ -1,5 +1,6 @@
 package org.keycloak.services.models.nosql.keycloak.data;
 
+import org.keycloak.services.models.nosql.api.AbstractNoSQLObject;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
@@ -9,22 +10,12 @@ import org.keycloak.services.models.nosql.api.NoSQLObject;
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
  */
 @NoSQLCollection(collectionName = "socialLinks")
-public class SocialLinkData implements NoSQLObject {
+public class SocialLinkData extends AbstractNoSQLObject {
 
-    private String id;
     private String socialUsername;
     private String socialProvider;
     private String userId;
 
-    @NoSQLId
-    public String getId() {
-        return id;
-    }
-
-    public void setId(String id) {
-        this.id = id;
-    }
-
     @NoSQLField
     public String getSocialUsername() {
         return socialUsername;
diff --git a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/UserData.java b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/UserData.java
index cde0b42..c2321fb 100644
--- a/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/UserData.java
+++ b/services/src/main/java/org/keycloak/services/models/nosql/keycloak/data/UserData.java
@@ -1,9 +1,12 @@
 package org.keycloak.services.models.nosql.keycloak.data;
 
 import org.keycloak.services.models.nosql.api.AbstractAttributedNoSQLObject;
+import org.keycloak.services.models.nosql.api.NoSQL;
 import org.keycloak.services.models.nosql.api.NoSQLCollection;
 import org.keycloak.services.models.nosql.api.NoSQLField;
 import org.keycloak.services.models.nosql.api.NoSQLId;
+import org.keycloak.services.models.nosql.api.query.NoSQLQuery;
+import org.keycloak.services.models.nosql.keycloak.data.credentials.PasswordData;
 
 /**
  * @author <a href="mailto:mposolda@redhat.com">Marek Posolda</a>
@@ -103,4 +106,15 @@ public class UserData extends AbstractAttributedNoSQLObject {
     public void setScopeIds(String[] scopeIds) {
         this.scopeIds = scopeIds;
     }
+
+    @Override
+    public void afterRemove(NoSQL noSQL) {
+        NoSQLQuery query = noSQL.createQueryBuilder()
+                .andCondition("userId", id)
+                .build();
+
+        // Remove social links and passwords of this user
+        noSQL.removeObjects(SocialLinkData.class, query);
+        noSQL.removeObjects(PasswordData.class, query);
+    }
 }
diff --git a/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java b/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
index dc8d711..dabba40 100755
--- a/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
+++ b/services/src/main/java/org/keycloak/services/resources/KeycloakApplication.java
@@ -62,9 +62,25 @@ public class KeycloakApplication extends Application {
     }
 
     public static KeycloakSessionFactory buildSessionFactory() {
-        // EntityManagerFactory emf = Persistence.createEntityManagerFactory("keycloak-identity-store");
-        // return new PicketlinkKeycloakSessionFactory(emf, buildPartitionManager());
-        return new MongoDBSessionFactory("localhost", 27017, "keycloak", true);
+        String sessionFactoryType = System.getProperty("keycloak.sessionFactory", "picketlink");
+        if ("mongo".equals(sessionFactoryType)) {
+            return buildMongoDBSessionFactory();
+        } else {
+            return buildPicketlinkSessionFactory();
+        }
+    }
+
+    private static KeycloakSessionFactory buildPicketlinkSessionFactory() {
+        EntityManagerFactory emf = Persistence.createEntityManagerFactory("keycloak-identity-store");
+        return new PicketlinkKeycloakSessionFactory(emf, buildPartitionManager());
+    }
+
+    private static KeycloakSessionFactory buildMongoDBSessionFactory() {
+        String host = System.getProperty("keycloak.mongodb.host", "localhost");
+        int port = Integer.parseInt(System.getProperty("keycloak.mongodb.port", "27017"));
+        String dbName = System.getProperty("keycloak.mongodb.databaseName", "keycloak");
+        boolean removeAllObjectsOnStartup = Boolean.parseBoolean(System.getProperty("keycloak.mongodb.removeAllObjectsOnStartup", "true"));
+        return new MongoDBSessionFactory(host, port, dbName, removeAllObjectsOnStartup);
     }
 
     public KeycloakSessionFactory getFactory() {