DATAMONGO-387 - Repository query execution for GeoPage results are now working correctly.

odrotbohm · odrotbohm · commit 91fe5480a803 · 2012-02-01T14:23:29.000+01:00
Added special handling of GeoPage return types for repository query methods.
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/AbstractMongoQuery.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/AbstractMongoQuery.java
@@ -15,15 +15,15 @@
  */
 package org.springframework.data.mongodb.repository.query;
 
-import static org.springframework.data.mongodb.repository.query.QueryUtils.applyPagination;
+import static org.springframework.data.mongodb.repository.query.QueryUtils.*;
 
 import java.util.List;
 
 import org.springframework.data.domain.PageImpl;
 import org.springframework.data.domain.Pageable;
-import org.springframework.data.mongodb.core.CollectionCallback;
 import org.springframework.data.mongodb.core.MongoOperations;
 import org.springframework.data.mongodb.core.geo.Distance;
+import org.springframework.data.mongodb.core.geo.GeoPage;
 import org.springframework.data.mongodb.core.geo.GeoResult;
 import org.springframework.data.mongodb.core.geo.GeoResults;
 import org.springframework.data.mongodb.core.geo.Point;
@@ -34,10 +34,6 @@
 import org.springframework.data.util.TypeInformation;
 import org.springframework.util.Assert;
 
-import com.mongodb.DBCollection;
-import com.mongodb.DBCursor;
-import com.mongodb.DBObject;
-
 /**
  * Base class for {@link RepositoryQuery} implementations for Mongo.
  * 
@@ -51,8 +47,8 @@ public abstract class AbstractMongoQuery implements RepositoryQuery {
 	/**
 	 * Creates a new {@link AbstractMongoQuery} from the given {@link MongoQueryMethod} and {@link MongoOperations}.
 	 * 
-	 * @param method
-	 * @param template
+	 * @param method must not be {@literal null}.
+	 * @param template must not be {@literal null}.
 	 */
 	public AbstractMongoQuery(MongoQueryMethod method, MongoOperations template) {
 
@@ -63,25 +59,30 @@ public AbstractMongoQuery(MongoQueryMethod method, MongoOperations template) {
 		this.mongoOperations = template;
 	}
 
-	/* (non-Javadoc)
+	/* 
+	 * (non-Javadoc)
 	 * @see org.springframework.data.repository.query.RepositoryQuery#getQueryMethod()
 	 */
 	public MongoQueryMethod getQueryMethod() {
-
 		return method;
 	}
 
 	/*
-	  * (non-Javadoc)
-	  *
-	  * @see org.springframework.data.repository.query.RepositoryQuery#execute(java .lang.Object[])
-	  */
+	 * (non-Javadoc)
+	 * @see org.springframework.data.repository.query.RepositoryQuery#execute(java.lang.Object[])
+	 */
 	public Object execute(Object[] parameters) {
 
 		MongoParameterAccessor accessor = new MongoParametersParameterAccessor(method, parameters);
 		Query query = createQuery(new ConvertingParameterAccessor(mongoOperations.getConverter(), accessor));
 
-		if (method.isGeoNearQuery()) {
+		if (method.isGeoNearQuery() && method.isPageQuery()) {
+
+			MongoParameterAccessor countAccessor = new MongoParametersParameterAccessor(method, parameters);
+			Query countQuery = createCountQuery(new ConvertingParameterAccessor(mongoOperations.getConverter(), countAccessor));
+
+			return new GeoNearExecution(accessor).execute(query, countQuery);
+		} else if (method.isGeoNearQuery()) {
 			return new GeoNearExecution(accessor).execute(query);
 		} else if (method.isCollectionQuery()) {
 			return new CollectionExecution().execute(query);
@@ -93,14 +94,25 @@ public Object execute(Object[] parameters) {
 	}
 
 	/**
-	 * Create a {@link Query} instance using the given {@link ParameterAccessor}
+	 * Creates a {@link Query} instance using the given {@link ParameterAccessor}
 	 * 
-	 * @param accessor
-	 * @param converter
+	 * @param accessor must not be {@literal null}.
 	 * @return
 	 */
 	protected abstract Query createQuery(ConvertingParameterAccessor accessor);
 
+	/**
+	 * Creates a {@link Query} instance using the given {@link ConvertingParameterAccessor}. Will delegate to
+	 * {@link #createQuery(ConvertingParameterAccessor)} by default but allows customization of the count query to be
+	 * triggered.
+	 * 
+	 * @param accessor must not be {@literal null}.
+	 * @return
+	 */
+	protected Query createCountQuery(ConvertingParameterAccessor accessor) {
+		return createQuery(accessor);
+	}
+
 	private abstract class Execution {
 
 		abstract Object execute(Query query);
@@ -122,13 +134,11 @@ protected List<?> readCollection(Query query) {
 	class CollectionExecution extends Execution {
 
 		/*
-		   * (non-Javadoc)
-		   *
-		   * @see org.springframework.data.mongodb.repository.MongoQuery.Execution #execute(com.mongodb.DBObject)
-		   */
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.repository.query.AbstractMongoQuery.Execution#execute(org.springframework.data.mongodb.core.query.Query)
+		 */
 		@Override
 		public Object execute(Query query) {
-
 			return readCollection(query);
 		}
 	}
@@ -162,24 +172,13 @@ public PagedExecution(Pageable pageable) {
 		Object execute(Query query) {
 
 			MongoEntityInformation<?, ?> metadata = method.getEntityInformation();
-			int count = getCollectionCursor(metadata.getCollectionName(), query.getQueryObject()).count();
+			long count = mongoOperations.count(query, metadata.getCollectionName());
 
 			List<?> result = mongoOperations.find(applyPagination(query, pageable), metadata.getJavaType(),
 					metadata.getCollectionName());
 
 			return new PageImpl(result, pageable, count);
 		}
-
-		private DBCursor getCollectionCursor(String collectionName, final DBObject query) {
-
-			return mongoOperations.execute(collectionName, new CollectionCallback<DBCursor>() {
-
-				public DBCursor doInCollection(DBCollection collection) {
-
-					return collection.find(query);
-				}
-			});
-		}
 	}
 
 	/**
@@ -221,6 +220,28 @@ public GeoNearExecution(MongoParameterAccessor accessor) {
 		@Override
 		Object execute(Query query) {
 
+			GeoResults<?> results = doExecuteQuery(query);
+			return isListOfGeoResult() ? results.getContent() : results;
+		}
+
+		/**
+		 * Executes the given {@link Query} to return a page.
+		 * 
+		 * @param query must not be {@literal null}.
+		 * @param countQuery must not be {@literal null}.
+		 * @return
+		 */
+		Object execute(Query query, Query countQuery) {
+
+			MongoEntityInformation<?, ?> information = method.getEntityInformation();
+			long count = mongoOperations.count(countQuery, information.getCollectionName());
+
+			return new GeoPage<Object>(doExecuteQuery(query), accessor.getPageable(), count);
+		}
+
+		@SuppressWarnings("unchecked")
+		private GeoResults<Object> doExecuteQuery(Query query) {
+
 			Point nearLocation = accessor.getGeoNearLocation();
 			NearQuery nearQuery = NearQuery.near(nearLocation);
 
@@ -234,10 +255,8 @@ Object execute(Query query) {
 			}
 
 			MongoEntityInformation<?, ?> entityInformation = method.getEntityInformation();
-			GeoResults<?> results = mongoOperations.geoNear(nearQuery, entityInformation.getJavaType(),
+			return (GeoResults<Object>) mongoOperations.geoNear(nearQuery, entityInformation.getJavaType(),
 					entityInformation.getCollectionName());
-
-			return isListOfGeoResult() ? results.getContent() : results;
 		}
 
 		private boolean isListOfGeoResult() {
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/MongoQueryCreator.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/MongoQueryCreator.java
@@ -203,7 +203,8 @@ private Criteria from(Type type, Criteria criteria, PotentiallyConvertingIterato
 		case NEAR:
 
 			Distance distance = accessor.getMaxDistance();
-			Point point = nextAs(parameters, Point.class);
+			Point point = accessor.getGeoNearLocation();
+			point = point == null ? nextAs(parameters, Point.class) : point;
 
 			if (distance == null) {
 				return criteria.near(point);
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/PartTreeMongoQuery.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/PartTreeMongoQuery.java
@@ -38,8 +38,8 @@ public class PartTreeMongoQuery extends AbstractMongoQuery {
 	/**
 	 * Creates a new {@link PartTreeMongoQuery} from the given {@link QueryMethod} and {@link MongoTemplate}.
 	 * 
-	 * @param method
-	 * @param template
+	 * @param method must not be {@literal null}.
+	 * @param template must not be {@literal null}.
 	 */
 	public PartTreeMongoQuery(MongoQueryMethod method, MongoOperations mongoOperations) {
 
@@ -50,23 +50,31 @@ public PartTreeMongoQuery(MongoQueryMethod method, MongoOperations mongoOperatio
 	}
 
 	/**
+	 * Return the {@link PartTree} backing the query.
+	 * 
 	 * @return the tree
 	 */
 	public PartTree getTree() {
 		return tree;
 	}
 
 	/*
-	  * (non-Javadoc)
-	  *
-	  * @see
-	  * org.springframework.data.mongodb.repository.AbstractMongoQuery#createQuery(org.springframework.data.
-	  * document.mongodb.repository.ConvertingParameterAccessor)
-	  */
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.repository.query.AbstractMongoQuery#createQuery(org.springframework.data.mongodb.repository.query.ConvertingParameterAccessor, boolean)
+	 */
 	@Override
 	protected Query createQuery(ConvertingParameterAccessor accessor) {
 
 		MongoQueryCreator creator = new MongoQueryCreator(tree, accessor, context, isGeoNearQuery);
 		return creator.createQuery();
 	}
+
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.data.mongodb.repository.query.AbstractMongoQuery#createCountQuery(org.springframework.data.mongodb.repository.query.ConvertingParameterAccessor)
+	 */
+	@Override
+	protected Query createCountQuery(ConvertingParameterAccessor accessor) {
+		return new MongoQueryCreator(tree, accessor, context, false).createQuery();
+	}
 }
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/AbstractPersonRepositoryIntegrationTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/AbstractPersonRepositoryIntegrationTests.java
@@ -35,6 +35,7 @@
 import org.springframework.data.mongodb.core.geo.Box;
 import org.springframework.data.mongodb.core.geo.Circle;
 import org.springframework.data.mongodb.core.geo.Distance;
+import org.springframework.data.mongodb.core.geo.GeoPage;
 import org.springframework.data.mongodb.core.geo.GeoResults;
 import org.springframework.data.mongodb.core.geo.Metrics;
 import org.springframework.data.mongodb.core.geo.Point;
@@ -366,6 +367,18 @@ public void executesGeoNearQueryForResultsCorrectly() {
 		assertThat(results.getContent().isEmpty(), is(false));
 	}
 
+	@Test
+	public void executesGeoPageQueryForResultsCorrectly() {
+
+		Point point = new Point(-73.99171, 40.738868);
+		dave.setLocation(point);
+		repository.save(dave);
+
+		GeoPage<Person> results = repository.findByLocationNear(new Point(-73.99, 40.73), new Distance(2000,
+				Metrics.KILOMETERS), new PageRequest(0, 20));
+		assertThat(results.getContent().isEmpty(), is(false));
+	}
+
 	/**
 	 * @see DATAMONGO-323
 	 */
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/PersonRepository.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/PersonRepository.java
@@ -24,6 +24,7 @@
 import org.springframework.data.mongodb.core.geo.Box;
 import org.springframework.data.mongodb.core.geo.Circle;
 import org.springframework.data.mongodb.core.geo.Distance;
+import org.springframework.data.mongodb.core.geo.GeoPage;
 import org.springframework.data.mongodb.core.geo.GeoResults;
 import org.springframework.data.mongodb.core.geo.Point;
 import org.springframework.data.mongodb.core.geo.Polygon;
@@ -148,4 +149,6 @@ public interface PersonRepository extends MongoRepository<Person, String>, Query
 	List<Person> findByNamedQuery(String firstname);
 
 	GeoResults<Person> findByLocationNear(Point point, Distance maxDistance);
+
+	GeoPage<Person> findByLocationNear(Point point, Distance maxDistance, Pageable pageable);
 }
