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.
This commit is contained in:
Christoph Strobl
2021-03-11 11:28:40 +01:00
committed by Mark Paluch
parent 193b7de2d9
commit af40f15a36
3 changed files with 63 additions and 1 deletions

View File

@@ -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;
}

View File

@@ -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) {

View File

@@ -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.