Fix reference sorting via explicit sort clause

christophstrobl · christophstrobl · commit 5f260fc79113 · 2021-05-19T13:53:48.000+02:00
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/ReferenceLookupDelegate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/ReferenceLookupDelegate.java
@@ -63,6 +63,12 @@ public final class ReferenceLookupDelegate {
 	private final SpELContext spELContext;
 	private final ParameterBindingDocumentCodec codec;
 
+	/**
+	 * Create a new {@link ReferenceLookupDelegate}.
+	 *
+	 * @param mappingContext must not be {@literal null}.
+	 * @param spELContext must not be {@literal null}.
+	 */
 	public ReferenceLookupDelegate(
 			MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext,
 			SpELContext spELContext) {
@@ -75,8 +81,18 @@ public ReferenceLookupDelegate(
 		this.codec = new ParameterBindingDocumentCodec();
 	}
 
+	/**
+	 * Read the reference expressed by the given property.
+	 *
+	 * @param property the reference defining property. Must not be {@literal null}. THe
+	 * @param value the source value identifying to the referenced entity. Must not be {@literal null}.
+	 * @param lookupFunction to execute a lookup query. Must not be {@literal null}.
+	 * @param entityReader the callback to convert raw source values into actual domain types. Must not be
+	 *          {@literal null}.
+	 * @return can be {@literal null}.
+	 */
 	@Nullable
-	Object readReference(MongoPersistentProperty property, Object value, LookupFunction lookupFunction,
+	public Object readReference(MongoPersistentProperty property, Object value, LookupFunction lookupFunction,
 			MongoEntityReader entityReader) {
 
 		DocumentReferenceQuery filter = computeFilter(property, value, spELContext);
@@ -98,10 +114,12 @@ Object readReference(MongoPersistentProperty property, Object value, LookupFunct
 	private ReferenceCollection computeReferenceContext(MongoPersistentProperty property, Object value,
 			SpELContext spELContext) {
 
+		// Use the first value as a reference for others in case of collection like
 		if (value instanceof Iterable) {
 			value = ((Iterable<?>) value).iterator().next();
 		}
 
+		// handle DBRef value
 		if (value instanceof DBRef) {
 			return ReferenceCollection.fromDBRef((DBRef) value);
 		}
@@ -110,21 +128,21 @@ private ReferenceCollection computeReferenceContext(MongoPersistentProperty prop
 
 		if (value instanceof Document) {
 
-			Document ref = (Document) value;
+			Document documentPointer = (Document) value;
 
 			if (property.isDocumentReference()) {
 
 				ParameterBindingContext bindingContext = bindingContext(property, value, spELContext);
 				DocumentReference documentReference = property.getDocumentReference();
 
 				String targetDatabase = parseValueOrGet(documentReference.db(), bindingContext,
-						() -> ref.get("db", String.class));
+						() -> documentPointer.get("db", String.class));
 				String targetCollection = parseValueOrGet(documentReference.collection(), bindingContext,
-						() -> ref.get("collection", collection));
+						() -> documentPointer.get("collection", collection));
 				return new ReferenceCollection(targetDatabase, targetCollection);
 			}
 
-			return new ReferenceCollection(ref.getString("db"), ref.get("collection", collection));
+			return new ReferenceCollection(documentPointer.getString("db"), documentPointer.get("collection", collection));
 		}
 
 		if (property.isDocumentReference()) {
@@ -141,19 +159,34 @@ private ReferenceCollection computeReferenceContext(MongoPersistentProperty prop
 		return new ReferenceCollection(null, collection);
 	}
 
+	/**
+	 * Use the given {@link ParameterBindingContext} to compute potential expressions against the value.
+	 *
+	 * @param value must not be {@literal null}.
+	 * @param bindingContext must not be {@literal null}.
+	 * @param defaultValue
+	 * @param <T>
+	 * @return can be {@literal null}.
+	 */
+	@Nullable
 	@SuppressWarnings("unchecked")
 	private <T> T parseValueOrGet(String value, ParameterBindingContext bindingContext, Supplier<T> defaultValue) {
 
 		if (!StringUtils.hasText(value)) {
 			return defaultValue.get();
 		}
 
+		// parameter binding requires a document, since we do not have one, construct it.
 		if (!BsonUtils.isJsonDocument(value) && value.contains("?#{")) {
 			String s = "{ 'target-value' : " + value + "}";
 			T evaluated = (T) codec.decode(s, bindingContext).get("target-value");
 			return evaluated != null ? evaluated : defaultValue.get();
 		}
 
+		if(BsonUtils.isJsonDocument(value)) {
+			return (T) codec.decode(value, bindingContext);
+		}
+
 		T evaluated = (T) bindingContext.evaluateExpression(value);
 		return evaluated != null ? evaluated : defaultValue.get();
 	}
@@ -165,8 +198,8 @@ ParameterBindingContext bindingContext(MongoPersistentProperty property, Object
 	}
 
 	ValueProvider valueProviderFor(Object source) {
-		return (index) -> {
 
+		return (index) -> {
 			if (source instanceof Document) {
 				return Streamable.of(((Document) source).values()).toList().get(index);
 			}
@@ -183,13 +216,22 @@ EvaluationContext evaluationContextFor(MongoPersistentProperty property, Object
 		return ctx;
 	}
 
+	/**
+	 * Compute the query to retrieve linked documents.
+	 *
+	 * @param property must not be {@literal null}.
+	 * @param value must not be {@literal null}.
+	 * @param spELContext must not be {@literal null}.
+	 * @return never {@literal null}.
+	 */
 	@SuppressWarnings("unchecked")
 	DocumentReferenceQuery computeFilter(MongoPersistentProperty property, Object value, SpELContext spELContext) {
 
 		DocumentReference documentReference = property.getDocumentReference();
 		String lookup = documentReference.lookup();
 
-		Document sort = parseValueOrGet(documentReference.sort(), bindingContext(property, value, spELContext), () -> null);
+		Document sort = parseValueOrGet(documentReference.sort(), bindingContext(property, value, spELContext),
+				() -> new Document());
 
 		if (property.isCollectionLike() && value instanceof Collection) {
 
@@ -219,45 +261,59 @@ DocumentReferenceQuery computeFilter(MongoPersistentProperty property, Object va
 		return new SingleDocumentReferenceQuery(codec.decode(lookup, bindingContext(property, value, spELContext)), sort);
 	}
 
+	/**
+	 * {@link DocumentReferenceQuery} implementation fetching a single {@link Document}.
+	 */
 	static class SingleDocumentReferenceQuery implements DocumentReferenceQuery {
 
-		Document filter;
-		Document sort;
+		private final Document query;
+		private final Document sort;
+
+		public SingleDocumentReferenceQuery(Document query, Document sort) {
 
-		public SingleDocumentReferenceQuery(Document filter, Document sort) {
-			this.filter = filter;
+			this.query = query;
 			this.sort = sort;
 		}
 
 		@Override
 		public Bson getQuery() {
-			return filter;
+			return query;
+		}
+
+		@Override
+		public Document getSort() {
+			return sort;
 		}
 
 		@Override
 		public Iterable<Document> apply(MongoCollection<Document> collection) {
 
-			Document result = collection.find(getQuery()).limit(1).first();
+			Document result = collection.find(getQuery()).sort(getSort()).limit(1).first();
 			return result != null ? Collections.singleton(result) : Collections.emptyList();
 		}
 	}
 
+	/**
+	 * {@link DocumentReferenceQuery} implementation to retrieve linked {@link Document documents} stored inside a
+	 * {@link Map} structure. Restores the original map order by matching individual query documents against the actual
+	 * values.
+	 */
 	static class MapDocumentReferenceQuery implements DocumentReferenceQuery {
 
-		private final Document filter;
+		private final Document query;
 		private final Document sort;
 		private final Map<Object, Document> filterOrderMap;
 
-		public MapDocumentReferenceQuery(Document filter, Document sort, Map<Object, Document> filterOrderMap) {
+		public MapDocumentReferenceQuery(Document query, Document sort, Map<Object, Document> filterOrderMap) {
 
-			this.filter = filter;
+			this.query = query;
 			this.sort = sort;
 			this.filterOrderMap = filterOrderMap;
 		}
 
 		@Override
 		public Bson getQuery() {
-			return filter;
+			return query;
 		}
 
 		@Override
@@ -283,33 +339,38 @@ public Iterable<Document> restoreOrder(Iterable<Document> documents) {
 		}
 	}
 
+	/**
+	 * {@link DocumentReferenceQuery} implementation to retrieve linked {@link Document documents} stored inside a
+	 * {@link Collection} like structure. Restores the original order by matching individual query documents against the
+	 * actual values.
+	 */
 	static class ListDocumentReferenceQuery implements DocumentReferenceQuery {
 
-		private final Document filter;
+		private final Document query;
 		private final Document sort;
 
-		public ListDocumentReferenceQuery(Document filter, Document sort) {
+		public ListDocumentReferenceQuery(Document query, Document sort) {
 
-			this.filter = filter;
+			this.query = query;
 			this.sort = sort;
 		}
 
 		@Override
 		public Iterable<Document> restoreOrder(Iterable<Document> documents) {
 
-			if (filter.containsKey("$or")) {
-				List<Document> ors = filter.get("$or", List.class);
-				List<Document> target = documents instanceof List ? (List<Document>) documents
-						: Streamable.of(documents).toList();
-				return target.stream().sorted((o1, o2) -> compareAgainstReferenceIndex(ors, o1, o2))
-						.collect(Collectors.toList());
+			List<Document> target = documents instanceof List ? (List<Document>) documents
+					: Streamable.of(documents).toList();
+
+			if (!sort.isEmpty() || !query.containsKey("$or")) {
+				return target;
 			}
 
-			return documents;
+			List<Document> ors = query.get("$or", List.class);
+			return target.stream().sorted((o1, o2) -> compareAgainstReferenceIndex(ors, o1, o2)).collect(Collectors.toList());
 		}
 
 		public Document getQuery() {
-			return filter;
+			return query;
 		}
 
 		@Override
@@ -333,9 +394,18 @@ int compareAgainstReferenceIndex(List<Document> referenceList, Document document
 		}
 	}
 
+	/**
+	 * The function that can execute a given {@link DocumentReferenceQuery} within the {@link ReferenceCollection} to
+	 * obtain raw results.
+	 */
 	@FunctionalInterface
 	interface LookupFunction {
 
+		/**
+		 * @param referenceQuery never {@literal null}.
+		 * @param referenceCollection never {@literal null}.
+		 * @return never {@literal null}.
+		 */
 		Iterable<Document> apply(DocumentReferenceQuery referenceQuery, ReferenceCollection referenceCollection);
 	}
 }
diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/BsonUtils.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/util/BsonUtils.java
@@ -349,7 +349,13 @@ public static String toJson(@Nullable Document source) {
 	 * @since 3.0
 	 */
 	public static boolean isJsonDocument(@Nullable String value) {
-		return StringUtils.hasText(value) && (value.startsWith("{") && value.endsWith("}"));
+
+		if(!StringUtils.hasText(value)) {
+			return false;
+		}
+
+		String potentialJson = value.trim();
+		return potentialJson.startsWith("{") && potentialJson.endsWith("}");
 	}
 
 	/**
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateDocumentReferenceTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateDocumentReferenceTests.java
@@ -380,6 +380,33 @@ void readCollectionObjectReferenceFromDocumentDeclaringCollectionName() {
 				new ObjectRefOfDocumentWithEmbeddedCollectionName("ref-1", "me-the-1-referenced-object"));
 	}
 
+	@Test // GH-3602
+	void useOrderFromAnnotatedSort() {
+
+		String rootCollectionName = template.getCollectionName(CollectionRefRoot.class);
+		String refCollectionName = template.getCollectionName(SimpleObjectRef.class);
+		Document refSource1 = new Document("_id", "ref-1").append("value", "me-the-1-referenced-object");
+		Document refSource2 = new Document("_id", "ref-2").append("value", "me-the-2-referenced-object");
+		Document refSource3 = new Document("_id", "ref-3").append("value", "me-the-3-referenced-object");
+		Document source = new Document("_id", "id-1").append("value", "v1").append("simpleSortedValueRef",
+				Arrays.asList("ref-1", "ref-3", "ref-2"));
+
+		template.execute(db -> {
+
+			db.getCollection(refCollectionName).insertOne(refSource1);
+			db.getCollection(refCollectionName).insertOne(refSource2);
+			db.getCollection(refCollectionName).insertOne(refSource3);
+			db.getCollection(rootCollectionName).insertOne(source);
+			return null;
+		});
+
+		CollectionRefRoot result = template.findOne(query(where("id").is("id-1")), CollectionRefRoot.class);
+		assertThat(result.getSimpleSortedValueRef()).containsExactly(
+				new SimpleObjectRef("ref-3", "me-the-3-referenced-object"),
+				new SimpleObjectRef("ref-2", "me-the-2-referenced-object"),
+				new SimpleObjectRef("ref-1", "me-the-1-referenced-object"));
+	}
+
 	@Test // GH-3602
 	void readObjectReferenceFromDocumentNotRelatingToTheIdProperty() {
 
@@ -857,7 +884,8 @@ void updateDerivedMappingFromLookup() {
 
 		template.save(book);
 
-		template.update(Book.class).matching(where("id").is(book.id)).apply(new Update().set("publisher", publisher)).first();
+		template.update(Book.class).matching(where("id").is(book.id)).apply(new Update().set("publisher", publisher))
+				.first();
 
 		Document target = template.execute(db -> {
 			return db.getCollection(template.getCollectionName(Book.class)).find(Filters.eq("_id", book.id)).first();
@@ -930,6 +958,9 @@ static class CollectionRefRoot {
 		@DocumentReference(lookup = "{ '_id' : '?#{#target}' }") //
 		List<SimpleObjectRef> simpleValueRef;
 
+		@DocumentReference(lookup = "{ '_id' : '?#{#target}' }", sort = "{ '_id' : -1 } ") //
+		List<SimpleObjectRef> simpleSortedValueRef;
+
 		@DocumentReference(lookup = "{ '_id' : '?#{#target}' }") //
 		Map<String, SimpleObjectRef> mapValueRef;
 
diff --git a/src/main/asciidoc/reference/mapping.adoc b/src/main/asciidoc/reference/mapping.adoc
@@ -902,6 +902,10 @@ It is possible to alter resolution defaults (listed below) via the attributes of
 | The single document lookup query evaluating placeholders via SpEL expressions using `#target` as the marker for a given source value. `Collection` like or `Map` properties combine individual lookups via an `$or` operator.
 | An `_id` field based query (`{ '_id' : ?#{#target} }`) using the loaded source value.
 
+| `sort`
+| Used for sorting result documents on server side.
+| None by default. Result order of `Collection` like properties is restored based on the used lookup query.
+
 | `lazy`
 | If set to `true` value resolution is delayed upon first access of the property.
 | Resolves properties eagerly by default.
@@ -1182,7 +1186,7 @@ We know it is tempting to use all kinds of MongoDB query operators in the lookup
 * Mind that resolution takes time and consider a lazy strategy.
 * A collection of document references is bulk loaded using an `$or` operator. +
 The original element order is restored in memory which cannot be done when using MongoDB query operators.
-In this case Results will be ordered as they are received from the store.
+In this case Results will be ordered as they are received from the store or via the provided `@DocumentReference(sort = ...)` attribute.
 
 And a few more general remarks:
 
