DATAMONGO-2029 - Encode collections of UUID and byte array query method arguments to their binary form.

We now convert collections that only contain UUID or byte array items to a BSON list that contains the encoded form of these items. Previously, we only converted single UUID and byte arrays into $binary so lists rendered to e.g. $uuid which does not work for queries.

Encoding is now encapsulated in strategy objects that implement the encoding only for their type. This allows to break up the conditional flow and improve organization of responsibilities.

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/query/ExpressionEvaluatingParameterBinder.java

@@ -16,16 +16,19 @@
 package org.springframework.data.mongodb.repository.query;
 
 import lombok.EqualsAndHashCode;
+import lombok.RequiredArgsConstructor;
 import lombok.Value;
 import lombok.experimental.UtilityClass;
 
 import java.io.StringWriter;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.UUID;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -205,6 +208,7 @@ class ExpressionEvaluatingParameterBinder {
 	 * @param binding must not be {@literal null}.
 	 * @return
 	 */
+	@SuppressWarnings("unchecked")
 	private String getParameterValueForBinding(MongoParameterAccessor accessor, MongoParameters parameters,
 			ParameterBinding binding) {
 
@@ -221,37 +225,7 @@ class ExpressionEvaluatingParameterBinder {
 			return binding.isExpression() ? JSON.serialize(value) : QuotedString.unquote(JSON.serialize(value));
 		}
 
-		if (value instanceof byte[]) {
-
-			if (binding.isQuoted()) {
-				return DatatypeConverter.printBase64Binary((byte[]) value);
-			}
-
-			return encode(new Binary((byte[]) value), BinaryCodec::new);
-		}
-
-		if (value instanceof UUID) {
-
-			if (binding.isQuoted()) {
-				return value.toString();
-			}
-
-			return encode((UUID) value, UuidCodec::new);
-		}
-
-		return JSON.serialize(value);
-	}
-
-	private <T> String encode(T value, Supplier<Codec<T>> defaultCodec) {
-
-		Codec<T> codec = codecRegistryProvider.getCodecFor((Class<T>) value.getClass()).orElseGet(defaultCodec);
-
-		StringWriter writer = new StringWriter();
-
-		codec.encode(new JsonWriter(writer), value, null);
-		writer.flush();
-
-		return writer.toString();
+		return EncodableValue.create(value).encode(codecRegistryProvider, binding.isQuoted());
 	}
 
 	/**
@@ -473,4 +447,221 @@ class ExpressionEvaluatingParameterBinder {
 			return quoted.substring(1, quoted.length() - 1);
 		}
 	}
+
+	/**
+	 * Value object encapsulating a bindable value, that can be encoded to be represented as JSON (BSON).
+	 *
+	 * @author Mark Paluch
+	 */
+	abstract static class EncodableValue {
+
+		/**
+		 * Obtain a {@link EncodableValue} given {@code value}.
+		 *
+		 * @param value the value to encode, may be {@literal null}.
+		 * @return the {@link EncodableValue} for {@code value}.
+		 */
+		@SuppressWarnings("unchecked")
+		public static EncodableValue create(@Nullable Object value) {
+
+			if (value instanceof byte[]) {
+				return new BinaryValue((byte[]) value);
+			}
+
+			if (value instanceof UUID) {
+				return new UuidValue((UUID) value);
+			}
+
+			if (value instanceof Collection) {
+
+				Collection<?> collection = (Collection<?>) value;
+				Class<?> commonElement = CollectionUtils.findCommonElementType(collection);
+
+				if (commonElement != null) {
+
+					if (UUID.class.isAssignableFrom(commonElement)) {
+						return new UuidCollection((Collection<UUID>) value);
+					}
+
+					if (byte[].class.isAssignableFrom(commonElement)) {
+						return new BinaryCollectionValue((Collection<byte[]>) value);
+					}
+				}
+			}
+
+			return new ObjectValue(value);
+		}
+
+		/**
+		 * Encode the encapsulated value.
+		 *
+		 * @param provider
+		 * @param quoted
+		 * @return
+		 */
+		public abstract String encode(CodecRegistryProvider provider, boolean quoted);
+
+		/**
+		 * Encode a {@code value} to JSON.
+		 *
+		 * @param provider
+		 * @param value
+		 * @param defaultCodec
+		 * @param <V>
+		 * @return
+		 */
+		protected <V> String encode(CodecRegistryProvider provider, V value, Supplier<Codec<V>> defaultCodec) {
+
+			StringWriter writer = new StringWriter();
+
+			doEncode(provider, writer, value, defaultCodec);
+
+			return writer.toString();
+		}
+
+		/**
+		 * Encode a {@link Collection} to JSON and potentially apply a {@link Function mapping function} before encoding.
+		 *
+		 * @param provider
+		 * @param value
+		 * @param mappingFunction
+		 * @param defaultCodec
+		 * @param <I> Input value type.
+		 * @param <V> Target type.
+		 * @return
+		 */
+		protected <I, V> String encodeCollection(CodecRegistryProvider provider, Iterable<I> value,
+				Function<I, V> mappingFunction, Supplier<Codec<V>> defaultCodec) {
+
+			StringWriter writer = new StringWriter();
+
+			writer.append("[");
+			value.forEach(it -> {
+
+				if (writer.getBuffer().length() > 1) {
+					writer.append(", ");
+				}
+
+				doEncode(provider, writer, mappingFunction.apply(it), defaultCodec);
+			});
+
+			writer.append("]");
+			writer.flush();
+
+			return writer.toString();
+		}
+
+		@SuppressWarnings("unchecked")
+		private <V> void doEncode(CodecRegistryProvider provider, StringWriter writer, V value,
+				Supplier<Codec<V>> defaultCodec) {
+
+			Codec<V> codec = provider.getCodecFor((Class<V>) value.getClass()).orElseGet(defaultCodec);
+
+			JsonWriter jsonWriter = new JsonWriter(writer);
+			codec.encode(jsonWriter, value, null);
+			jsonWriter.flush();
+		}
+	}
+
+	/**
+	 * {@link EncodableValue} for {@code byte[]} to render to {@literal $binary}.
+	 */
+	@RequiredArgsConstructor
+	static class BinaryValue extends EncodableValue {
+
+		private final byte[] value;
+
+		/* 
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.repository.query.ExpressionEvaluatingParameterBinder.EncodableValue#encode(org.springframework.data.mongodb.CodecRegistryProvider, boolean)
+		 */
+		@Override
+		public String encode(CodecRegistryProvider provider, boolean quoted) {
+
+			if (quoted) {
+				return DatatypeConverter.printBase64Binary(this.value);
+			}
+
+			return encode(provider, new Binary(this.value), BinaryCodec::new);
+		}
+	}
+
+	/**
+	 * {@link EncodableValue} for {@link Collection} containing only {@code byte[]} items to render to a BSON list
+	 * containing {@literal $binary}.
+	 */
+	@RequiredArgsConstructor
+	static class BinaryCollectionValue extends EncodableValue {
+
+		private final Collection<byte[]> value;
+
+		/* 
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.repository.query.ExpressionEvaluatingParameterBinder.EncodableValue#encode(org.springframework.data.mongodb.CodecRegistryProvider, boolean)
+		 */
+		@Override
+		public String encode(CodecRegistryProvider provider, boolean quoted) {
+			return encodeCollection(provider, this.value, Binary::new, BinaryCodec::new);
+		}
+	}
+
+	/**
+	 * {@link EncodableValue} for {@link UUID} to render to {@literal $binary}.
+	 */
+	@RequiredArgsConstructor
+	static class UuidValue extends EncodableValue {
+
+		private final UUID value;
+
+		/* 
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.repository.query.ExpressionEvaluatingParameterBinder.EncodableValue#encode(org.springframework.data.mongodb.CodecRegistryProvider, boolean)
+		 */
+		@Override
+		public String encode(CodecRegistryProvider provider, boolean quoted) {
+
+			if (quoted) {
+				return this.value.toString();
+			}
+
+			return encode(provider, this.value, UuidCodec::new);
+		}
+	}
+
+	/**
+	 * {@link EncodableValue} for {@link Collection} containing only {@link UUID} items to render to a BSON list
+	 * containing {@literal $binary}.
+	 */
+	@RequiredArgsConstructor
+	static class UuidCollection extends EncodableValue {
+
+		private final Collection<UUID> value;
+
+		/* 
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.repository.query.ExpressionEvaluatingParameterBinder.EncodableValue#encode(org.springframework.data.mongodb.CodecRegistryProvider, boolean)
+		 */
+		@Override
+		public String encode(CodecRegistryProvider provider, boolean quoted) {
+			return encodeCollection(provider, this.value, Function.identity(), UuidCodec::new);
+		}
+	}
+
+	/**
+	 * Fallback-{@link EncodableValue} for {@link Object}-typed values.
+	 */
+	@RequiredArgsConstructor
+	static class ObjectValue extends EncodableValue {
+
+		private final @Nullable Object value;
+
+		/* 
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mongodb.repository.query.ExpressionEvaluatingParameterBinder.EncodableValue#encode(org.springframework.data.mongodb.CodecRegistryProvider, boolean)
+		 */
+		@Override
+		public String encode(CodecRegistryProvider provider, boolean quoted) {
+			return JSON.serialize(this.value);
+		}
+	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/query/StringBasedMongoQueryUnitTests.java

@@ -23,6 +23,7 @@ import static org.mockito.Mockito.*;
 import java.lang.reflect.Method;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -322,6 +323,21 @@ public class StringBasedMongoQueryUnitTests {
 		assertThat(query.getQueryObject().toJson(), is(reference.getQueryObject().toJson()));
 	}
 
+	@Test // DATAMONGO-2029
+	public void shouldSupportNonQuotedBinaryCollectionDataReplacement() {
+
+		byte[] binaryData = "Matthews".getBytes(StandardCharsets.UTF_8);
+		ConvertingParameterAccessor accessor = StubParameterAccessor.getAccessor(converter,
+				(Object) Arrays.asList(binaryData));
+		StringBasedMongoQuery mongoQuery = createQueryForMethod("findByLastnameAsBinaryIn", List.class);
+
+		org.springframework.data.mongodb.core.query.Query query = mongoQuery.createQuery(accessor);
+		org.springframework.data.mongodb.core.query.Query reference = new BasicQuery("{'lastname' : { $in: [{'$binary' : '"
+				+ DatatypeConverter.printBase64Binary(binaryData) + "', '$type' : '" + BSON.B_GENERAL + "'}] }}");
+
+		assertThat(query.getQueryObject().toJson(), is(reference.getQueryObject().toJson()));
+	}
+
 	@Test // DATAMONGO-1911
 	public void shouldSupportNonQuotedUUIDReplacement() {
 
@@ -336,6 +352,23 @@ public class StringBasedMongoQueryUnitTests {
 		assertThat(query.getQueryObject().toJson(), is(reference.getQueryObject().toJson()));
 	}
 
+	@Test // DATAMONGO-2029
+	public void shouldSupportNonQuotedUUIDCollectionReplacement() {
+
+		UUID uuid1 = UUID.fromString("864de43b-e3ea-f1e4-3663-fb8240b659b9");
+		UUID uuid2 = UUID.fromString("864de43b-cafe-f1e4-3663-fb8240b659b9");
+
+		ConvertingParameterAccessor accessor = StubParameterAccessor.getAccessor(converter,
+				(Object) Arrays.asList(uuid1, uuid2));
+		StringBasedMongoQuery mongoQuery = createQueryForMethod("findByLastnameAsUUIDIn", List.class);
+
+		org.springframework.data.mongodb.core.query.Query query = mongoQuery.createQuery(accessor);
+		org.springframework.data.mongodb.core.query.Query reference = new BasicQuery(
+				"{'lastname' : { $in: [{ $binary : \"5PHq4zvkTYa5WbZAgvtjNg==\", $type : \"03\" }, { $binary : \"5PH+yjvkTYa5WbZAgvtjNg==\", $type : \"03\" }]}}");
+
+		assertThat(query.getQueryObject().toJson(), is(reference.getQueryObject().toJson()));
+	}
+
 	@Test // DATAMONGO-1911
 	public void shouldSupportQuotedUUIDReplacement() {
 
@@ -580,9 +613,15 @@ public class StringBasedMongoQueryUnitTests {
 		@Query("{ 'lastname' : ?0 }")
 		Person findByLastnameAsBinary(byte[] lastname);
 
+		@Query("{ 'lastname' : { $in: ?0} }")
+		Person findByLastnameAsBinaryIn(List<byte[]> lastname);
+
 		@Query("{ 'lastname' : ?0 }")
 		Person findByLastnameAsUUID(UUID lastname);
 
+		@Query("{ 'lastname' : { $in : ?0} }")
+		Person findByLastnameAsUUIDIn(List<UUID> lastname);
+
 		@Query("{ 'lastname' : '?0' }")
 		Person findByLastnameAsStringUUID(UUID lastname);