Skip to content
This repository was archived by the owner on Aug 29, 2024. It is now read-only.

Commit 6c20721

Browse files
kushagraThaparkushagraThapar
authored
Added new feature API for find by id and partition using read document() call (#433)
1 parent 4f43dcf commit 6c20721

13 files changed

+170
-11
lines changed

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbOperations.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ public interface DocumentDbOperations {
3131

3232
<T> T findById(String collectionName, Object id, Class<T> entityClass);
3333

34+
<T> T findById(Object id, Class<T> entityClass, PartitionKey partitionKey);
35+
3436
<T> T insert(T objectToSave, PartitionKey partitionKey);
3537

3638
<T> T insert(String collectionName, T objectToSave, PartitionKey partitionKey);

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplate.java

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,32 @@ public <T> T findById(String collectionName, Object id, Class<T> domainClass) {
153153
}
154154
}
155155

156+
@Override
157+
public <T> T findById(Object id, Class<T> entityClass, PartitionKey partitionKey) {
158+
Assert.notNull(entityClass, "entityClass should not be null");
159+
Assert.notNull(partitionKey, "partitionKey should not be null");
160+
assertValidId(id);
161+
162+
final com.azure.data.cosmos.PartitionKey pk = toCosmosPartitionKey(partitionKey);
163+
164+
try {
165+
final String collectionName = getCollectionName(entityClass);
166+
return cosmosClient
167+
.getDatabase(databaseName)
168+
.getContainer(collectionName)
169+
.getItem(id.toString(), pk)
170+
.read()
171+
.flatMap(cosmosItemResponse -> Mono.justOrEmpty(toDomainObject(entityClass,
172+
cosmosItemResponse.properties())))
173+
.onErrorResume(Mono::error)
174+
.block();
175+
176+
} catch (Exception e) {
177+
throw new DocumentDBAccessException("findById exception", e);
178+
}
179+
180+
}
181+
156182
public <T> void upsert(T object, PartitionKey partitionKey) {
157183
Assert.notNull(object, "Upsert object should not be null");
158184

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/ReactiveCosmosOperations.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
import com.azure.data.cosmos.CosmosContainerResponse;
1010
import com.azure.data.cosmos.PartitionKey;
11+
import com.microsoft.azure.spring.data.cosmosdb.core.convert.MappingDocumentDbConverter;
1112
import com.microsoft.azure.spring.data.cosmosdb.core.query.DocumentQuery;
1213
import com.microsoft.azure.spring.data.cosmosdb.repository.support.DocumentDbEntityInformation;
1314
import reactor.core.publisher.Flux;
@@ -25,6 +26,8 @@ public interface ReactiveCosmosOperations {
2526

2627
<T> Mono<T> findById(String collectionName, Object id, Class<T> entityClass);
2728

29+
<T> Mono<T> findById(Object id, Class<T> entityClass, PartitionKey partitionKey);
30+
2831
<T> Mono<T> insert(T objectToSave, PartitionKey partitionKey);
2932

3033
<T> Mono<T> insert(String collectionName, Object objectToSave, PartitionKey partitionKey);
@@ -50,4 +53,6 @@ public interface ReactiveCosmosOperations {
5053
Mono<Long> count(String collectionName);
5154

5255
Mono<Long> count(DocumentQuery query, String containerName);
56+
57+
MappingDocumentDbConverter getConverter();
5358
}

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/core/ReactiveCosmosTemplate.java

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica
4343

4444
private final String databaseName;
4545

46+
private final MappingDocumentDbConverter mappingDocumentDbConverter;
47+
4648
private final CosmosClient cosmosClient;
4749

4850
private final List<String> collectionCache;
@@ -62,6 +64,7 @@ public ReactiveCosmosTemplate(CosmosDbFactory cosmosDbFactory,
6264

6365
this.databaseName = dbName;
6466
this.collectionCache = new ArrayList<>();
67+
this.mappingDocumentDbConverter = mappingDocumentDbConverter;
6568

6669
this.cosmosClient = cosmosDbFactory.getCosmosClient();
6770
}
@@ -74,6 +77,11 @@ public void setApplicationContext(@NonNull ApplicationContext applicationContext
7477
// NOTE: When application context instance variable gets introduced, assign it here.
7578
}
7679

80+
@Override
81+
public MappingDocumentDbConverter getConverter() {
82+
return this.mappingDocumentDbConverter;
83+
}
84+
7785
/**
7886
* Creates a collection if it doesn't already exist
7987
*
@@ -157,12 +165,36 @@ public <T> Mono<T> findById(String containerName, Object id, Class<T> entityClas
157165
.flatMap(cosmosItemFeedResponse -> Mono.justOrEmpty(cosmosItemFeedResponse
158166
.results()
159167
.stream()
160-
.map(cosmosItem -> cosmosItem.toObject(entityClass))
168+
.map(cosmosItem -> toDomainObject(entityClass, cosmosItem))
161169
.findFirst()))
162170
.onErrorResume(Mono::error)
163171
.next();
164172
}
165173

174+
/**
175+
* Find by id
176+
*
177+
* @param id the id
178+
* @param entityClass the entity class
179+
* @param partitionKey partition Key
180+
* @return Mono with the item or error
181+
*/
182+
@Override
183+
public <T> Mono<T> findById(Object id, Class<T> entityClass, PartitionKey partitionKey) {
184+
Assert.notNull(entityClass, "entityClass should not be null");
185+
assertValidId(id);
186+
187+
final String containerName = getContainerName(entityClass);
188+
return cosmosClient.getDatabase(databaseName)
189+
.getContainer(containerName)
190+
.getItem(id.toString(), partitionKey)
191+
.read()
192+
.flatMap(cosmosItemResponse -> Mono.justOrEmpty(toDomainObject(entityClass,
193+
cosmosItemResponse.properties())))
194+
.onErrorResume(Mono::error);
195+
}
196+
197+
166198
/**
167199
* Insert
168200
*
@@ -190,7 +222,7 @@ public <T> Mono<T> insert(T objectToSave) {
190222
.getContainer(getContainerName(objectToSave.getClass()))
191223
.createItem(objectToSave, new CosmosItemRequestOptions())
192224
.onErrorResume(Mono::error)
193-
.flatMap(cosmosItemResponse -> Mono.just(cosmosItemResponse.properties().toObject(domainClass)));
225+
.flatMap(cosmosItemResponse -> Mono.just(toDomainObject(domainClass, cosmosItemResponse.properties())));
194226
}
195227

196228
/**
@@ -215,7 +247,7 @@ public <T> Mono<T> insert(String containerName, Object objectToSave, PartitionKe
215247
.getContainer(containerName)
216248
.createItem(objectToSave, options)
217249
.onErrorResume(Mono::error)
218-
.flatMap(cosmosItemResponse -> Mono.just(cosmosItemResponse.properties().toObject(domainClass)));
250+
.flatMap(cosmosItemResponse -> Mono.just(toDomainObject(domainClass, cosmosItemResponse.properties())));
219251
}
220252

221253
/**
@@ -249,7 +281,7 @@ public <T> Mono<T> upsert(String containerName, T object, PartitionKey partition
249281
return cosmosClient.getDatabase(this.databaseName)
250282
.getContainer(containerName)
251283
.upsertItem(object, options)
252-
.flatMap(cosmosItemResponse -> Mono.just(cosmosItemResponse.properties().toObject(domainClass)))
284+
.flatMap(cosmosItemResponse -> Mono.just(toDomainObject(domainClass, cosmosItemResponse.properties())))
253285
.onErrorResume(Mono::error);
254286
}
255287

@@ -341,7 +373,7 @@ public <T> Mono<T> delete(DocumentQuery query, Class<T> entityClass, String cont
341373
@Override
342374
public <T> Flux<T> find(DocumentQuery query, Class<T> entityClass, String containerName) {
343375
return findDocuments(query, entityClass, containerName)
344-
.map(cosmosItemProperties -> cosmosItemProperties.toObject(entityClass));
376+
.map(cosmosItemProperties -> toDomainObject(entityClass, cosmosItemProperties));
345377
}
346378

347379
/**
@@ -498,4 +530,9 @@ private Mono<CosmosItemProperties> deleteDocument(@NonNull CosmosItemProperties
498530
.map(cosmosItemResponse -> cosmosItemProperties);
499531
}
500532

533+
private <T> T toDomainObject(@NonNull Class<T> domainClass, CosmosItemProperties cosmosItemProperties) {
534+
return this.mappingDocumentDbConverter.read(domainClass, cosmosItemProperties);
535+
}
536+
537+
501538
}

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/DocumentDbRepository.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,25 @@
66

77
package com.microsoft.azure.spring.data.cosmosdb.repository;
88

9+
import com.microsoft.azure.documentdb.PartitionKey;
910
import org.springframework.data.repository.NoRepositoryBean;
1011
import org.springframework.data.repository.PagingAndSortingRepository;
1112

1213
import java.io.Serializable;
14+
import java.util.Optional;
1315

1416
@NoRepositoryBean
1517
public interface DocumentDbRepository<T, ID extends Serializable> extends PagingAndSortingRepository<T, ID> {
18+
19+
/**
20+
* Retrieves an entity by its id.
21+
*
22+
* @param id must not be {@literal null}.
23+
* @param partitionKey partition key value of entity
24+
* @return the entity with the given id or {@literal Optional#empty()} if none found
25+
* @throws IllegalArgumentException if {@code id} is {@literal null}.
26+
*/
27+
Optional<T> findById(ID id, PartitionKey partitionKey);
28+
1629
}
1730

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/ReactiveCosmosRepository.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,20 @@
55
*/
66
package com.microsoft.azure.spring.data.cosmosdb.repository;
77

8+
import com.azure.data.cosmos.PartitionKey;
89
import org.springframework.data.repository.NoRepositoryBean;
910
import org.springframework.data.repository.reactive.ReactiveSortingRepository;
11+
import reactor.core.publisher.Mono;
1012

1113
@NoRepositoryBean
1214
public interface ReactiveCosmosRepository<T, K> extends ReactiveSortingRepository<T, K> {
15+
16+
/**
17+
* Retrieves an entity by its id and partition key.
18+
* @param id must not be {@literal null}.
19+
* @param partitionKey partition key value of the entity.
20+
* @return {@link Mono} emitting the entity with the given id or {@link Mono#empty()} if none found.
21+
* @throws IllegalArgumentException in case the given {@code id} is {@literal null}.
22+
*/
23+
Mono<T> findById(K id, PartitionKey partitionKey);
1324
}

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleDocumentDbRepository.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,24 @@ public Optional<T> findById(ID id) {
141141
return Optional.ofNullable(operation.findById(information.getCollectionName(), id, information.getJavaType()));
142142
}
143143

144+
/**
145+
* find one entity per id without partitions
146+
*
147+
* @param id
148+
* @return
149+
*/
150+
@Override
151+
public Optional<T> findById(ID id, PartitionKey partitionKey) {
152+
Assert.notNull(id, "id must not be null");
153+
154+
if (id instanceof String && !StringUtils.hasText((String) id)) {
155+
return Optional.empty();
156+
}
157+
158+
return Optional.ofNullable(operation.findById(id, information.getJavaType(), partitionKey));
159+
}
160+
161+
144162
/**
145163
* return count of documents in one collection without partitions
146164
*

‎src/main/java/com/microsoft/azure/spring/data/cosmosdb/repository/support/SimpleReactiveCosmosRepository.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,13 @@ public Mono<T> findById(Publisher<K> publisher) {
8383
id, entityInformation.getJavaType()));
8484
}
8585

86+
@Override
87+
public Mono<T> findById(K id, PartitionKey partitionKey) {
88+
Assert.notNull(id, "The given id must not be null!");
89+
return cosmosOperations.findById(id,
90+
entityInformation.getJavaType(), partitionKey);
91+
}
92+
8693
@Override
8794
public Mono<Boolean> existsById(K id) {
8895
Assert.notNull(id, "The given id must not be null!");

‎src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/DocumentDbTemplatePartitionIT.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,15 @@ public void testFindWithPartition() {
111111
assertEquals(TEST_PERSON, result.get(0));
112112
}
113113

114+
@Test
115+
public void testFindByIdWithPartition() {
116+
final PartitionPerson partitionPersonById = dbTemplate.findById(TEST_PERSON.getId(),
117+
PartitionPerson.class,
118+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON)));
119+
120+
assertEquals(TEST_PERSON, partitionPersonById);
121+
}
122+
114123
@Test
115124
public void testFindByNonExistIdWithPartition() {
116125
final Criteria criteria = Criteria.getInstance(IS_EQUAL, PROPERTY_ID, Arrays.asList(NOT_EXIST_ID));

‎src/test/java/com/microsoft/azure/spring/data/cosmosdb/core/ReactiveCosmosTemplatePartitionIT.java

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
*/
66
package com.microsoft.azure.spring.data.cosmosdb.core;
77

8-
import com.fasterxml.jackson.databind.ObjectMapper;
98
import com.azure.data.cosmos.PartitionKey;
109
import com.microsoft.azure.spring.data.cosmosdb.CosmosDbFactory;
1110
import com.microsoft.azure.spring.data.cosmosdb.config.DocumentDBConfig;
@@ -80,14 +79,13 @@ public void setUp() throws ClassNotFoundException {
8079
final CosmosDbFactory dbFactory = new CosmosDbFactory(dbConfig);
8180

8281
final DocumentDbMappingContext mappingContext = new DocumentDbMappingContext();
83-
final ObjectMapper objectMapper = new ObjectMapper();
8482

8583
personInfo = new DocumentDbEntityInformation<>(PartitionPerson.class);
8684
containerName = personInfo.getCollectionName();
8785

8886
mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
8987

90-
final MappingDocumentDbConverter dbConverter = new MappingDocumentDbConverter(mappingContext, objectMapper);
88+
final MappingDocumentDbConverter dbConverter = new MappingDocumentDbConverter(mappingContext, null);
9189
cosmosTemplate = new ReactiveCosmosTemplate(dbFactory, dbConverter, DB_NAME);
9290
cosmosTemplate.createCollectionIfNotExists(personInfo).block();
9391

@@ -113,6 +111,17 @@ public void testFindWithPartition() {
113111
}).verifyComplete();
114112
}
115113

114+
@Test
115+
public void testFindByIdWithPartition() {
116+
final Mono<PartitionPerson> partitionPersonMono = cosmosTemplate.findById(TEST_PERSON.getId(),
117+
PartitionPerson.class,
118+
new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON)));
119+
StepVerifier.create(partitionPersonMono).consumeNextWith(actual -> {
120+
Assert.assertThat(actual.getFirstName(), is(equalTo(TEST_PERSON.getFirstName())));
121+
Assert.assertThat(actual.getLastName(), is(equalTo(TEST_PERSON.getLastName())));
122+
}).verifyComplete();
123+
}
124+
116125
// @Test
117126
// public void testFindByNonExistIdWithPartition() {
118127
//

‎src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/SimpleDocumentDbRepositoryUnitTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ public void testSave() {
6767

6868
@Test
6969
public void testFindOne() {
70-
when(dbOperations.findById(anyString(), any(), any())).thenReturn(TEST_PERSON);
70+
when(dbOperations.findById(anyString(), anyString(), any())).thenReturn(TEST_PERSON);
7171

7272
repository.save(TEST_PERSON);
7373

@@ -82,7 +82,7 @@ public void testFindOneExceptionForPartitioned() {
8282

8383
repository.save(TEST_PERSON);
8484

85-
when(dbOperations.findById(anyString(), any(), any()))
85+
when(dbOperations.findById(anyString(), anyString(), any()))
8686
.thenThrow(new UnsupportedOperationException(PARTITION_VALUE_REQUIRED_MSG));
8787

8888
final Person result = repository.findById(TEST_PERSON.getId()).get();
@@ -98,7 +98,7 @@ public void testUpdate() {
9898
TestConstants.UPDATED_HOBBIES, updatedAddress);
9999
repository.save(updatedPerson);
100100

101-
when(dbOperations.findById(anyString(), any(), any())).thenReturn(updatedPerson);
101+
when(dbOperations.findById(anyString(), anyString(), any())).thenReturn(updatedPerson);
102102

103103
final Person result = repository.findById(TEST_PERSON.getId()).get();
104104
assertEquals(updatedPerson, result);

‎src/test/java/com/microsoft/azure/spring/data/cosmosdb/repository/integration/AddressRepositoryIT.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
*/
66
package com.microsoft.azure.spring.data.cosmosdb.repository.integration;
77

8+
import com.microsoft.azure.documentdb.PartitionKey;
89
import com.microsoft.azure.spring.data.cosmosdb.common.TestConstants;
910
import com.microsoft.azure.spring.data.cosmosdb.common.TestUtils;
1011
import com.microsoft.azure.spring.data.cosmosdb.core.DocumentDbTemplate;
@@ -23,6 +24,7 @@
2324
import java.util.Arrays;
2425
import java.util.Comparator;
2526
import java.util.List;
27+
import java.util.Optional;
2628

2729
import static org.assertj.core.api.Assertions.assertThat;
2830

@@ -77,6 +79,14 @@ public void testFindAll() {
7779
assertThat(result.size()).isEqualTo(4);
7880
}
7981

82+
@Test
83+
public void testFindByIdWithPartitionKey() {
84+
final Optional<Address> addressById = repository.findById(TEST_ADDRESS1_PARTITION1.getPostalCode(),
85+
new PartitionKey(entityInformation.getPartitionKeyFieldValue(TEST_ADDRESS1_PARTITION1)));
86+
87+
assertThat(addressById.equals(TEST_ADDRESS1_PARTITION1));
88+
}
89+
8090
@Test
8191
public void testFindByIdForPartitionedCollection() {
8292
final List<Address> addresses = repository.findByPostalCode(TestConstants.POSTAL_CODE);

0 commit comments

Comments
 (0)