spring-projectsGH-3726: Add support for $sampleRate.

Closes spring-projects#3726

Author: JamesMcNee

Diff for: spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Criteria.java

@@ -15,19 +15,7 @@
  */
 package org.springframework.data.mongodb.core.query;
 
-import static org.springframework.util.ObjectUtils.*;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
-import java.util.regex.Pattern;
-import java.util.stream.Collectors;
-
+import com.mongodb.BasicDBList;
 import org.bson.BsonRegularExpression;
 import org.bson.BsonType;
 import org.bson.Document;
@@ -50,7 +38,18 @@
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 
-import com.mongodb.BasicDBList;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+import static org.springframework.util.ObjectUtils.*;
 
 /**
  * Central class for creating queries. It follows a fluent API style so that you can easily chain together multiple
@@ -64,6 +63,7 @@
  * @author Andreas Zink
  * @author Ziemowit Stolarczyk
  * @author Clément Petit
+ * @author James McNee
  */
 public class Criteria implements CriteriaDefinition {
 
@@ -390,6 +390,21 @@ public Criteria exists(boolean value) {
 		return this;
 	}
 
+	/**
+	 * Creates a criterion using the {@literal $sampleRate} operator.
+	 *
+	 * @param sampleRate sample rate to determine number of documents to be randomly selected from the input.
+	 * @return this.
+	 * @see <a href="https://docs.mongodb.com/manual/reference/operator/aggregation/sampleRate/">MongoDB Query operator: $sampleRate</a>
+	 */
+	public Criteria sampleRate(double sampleRate) {
+		Assert.isTrue(sampleRate >= 0, "The sample rate must be greater than zero!");
+		Assert.isTrue(sampleRate <= 1, "The sample rate must not be greater than one!");
+
+		criteria.put("$sampleRate", sampleRate);
+		return this;
+	}
+
 	/**
 	 * Creates a criterion using the {@literal $type} operator.
 	 *

Diff for: spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/CriteriaUnitTests.java

@@ -15,21 +15,20 @@
  */
 package org.springframework.data.mongodb.core.query;
 
-import static org.springframework.data.mongodb.test.util.Assertions.*;
-
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-
 import org.bson.Document;
 import org.junit.Test;
-
 import org.springframework.data.geo.Point;
 import org.springframework.data.mongodb.InvalidMongoDbApiUsageException;
 import org.springframework.data.mongodb.core.geo.GeoJsonLineString;
 import org.springframework.data.mongodb.core.geo.GeoJsonPoint;
 import org.springframework.data.mongodb.core.schema.MongoJsonSchema;
 
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+
+import static org.springframework.data.mongodb.test.util.Assertions.*;
+
 /**
  * Unit tests for {@link Criteria}.
  *
@@ -156,6 +155,22 @@ public void shouldNegateFollowingSimpleExpression() {
 		assertThat(co).isEqualTo(Document.parse("{ \"age\" : { \"$not\" : { \"$gt\" : 18}} , \"status\" : \"student\"}"));
 	}
 
+	@Test // GH-3726
+	public void shouldBuildCorrectSampleRateOperation() {
+		Criteria c = new Criteria().sampleRate(0.4);
+		assertThat(c.getCriteriaObject()).isEqualTo(Document.parse("{ \"$sampleRate\" : 0.4 }"));
+	}
+
+	@Test // GH-3726
+	public void shouldThrowExceptionWhenSampleRateIsNegative() {
+		assertThatIllegalArgumentException().isThrownBy(() -> new Criteria().sampleRate(-1));
+	}
+
+	@Test // GH-3726
+	public void shouldThrowExceptionWhenSampleRateIsGreatedThanOne() {
+		assertThatIllegalArgumentException().isThrownBy(() -> new Criteria().sampleRate(1.01));
+	}
+
 	@Test // DATAMONGO-1068
 	public void getCriteriaObjectShouldReturnEmptyDocumentWhenNoCriteriaSpecified() {
 

Diff for: src/main/asciidoc/reference/mongo-repositories.adoc

@@ -281,6 +281,10 @@ lower / upper bounds (`$gt` / `$gte` & `$lt` / `$lte`) according to `Range`
 | `Exists`
 | `findByLocationExists(boolean exists)`
 | `{"location" : {"$exists" : exists }}`
+
+| `SampleRate`
+| `sampleRate(double sampleRate)`
+| `{"$sampleRate" : sampleRate }`
 |===
 
 NOTE: If the property criterion compares a document, the order of the fields and exact equality in the document matters.