Added unique key policy support to spring data SDK (#27270)

* Added unique key policy support to spring data SDK

* Added empty line at the end

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/UniqueIndexTest.java

@@ -166,8 +166,9 @@ public void uniqueKeySerializationDeserialization() {
 
         CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(UUID.randomUUID().toString(), partitionKeyDef);
         UniqueKeyPolicy uniqueKeyPolicy = new UniqueKeyPolicy();
-        UniqueKey uniqueKey = new UniqueKey(ImmutableList.of("/name", "/description"));
-        uniqueKeyPolicy.setUniqueKeys(Lists.newArrayList(uniqueKey));
+        UniqueKey uniqueKey = new UniqueKey(ImmutableList.of("/name"));
+        UniqueKey uniqueKey1 = new UniqueKey(ImmutableList.of("/description"));
+        uniqueKeyPolicy.setUniqueKeys(Lists.newArrayList(uniqueKey, uniqueKey1));
         collectionDefinition.setUniqueKeyPolicy(uniqueKeyPolicy);
 
         IndexingPolicy indexingPolicy = new IndexingPolicy();
@@ -191,10 +192,10 @@ public void uniqueKeySerializationDeserialization() {
         assertThat(collection.getUniqueKeyPolicy().getUniqueKeys()).isNotNull();
         assertThat(collection.getUniqueKeyPolicy().getUniqueKeys())
                 .hasSameSizeAs(collectionDefinition.getUniqueKeyPolicy().getUniqueKeys());
-        assertThat(collection.getUniqueKeyPolicy().getUniqueKeys()
-                .stream().map(ui -> ui.getPaths()).collect(Collectors.toList()))
-                .containsExactlyElementsOf(
-                        ImmutableList.of(ImmutableList.of("/name", "/description")));
+        //  check first unique key policy
+        assertThat(collection.getUniqueKeyPolicy().getUniqueKeys().get(0).getPaths().get(0)).isEqualTo("/name");
+        //  check second unique key policy
+        assertThat(collection.getUniqueKeyPolicy().getUniqueKeys().get(1).getPaths().get(0)).isEqualTo("/description");
     }
 
     private CosmosException getDocumentClientException(RuntimeException e) {

sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/common/PageTestUtils.java

@@ -2,16 +2,21 @@
 // Licensed under the MIT License.
 package com.azure.spring.data.cosmos.common;
 
+import com.azure.spring.data.cosmos.core.convert.ObjectMapperFactory;
 import com.azure.spring.data.cosmos.core.query.CosmosPageRequest;
-import org.json.JSONException;
-import org.json.JSONObject;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.Pageable;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.Assert.assertTrue;
 
 public class PageTestUtils {
+
+    private static final ObjectMapper objectMapper = ObjectMapperFactory.getObjectMapper();
+
     public static void validateLastPage(Page<?> page, int pageSize) {
         final Pageable pageable = page.getPageable();
 
@@ -36,15 +41,12 @@ private static boolean continuationTokenIsNull(CosmosPageRequest pageRequest) {
         if (tokenJson == null) {
             return true;
         }
-
-        final JSONObject jsonObject;
         try {
-            jsonObject = new JSONObject(tokenJson);
-            return jsonObject.isNull("compositeToken");
-        } catch (JSONException e) {
+            JsonNode jsonNode = objectMapper.readTree(tokenJson);
+            return jsonNode.get("compositeToken") == null;
+        } catch (JsonProcessingException e) {
             e.printStackTrace();
             return false;
         }
-
     }
 }

sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/domain/UniqueKeyPolicyEntity.java

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.domain;
+
+import com.azure.spring.data.cosmos.core.mapping.Container;
+import com.azure.spring.data.cosmos.core.mapping.CosmosUniqueKey;
+import com.azure.spring.data.cosmos.core.mapping.CosmosUniqueKeyPolicy;
+import com.azure.spring.data.cosmos.core.mapping.PartitionKey;
+import org.springframework.data.annotation.Id;
+
+@Container
+@CosmosUniqueKeyPolicy(uniqueKeys = {
+    @CosmosUniqueKey(paths = {"/lastName", "/zipCode"}),
+    @CosmosUniqueKey(paths = {"/city"})
+})
+public class UniqueKeyPolicyEntity {
+
+    @Id
+    String id;
+
+    @PartitionKey
+    String firstName;
+
+    String lastName;
+    String zipCode;
+    String city;
+
+    public UniqueKeyPolicyEntity(String id, String firstName, String lastName, String zipCode, String city) {
+        this.id = id;
+        this.firstName = firstName;
+        this.lastName = lastName;
+        this.zipCode = zipCode;
+        this.city = city;
+    }
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(String id) {
+        this.id = id;
+    }
+
+    public String getFirstName() {
+        return firstName;
+    }
+
+    public void setFirstName(String firstName) {
+        this.firstName = firstName;
+    }
+
+    public String getZipCode() {
+        return zipCode;
+    }
+
+    public void setZipCode(String zipCode) {
+        this.zipCode = zipCode;
+    }
+
+    public String getLastName() {
+        return lastName;
+    }
+
+    public void setLastName(String lastName) {
+        this.lastName = lastName;
+    }
+
+    public String getCity() {
+        return city;
+    }
+
+    public void setCity(String city) {
+        this.city = city;
+    }
+}

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

@@ -0,0 +1,162 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.spring.data.cosmos.repository.integration;
+
+import com.azure.cosmos.models.CosmosContainerProperties;
+import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.UniqueKey;
+import com.azure.cosmos.models.UniqueKeyPolicy;
+import com.azure.spring.data.cosmos.IntegrationTestCollectionManager;
+import com.azure.spring.data.cosmos.core.CosmosTemplate;
+import com.azure.spring.data.cosmos.core.ReactiveCosmosTemplate;
+import com.azure.spring.data.cosmos.domain.CompositeIndexEntity;
+import com.azure.spring.data.cosmos.domain.UniqueKeyPolicyEntity;
+import com.azure.spring.data.cosmos.exception.CosmosAccessException;
+import com.azure.spring.data.cosmos.repository.TestRepositoryConfig;
+import com.azure.spring.data.cosmos.repository.repository.UniqueKeyPolicyEntityRepository;
+import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation;
+import com.azure.spring.data.cosmos.repository.support.SimpleCosmosRepository;
+import com.azure.spring.data.cosmos.repository.support.SimpleReactiveCosmosRepository;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import java.util.Arrays;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.fail;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(classes = TestRepositoryConfig.class)
+public class UniqueKeyPolicyIT {
+
+    @ClassRule
+    public static final IntegrationTestCollectionManager collectionManager = new IntegrationTestCollectionManager();
+
+    private static final UniqueKeyPolicyEntity ENTITY_1 = new UniqueKeyPolicyEntity("id-1", "firstName-1", "lastName"
+        + "-1", "100", "city-1");
+    private static final UniqueKeyPolicyEntity ENTITY_2 = new UniqueKeyPolicyEntity("id-2", "firstName-1", "lastName"
+        + "-2", "100", "city-2");
+    private static final UniqueKeyPolicyEntity ENTITY_3 = new UniqueKeyPolicyEntity("id-3", "firstName-2", "lastName"
+        + "-3", "100", "city-1");
+    private static final UniqueKeyPolicyEntity ENTITY_4 = new UniqueKeyPolicyEntity("id-4", "firstName-2", "lastName"
+        + "-4", "100", "city-2");
+    private static final UniqueKeyPolicyEntity ENTITY_5 = new UniqueKeyPolicyEntity("id-5", "firstName-3", "lastName"
+        + "-5", "100", "city-3");
+
+    @Autowired
+    UniqueKeyPolicyEntityRepository repository;
+
+    @Autowired
+    CosmosTemplate template;
+
+    @Autowired
+    ReactiveCosmosTemplate reactiveTemplate;
+
+    CosmosEntityInformation<UniqueKeyPolicyEntity, String> information =
+        new CosmosEntityInformation<>(UniqueKeyPolicyEntity.class);
+
+    @Before
+    public void setup() {
+        collectionManager.ensureContainersCreatedAndEmpty(template, CompositeIndexEntity.class);
+        repository.saveAll(Arrays.asList(ENTITY_1, ENTITY_2, ENTITY_3, ENTITY_4, ENTITY_5));
+    }
+
+    @Test
+    public void canSetUniqueKeyPolicy() {
+        new SimpleCosmosRepository<>(information, template);
+        CosmosContainerProperties properties = template.getContainerProperties(information.getContainerName());
+        UniqueKeyPolicy uniqueKeyPolicy = properties.getUniqueKeyPolicy();
+        List<UniqueKey> uniqueKeys = uniqueKeyPolicy.getUniqueKeys();
+
+        assertThat(uniqueKeys.size()).isEqualTo(2);
+
+        assertThat(uniqueKeys.get(0).getPaths().get(0)).isEqualTo("/lastName");
+        assertThat(uniqueKeys.get(0).getPaths().get(1)).isEqualTo("/zipCode");
+
+        assertThat(uniqueKeys.get(1).getPaths().get(0)).isEqualTo("/city");
+    }
+
+    @Test
+    public void canSetUniqueKeyPolicyReactive() {
+        new SimpleReactiveCosmosRepository<>(information, reactiveTemplate);
+        CosmosContainerProperties properties =
+            reactiveTemplate.getContainerProperties(information.getContainerName()).block();
+        UniqueKeyPolicy uniqueKeyPolicy = properties.getUniqueKeyPolicy();
+        List<UniqueKey> uniqueKeys = uniqueKeyPolicy.getUniqueKeys();
+
+        assertThat(uniqueKeys.size()).isEqualTo(2);
+
+        assertThat(uniqueKeys.get(0).getPaths().get(0)).isEqualTo("/lastName");
+        assertThat(uniqueKeys.get(0).getPaths().get(1)).isEqualTo("/zipCode");
+
+        assertThat(uniqueKeys.get(1).getPaths().get(0)).isEqualTo("/city");
+    }
+
+    @Test
+    public void canSaveNewEntityWithDifferentUniqueKeys() {
+        long count = repository.count();
+        assertThat(count).isEqualTo(5);
+        UniqueKeyPolicyEntity entity = new UniqueKeyPolicyEntity("id-6", "firstName-3", "lastName-6",
+            "100", "city-1");
+        repository.save(entity);
+        count = repository.count();
+        assertThat(count).isEqualTo(6);
+        repository.deleteById("id-6", new PartitionKey("firstName-3"));
+        count = repository.count();
+        assertThat(count).isEqualTo(5);
+    }
+
+    @Test
+    public void cannotSaveNewEntityWithUniqueKeysLastNameAndZipCode() {
+        //  save with same lastName and zip code (which already exists in the same logical partition), though with a new id
+        UniqueKeyPolicyEntity entity = new UniqueKeyPolicyEntity("id-6", "firstName-3", "lastName-5",
+            "100", "city-6");
+        try {
+            repository.save(entity);
+            fail("Save call should have failed with unique constraints exception");
+        } catch (CosmosAccessException cosmosAccessException) {
+            assertThat(cosmosAccessException.getCosmosException().getStatusCode()).isEqualTo(409);
+            assertThat(cosmosAccessException.getCosmosException().getMessage()).contains("Unique index constraint "
+                + "violation.");
+        }
+        //  change logical partition, now the entity should be saved
+        entity.setFirstName("firstName-2");
+        repository.save(entity);
+        long count = repository.count();
+        assertThat(count).isEqualTo(6);
+        repository.deleteById("id-6", new PartitionKey("firstName-2"));
+        count = repository.count();
+        assertThat(count).isEqualTo(5);
+    }
+
+    @Test
+    public void cannotSaveNewEntityWithUniqueKeysCity() {
+        //  save with same city (which already exists in the same logical partition), though with a new id
+        UniqueKeyPolicyEntity entity = new UniqueKeyPolicyEntity("id-6", "firstName-3", "lastName-6",
+            "100", "city-3");
+        try {
+            repository.save(entity);
+            fail("Save call should have failed with unique constraints exception");
+        } catch (CosmosAccessException cosmosAccessException) {
+            assertThat(cosmosAccessException.getCosmosException().getStatusCode()).isEqualTo(409);
+            assertThat(cosmosAccessException.getCosmosException().getMessage()).contains("Unique index constraint "
+                + "violation.");
+        }
+        //  change logical partition, now the entity should be saved
+        entity.setFirstName("firstName-2");
+        repository.save(entity);
+        long count = repository.count();
+        assertThat(count).isEqualTo(6);
+        repository.deleteById("id-6", new PartitionKey("firstName-2"));
+        count = repository.count();
+        assertThat(count).isEqualTo(5);
+    }
+
+}

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

@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.repository.repository;
+
+import com.azure.spring.data.cosmos.domain.UniqueKeyPolicyEntity;
+import com.azure.spring.data.cosmos.repository.CosmosRepository;
+import org.springframework.stereotype.Repository;
+
+@Repository
+public interface UniqueKeyPolicyEntityRepository extends CosmosRepository<UniqueKeyPolicyEntity, String> {
+
+}

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

@@ -402,12 +402,13 @@ public class SampleApplication implements CommandLineRunner {
     - delete entity
 
 ### Spring Data Annotations
-- Spring Data [@Id annotation][spring_data_commons_id_annotation].
+#### Spring Data [@Id annotation][spring_data_commons_id_annotation]
   There are 2 ways to map a field in domain class to `id` field of Azure Cosmos DB Item.
   - annotate a field in domain class with `@Id`, this field will be mapped to Item `id` in Cosmos DB.
   - set name of this field to `id`, this field will be mapped to Item `id` in Azure Cosmos DB.
-- Supports auto generation of string type UUIDs using the @GeneratedValue annotation. The id field of an entity with a string
- type id can be annotated with `@GeneratedValue` to automatically generate a random UUID prior to insertion.
+#### Id auto generation
+  - Supports auto generation of string type UUIDs using the @GeneratedValue annotation. The id field of an entity with a string
+   type id can be annotated with `@GeneratedValue` to automatically generate a random UUID prior to insertion.
  ```java readme-sample-GeneratedIdEntity
  public class GeneratedIdEntity {
 
@@ -417,7 +418,7 @@ public class SampleApplication implements CommandLineRunner {
 
  }
  ```
-- SpEL Expression and Custom Container Name.
+#### SpEL Expression and Custom Container Name.
   - By default, container name will be class name of user domain class. To customize it, add the `@Container(containerName="myCustomContainerName")` annotation to the domain class. The container field also supports SpEL expressions (eg. `container = "${dynamic.container.name}"` or `container = "#{@someBean.getContainerName()}"`) in order to provide container names programmatically/via configuration properties.
   - In order for SpEL expressions to work properly, you need to add `@DependsOn("expressionResolver")` on top of Spring Application class.
 ```java
@@ -427,8 +428,8 @@ public class SampleApplication {
 
 }
 ```
-- Custom IndexingPolicy
-  By default, IndexingPolicy will be set by azure service. To customize it add annotation `@CosmosIndexingPolicy` to domain class. This annotation has 4 attributes to customize, see following:
+#### Indexing Policy
+- By default, IndexingPolicy will be set by azure service. To customize it add annotation `@CosmosIndexingPolicy` to domain class. This annotation has 4 attributes to customize, see following:
 ```java readme-sample-CosmosIndexingPolicyCodeSnippet
 // Indicate if indexing policy use automatic or not
 // Default value is true
@@ -443,7 +444,27 @@ String[] includePaths() default {};
 // Excluded paths for indexing
 String[] excludePaths() default {};
 ```
+#### Unique Key Policy
+ - Spring Data Cosmos SDK supports setting `UniqueKeyPolicy` on container by adding the annotation `@CosmosUniqueKeyPolicy` to domain class. This annotation has the following attributes:
+```java readme-sample-CosmosUniqueKeyPolicyCodeSnippet
+@Container
+@CosmosUniqueKeyPolicy(uniqueKeys = {
+    @CosmosUniqueKey(paths = {"/lastName", "/zipCode"}),
+    @CosmosUniqueKey(paths = {"/city"})
+})
+public class CosmosUniqueKeyPolicyCodeSnippet {
 
+    @Id
+    String id;
+
+    @PartitionKey
+    String firstName;
+
+    String lastName;
+    String zipCode;
+    String city;
+}
+```
 ### Azure Cosmos DB Partition
 - Azure-spring-data-cosmos supports [Azure Cosmos DB partition][azure_cosmos_db_partition].
 - To specify a field of domain class to be partition key field, just annotate it with `@PartitionKey`.