Add support for keyset extraction of nested property paths.

Closes #4326
Original Pull Request: #4317

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

@@ -124,7 +124,7 @@ class EntityOperations {
 			return new SimpleMappedEntity((Map<String, Object>) entity);
 		}
 
-		return MappedEntity.of(entity, context);
+		return MappedEntity.of(entity, context, this);
 	}
 
 	/**
@@ -148,7 +148,7 @@ class EntityOperations {
 			return new SimpleMappedEntity((Map<String, Object>) entity);
 		}
 
-		return AdaptibleMappedEntity.of(entity, context, conversionService);
+		return AdaptibleMappedEntity.of(entity, context, conversionService, this);
 	}
 
 	/**
@@ -386,6 +386,16 @@ class EntityOperations {
 		 */
 		Object getId();
 
+		/**
+		 * Returns the property value for {@code key}.
+		 *
+		 * @param key
+		 * @return
+		 * @since 4.1
+		 */
+		@Nullable
+		Object getPropertyValue(String key);
+
 		/**
 		 * Returns the {@link Query} to find the entity by its identifier.
 		 *
@@ -457,6 +467,11 @@ class EntityOperations {
 		 */
 		boolean isNew();
 
+		/**
+		 * @param sortObject
+		 * @return
+		 * @since 3.1
+		 */
 		Map<String, Object> extractKeys(Document sortObject);
 
 	}
@@ -518,7 +533,12 @@ class EntityOperations {
 
 		@Override
 		public Object getId() {
-			return map.get(ID_FIELD);
+			return getPropertyValue(ID_FIELD);
+		}
+
+		@Override
+		public Object getPropertyValue(String key) {
+			return map.get(key);
 		}
 
 		@Override
@@ -613,23 +633,26 @@ class EntityOperations {
 		private final MongoPersistentEntity<?> entity;
 		private final IdentifierAccessor idAccessor;
 		private final PersistentPropertyAccessor<T> propertyAccessor;
+		private final EntityOperations entityOperations;
 
 		protected MappedEntity(MongoPersistentEntity<?> entity, IdentifierAccessor idAccessor,
-				PersistentPropertyAccessor<T> propertyAccessor) {
+				PersistentPropertyAccessor<T> propertyAccessor, EntityOperations entityOperations) {
 
 			this.entity = entity;
 			this.idAccessor = idAccessor;
 			this.propertyAccessor = propertyAccessor;
+			this.entityOperations = entityOperations;
 		}
 
 		private static <T> MappedEntity<T> of(T bean,
-				MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> context) {
+				MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> context,
+				EntityOperations entityOperations) {
 
 			MongoPersistentEntity<?> entity = context.getRequiredPersistentEntity(bean.getClass());
 			IdentifierAccessor identifierAccessor = entity.getIdentifierAccessor(bean);
 			PersistentPropertyAccessor<T> propertyAccessor = entity.getPropertyAccessor(bean);
 
-			return new MappedEntity<>(entity, identifierAccessor, propertyAccessor);
+			return new MappedEntity<>(entity, identifierAccessor, propertyAccessor, entityOperations);
 		}
 
 		@Override
@@ -642,6 +665,11 @@ class EntityOperations {
 			return idAccessor.getRequiredIdentifier();
 		}
 
+		@Override
+		public Object getPropertyValue(String key) {
+			return propertyAccessor.getProperty(entity.getRequiredPersistentProperty(key));
+		}
+
 		@Override
 		public Query getByIdQuery() {
 
@@ -728,13 +756,38 @@ class EntityOperations {
 
 			for (String key : sortObject.keySet()) {
 
-				// TODO: make this work for nested properties
+				if (key.indexOf('.') != -1) {
-				MongoPersistentProperty persistentProperty = entity.getRequiredPersistentProperty(key);
+
-				keyset.put(key, propertyAccessor.getProperty(persistentProperty));
+					// follow the path across nested levels.
+					// TODO: We should have a MongoDB-specific property path abstraction to allow diving into Document.
+					keyset.put(key, getNestedPropertyValue(key));
+				} else {
+					keyset.put(key, getPropertyValue(key));
+				}
 			}
 
 			return keyset;
 		}
+
+		@Nullable
+		private Object getNestedPropertyValue(String key) {
+
+			String[] segments = key.split("\\.");
+			Entity<?> currentEntity = this;
+			Object currentValue = null;
+
+			for (int i = 0; i < segments.length; i++) {
+
+				String segment = segments[i];
+				currentValue = currentEntity.getPropertyValue(segment);
+
+				if (i < segments.length - 1) {
+					currentEntity = entityOperations.forEntity(currentValue);
+				}
+			}
+
+			return currentValue;
+		}
 	}
 
 	private static class AdaptibleMappedEntity<T> extends MappedEntity<T> implements AdaptibleEntity<T> {
@@ -744,9 +797,9 @@ class EntityOperations {
 		private final IdentifierAccessor identifierAccessor;
 
 		private AdaptibleMappedEntity(MongoPersistentEntity<?> entity, IdentifierAccessor identifierAccessor,
-				ConvertingPropertyAccessor<T> propertyAccessor) {
+				ConvertingPropertyAccessor<T> propertyAccessor, EntityOperations entityOperations) {
 
-			super(entity, identifierAccessor, propertyAccessor);
+			super(entity, identifierAccessor, propertyAccessor, entityOperations);
 
 			this.entity = entity;
 			this.propertyAccessor = propertyAccessor;
@@ -755,14 +808,14 @@ class EntityOperations {
 
 		private static <T> AdaptibleEntity<T> of(T bean,
 				MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> context,
-				ConversionService conversionService) {
+				ConversionService conversionService, EntityOperations entityOperations) {
 
 			MongoPersistentEntity<?> entity = context.getRequiredPersistentEntity(bean.getClass());
 			IdentifierAccessor identifierAccessor = entity.getIdentifierAccessor(bean);
 			PersistentPropertyAccessor<T> propertyAccessor = entity.getPropertyAccessor(bean);
 
 			return new AdaptibleMappedEntity<>(entity, identifierAccessor,
-					new ConvertingPropertyAccessor<>(propertyAccessor, conversionService));
+					new ConvertingPropertyAccessor<>(propertyAccessor, conversionService), entityOperations);
 		}
 
 		@Nullable

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/EntityOperationsUnitTests.java

@@ -17,10 +17,14 @@ package org.springframework.data.mongodb.core;
 
 import static org.assertj.core.api.Assertions.*;
 
+import lombok.AllArgsConstructor;
+import lombok.NoArgsConstructor;
+
 import java.time.Instant;
+import java.util.Map;
 
+import org.bson.Document;
 import org.junit.jupiter.api.Test;
-
 import org.springframework.core.convert.ConversionService;
 import org.springframework.core.convert.support.DefaultConversionService;
 import org.springframework.data.annotation.Id;
@@ -61,6 +65,57 @@ class EntityOperationsUnitTests {
 		assertThat(initAdaptibleEntity(new DomainTypeWithIdProperty()).populateIdIfNecessary(null)).isNotNull();
 	}
 
+	@Test // GH-4308
+	void shouldExtractKeysFromEntity() {
+
+		WithNestedDocument object = new WithNestedDocument("foo");
+
+		Map<String, Object> keys = operations.forEntity(object).extractKeys(new Document("id", 1));
+
+		assertThat(keys).containsEntry("id", "foo");
+	}
+
+	@Test // GH-4308
+	void shouldExtractKeysFromDocument() {
+
+		Document object = new Document("id", "foo");
+
+		Map<String, Object> keys = operations.forEntity(object).extractKeys(new Document("id", 1));
+
+		assertThat(keys).containsEntry("id", "foo");
+	}
+
+	@Test // GH-4308
+	void shouldExtractKeysFromNestedEntity() {
+
+		WithNestedDocument object = new WithNestedDocument("foo", new WithNestedDocument("bar"), null);
+
+		Map<String, Object> keys = operations.forEntity(object).extractKeys(new Document("nested.id", 1));
+
+		assertThat(keys).containsEntry("nested.id", "bar");
+	}
+
+	@Test // GH-4308
+	void shouldExtractKeysFromNestedEntityDocument() {
+
+		WithNestedDocument object = new WithNestedDocument("foo", new WithNestedDocument("bar"),
+				new Document("john", "doe"));
+
+		Map<String, Object> keys = operations.forEntity(object).extractKeys(new Document("document.john", 1));
+
+		assertThat(keys).containsEntry("document.john", "doe");
+	}
+
+	@Test // GH-4308
+	void shouldExtractKeysFromNestedDocument() {
+
+		Document object = new Document("document", new Document("john", "doe"));
+
+		Map<String, Object> keys = operations.forEntity(object).extractKeys(new Document("document.john", 1));
+
+		assertThat(keys).containsEntry("document.john", "doe");
+	}
+
 	<T> EntityOperations.AdaptibleEntity<T> initAdaptibleEntity(T source) {
 		return operations.forEntity(source, conversionService);
 	}
@@ -80,4 +135,19 @@ class EntityOperationsUnitTests {
 	static class InvalidMetaField {
 		Instant time;
 	}
+
+	@AllArgsConstructor
+	@NoArgsConstructor
+	class WithNestedDocument {
+
+		String id;
+
+		WithNestedDocument nested;
+
+		Document document;
+
+		public WithNestedDocument(String id) {
+			this.id = id;
+		}
+	}
 }

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateScrollTests.java

@@ -18,18 +18,23 @@ package org.springframework.data.mongodb.core;
 import static org.springframework.data.mongodb.core.query.Criteria.*;
 import static org.springframework.data.mongodb.test.util.Assertions.*;
 
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
 import java.util.Arrays;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
 import org.bson.Document;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.data.annotation.PersistenceCreator;
 import org.springframework.data.auditing.IsNewAwareAuditingHandler;
 import org.springframework.data.domain.KeysetScrollPosition;
 import org.springframework.data.domain.OffsetScrollPosition;
@@ -69,7 +74,6 @@ class MongoTemplateScrollTests {
 
 		cfg.configureMappingContext(it -> {
 			it.autocreateIndex(false);
-			it.initialEntitySet(AuditablePerson.class);
 		});
 
 		cfg.configureApplicationContext(it -> {
@@ -87,6 +91,39 @@ class MongoTemplateScrollTests {
 	@BeforeEach
 	void setUp() {
 		template.remove(Person.class).all();
+		template.remove(WithNestedDocument.class).all();
+	}
+
+	@Test
+	void shouldUseKeysetScrollingWithNestedSort() {
+
+		WithNestedDocument john20 = new WithNestedDocument(null, "John", 120, new WithNestedDocument("John", 20),
+				new Document("name", "bar"));
+		WithNestedDocument john40 = new WithNestedDocument(null, "John", 140, new WithNestedDocument("John", 40),
+				new Document("name", "baz"));
+		WithNestedDocument john41 = new WithNestedDocument(null, "John", 141, new WithNestedDocument("John", 41),
+				new Document("name", "foo"));
+
+		template.insertAll(Arrays.asList(john20, john40, john41));
+
+		Query q = new Query(where("name").regex("J.*")).with(Sort.by("nested.name", "nested.age", "document.name"))
+				.limit(2);
+		q.with(KeysetScrollPosition.initial());
+
+		Scroll<WithNestedDocument> scroll = template.scroll(q, WithNestedDocument.class);
+
+		assertThat(scroll.hasNext()).isTrue();
+		assertThat(scroll.isLast()).isFalse();
+		assertThat(scroll).hasSize(2);
+		assertThat(scroll).containsOnly(john20, john40);
+
+		scroll = template.scroll(q.with(scroll.lastPosition()), WithNestedDocument.class);
+
+		assertThat(scroll.hasNext()).isFalse();
+		assertThat(scroll.isLast()).isTrue();
+		assertThat(scroll).hasSize(1);
+		assertThat(scroll).containsOnly(john41);
+
 	}
 
 	@ParameterizedTest // GH-4308
@@ -144,4 +181,32 @@ class MongoTemplateScrollTests {
 		return new Document("_class", person.getClass().getName()).append("_id", person.getId()).append("active", true)
 				.append("firstName", person.getFirstName()).append("age", person.getAge());
 	}
+
+	@NoArgsConstructor
+	@Data
+	class WithNestedDocument {
+
+		String id;
+		String name;
+
+		int age;
+
+		WithNestedDocument nested;
+
+		Document document;
+
+		public WithNestedDocument(String name, int age) {
+			this.name = name;
+			this.age = age;
+		}
+
+		@PersistenceCreator
+		public WithNestedDocument(String id, String name, int age, WithNestedDocument nested, Document document) {
+			this.id = id;
+			this.name = name;
+			this.age = age;
+			this.nested = nested;
+			this.document = document;
+		}
+	}
 }