Fixing Nested PK functionality in CosmosEntityInformation (#38274)

sdk/spring/azure-spring-data-cosmos/CHANGELOG.md

@@ -7,9 +7,11 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed a bug with auto generated id's when using bulk `saveAll` - See [PR 38356](https://github.com/Azure/azure-sdk-for-java/pull/38356).
 
 #### Other Changes
 * Implemented a custom scheduler for `azure-spring-data-cosmos` - See [PR 38029](https://github.com/Azure/azure-sdk-for-java/pull/38029).
+* Optimized querying entities with nested partition keys by passing the nested partition key in `CosmosQueryRequestOptions` - See [PR 38356](https://github.com/Azure/azure-sdk-for-java/pull/38356).
 
 ### 5.8.0 (2023-12-14)
 

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java

@@ -260,6 +260,7 @@ public <S extends T, T> Iterable<S> insertAll(CosmosEntityInformation<T, ?> info
 
         List<CosmosItemOperation> cosmosItemOperations = new ArrayList<>();
         entities.forEach(entity -> {
+            generateIdIfNullAndAutoGenerationEnabled(entity, domainType);
             JsonNode originalItem = mappingCosmosConverter.writeJsonNode(entity);
             PartitionKey partitionKey = new PartitionKey(information.getPartitionKeyFieldValue(entity));
             final CosmosBulkItemRequestOptions options = new CosmosBulkItemRequestOptions();
@@ -750,7 +751,6 @@ public <S extends T, T> void deleteEntities(CosmosEntityInformation<T, ?> inform
         List<CosmosItemOperation> cosmosItemOperations = new ArrayList<>();
         entities.forEach(entity -> {
             JsonNode originalItem = mappingCosmosConverter.writeJsonNode(entity);
-            PartitionKey partitionKey = new PartitionKey(information.getPartitionKeyFieldValue(entity));
             final CosmosBulkItemRequestOptions options = new CosmosBulkItemRequestOptions();
             applyBulkVersioning(domainType, originalItem, options);
             cosmosItemOperations.add(CosmosBulkOperations.getDeleteItemOperation(String.valueOf(information.getId(entity)),

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java

@@ -497,6 +497,7 @@ public <S extends T, T> Flux<S> insertAll(CosmosEntityInformation<T, ?> entityIn
         Class<T> domainType = entityInformation.getJavaType();
 
         Flux<CosmosItemOperation> cosmosItemOperationsFlux = entities.map(entity -> {
+            generateIdIfNullAndAutoGenerationEnabled(entity, domainType);
             JsonNode originalItem = mappingCosmosConverter.writeJsonNode(entity);
             PartitionKey partitionKey = new PartitionKey(entityInformation.getPartitionKeyFieldValue(entity));
             final CosmosBulkItemRequestOptions options = new CosmosBulkItemRequestOptions();

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/repository/support/CosmosEntityInformation.java

@@ -35,6 +35,7 @@
 import java.util.List;
 import java.util.UUID;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static com.azure.spring.data.cosmos.common.ExpressionResolver.resolveExpression;
 
@@ -264,16 +265,39 @@ public String getVersionFieldValue(Object entity) {
      *
      * @param entity the target object from which to get the field
      * @return partition key field
+     * @throws RuntimeException thrown if field is not found
      */
     public Object getPartitionKeyFieldValue(T entity) {
-        return partitionKeyField == null ? null : ReflectionUtils.getField(partitionKeyField, entity);
+        if (partitionKeyField == null && partitionKeyPath != null) {
+            List<String> parts = Arrays.stream(partitionKeyPath.split("/")).collect(Collectors.toList());
+            final Object[] currentObject = {entity};
+            parts.forEach(part -> {
+                if (!part.isEmpty()) {
+                    Field f = null;
+                    try {
+                        f = currentObject[0].getClass().getDeclaredField(part);
+                    } catch (NoSuchFieldException e) {
+                        throw new RuntimeException(e);
+                    }
+                    ReflectionUtils.makeAccessible(f);
+                    currentObject[0] = ReflectionUtils.getField(f, currentObject[0]);
+                }
+            });
+            return currentObject[0];
+        } else {
+            return partitionKeyField == null ? null : ReflectionUtils.getField(partitionKeyField, entity);
+        }
     }
 
     /**
      * @return the partition key field name
      */
     public String getPartitionKeyFieldName() {
-        return partitionKeyField == null ? null : partitionKeyField.getName();
+        if (partitionKeyField == null && partitionKeyPath != null) {
+            return partitionKeyPath.substring(1).replace("/", ".");
+        } else {
+            return partitionKeyField == null ? null : partitionKeyField.getName();
+        }
     }
 
     /**

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/integration/NestedPartitionKeyRepositoryIT.java

@@ -113,4 +113,17 @@ public void testDeleteByEntity() {
         assertThat(nestedPartitionKeyEntityWithGeneratedValues.get(0)).isEqualTo(NESTED_ENTITY_2);
     }
 
+    @Test
+    public void testIdIsAutoGenerated() {
+        repository.saveAll(Arrays.asList(NESTED_ENTITY_1, NESTED_ENTITY_2));
+
+        Iterable<NestedPartitionKeyEntityWithGeneratedValue> iterable = repository.findAll();
+        List<NestedPartitionKeyEntityWithGeneratedValue> nestedPartitionKeyEntityWithGeneratedValues =
+            TestUtils.toList(iterable);
+
+        String pattern = "^[a-zA-Z0-9]{8}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{12}$";
+        assertThat(nestedPartitionKeyEntityWithGeneratedValues.get(0).getId()).matches(pattern);
+        assertThat(nestedPartitionKeyEntityWithGeneratedValues.get(1).getId()).matches(pattern);
+    }
+
 }

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/integration/ReactiveNestedPartitionKeyRepositoryIT.java

@@ -22,6 +22,10 @@
 import reactor.test.StepVerifier;
 
 import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
 
 @RunWith(SpringJUnit4ClassRunner.class)
 @ContextConfiguration(classes = TestRepositoryConfig.class)
@@ -119,4 +123,20 @@ public void testDeleteByEntity() {
         StepVerifier.create(findFlux).expectNext(NESTED_ENTITY_2).verifyComplete();
     }
 
+    @Test
+    public void testIdIsAutoGenerated() {
+        Flux<NestedPartitionKeyEntity> nestedPartitionKeyEntityFlux =
+            repository.saveAll(Arrays.asList(NESTED_ENTITY_1, NESTED_ENTITY_2));
+
+        StepVerifier.create(nestedPartitionKeyEntityFlux).expectNextCount(2).verifyComplete();
+
+        Flux<NestedPartitionKeyEntity> fluxEntities = repository.findAll();
+        StepVerifier.create(fluxEntities).expectNext(NESTED_ENTITY_1).expectNext(NESTED_ENTITY_2).verifyComplete();
+
+        List<NestedPartitionKeyEntity> results = fluxEntities.toStream().collect(Collectors.toList());
+        String pattern = "^[a-zA-Z0-9]{8}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{4}-[a-zA-Z0-9]{12}$";
+        assertThat(results.get(0).getId()).matches(pattern);
+        assertThat(results.get(0).getId()).matches(pattern);
+    }
+
 }

sdk/spring/azure-spring-data-cosmos/src/test/java/com/azure/spring/data/cosmos/repository/support/CosmosEntityInformationUnitTest.java

@@ -66,7 +66,7 @@ public void testCustomContainerName() {
     }
 
     @Test
-    public void testGetPartitionKeyName() {
+    public void testGetPartitionKeyPath() {
         final CosmosEntityInformation<VolunteerWithPartitionKey, String> entityInformation =
                 new CosmosEntityInformation<>(VolunteerWithPartitionKey.class);
 
@@ -75,7 +75,7 @@ public void testGetPartitionKeyName() {
     }
 
     @Test
-    public void testNullPartitionKeyName() {
+    public void testNullPartitionKeyPath() {
         final CosmosEntityInformation<Volunteer, String> entityInformation =
                 new CosmosEntityInformation<>(Volunteer.class);
 
@@ -84,7 +84,7 @@ public void testNullPartitionKeyName() {
     }
 
     @Test
-    public void testCustomPartitionKeyName() {
+    public void testCustomPartitionKeyPath() {
         final CosmosEntityInformation<VolunteerWithCustomPartitionKey, String> entityInformation =
                 new CosmosEntityInformation<>(VolunteerWithCustomPartitionKey.class);
 
@@ -98,7 +98,7 @@ public void testPartitionKeyPathAnnotation() {
             new CosmosEntityInformation<>(VolunteerWithPartitionKeyPath.class);
 
         final String partitionKeyPath = entityInformation.getPartitionKeyPath();
-        assertThat(partitionKeyPath).isEqualTo("/partitionKeyPath");
+        assertThat(partitionKeyPath).isEqualTo("/volunteerWithPartitionKey/name");
     }
 
     @Test
@@ -110,6 +110,92 @@ public void testPartitionKeyPathAndPartitionKeyAnnotation() {
         assertThat(partitionKeyPath).isEqualTo("/name");
     }
 
+    @Test
+    public void testGetPartitionKeyName() {
+        final CosmosEntityInformation<VolunteerWithPartitionKey, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithPartitionKey.class);
+
+        final String partitionKeyName = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyName).isEqualTo("name");
+    }
+
+    @Test
+    public void testNullPartitionKeyName() {
+        final CosmosEntityInformation<Volunteer, String> entityInformation =
+            new CosmosEntityInformation<>(Volunteer.class);
+
+        final String partitionKeyName = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyName).isEqualTo(null);
+    }
+
+    @Test
+    public void testCustomPartitionKeyName() {
+        final CosmosEntityInformation<VolunteerWithCustomPartitionKey, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithCustomPartitionKey.class);
+
+        final String partitionKeyName = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyName).isEqualTo("name");
+    }
+
+    @Test
+    public void testPartitionKeyPathAnnotationGetFieldName() {
+        final CosmosEntityInformation<VolunteerWithPartitionKeyPath, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithPartitionKeyPath.class);
+
+        final String partitionKeyPath = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyPath).isEqualTo("volunteerWithPartitionKey.name");
+    }
+
+    @Test
+    public void testPartitionKeyPathAndPartitionKeyAnnotationGetFieldName() {
+        final CosmosEntityInformation<VolunteerWithPartitionKeyPathAndPartitionKey, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithPartitionKeyPathAndPartitionKey.class);
+
+        final String partitionKeyPath = entityInformation.getPartitionKeyFieldName();
+        assertThat(partitionKeyPath).isEqualTo("name");
+    }
+
+    @Test
+    public void testGetPartitionKeyValue() {
+        final CosmosEntityInformation<VolunteerWithPartitionKey, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithPartitionKey.class);
+
+        VolunteerWithPartitionKey entity = new VolunteerWithPartitionKey();
+        entity.setName("MyName");
+
+        final String partitionKeyName = String.valueOf(entityInformation.getPartitionKeyFieldValue(entity));
+        assertThat(partitionKeyName).isEqualTo("MyName");
+    }
+
+    @Test
+    public void testPartitionKeyPathAnnotationGetFieldValue() {
+        final CosmosEntityInformation<VolunteerWithPartitionKeyPath, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithPartitionKeyPath.class);
+
+        VolunteerWithPartitionKeyPath entity = new VolunteerWithPartitionKeyPath();
+        VolunteerWithPartitionKey nestedEntity = new VolunteerWithPartitionKey();
+        nestedEntity.setName("NestedMyName");
+        entity.setVolunteerWithPartitionKey(nestedEntity);
+
+        final String partitionKeyPath = String.valueOf(entityInformation.getPartitionKeyFieldValue(entity));
+        assertThat(partitionKeyPath).isEqualTo("NestedMyName");
+    }
+
+    @Test
+    public void testPartitionKeyPathAndPartitionKeyAnnotationGetFieldValue() {
+        final CosmosEntityInformation<VolunteerWithPartitionKeyPathAndPartitionKey, String> entityInformation =
+            new CosmosEntityInformation<>(VolunteerWithPartitionKeyPathAndPartitionKey.class);
+
+        VolunteerWithPartitionKeyPathAndPartitionKey entity = new VolunteerWithPartitionKeyPathAndPartitionKey();
+        entity.setName("ActualName");
+        VolunteerWithPartitionKey nestedEntity = new VolunteerWithPartitionKey();
+        nestedEntity.setName("NestedMyName");
+        entity.setVolunteerWithPartitionKey(nestedEntity);
+
+        final String partitionKeyPath = String.valueOf(entityInformation.getPartitionKeyFieldValue(entity));
+        assertThat(partitionKeyPath).isEqualTo("ActualName");
+    }
+
     @Test
     public void testVersionedEntity() {
         final CosmosEntityInformation<VersionedVolunteer, String> entityInformation =
@@ -197,17 +283,31 @@ public void setName(String name) {
         }
     }
 
-    @Container(partitionKeyPath = "/partitionKeyPath")
+    @Container(partitionKeyPath = "/volunteerWithPartitionKey/name")
     private static class VolunteerWithPartitionKeyPath {
         private String id;
         private String name;
+        private VolunteerWithPartitionKey volunteerWithPartitionKey;
+
+        public void setVolunteerWithPartitionKey(VolunteerWithPartitionKey volunteerWithPartitionKey) {
+            this.volunteerWithPartitionKey = volunteerWithPartitionKey;
+        }
     }
 
-    @Container(partitionKeyPath = "/partitionKeyPath")
+    @Container(partitionKeyPath = "/volunteerWithPartitionKey/name")
     private static class VolunteerWithPartitionKeyPathAndPartitionKey {
         private String id;
         @PartitionKey
         private String name;
+        private VolunteerWithPartitionKey volunteerWithPartitionKey;
+
+        public void setVolunteerWithPartitionKey(VolunteerWithPartitionKey volunteerWithPartitionKey) {
+            this.volunteerWithPartitionKey = volunteerWithPartitionKey;
+        }
+
+        public void setName(String name) {
+            this.name = name;
+        }
     }
 
     @Container(containerName = "testContainer")