Add autoscale RU support for azure-spring-data-cosmos (#21851)

* Add autoscale RU support for azure-spring-data-cosmos

Resolves #12711

* Added sample to read me for auto scale throughput

* Fixed readme link

Co-authored-by: Kushagra Thapar <kuthapar@microsoft.com>

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

@@ -10,6 +10,7 @@
 
 public final class TestConstants {
 
+    public static final String AUTOSCALE_MAX_THROUGHPUT = "4000";
     private static final Address ADDRESS_1 = new Address("201107", "Zixing Road", "Shanghai");
     private static final Address ADDRESS_2 = new Address("200000", "Xuhui", "Shanghai");
     public static final String HOBBY1 = "photography";

sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java

@@ -6,8 +6,10 @@
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.implementation.ConflictException;
+import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.SqlQuerySpec;
+import com.azure.cosmos.models.ThroughputResponse;
 import com.azure.spring.data.cosmos.CosmosFactory;
 import com.azure.spring.data.cosmos.IntegrationTestCollectionManager;
 import com.azure.spring.data.cosmos.common.PageTestUtils;
@@ -23,6 +25,7 @@
 import com.azure.spring.data.cosmos.core.query.Criteria;
 import com.azure.spring.data.cosmos.core.query.CriteriaType;
 import com.azure.spring.data.cosmos.domain.AuditableEntity;
+import com.azure.spring.data.cosmos.domain.AutoScaleSample;
 import com.azure.spring.data.cosmos.domain.GenIdEntity;
 import com.azure.spring.data.cosmos.domain.Person;
 import com.azure.spring.data.cosmos.exception.CosmosAccessException;
@@ -93,6 +96,7 @@ public class CosmosTemplateIT {
     @ClassRule
     public static final IntegrationTestCollectionManager collectionManager = new IntegrationTestCollectionManager();
 
+    private static CosmosAsyncClient client;
     private static CosmosTemplate cosmosTemplate;
     private static CosmosEntityInformation<Person, String> personInfo;
     private static String containerName;
@@ -113,7 +117,7 @@ public class CosmosTemplateIT {
     @Before
     public void setUp() throws ClassNotFoundException {
         if (cosmosTemplate == null) {
-            CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
+            client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
             final CosmosFactory cosmosFactory = new CosmosFactory(client, TestConstants.DB_NAME);
 
             final CosmosMappingContext mappingContext = new CosmosMappingContext();
@@ -609,4 +613,19 @@ public void testSliceQuery() {
         assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics()).isNotNull();
         assertThat(responseDiagnosticsTestUtils.getCosmosResponseStatistics().getRequestCharge()).isGreaterThan(0);
     }
+
+    @Test
+    public void createWithAutoscale() throws ClassNotFoundException {
+        final CosmosEntityInformation<AutoScaleSample, String> autoScaleSampleInfo =
+            new CosmosEntityInformation<>(AutoScaleSample.class);
+        CosmosContainerProperties containerProperties = cosmosTemplate.createContainerIfNotExists(autoScaleSampleInfo);
+        assertNotNull(containerProperties);
+        ThroughputResponse throughput = client.getDatabase(TestConstants.DB_NAME)
+            .getContainer(autoScaleSampleInfo.getContainerName())
+            .readThroughput()
+            .block();
+        assertNotNull(throughput);
+        assertEquals(Integer.parseInt(TestConstants.AUTOSCALE_MAX_THROUGHPUT),
+            throughput.getProperties().getAutoscaleMaxThroughput());
+    }
 }

sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplateIT.java

@@ -7,8 +7,11 @@
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.implementation.ConflictException;
+import com.azure.cosmos.models.CosmosContainerProperties;
+import com.azure.cosmos.models.CosmosContainerResponse;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.SqlQuerySpec;
+import com.azure.cosmos.models.ThroughputResponse;
 import com.azure.spring.data.cosmos.CosmosFactory;
 import com.azure.spring.data.cosmos.ReactiveIntegrationTestCollectionManager;
 import com.azure.spring.data.cosmos.common.ResponseDiagnosticsTestUtils;
@@ -21,6 +24,7 @@
 import com.azure.spring.data.cosmos.core.query.Criteria;
 import com.azure.spring.data.cosmos.core.query.CriteriaType;
 import com.azure.spring.data.cosmos.domain.AuditableEntity;
+import com.azure.spring.data.cosmos.domain.AutoScaleSample;
 import com.azure.spring.data.cosmos.domain.GenIdEntity;
 import com.azure.spring.data.cosmos.domain.Person;
 import com.azure.spring.data.cosmos.exception.CosmosAccessException;
@@ -59,6 +63,8 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 
 @RunWith(SpringJUnit4ClassRunner.class)
@@ -92,6 +98,7 @@ public class ReactiveCosmosTemplateIT {
     @Value("${cosmos.key}")
     private String cosmosDbKey;
 
+    private static CosmosAsyncClient client;
     private static ReactiveCosmosTemplate cosmosTemplate;
     private static String containerName;
     private static CosmosEntityInformation<Person, String> personInfo;
@@ -115,7 +122,7 @@ public void setUp() throws ClassNotFoundException {
         if (cosmosTemplate == null) {
             azureKeyCredential = new AzureKeyCredential(cosmosDbKey);
             cosmosClientBuilder.credential(azureKeyCredential);
-            CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
+            client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
             final CosmosFactory dbFactory = new CosmosFactory(client, TestConstants.DB_NAME);
 
             final CosmosMappingContext mappingContext = new CosmosMappingContext();
@@ -501,4 +508,20 @@ public void testRunQueryWithReturnTypeContainingLocalDateTime() {
         StepVerifier.create(flux).expectNextCount(1).verifyComplete();
     }
 
+    @Test
+    public void createWithAutoscale() {
+        final CosmosEntityInformation<AutoScaleSample, String> autoScaleSampleInfo =
+            new CosmosEntityInformation<>(AutoScaleSample.class);
+        CosmosContainerResponse containerResponse = cosmosTemplate
+            .createContainerIfNotExists(autoScaleSampleInfo)
+            .block();
+        assertNotNull(containerResponse);
+        ThroughputResponse throughput = client.getDatabase(TestConstants.DB_NAME)
+            .getContainer(autoScaleSampleInfo.getContainerName())
+            .readThroughput()
+            .block();
+        assertNotNull(throughput);
+        assertEquals(Integer.parseInt(TestConstants.AUTOSCALE_MAX_THROUGHPUT),
+            throughput.getProperties().getAutoscaleMaxThroughput());
+    }
 }

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

@@ -0,0 +1,27 @@
+// 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.common.TestConstants;
+import com.azure.spring.data.cosmos.core.mapping.Container;
+
+@Container(ru = TestConstants.AUTOSCALE_MAX_THROUGHPUT, autoScale = true)
+public class AutoScaleSample {
+    private String id;
+
+    public AutoScaleSample() {
+    }
+
+    public AutoScaleSample(String id) {
+        this.id = id;
+    }
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(String id) {
+        this.id = id;
+    }
+
+}

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

@@ -244,9 +244,13 @@ public class User {
 
 - Annotation `@Container(containerName="myContainer")` specifies container name in Azure Cosmos DB.
 - Annotation `@PartitionKey` on `lastName` field specifies this field as partition key in Azure Cosmos DB.
+
+#### Creating Containers with autoscale throughput
+- Annotation `autoScale` field specifies container to be created with autoscale throughput if set to true. Default is false, which means containers are created with manual throughput.
+- Read more about autoscale throughput [here][autoscale-throughput]
 <!-- embedme src/samples/java/com/azure/spring/data/cosmos/UserSample.java#L14-L19 -->
 ```java
-@Container(containerName = "myContainer")
+@Container(containerName = "myContainer", autoScale = true, ru = "4000")
 public class UserSample {
     @Id
     private String emailAddress;
@@ -935,5 +939,6 @@ or contact [opencode@microsoft.com][coc_contact] with any additional questions o
 [spring_data_custom_query]: https://docs.spring.io/spring-data/commons/docs/current/reference/html/#repositories.query-methods.details
 [sql_queries_in_cosmos]: https://docs.microsoft.com/azure/cosmos-db/tutorial-query-sql-api
 [sql_queries_getting_started]: https://docs.microsoft.com/azure/cosmos-db/sql-query-getting-started
+[autoscale-throughput]: https://docs.microsoft.com/azure/cosmos-db/provision-throughput-autoscale
 
 ![Impressions](https://azure-sdk-impressions.azurewebsites.net/api/impressions/azure-sdk-for-java%2Fsdk%2Fcosmos%2F%2Fazure-spring-data-cosmos%2FREADME.png)

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

@@ -15,6 +15,7 @@ public final class Constants {
     public static final String DEFAULT_REPOSITORY_IMPLEMENT_POSTFIX = "Impl";
     public static final int DEFAULT_TIME_TO_LIVE = -1; // Indicates never expire
     public static final boolean DEFAULT_AUTO_CREATE_CONTAINER = true;
+    public static final boolean DEFAULT_AUTO_SCALE = false;
 
     public static final String ID_PROPERTY_NAME = "id";
     public static final String ETAG_PROPERTY_DEFAULT_NAME = "_etag";

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

@@ -474,8 +474,9 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat
                     cosmosContainerResponseMono =
                         cosmosAsyncDatabase.createContainerIfNotExists(cosmosContainerProperties);
                 } else {
-                    ThroughputProperties throughputProperties =
-                        ThroughputProperties.createManualThroughput(information.getRequestUnit());
+                    ThroughputProperties throughputProperties = information.isAutoScale()
+                        ? ThroughputProperties.createAutoscaledThroughput(information.getRequestUnit())
+                        : ThroughputProperties.createManualThroughput(information.getRequestUnit());
                     cosmosContainerResponseMono =
                         cosmosAsyncDatabase.createContainerIfNotExists(cosmosContainerProperties,
                             throughputProperties);

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

@@ -158,8 +158,9 @@ public Mono<CosmosContainerResponse> createContainerIfNotExists(CosmosEntityInfo
                     cosmosContainerResponseMono =
                         database.createContainerIfNotExists(cosmosContainerProperties);
                 } else {
-                    ThroughputProperties throughputProperties =
-                        ThroughputProperties.createManualThroughput(information.getRequestUnit());
+                    ThroughputProperties throughputProperties = information.isAutoScale()
+                        ? ThroughputProperties.createAutoscaledThroughput(information.getRequestUnit())
+                        : ThroughputProperties.createManualThroughput(information.getRequestUnit());
                     cosmosContainerResponseMono =
                         database.createContainerIfNotExists(cosmosContainerProperties,
                             throughputProperties);

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

@@ -53,4 +53,10 @@
      * @return partition key path
      */
     String partitionKeyPath() default "";
+
+    /**
+     * To enable auto scale for container RU limit
+     * @return default as false
+     */
+    boolean autoScale() default Constants.DEFAULT_AUTO_SCALE;
 }

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

@@ -65,6 +65,7 @@ public class CosmosEntityInformation<T, ID> extends AbstractEntityInformation<T,
     private final boolean autoCreateContainer;
     private final boolean autoGenerateId;
     private final boolean persitable;
+    private final boolean autoScale;
 
 
     /**
@@ -98,6 +99,7 @@ public CosmosEntityInformation(Class<T> domainType) {
         this.indexingPolicy = getIndexingPolicy(domainType);
         this.autoCreateContainer = getIsAutoCreateContainer(domainType);
         this.persitable = Persistable.class.isAssignableFrom(domainType);
+        this.autoScale = getIsAutoScale(domainType);
     }
 
     @Override
@@ -257,6 +259,15 @@ public boolean isAutoCreateContainer() {
         return autoCreateContainer;
     }
 
+    /**
+     * Check if container should use autoscale for resource units
+     *
+     * @return boolean
+     */
+    public boolean isAutoScale() {
+        return autoScale;
+    }
+
     private IndexingPolicy getIndexingPolicy(Class<?> domainType) {
         final IndexingPolicy policy = new IndexingPolicy();
 
@@ -469,5 +480,16 @@ private boolean getIsAutoCreateContainer(Class<T> domainType) {
 
         return autoCreateContainer;
     }
+
+    private boolean getIsAutoScale(Class<T> domainType) {
+        final Container annotation = domainType.getAnnotation(Container.class);
+
+        boolean autoScale = Constants.DEFAULT_AUTO_SCALE;
+        if (annotation != null) {
+            autoScale = annotation.autoScale();
+        }
+
+        return autoScale;
+    }
 }
 

sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/UserSample.java

@@ -11,7 +11,7 @@
 import com.azure.spring.data.cosmos.core.mapping.Container;
 import org.springframework.data.annotation.Id;
 
-@Container(containerName = "myContainer")
+@Container(containerName = "myContainer", autoScale = true, ru = "4000")
 public class UserSample {
     @Id
     private String emailAddress;