Skip to content

Commit

Permalink
Fix ShardKey lookup for nested paths.
Browse files Browse the repository at this point in the history
This commit fixes the lookup of shard key values for nested paths using the dot (.) notation.

Closes: #3590
Original pull request: #3591.
  • Loading branch information
christophstrobl authored and mp911de committed Mar 15, 2021
1 parent 193b7de commit af40f15
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -658,7 +658,8 @@ <T> Document applyShardKey(MongoPersistentEntity<T> domainType, Document filter,
: 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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,41 @@ public static Document parse(String json, @Nullable CodecRegistryProvider codecR
.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) {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@
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;
Expand Down Expand Up @@ -1910,6 +1911,24 @@ void saveShouldAppendDefaultShardKeyIfNotPresentInFilter() {
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() {

Expand Down Expand Up @@ -2246,6 +2265,13 @@ static class Sith {
@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.
Expand Down

0 comments on commit af40f15

Please sign in to comment.