GH-3726: Add support for $sampleRate match expression.

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

@@ -42,6 +42,7 @@
 import org.springframework.data.mongodb.core.schema.JsonSchemaObject.Type;
 import org.springframework.data.mongodb.core.schema.JsonSchemaProperty;
 import org.springframework.data.mongodb.core.schema.MongoJsonSchema;
+import org.springframework.data.mongodb.util.RegexFlags;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.Base64Utils;
@@ -63,6 +64,7 @@
  * @author Andreas Zink
  * @author Ziemowit Stolarczyk
  * @author Clément Petit
+ * @author James McNee
  */
 public class Criteria implements CriteriaDefinition {
 
@@ -71,20 +73,6 @@ public class Criteria implements CriteriaDefinition {
 	 */
 	private static final Object NOT_SET = new Object();
 
-	private static final int[] FLAG_LOOKUP = new int[Character.MAX_VALUE];
-
-	static {
-		FLAG_LOOKUP['g'] = 256;
-		FLAG_LOOKUP['i'] = Pattern.CASE_INSENSITIVE;
-		FLAG_LOOKUP['m'] = Pattern.MULTILINE;
-		FLAG_LOOKUP['s'] = Pattern.DOTALL;
-		FLAG_LOOKUP['c'] = Pattern.CANON_EQ;
-		FLAG_LOOKUP['x'] = Pattern.COMMENTS;
-		FLAG_LOOKUP['d'] = Pattern.UNIX_LINES;
-		FLAG_LOOKUP['t'] = Pattern.LITERAL;
-		FLAG_LOOKUP['u'] = Pattern.UNICODE_CASE;
-	}
-
 	private @Nullable String key;
 	private List<Criteria> criteriaChain;
 	private LinkedHashMap<String, Object> criteria = new LinkedHashMap<String, Object>();
@@ -403,6 +391,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.
 	 *
@@ -530,7 +533,7 @@ private Pattern toPattern(String regex, @Nullable String options) {
 
 		Assert.notNull(regex, "Regex string must not be null!");
 
-		return Pattern.compile(regex, regexFlags(options));
+		return Pattern.compile(regex, RegexFlags.toRegexFlags(options));
 	}
 
 	/**
@@ -1099,47 +1102,6 @@ private static boolean requiresGeoJsonFormat(Object value) {
 				|| (value instanceof GeoCommand && ((GeoCommand) value).getShape() instanceof GeoJson);
 	}
 
-	/**
-	 * Lookup the MongoDB specific flags for a given regex option string.
-	 *
-	 * @param s the Regex option/flag to look up. Can be {@literal null}.
-	 * @return zero if given {@link String} is {@literal null} or empty.
-	 * @since 2.2
-	 */
-	private static int regexFlags(@Nullable String s) {
-
-		int flags = 0;
-
-		if (s == null) {
-			return flags;
-		}
-
-		for (final char f : s.toLowerCase().toCharArray()) {
-			flags |= regexFlag(f);
-		}
-
-		return flags;
-	}
-
-	/**
-	 * Lookup the MongoDB specific flags for a given character.
-	 *
-	 * @param c the Regex option/flag to look up.
-	 * @return
-	 * @throws IllegalArgumentException for unknown flags
-	 * @since 2.2
-	 */
-	private static int regexFlag(char c) {
-
-		int flag = FLAG_LOOKUP[c];
-
-		if (flag == 0) {
-			throw new IllegalArgumentException(String.format("Unrecognized flag [%c]", c));
-		}
-
-		return flag;
-	}
-
 	/**
 	 * MongoDB specific <a href="https://docs.mongodb.com/manual/reference/operator/query-bitwise/">bitwise query
 	 * operators</a> like {@code $bitsAllClear, $bitsAllSet,...} for usage with {@link Criteria#bits()} and {@link Query}.

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/RegexFlags.java

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2021 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.util;
+
+import java.util.regex.Pattern;
+
+/**
+ * Utility to translate {@link Pattern#flags() regex flags} to MongoDB regex options and vice versa.
+ *
+ * @author Mark Paluch
+ * @since 3.3
+ */
+public abstract class RegexFlags {
+
+	private static final int[] FLAG_LOOKUP = new int[Character.MAX_VALUE];
+
+	static {
+		FLAG_LOOKUP['g'] = 256;
+		FLAG_LOOKUP['i'] = Pattern.CASE_INSENSITIVE;
+		FLAG_LOOKUP['m'] = Pattern.MULTILINE;
+		FLAG_LOOKUP['s'] = Pattern.DOTALL;
+		FLAG_LOOKUP['c'] = Pattern.CANON_EQ;
+		FLAG_LOOKUP['x'] = Pattern.COMMENTS;
+		FLAG_LOOKUP['d'] = Pattern.UNIX_LINES;
+		FLAG_LOOKUP['t'] = Pattern.LITERAL;
+		FLAG_LOOKUP['u'] = Pattern.UNICODE_CASE;
+	}
+
+	private RegexFlags() {
+
+	}
+
+	/**
+	 * Lookup the MongoDB specific options from given {@link Pattern#flags() flags}.
+	 *
+	 * @param flags the Regex flags to look up.
+	 * @return the options string. May be empty.
+	 */
+	public static String toRegexOptions(int flags) {
+
+		if (flags == 0) {
+			return "";
+		}
+
+		StringBuilder buf = new StringBuilder();
+
+		for (int i = 'a'; i < 'z'; i++) {
+
+			if (FLAG_LOOKUP[i] == 0) {
+				continue;
+			}
+
+			if ((flags & FLAG_LOOKUP[i]) > 0) {
+				buf.append((char) i);
+			}
+		}
+
+		return buf.toString();
+	}
+
+	/**
+	 * Lookup the MongoDB specific flags for a given regex option string.
+	 *
+	 * @param s the Regex option/flag to look up. Can be {@literal null}.
+	 * @return zero if given {@link String} is {@literal null} or empty.
+	 * @since 2.2
+	 */
+	public static int toRegexFlags(String s) {
+
+		int flags = 0;
+
+		for (char f : s.toLowerCase().toCharArray()) {
+			flags |= toRegexFlag(f);
+		}
+
+		return flags;
+	}
+
+	/**
+	 * Lookup the MongoDB specific flags for a given character.
+	 *
+	 * @param c the Regex option/flag to look up.
+	 * @return
+	 * @throws IllegalArgumentException for unknown flags
+	 * @since 2.2
+	 */
+	public static int toRegexFlag(char c) {
+
+		int flag = FLAG_LOOKUP[c];
+
+		if (flag == 0) {
+			throw new IllegalArgumentException(String.format("Unrecognized flag [%c]", c));
+		}
+
+		return flag;
+	}
+}