DATAMONGO-2050 - Allow to specify the index to use for $geoNear aggregation operation.

Original pull request: #596.

Author: christophstrobl

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2015 the original author or authors.
+ * Copyright 2013-2018 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 @@
 
 import org.springframework.data.mongodb.core.query.NearQuery;
 import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
 
 import com.mongodb.BasicDBObject;
 import com.mongodb.DBObject;
@@ -26,32 +27,58 @@
  * <p>
  * We recommend to use the static factory method {@link Aggregation#geoNear(NearQuery, String)} instead of creating
  * instances of this class directly.
- * 
+ *
  * @author Thomas Darimont
  * @since 1.3
  */
 public class GeoNearOperation implements AggregationOperation {
 
 	private final NearQuery nearQuery;
 	private final String distanceField;
+	private final String indexKey;
 
 	/**
 	 * Creates a new {@link GeoNearOperation} from the given {@link NearQuery} and the given distance field. The
 	 * {@code distanceField} defines output field that contains the calculated distance.
-	 * 
-	 * @param query must not be {@literal null}.
+	 *
+	 * @param nearQuery must not be {@literal null}.
 	 * @param distanceField must not be {@literal null}.
 	 */
 	public GeoNearOperation(NearQuery nearQuery, String distanceField) {
+		this(nearQuery, distanceField, null);
+	}
+
+	/**
+	 * Creates a new {@link GeoNearOperation} from the given {@link NearQuery} and the given distance field. The
+	 * {@code distanceField} defines output field that contains the calculated distance.
+	 *
+	 * @param nearQuery must not be {@literal null}.
+	 * @param distanceField must not be {@literal null}.
+	 * @param indexKey can be {@literal null};
+	 */
+	private GeoNearOperation(NearQuery nearQuery, String distanceField, String indexKey) {
 
 		Assert.notNull(nearQuery, "NearQuery must not be null.");
 		Assert.hasLength(distanceField, "Distance field must not be null or empty.");
 
 		this.nearQuery = nearQuery;
 		this.distanceField = distanceField;
+		this.indexKey = indexKey;
+	}
+
+	/**
+	 * Optionally specify the geospatial index to use via the field to use in the calculation. <br />
+	 * <strong>NOTE:</strong> Requires MongoDB 4.0 or later.
+	 *
+	 * @param key the geospatial index field to use when calculating the distance.
+	 * @return new instance of {@link GeoNearOperation}.
+	 * @since 2.1
+	 */
+	public GeoNearOperation useIndex(String key) {
+		return new GeoNearOperation(nearQuery, distanceField, key);
 	}
 
-	/* 
+	/*
 	 * (non-Javadoc)
 	 * @see org.springframework.data.mongodb.core.aggregation.AggregationOperation#toDBObject(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext)
 	 */
@@ -61,6 +88,10 @@ public DBObject toDBObject(AggregationOperationContext context) {
 		BasicDBObject command = (BasicDBObject) context.getMappedObject(nearQuery.toDBObject());
 		command.put("distanceField", distanceField);
 
+		if (StringUtils.hasText(indexKey)) {
+			command.put("key", indexKey);
+		}
+
 		return new BasicDBObject("$geoNear", command);
 	}
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2017 the original author or authors.
+ * Copyright 2013-2018 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.
@@ -27,9 +27,10 @@
 
 /**
  * Unit tests for {@link GeoNearOperation}.
- * 
+ *
  * @author Oliver Gierke
  * @author Thomas Darimont
+ * @author Christoph Strobl
  */
 public class GeoNearOperationUnitTests {
 
@@ -42,7 +43,21 @@ public void rendersNearQueryAsAggregationOperation() {
 
 		DBObject nearClause = DBObjectTestUtils.getAsDBObject(dbObject, "$geoNear");
 
-		DBObject expected = (DBObject) new BasicDBObject(query.toDBObject().toMap()).append("distanceField", "distance");
+		DBObject expected = new BasicDBObject(query.toDBObject().toMap()).append("distanceField", "distance");
+		assertThat(nearClause, is(expected));
+	}
+
+	@Test // DATAMONGO-2050
+	public void rendersNearQueryWithKeyCorrectly() {
+
+		NearQuery query = NearQuery.near(10.0, 10.0);
+		GeoNearOperation operation = new GeoNearOperation(query, "distance").useIndex("geo-index-1");
+		DBObject dbObject = operation.toDBObject(Aggregation.DEFAULT_CONTEXT);
+
+		DBObject nearClause = DBObjectTestUtils.getAsDBObject(dbObject, "$geoNear");
+
+		DBObject expected = new BasicDBObject(query.toDBObject().toMap()).append("distanceField", "distance").append("key",
+				"geo-index-1");
 		assertThat(nearClause, is(expected));
 	}
 }