From af40f15a360c25967597b986e23641f473df5955 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@vmware.com>
Date: Thu, 11 Mar 2021 11:28:40 +0100
Subject: [PATCH] Fix ShardKey lookup for nested paths.

This commit fixes the lookup of shard key values for nested paths using the dot (.) notation.

Closes: #3590
Original pull request: #3591.
---
 .../data/mongodb/core/QueryOperations.java    |  3 +-
 .../data/mongodb/util/BsonUtils.java          | 35 +++++++++++++++++++
 .../mongodb/core/MongoTemplateUnitTests.java  | 26 ++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java
index 8376a182c..9d53cecb5 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/QueryOperations.java
@@ -658,7 +658,8 @@ class QueryOperations {
 					: mappedDocument != null ? mappedDocument.getDocument() : getMappedUpdate(domainType);
 
 			Document filterWithShardKey = new Document(filter);
-			getMappedShardKeyFields(domainType).forEach(key -> filterWithShardKey.putIfAbsent(key, shardKeySource.get(key)));
+			getMappedShardKeyFields(domainType)
+					.forEach(key -> filterWithShardKey.putIfAbsent(key, BsonUtils.resolveValue(shardKeySource, key)));
 
 			return filterWithShardKey;
 		}
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
index 5b036072b..2ac6744db 100644
--- 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
@@ -282,6 +282,41 @@ public class BsonUtils {
 				.orElseGet(() -> new DocumentCodec(codecRegistryProvider.getCodecRegistry())));
 	}
 
+	/**
+	 * Resolve a the value for a given key. If the given {@link Bson} value contains the key the value is immediately
+	 * returned. If not and the key contains a path using the dot ({@code .}) notation it will try to resolve the path by
+	 * inspecting the individual parts. If one of the intermediate ones is {@literal null} or cannot be inspected further
+	 * (wrong) type, {@literal null} is returned.
+	 *
+	 * @param bson the source to inspect. Must not be {@literal null}.
+	 * @param key the key to lookup. Must not be {@literal null}.
+	 * @return can be {@literal null}.
+	 */
+	@Nullable
+	public static Object resolveValue(Bson bson, String key) {
+
+		Map<String, Object> source = asMap(bson);
+
+		if (source.containsKey(key) || !key.contains(".")) {
+			return source.get(key);
+		}
+
+		String[] parts = key.split("\\.");
+
+		for (int i = 1; i < parts.length; i++) {
+
+			Object result = source.get(parts[i - 1]);
+
+			if (result == null || !(result instanceof Bson)) {
+				return null;
+			}
+
+			source = asMap((Bson) result);
+		}
+
+		return source.get(parts[parts.length - 1]);
+	}
+
 	@Nullable
 	private static String toJson(@Nullable Object value) {
 
diff --git a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
index 0ef616c48..f1ce4005a 100644
--- a/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
+++ b/spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java
@@ -84,6 +84,7 @@ import org.springframework.data.mongodb.core.geo.GeoJsonPoint;
 import org.springframework.data.mongodb.core.index.MongoPersistentEntityIndexCreator;
 import org.springframework.data.mongodb.core.mapping.Field;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.data.mongodb.core.mapping.Sharded;
 import org.springframework.data.mongodb.core.mapping.event.AbstractMongoEventListener;
 import org.springframework.data.mongodb.core.mapping.event.AfterConvertCallback;
 import org.springframework.data.mongodb.core.mapping.event.AfterSaveCallback;
@@ -1910,6 +1911,24 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		verify(findIterable, never()).first();
 	}
 
+	@Test // GH-3590
+	void shouldIncludeValueFromNestedShardKeyPath() {
+
+		WithShardKeyPoitingToNested source = new WithShardKeyPoitingToNested();
+		source.id = "id-1";
+		source.value = "v1";
+		source.nested = new WithNamedFields();
+		source.nested.customName = "cname";
+		source.nested.name = "name";
+
+		template.save(source);
+
+		ArgumentCaptor<Bson> filter = ArgumentCaptor.forClass(Bson.class);
+		verify(collection).replaceOne(filter.capture(), any(), any());
+
+		assertThat(filter.getValue()).isEqualTo(new Document("_id", "id-1").append("value", "v1").append("nested.custom-named-field", "cname"));
+	}
+
 	@Test // DATAMONGO-2341
 	void saveShouldProjectOnShardKeyWhenLoadingExistingDocument() {
 
@@ -2246,6 +2265,13 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		@Field("firstname") String name;
 	}
 
+	@Sharded(shardKey = {"value", "nested.customName"})
+	static class WithShardKeyPoitingToNested {
+		String id;
+		String value;
+		WithNamedFields nested;
+	}
+
 	/**
 	 * Mocks out the {@link MongoTemplate#getDb()} method to return the {@link DB} mock instead of executing the actual
 	 * behaviour.