From 62eb719b1e9adb050c33be254cbe9bf7a527415e Mon Sep 17 00:00:00 2001
From: James McNee <james@jamesmcnee.co.uk>
Date: Thu, 26 Aug 2021 21:25:36 +0100
Subject: [PATCH] Add support for `$sampleRate` criteria.

Closes #3726
Original pull request: #3765.
---
 .../data/mongodb/core/query/Criteria.java        | 16 ++++++++++++++++
 .../mongodb/core/query/CriteriaUnitTests.java    | 16 ++++++++++++++++
 .../asciidoc/reference/mongo-repositories.adoc   |  4 ++++
 3 files changed, 36 insertions(+)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Criteria.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Criteria.java
index f9a354c38..3ec4caf3c 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Criteria.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/query/Criteria.java
@@ -64,6 +64,7 @@ import com.mongodb.BasicDBList;
  * @author Andreas Zink
  * @author Ziemowit Stolarczyk
  * @author Clément Petit
+ * @author James McNee
  */
 public class Criteria implements CriteriaDefinition {
 
@@ -390,6 +391,21 @@ public class Criteria implements CriteriaDefinition {
 		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 --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/CriteriaUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/CriteriaUnitTests.java
index 9edf3c43f..e24fc34be 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/CriteriaUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/query/CriteriaUnitTests.java
@@ -156,6 +156,22 @@ public class CriteriaUnitTests {
 		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 --git a/src/main/asciidoc/reference/mongo-repositories.adoc b/src/main/asciidoc/reference/mongo-repositories.adoc
index 328a547b5..b847174c6 100644
--- a/src/main/asciidoc/reference/mongo-repositories.adoc
+++ b/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.