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

Original pull request: #596.

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2016 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,14 +17,16 @@ package org.springframework.data.mongodb.core.aggregation;
 
 import org.bson.Document;
 import org.springframework.data.mongodb.core.query.NearQuery;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
 
 /**
  * Represents a {@code geoNear} aggregation operation.
  * <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
  * @author Christoph Strobl
  * @since 1.3
@@ -33,24 +35,50 @@ public class GeoNearOperation implements AggregationOperation {
 
 	private final NearQuery nearQuery;
 	private final String distanceField;
+	private final @Nullable 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, @Nullable 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#toDocument(org.springframework.data.mongodb.core.aggregation.AggregationOperationContext)
 	 */
@@ -60,6 +88,10 @@ public class GeoNearOperation implements AggregationOperation {
 		Document command = context.getMappedObject(nearQuery.toDocument());
 		command.put("distanceField", distanceField);
 
+		if (StringUtils.hasText(indexKey)) {
+			command.put("key", indexKey);
+		}
+
 		return new Document("$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.
@@ -25,9 +25,10 @@ import org.springframework.data.mongodb.core.query.NearQuery;
 
 /**
  * Unit tests for {@link GeoNearOperation}.
- * 
+ *
  * @author Oliver Gierke
  * @author Thomas Darimont
+ * @author Christoph Strobl
  */
 public class GeoNearOperationUnitTests {
 
@@ -43,4 +44,17 @@ public class GeoNearOperationUnitTests {
 		Document expected = new Document(query.toDocument()).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");
+		Document document = operation.toDocument(Aggregation.DEFAULT_CONTEXT);
+
+		Document nearClause = DocumentTestUtils.getAsDocument(document, "$geoNear");
+
+		Document expected = new Document(query.toDocument()).append("distanceField", "distance").append("key", "geo-index-1");
+		assertThat(nearClause, is(expected));
+	}
 }