DATAMONGO-709 - Added support for restricting results by document types.

Added restrict(…) method to the Query API that generates appropriate filter criteria to restrict the result to certain types only. 

Type restrictions in query expressions are now applied in QueryMapper via a MongoTypeMapper based on information passed in through the query object in a "special" key. Exposed MongoTypeMapper in MongoConverter and MappingMongoConverter. Merged DefaultTypeMapper and DefaultMongoTypeMapper.

Original pull request: #53.

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/DefaultMongoTypeMapper.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2011 the original author or authors.
+ * Copyright 2011-2013 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,16 +18,19 @@ package org.springframework.data.mongodb.core.convert;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
-import org.springframework.data.convert.SimpleTypeInformationMapper;
 import org.springframework.data.convert.DefaultTypeMapper;
+import org.springframework.data.convert.SimpleTypeInformationMapper;
 import org.springframework.data.convert.TypeAliasAccessor;
 import org.springframework.data.convert.TypeInformationMapper;
 import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.util.ClassTypeInformation;
 import org.springframework.data.util.TypeInformation;
+
 import com.mongodb.BasicDBList;
+import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
 
 /**
@@ -37,33 +40,43 @@ import com.mongodb.DBObject;
  * respectively.
  * 
  * @author Oliver Gierke
+ * @author Thomas Darimont
  */
 public class DefaultMongoTypeMapper extends DefaultTypeMapper<DBObject> implements MongoTypeMapper {
 
 	public static final String DEFAULT_TYPE_KEY = "_class";
-	@SuppressWarnings("rawtypes")
+	@SuppressWarnings("rawtypes")//
 	private static final TypeInformation<List> LIST_TYPE_INFO = ClassTypeInformation.from(List.class);
-	@SuppressWarnings("rawtypes")
+	@SuppressWarnings("rawtypes")//
 	private static final TypeInformation<Map> MAP_TYPE_INFO = ClassTypeInformation.from(Map.class);
-	private String typeKey = DEFAULT_TYPE_KEY;
+
+	private final TypeAliasAccessor<DBObject> accessor;
+	private final String typeKey;
 
 	public DefaultMongoTypeMapper() {
-		this(DEFAULT_TYPE_KEY, Arrays.asList(SimpleTypeInformationMapper.INSTANCE));
+		this(DEFAULT_TYPE_KEY);
 	}
 
 	public DefaultMongoTypeMapper(String typeKey) {
-		super(new DBObjectTypeAliasAccessor(typeKey));
-		this.typeKey = typeKey;
+		this(typeKey, Arrays.asList(SimpleTypeInformationMapper.INSTANCE));
 	}
 
 	public DefaultMongoTypeMapper(String typeKey, MappingContext<? extends PersistentEntity<?, ?>, ?> mappingContext) {
-		super(new DBObjectTypeAliasAccessor(typeKey), mappingContext, Arrays.asList(SimpleTypeInformationMapper.INSTANCE));
-		this.typeKey = typeKey;
+		this(typeKey, new DBObjectTypeAliasAccessor(typeKey), mappingContext, Arrays
+				.asList(SimpleTypeInformationMapper.INSTANCE));
 	}
 
 	public DefaultMongoTypeMapper(String typeKey, List<? extends TypeInformationMapper> mappers) {
-		super(new DBObjectTypeAliasAccessor(typeKey), mappers);
+		this(typeKey, new DBObjectTypeAliasAccessor(typeKey), null, mappers);
+	}
+
+	private DefaultMongoTypeMapper(String typeKey, TypeAliasAccessor<DBObject> accessor,
+			MappingContext<? extends PersistentEntity<?, ?>, ?> mappingContext, List<? extends TypeInformationMapper> mappers) {
+
+		super(accessor, mappingContext, mappers);
+
 		this.typeKey = typeKey;
+		this.accessor = accessor;
 	}
 
 	/*
@@ -74,6 +87,31 @@ public class DefaultMongoTypeMapper extends DefaultTypeMapper<DBObject> implemen
 		return typeKey == null ? false : typeKey.equals(key);
 	}
 
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.convert.MongoTypeMapper#writeTypeRestrictions(java.util.Set)
+	 */
+	@Override
+	public void writeTypeRestrictions(DBObject result, Set<Class<?>> restrictedTypes) {
+
+		if (restrictedTypes == null || restrictedTypes.isEmpty()) {
+			return;
+		}
+
+		BasicDBList restrictedMappedTypes = new BasicDBList();
+
+		for (Class<?> restrictedType : restrictedTypes) {
+
+			Object typeAlias = getAliasFor(ClassTypeInformation.from(restrictedType));
+
+			if (typeAlias != null) {
+				restrictedMappedTypes.add(typeAlias);
+			}
+		}
+
+		accessor.writeTypeTo(result, new BasicDBObject("$in", restrictedMappedTypes));
+	}
+
 	/* (non-Javadoc)
 	 * @see org.springframework.data.convert.DefaultTypeMapper#getFallbackTypeFor(java.lang.Object)
 	 */
@@ -83,6 +121,7 @@ public class DefaultMongoTypeMapper extends DefaultTypeMapper<DBObject> implemen
 	}
 
 	/**
+	 * {@link TypeAliasAccessor} to store aliases in a {@link DBObject}.
 	 * 
 	 * @author Oliver Gierke
 	 */

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java

@@ -72,6 +72,7 @@ import com.mongodb.DBRef;
  * @author Oliver Gierke
  * @author Jon Brisbin
  * @author Patrik Wasik
+ * @author Thomas Darimont
  */
 public class MappingMongoConverter extends AbstractMongoConverter implements ApplicationContextAware {
 
@@ -124,6 +125,15 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 				mappingContext) : typeMapper;
 	}
 
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.core.convert.MongoConverter#getTypeMapper()
+	 */
+	@Override
+	public MongoTypeMapper getTypeMapper() {
+		return this.typeMapper;
+	}
+
 	/**
 	 * Configure the characters dots potentially contained in a {@link Map} shall be replaced with. By default we don't do
 	 * any translation but rather reject a {@link Map} with keys containing dots causing the conversion for the entire
@@ -356,8 +366,7 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 			try {
 				Object id = wrapper.getProperty(idProperty, Object.class, fieldAccessOnly);
 				dbo.put("_id", idMapper.convertId(id));
-			} catch (ConversionException ignored) {
-			}
+			} catch (ConversionException ignored) {}
 		}
 
 		// Write the properties

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MongoConverter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2011 the original author or authors.
+ * Copyright 2010-2013 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@ package org.springframework.data.mongodb.core.convert;
 
 import org.springframework.data.convert.EntityConverter;
 import org.springframework.data.convert.EntityReader;
+import org.springframework.data.convert.TypeMapper;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 
@@ -26,9 +27,17 @@ import com.mongodb.DBObject;
  * Central Mongo specific converter interface which combines {@link MongoWriter} and {@link MongoReader}.
  * 
  * @author Oliver Gierke
+ * @author Thomas Darimont
  */
 public interface MongoConverter extends
 		EntityConverter<MongoPersistentEntity<?>, MongoPersistentProperty, Object, DBObject>, MongoWriter<Object>,
 		EntityReader<Object, DBObject> {
 
+	/**
+	 * Returns thw {@link TypeMapper} being used to write type information into {@link DBObject}s created with that
+	 * converter.
+	 * 
+	 * @return will never be {@literal null}.
+	 */
+	MongoTypeMapper getTypeMapper();
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MongoTypeMapper.java

@@ -15,6 +15,8 @@
  */
 package org.springframework.data.mongodb.core.convert;
 
+import java.util.Set;
+
 import org.springframework.data.convert.TypeMapper;
 
 import com.mongodb.DBObject;
@@ -32,4 +34,14 @@ public interface MongoTypeMapper extends TypeMapper<DBObject> {
 	 * @return
 	 */
 	boolean isTypeKey(String key);
+
+	/**
+	 * Writes type restrictions to the given {@link DBObject}. This usually results in an {@code $in}-clause to be
+	 * generated that restricts the type-key (e.g. {@code _class}) to be in the set of type aliases for the given
+	 * {@code restrictedTypes}.
+	 * 
+	 * @param result must not be {@literal null}
+	 * @param restrictedTypes must not be {@literal null}
+	 */
+	void writeTypeRestrictions(DBObject result, Set<Class<?>> restrictedTypes);
 }

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/QueryMapper.java

@@ -30,6 +30,7 @@ import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mapping.context.PersistentPropertyPath;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
+import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.util.Assert;
 
 import com.mongodb.BasicDBList;
@@ -75,6 +76,7 @@ public class QueryMapper {
 	 * @param entity can be {@literal null}.
 	 * @return
 	 */
+	@SuppressWarnings("deprecation")
 	public DBObject getMappedObject(DBObject query, MongoPersistentEntity<?> entity) {
 
 		if (Keyword.isKeyword(query)) {
@@ -85,6 +87,16 @@ public class QueryMapper {
 
 		for (String key : query.keySet()) {
 
+			// TODO: remove one once QueryMapper can work with Query instances directly
+			if (Query.isRestrictedTypeKey(key)) {
+
+				@SuppressWarnings("unchecked")
+				Set<Class<?>> restrictedTypes = (Set<Class<?>>) query.get(key);
+				this.converter.getTypeMapper().writeTypeRestrictions(result, restrictedTypes);
+
+				continue;
+			}
+
 			if (Keyword.isKeyword(key)) {
 				result.putAll(getMappedKeyword(new Keyword(query, key), entity));
 				continue;

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Query.java

@@ -19,8 +19,11 @@ import static org.springframework.data.mongodb.core.query.SerializationUtils.*;
 import static org.springframework.util.ObjectUtils.*;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Set;
 
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Sort;
@@ -34,9 +37,13 @@ import com.mongodb.DBObject;
 /**
  * @author Thomas Risberg
  * @author Oliver Gierke
+ * @author Thomas Darimont
  */
 public class Query {
 
+	private final static String RESTRICTED_TYPES_KEY = "_$RESTRICTED_TYPES";
+
+	private final Set<Class<?>> restrictedTypes = new HashSet<Class<?>>();
 	private LinkedHashMap<String, Criteria> criteria = new LinkedHashMap<String, Criteria>();
 	private Field fieldSpec;
 	private Sort sort;
@@ -54,8 +61,7 @@ public class Query {
 		return new Query(criteria);
 	}
 
-	public Query() {
-	}
+	public Query() {}
 
 	/**
 	 * Creates a new {@link Query} using the given {@link Criteria}.
@@ -161,13 +167,46 @@ public class Query {
 		return this;
 	}
 
+	/**
+	 * @return the restrictedTypes
+	 */
+	public Set<Class<?>> getRestrictedTypes() {
+		return restrictedTypes == null ? Collections.<Class<?>> emptySet() : restrictedTypes;
+	}
+
+	/**
+	 * Restricts the query to only return documents instances that are exactly of the given types.
+	 * 
+	 * @param type may not be {@literal null}
+	 * @param additionalTypes may not be {@literal null}
+	 * @return
+	 */
+	public Query restrict(Class<?> type, Class<?>... additionalTypes) {
+
+		Assert.notNull(type, "Type must not be null!");
+		Assert.notNull(additionalTypes, "AdditionalTypes must not be null");
+
+		restrictedTypes.add(type);
+		for (Class<?> additionalType : additionalTypes) {
+			restrictedTypes.add(additionalType);
+		}
+
+		return this;
+	}
+
 	public DBObject getQueryObject() {
+
 		DBObject dbo = new BasicDBObject();
 		for (String k : criteria.keySet()) {
 			CriteriaDefinition c = criteria.get(k);
 			DBObject cl = c.getCriteriaObject();
 			dbo.putAll(cl);
 		}
+
+		if (!restrictedTypes.isEmpty()) {
+			dbo.put(RESTRICTED_TYPES_KEY, getRestrictedTypes());
+		}
+
 		return dbo;
 	}
 
@@ -266,4 +305,17 @@ public class Query {
 
 		return result;
 	}
+
+	/**
+	 * Returns whether the given key is the one used to hold the type restriction information.
+	 * 
+	 * @deprecated don't call this method as the restricted type handling will undergo some significant changes going
+	 *             forward.
+	 * @param key
+	 * @return
+	 */
+	@Deprecated
+	public static boolean isRestrictedTypeKey(String key) {
+		return RESTRICTED_TYPES_KEY.equals(key);
+	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/BaseDoc.java

@@ -0,0 +1,8 @@
+package org.springframework.data.mongodb.core;
+
+import org.springframework.data.annotation.Id;
+
+public class BaseDoc {
+	@Id String id;
+	String value;
+}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoOperationsUnitTests.java

@@ -30,6 +30,7 @@ import org.springframework.dao.DataAccessException;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.convert.AbstractMongoConverter;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
+import org.springframework.data.mongodb.core.convert.MongoTypeMapper;
 import org.springframework.data.mongodb.core.geo.Point;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
@@ -45,14 +46,13 @@ import com.mongodb.DBRef;
  * instances of their implementation and thus can see if it correctly implements the {@link MongoOperations} interface.
  * 
  * @author Oliver Gierke
+ * @author Thomas Darimont
  */
 @RunWith(MockitoJUnitRunner.class)
 public abstract class MongoOperationsUnitTests {
 
-	@Mock
-	CollectionCallback<Object> collectionCallback;
-	@Mock
-	DbCallback<Object> dbCallback;
+	@Mock CollectionCallback<Object> collectionCallback;
+	@Mock DbCallback<Object> dbCallback;
 
 	MongoConverter converter;
 	Person person;
@@ -86,6 +86,11 @@ public abstract class MongoOperationsUnitTests {
 			public DBRef toDBRef(Object object, MongoPersistentProperty referingProperty) {
 				return null;
 			}
+
+			@Override
+			public MongoTypeMapper getTypeMapper() {
+				return null;
+			}
 		};
 	}
 

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateTests.java

@@ -156,6 +156,7 @@ public class MongoTemplateTests {
 		template.dropCollection(Document.class);
 		template.dropCollection(ObjectWith3AliasedFields.class);
 		template.dropCollection(ObjectWith3AliasedFieldsAndNestedAddress.class);
+		template.dropCollection(BaseDoc.class);
 	}
 
 	@Test
@@ -1843,6 +1844,96 @@ public class MongoTemplateTests {
 		assertThat(result.address.state, is(stateValue));
 	}
 
+	/**
+	 * @see DATAMONGO-709
+	 */
+	@Test
+	public void aQueryRestrictedWithOneRestrictedResultTypeShouldReturnOnlyInstancesOfTheRestrictedType() {
+
+		BaseDoc doc0 = new BaseDoc();
+		doc0.value = "foo";
+		SpecialDoc doc1 = new SpecialDoc();
+		doc1.value = "foo";
+		doc1.specialValue = "specialfoo";
+		VerySpecialDoc doc2 = new VerySpecialDoc();
+		doc2.value = "foo";
+		doc2.specialValue = "specialfoo";
+		doc2.verySpecialValue = 4711;
+
+		String collectionName = template.getCollectionName(BaseDoc.class);
+		template.insert(doc0, collectionName);
+		template.insert(doc1, collectionName);
+		template.insert(doc2, collectionName);
+
+		Query query = Query.query(where("value").is("foo")).restrict(SpecialDoc.class);
+		List<BaseDoc> result = template.find(query, BaseDoc.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.size(), is(1));
+		assertThat(result.get(0), is(instanceOf(SpecialDoc.class)));
+	}
+
+	/**
+	 * @see DATAMONGO-709
+	 */
+	@Test
+	public void aQueryRestrictedWithMultipleRestrictedResultTypesShouldReturnOnlyInstancesOfTheRestrictedTypes() {
+
+		BaseDoc doc0 = new BaseDoc();
+		doc0.value = "foo";
+		SpecialDoc doc1 = new SpecialDoc();
+		doc1.value = "foo";
+		doc1.specialValue = "specialfoo";
+		VerySpecialDoc doc2 = new VerySpecialDoc();
+		doc2.value = "foo";
+		doc2.specialValue = "specialfoo";
+		doc2.verySpecialValue = 4711;
+
+		String collectionName = template.getCollectionName(BaseDoc.class);
+		template.insert(doc0, collectionName);
+		template.insert(doc1, collectionName);
+		template.insert(doc2, collectionName);
+
+		Query query = Query.query(where("value").is("foo")).restrict(BaseDoc.class, VerySpecialDoc.class);
+		List<BaseDoc> result = template.find(query, BaseDoc.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.size(), is(2));
+		assertThat(result.get(0).getClass(), is((Object) BaseDoc.class));
+		assertThat(result.get(1).getClass(), is((Object) VerySpecialDoc.class));
+	}
+
+	/**
+	 * @see DATAMONGO-709
+	 */
+	@Test
+	public void aQueryWithNoRestrictedResultTypesShouldReturnAllInstancesWithinTheGivenCollection() {
+
+		BaseDoc doc0 = new BaseDoc();
+		doc0.value = "foo";
+		SpecialDoc doc1 = new SpecialDoc();
+		doc1.value = "foo";
+		doc1.specialValue = "specialfoo";
+		VerySpecialDoc doc2 = new VerySpecialDoc();
+		doc2.value = "foo";
+		doc2.specialValue = "specialfoo";
+		doc2.verySpecialValue = 4711;
+
+		String collectionName = template.getCollectionName(BaseDoc.class);
+		template.insert(doc0, collectionName);
+		template.insert(doc1, collectionName);
+		template.insert(doc2, collectionName);
+
+		Query query = Query.query(where("value").is("foo"));
+		List<BaseDoc> result = template.find(query, BaseDoc.class);
+
+		assertThat(result, is(notNullValue()));
+		assertThat(result.size(), is(3));
+		assertThat(result.get(0).getClass(), is((Object) BaseDoc.class));
+		assertThat(result.get(1).getClass(), is((Object) SpecialDoc.class));
+		assertThat(result.get(2).getClass(), is((Object) VerySpecialDoc.class));
+	}
+
 	static interface Model {
 		String value();
 

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/SpecialDoc.java

@@ -0,0 +1,5 @@
+package org.springframework.data.mongodb.core;
+
+public class SpecialDoc extends BaseDoc {
+	String specialValue;
+}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/VerySpecialDoc.java

@@ -0,0 +1,5 @@
+package org.springframework.data.mongodb.core;
+
+public class VerySpecialDoc extends SpecialDoc {
+	int verySpecialValue;
+}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DefaultMongoTypeMapperUnitTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2011 the original author or authors.
+ * Copyright 2011-2013 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,17 +20,21 @@ import static org.junit.Assert.*;
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.List;
+
 import org.junit.Before;
 import org.junit.Test;
-import org.springframework.data.convert.SimpleTypeInformationMapper;
 import org.springframework.data.convert.ConfigurableTypeInformationMapper;
+import org.springframework.data.convert.SimpleTypeInformationMapper;
+import org.springframework.data.mongodb.core.DBObjectUtils;
 import org.springframework.data.util.TypeInformation;
 
+import com.mongodb.BasicDBList;
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
 
 /**
- * Unit tests for {@link ConfigurableTypeMapper}.
+ * Unit tests for {@link DefaultMongoTypeMapper}.
  * 
  * @author Oliver Gierke
  */
@@ -48,26 +52,28 @@ public class DefaultMongoTypeMapperUnitTests {
 				"1"));
 		simpleTypeInformationMapper = SimpleTypeInformationMapper.INSTANCE;
 
-		typeMapper = new DefaultMongoTypeMapper(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY,
-				Arrays.asList(configurableTypeInformationMapper));
+		typeMapper = new DefaultMongoTypeMapper();
 	}
 
 	@Test
 	public void defaultInstanceWritesClasses() {
 
-		typeMapper = new DefaultMongoTypeMapper();
 		writesTypeToField(new BasicDBObject(), String.class, String.class.getName());
 	}
 
 	@Test
 	public void defaultInstanceReadsClasses() {
-		typeMapper = new DefaultMongoTypeMapper();
+
 		DBObject dbObject = new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, String.class.getName());
 		readsTypeFromField(dbObject, String.class);
 	}
 
 	@Test
 	public void writesMapKeyForType() {
+
+		typeMapper = new DefaultMongoTypeMapper(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY,
+				Arrays.asList(configurableTypeInformationMapper));
+
 		writesTypeToField(new BasicDBObject(), String.class, "1");
 		writesTypeToField(new BasicDBObject(), Object.class, null);
 	}
@@ -84,6 +90,10 @@ public class DefaultMongoTypeMapperUnitTests {
 
 	@Test
 	public void readsTypeForMapKey() {
+
+		typeMapper = new DefaultMongoTypeMapper(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY,
+				Arrays.asList(configurableTypeInformationMapper));
+
 		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, "1"), String.class);
 		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, "unmapped"), null);
 	}
@@ -98,6 +108,88 @@ public class DefaultMongoTypeMapperUnitTests {
 		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, Object.class.getName()), Object.class);
 	}
 
+	/**
+	 * @see DATAMONGO-709
+	 */
+	@Test
+	public void writesTypeRestrictionsCorrectly() {
+
+		DBObject result = new BasicDBObject();
+
+		typeMapper = new DefaultMongoTypeMapper();
+		typeMapper.writeTypeRestrictions(result, Collections.<Class<?>> singleton(String.class));
+
+		DBObject typeInfo = DBObjectUtils.getAsDBObject(result, DefaultMongoTypeMapper.DEFAULT_TYPE_KEY);
+		List<Object> aliases = DBObjectUtils.getAsDBList(typeInfo, "$in");
+		assertThat(aliases, hasSize(1));
+		assertThat(aliases.get(0), is((Object) String.class.getName()));
+	}
+
+	@Test
+	public void addsFullyQualifiedClassNameUnderDefaultKeyByDefault() {
+		writesTypeToField(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, new BasicDBObject(), String.class);
+	}
+
+	@Test
+	public void writesTypeToCustomFieldIfConfigured() {
+		typeMapper = new DefaultMongoTypeMapper("_custom");
+		writesTypeToField("_custom", new BasicDBObject(), String.class);
+	}
+
+	@Test
+	public void doesNotWriteTypeInformationInCaseKeyIsSetToNull() {
+		typeMapper = new DefaultMongoTypeMapper(null);
+		writesTypeToField(null, new BasicDBObject(), String.class);
+	}
+
+	@Test
+	public void readsTypeFromDefaultKeyByDefault() {
+		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, String.class.getName()), String.class);
+	}
+
+	@Test
+	public void readsTypeFromCustomFieldConfigured() {
+
+		typeMapper = new DefaultMongoTypeMapper("_custom");
+		readsTypeFromField(new BasicDBObject("_custom", String.class.getName()), String.class);
+	}
+
+	@Test
+	public void returnsListForBasicDBLists() {
+		readsTypeFromField(new BasicDBList(), null);
+	}
+
+	@Test
+	public void returnsNullIfNoTypeInfoInDBObject() {
+		readsTypeFromField(new BasicDBObject(), null);
+		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, ""), null);
+	}
+
+	@Test
+	public void returnsNullIfClassCannotBeLoaded() {
+		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, "fooBar"), null);
+	}
+
+	@Test
+	public void returnsNullIfTypeKeySetToNull() {
+		typeMapper = new DefaultMongoTypeMapper(null);
+		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, String.class), null);
+	}
+
+	@Test
+	public void returnsCorrectTypeKey() {
+
+		assertThat(typeMapper.isTypeKey(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY), is(true));
+
+		typeMapper = new DefaultMongoTypeMapper("_custom");
+		assertThat(typeMapper.isTypeKey("_custom"), is(true));
+		assertThat(typeMapper.isTypeKey(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY), is(false));
+
+		typeMapper = new DefaultMongoTypeMapper(null);
+		assertThat(typeMapper.isTypeKey("_custom"), is(false));
+		assertThat(typeMapper.isTypeKey(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY), is(false));
+	}
+
 	private void readsTypeFromField(DBObject dbObject, Class<?> type) {
 
 		TypeInformation<?> typeInfo = typeMapper.readType(dbObject);
@@ -110,6 +202,18 @@ public class DefaultMongoTypeMapperUnitTests {
 		}
 	}
 
+	private void writesTypeToField(String field, DBObject dbObject, Class<?> type) {
+
+		typeMapper.writeType(type, dbObject);
+
+		if (field == null) {
+			assertThat(dbObject.keySet().isEmpty(), is(true));
+		} else {
+			assertThat(dbObject.containsField(field), is(true));
+			assertThat(dbObject.get(field), is((Object) type.getName()));
+		}
+	}
+
 	private void writesTypeToField(DBObject dbObject, Class<?> type, Object value) {
 
 		typeMapper.writeType(type, dbObject);

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/DefaultTypeMapperUnitTests.java

@@ -1,130 +0,0 @@
-/*
- * Copyright 2011 the original author or authors.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.springframework.data.mongodb.core.convert;
-
-import static org.hamcrest.Matchers.*;
-import static org.junit.Assert.*;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.springframework.data.util.TypeInformation;
-
-import com.mongodb.BasicDBList;
-import com.mongodb.BasicDBObject;
-import com.mongodb.DBObject;
-
-/**
- * Unit tests for {@link DefaultMongoTypeMapper}.
- * 
- * @author Oliver Gierke
- */
-public class DefaultTypeMapperUnitTests {
-
-	DefaultMongoTypeMapper mapper;
-
-	@Before
-	public void setUp() {
-		mapper = new DefaultMongoTypeMapper();
-	}
-
-	@Test
-	public void addsFullyQualifiedClassNameUnderDefaultKeyByDefault() {
-		writesTypeToField(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, new BasicDBObject(), String.class);
-	}
-
-	@Test
-	public void writesTypeToCustomFieldIfConfigured() {
-		mapper = new DefaultMongoTypeMapper("_custom");
-		writesTypeToField("_custom", new BasicDBObject(), String.class);
-	}
-
-	@Test
-	public void doesNotWriteTypeInformationInCaseKeyIsSetToNull() {
-		mapper = new DefaultMongoTypeMapper(null);
-		writesTypeToField(null, new BasicDBObject(), String.class);
-	}
-
-	@Test
-	public void readsTypeFromDefaultKeyByDefault() {
-		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, String.class.getName()), String.class);
-	}
-
-	@Test
-	public void readsTypeFromCustomFieldConfigured() {
-		mapper = new DefaultMongoTypeMapper("_custom");
-		readsTypeFromField(new BasicDBObject("_custom", String.class.getName()), String.class);
-	}
-
-	@Test
-	public void returnsListForBasicDBLists() {
-		readsTypeFromField(new BasicDBList(), null);
-	}
-
-	@Test
-	public void returnsNullIfNoTypeInfoInDBObject() {
-		readsTypeFromField(new BasicDBObject(), null);
-		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, ""), null);
-	}
-
-	@Test
-	public void returnsNullIfClassCannotBeLoaded() {
-		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, "fooBar"), null);
-	}
-
-	@Test
-	public void returnsNullIfTypeKeySetToNull() {
-		mapper = new DefaultMongoTypeMapper(null);
-		readsTypeFromField(new BasicDBObject(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY, String.class), null);
-	}
-
-	@Test
-	public void returnsCorrectTypeKey() {
-
-		assertThat(mapper.isTypeKey(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY), is(true));
-
-		mapper = new DefaultMongoTypeMapper("_custom");
-		assertThat(mapper.isTypeKey("_custom"), is(true));
-		assertThat(mapper.isTypeKey(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY), is(false));
-
-		mapper = new DefaultMongoTypeMapper(null);
-		assertThat(mapper.isTypeKey("_custom"), is(false));
-		assertThat(mapper.isTypeKey(DefaultMongoTypeMapper.DEFAULT_TYPE_KEY), is(false));
-	}
-
-	private void readsTypeFromField(DBObject dbObject, Class<?> type) {
-
-		TypeInformation<?> typeInfo = mapper.readType(dbObject);
-
-		if (type != null) {
-			assertThat(typeInfo, is(notNullValue()));
-			assertThat(typeInfo.getType(), is(typeCompatibleWith(type)));
-		} else {
-			assertThat(typeInfo, is(nullValue()));
-		}
-	}
-
-	private void writesTypeToField(String field, DBObject dbObject, Class<?> type) {
-
-		mapper.writeType(type, dbObject);
-
-		if (field == null) {
-			assertThat(dbObject.keySet().isEmpty(), is(true));
-		} else {
-			assertThat(dbObject.containsField(field), is(true));
-			assertThat(dbObject.get(field), is((Object) type.getName()));
-		}
-	}
-}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/QueryTests.java

@@ -20,6 +20,8 @@ import static org.junit.Assert.*;
 import static org.springframework.data.mongodb.core.query.Criteria.*;
 import static org.springframework.data.mongodb.core.query.Query.*;
 
+import java.util.Arrays;
+
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
@@ -27,6 +29,7 @@ import org.junit.rules.ExpectedException;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.mongodb.InvalidMongoDbApiUsageException;
+import org.springframework.data.mongodb.core.SpecialDoc;
 
 /**
  * Unit tests for {@link Query}.
@@ -34,11 +37,11 @@ import org.springframework.data.mongodb.InvalidMongoDbApiUsageException;
  * @author Thomas Risberg
  * @author Oliver Gierke
  * @author Patryk Wasik
+ * @author Thomas Darimont
  */
 public class QueryTests {
 
-	@Rule
-	public ExpectedException exception = ExpectedException.none();
+	@Rule public ExpectedException exception = ExpectedException.none();
 
 	@Test
 	public void testSimpleQuery() {
@@ -59,8 +62,7 @@ public class QueryTests {
 		try {
 			new Query(where("name").not().is("Thomas"));
 			Assert.fail("This should have caused an InvalidDocumentStoreApiUsageException");
-		} catch (InvalidMongoDbApiUsageException e) {
-		}
+		} catch (InvalidMongoDbApiUsageException e) {}
 	}
 
 	@Test
@@ -195,4 +197,20 @@ public class QueryTests {
 
 		new Query().with(new Sort(new Sort.Order("foo").ignoreCase()));
 	}
+
+	/**
+	 * @see DATAMONGO-709
+	 */
+	@Test
+	@SuppressWarnings("unchecked")
+	public void shouldReturnClassHierarchyOfRestrictedTypes() {
+
+		Query query = new Query(where("name").is("foo")).restrict(SpecialDoc.class);
+		assertThat(
+				query.toString(),
+				is("Query: { \"name\" : \"foo\", \"_$RESTRICTED_TYPES\" : [ { $java : class org.springframework.data.mongodb.core.SpecialDoc } ] }, Fields: null, Sort: null"));
+		assertThat(query.getRestrictedTypes(), is(notNullValue()));
+		assertThat(query.getRestrictedTypes().size(), is(1));
+		assertThat(query.getRestrictedTypes(), hasItems(Arrays.asList(SpecialDoc.class).toArray(new Class<?>[0])));
+	}
 }