DATADOC-183 - Added count(…) methods to MongoOperations.

MongoTemplate is now able to count documents using either an entity class or collection name.

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

@@ -79,22 +79,22 @@ public interface MongoOperations {
 	CommandResult executeCommand(DBObject command, int options);
 	
 	/**
-	 * Execute a MongoDB query and iterate over the query results on a per-document basis with a DocumentCallbackHandler. 
+	 * Execute a MongoDB query and iterate over the query results on a per-document basis with a DocumentCallbackHandler.
 	 * 
 	 * @param query the query class that specifies the criteria used to find a record and also an optional fields
 	 *          specification
-	 * @param collectionName name of the collection to retrieve the objects from         
-	 * @param dch the handler that will extract results, one document at a time 
+	 * @param collectionName name of the collection to retrieve the objects from
+	 * @param dch the handler that will extract results, one document at a time
 	 */
 	void executeQuery(Query query, String collectionName, DocumentCallbackHandler dch);
 	
 	/**
-	 * Execute a MongoDB query and iterate over the query results on a per-document basis with a DocumentCallbackHandler using the 
+	 * Execute a MongoDB query and iterate over the query results on a per-document basis with a DocumentCallbackHandler using the
 	 * provided CursorPreparer.
 	 * @param query the query class that specifies the criteria used to find a record and also an optional fields
 	 *          specification
-	 * @param collectionName name of the collection to retrieve the objects from         
-	 * @param dch the handler that will extract results, one document at a time 
+	 * @param collectionName name of the collection to retrieve the objects from
+	 * @param dch the handler that will extract results, one document at a time
 	 * @param preparer allows for customization of the DBCursor used when iterating over the result set, (apply limits,
 	 *          skips and so on).
 	 */
@@ -504,6 +504,24 @@ public interface MongoOperations {
 	 */
 	<T> T findAndRemove(Query query, Class<T> entityClass, String collectionName);
 
+	/**
+	 * Returns the number of documents for the given {@link Query} by querying the collection of the given entity class.
+	 * 
+	 * @param query
+	 * @param entityClass must not be {@literal null}.
+	 * @return
+	 */
+	long count(Query query, Class<?> entityClass);
+	
+	/**
+	 * Returns the number of documents for the given {@link Query} querying the given collection.
+	 * 
+	 * @param query
+	 * @param collectionName must not be {@literal null} or empty.
+	 * @return
+	 */
+	long count(Query query, String collectionName);
+	
 	/**
 	 * Insert the object into the collection for the entity type of the object to save.
 	 * <p/>

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

@@ -512,6 +512,28 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 				entityClass);
 	}
 
+	public long count(Query query, Class<?> entityClass) {
+		Assert.notNull(entityClass);
+		return count(query, entityClass, determineCollectionName(entityClass));
+	}
+
+	public long count(final Query query, String collectionName) {
+		return count(query, null, collectionName);
+	}
+	
+	private long count(Query query, Class<?> entityClass, String collectionName) {
+		
+		Assert.hasText(collectionName);
+		final DBObject dbObject = query == null ? null : mapper.getMappedObject(query.getQueryObject(),
+				entityClass == null ? null : mappingContext.getPersistentEntity(entityClass));
+		
+		return execute(collectionName, new CollectionCallback<Long>() {
+			public Long doInCollection(DBCollection collection) throws MongoException, DataAccessException {
+				return collection.count(dbObject);
+			}
+		});
+	}
+
 	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.MongoOperations#insert(java.lang.Object)
@@ -832,8 +854,10 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 				entityClass), collectionName);
 	}
 
-	public <T> MapReduceResults<T> mapReduce(String inputCollectionName, String mapFunction, String reduceFunction, Class<T> entityClass) {
-		return mapReduce(null, inputCollectionName, mapFunction, reduceFunction, new MapReduceOptions().outputTypeInline(), entityClass);
+	public <T> MapReduceResults<T> mapReduce(String inputCollectionName, String mapFunction, String reduceFunction,
+			Class<T> entityClass) {
+		return mapReduce(null, inputCollectionName, mapFunction, reduceFunction, new MapReduceOptions().outputTypeInline(),
+				entityClass);
 	}
 
 	public <T> MapReduceResults<T> mapReduce(String inputCollectionName, String mapFunction, String reduceFunction,
@@ -841,12 +865,14 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 		return mapReduce(null, inputCollectionName, mapFunction, reduceFunction, mapReduceOptions, entityClass);
 	}
 
-	public <T> MapReduceResults<T> mapReduce(Query query, String inputCollectionName, String mapFunction, String reduceFunction, Class<T> entityClass) {
-		return mapReduce(query, inputCollectionName, mapFunction, reduceFunction, new MapReduceOptions().outputTypeInline(), entityClass);
+	public <T> MapReduceResults<T> mapReduce(Query query, String inputCollectionName, String mapFunction,
+			String reduceFunction, Class<T> entityClass) {
+		return mapReduce(query, inputCollectionName, mapFunction, reduceFunction,
+				new MapReduceOptions().outputTypeInline(), entityClass);
 	}
 
-	public <T> MapReduceResults<T> mapReduce(Query query, String inputCollectionName, String mapFunction, String reduceFunction,
-			MapReduceOptions mapReduceOptions, Class<T> entityClass) {
+	public <T> MapReduceResults<T> mapReduce(Query query, String inputCollectionName, String mapFunction,
+			String reduceFunction, MapReduceOptions mapReduceOptions, Class<T> entityClass) {
 		String mapFunc = replaceWithResourceIfNecessary(mapFunction);
 		String reduceFunc = replaceWithResourceIfNecessary(reduceFunction);
 		DBCollection inputCollection = getCollection(inputCollectionName);
@@ -872,9 +898,10 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware {
 		}
 		String error = commandResult.getErrorMessage();
 		if (error != null) {
-			throw new InvalidDataAccessApiUsageException("Command execution failed:  Error [" + error + "], Command = " + commandObject);
+			throw new InvalidDataAccessApiUsageException("Command execution failed:  Error [" + error + "], Command = "
+					+ commandObject);
 		}
-			
+
 		if (LOGGER.isDebugEnabled()) {
 			LOGGER.debug("MapReduce command result = [" + commandResult + "]");
 		}

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

@@ -94,7 +94,7 @@ public class MongoTemplateTests {
 		this.mappingTemplate = new MongoTemplate(factory, mappingConverter);
 	}
 
-	@Before	
+	@Before
 	public void setUp() {
 		cleanDb();
 	}
@@ -847,4 +847,46 @@ public class MongoTemplateTests {
 		assertEquals(1, names.size());
 		//template.remove(new Query(), Person.class);
 	}
+
+	/**
+	 * @see DATADOC-183
+	 */
+	@Test
+	public void countsDocumentsCorrectly() {
+		
+		assertThat(template.count(new Query(), Person.class), is(0L));
+		
+		Person dave = new Person("Dave");
+		Person carter = new Person("Carter");
+		
+		template.save(dave);
+		template.save(carter);
+		
+		assertThat(template.count(null, Person.class), is(2L));
+		assertThat(template.count(query(where("firstName").is("Carter")), Person.class), is(1L));
+	}
+	
+	/**
+	 * @see DATADOC-183
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void countRejectsNullEntityClass() {
+		template.count(null, (Class<?>) null);
+	}
+	
+	/**
+	 * @see DATADOC-183
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void countRejectsEmptyCollectionName() {
+		template.count(null, "");
+	}
+	
+	/**
+	 * @see DATADOC-183
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void countRejectsNullCollectionName() {
+		template.count(null, (String) null);
+	}
 }