add support for setting throughput on database creation (#24456)

* add support for setting throughput on database creation

* added section to readme

* removed locale from links

* fix checkstyle issues

* do not overwrite cosmosTemplate

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

@@ -3,6 +3,7 @@
 package com.azure.spring.data.cosmos.core;
 
 import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncDatabase;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.implementation.ConflictException;
@@ -117,18 +118,9 @@ public class CosmosTemplateIT {
     public void setUp() throws ClassNotFoundException {
         if (cosmosTemplate == null) {
             client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
-            final CosmosFactory cosmosFactory = new CosmosFactory(client, TestConstants.DB_NAME);
-
-            final CosmosMappingContext mappingContext = new CosmosMappingContext();
             personInfo = new CosmosEntityInformation<>(Person.class);
             containerName = personInfo.getContainerName();
-
-            mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
-
-            final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext,
-                null);
-
-            cosmosTemplate = new CosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter);
+            cosmosTemplate = createCosmosTemplate(cosmosConfig, TestConstants.DB_NAME);
         }
 
         collectionManager.ensureContainersCreatedAndEmpty(cosmosTemplate, Person.class,
@@ -137,6 +129,14 @@ public void setUp() throws ClassNotFoundException {
             new PartitionKey(TEST_PERSON.getLastName()));
     }
 
+    private CosmosTemplate createCosmosTemplate(CosmosConfig config, String dbName) throws ClassNotFoundException {
+        final CosmosFactory cosmosFactory = new CosmosFactory(client, dbName);
+        final CosmosMappingContext mappingContext = new CosmosMappingContext();
+        mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
+        final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null);
+        return new CosmosTemplate(cosmosFactory, config, cosmosConverter);
+    }
+
     private void insertPerson(Person person) {
         cosmosTemplate.insert(person,
             new PartitionKey(personInfo.getPartitionKeyFieldValue(person)));
@@ -660,4 +660,34 @@ public void createWithAutoscale() throws ClassNotFoundException {
         assertEquals(Integer.parseInt(TestConstants.AUTOSCALE_MAX_THROUGHPUT),
             throughput.getProperties().getAutoscaleMaxThroughput());
     }
+
+    @Test
+    public void createDatabaseWithThroughput() throws ClassNotFoundException {
+        final String configuredThroughputDbName = TestConstants.DB_NAME + "-configured-throughput";
+        deleteDatabaseIfExists(configuredThroughputDbName);
+
+        Integer expectedRequestUnits = 700;
+        final CosmosConfig config = CosmosConfig.builder()
+            .enableDatabaseThroughput(false, expectedRequestUnits)
+            .build();
+        final CosmosTemplate configuredThroughputCosmosTemplate = createCosmosTemplate(config, configuredThroughputDbName);
+
+        final CosmosEntityInformation<Person, String> personInfo =
+            new CosmosEntityInformation<>(Person.class);
+        configuredThroughputCosmosTemplate.createContainerIfNotExists(personInfo);
+
+        final CosmosAsyncDatabase database = client.getDatabase(configuredThroughputDbName);
+        final ThroughputResponse response = database.readThroughput().block();
+        assertEquals(expectedRequestUnits, response.getProperties().getManualThroughput());
+    }
+
+    private void deleteDatabaseIfExists(String dbName) {
+        CosmosAsyncDatabase database = client.getDatabase(dbName);
+        try {
+            database.delete().block();
+        } catch (CosmosException ex) {
+            assertEquals(ex.getStatusCode(), 404);
+        }
+    }
+
 }

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

@@ -4,6 +4,7 @@
 
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncDatabase;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.implementation.ConflictException;
@@ -119,17 +120,9 @@ public void setUp() throws ClassNotFoundException {
             azureKeyCredential = new AzureKeyCredential(cosmosDbKey);
             cosmosClientBuilder.credential(azureKeyCredential);
             client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
-            final CosmosFactory dbFactory = new CosmosFactory(client, TestConstants.DB_NAME);
-
-            final CosmosMappingContext mappingContext = new CosmosMappingContext();
             personInfo = new CosmosEntityInformation<>(Person.class);
             containerName = personInfo.getContainerName();
-
-            mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
-
-            final MappingCosmosConverter dbConverter =
-                new MappingCosmosConverter(mappingContext, null);
-            cosmosTemplate = new ReactiveCosmosTemplate(dbFactory, cosmosConfig, dbConverter);
+            cosmosTemplate = createReactiveCosmosTemplate(cosmosConfig, TestConstants.DB_NAME);
         }
 
         collectionManager.ensureContainersCreatedAndEmpty(cosmosTemplate, Person.class, GenIdEntity.class, AuditableEntity.class);
@@ -138,6 +131,14 @@ public void setUp() throws ClassNotFoundException {
             new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON))).block();
     }
 
+    private ReactiveCosmosTemplate createReactiveCosmosTemplate(CosmosConfig config, String dbName) throws ClassNotFoundException {
+        final CosmosFactory cosmosFactory = new CosmosFactory(client, dbName);
+        final CosmosMappingContext mappingContext = new CosmosMappingContext();
+        mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class));
+        final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null);
+        return new ReactiveCosmosTemplate(cosmosFactory, config, cosmosConverter);
+    }
+
     @After
     public void cleanup() {
         //  Reset master key
@@ -545,4 +546,34 @@ public void createWithAutoscale() {
         assertEquals(Integer.parseInt(TestConstants.AUTOSCALE_MAX_THROUGHPUT),
             throughput.getProperties().getAutoscaleMaxThroughput());
     }
+
+    @Test
+    public void createDatabaseWithThroughput() throws ClassNotFoundException {
+        final String configuredThroughputDbName = TestConstants.DB_NAME + "-other";
+        deleteDatabaseIfExists(configuredThroughputDbName);
+
+        Integer expectedRequestUnits = 700;
+        final CosmosConfig config = CosmosConfig.builder()
+            .enableDatabaseThroughput(false, expectedRequestUnits)
+            .build();
+        final ReactiveCosmosTemplate configuredThroughputCosmosTemplate = createReactiveCosmosTemplate(config, configuredThroughputDbName);
+
+        final CosmosEntityInformation<Person, String> personInfo =
+            new CosmosEntityInformation<>(Person.class);
+        configuredThroughputCosmosTemplate.createContainerIfNotExists(personInfo).block();
+
+        final CosmosAsyncDatabase database = client.getDatabase(configuredThroughputDbName);
+        final ThroughputResponse response = database.readThroughput().block();
+        assertEquals(expectedRequestUnits, response.getProperties().getManualThroughput());
+    }
+
+    private void deleteDatabaseIfExists(String dbName) {
+        CosmosAsyncDatabase database = client.getDatabase(dbName);
+        try {
+            database.delete().block();
+        } catch (CosmosException ex) {
+            assertEquals(ex.getStatusCode(), 404);
+        }
+    }
+
 }

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

@@ -180,6 +180,24 @@ public CosmosConfig cosmosConfig() {
 By default, `@EnableCosmosRepositories` will scan the current package for any interfaces that extend one of Spring Data's repository interfaces.
 Use it to annotate your Configuration class to scan a different root package by `@EnableCosmosRepositories(basePackageClass=UserRepository.class)` if your project layout has multiple projects.
 
+#### Using database provisioned throughput
+
+Cosmos supports both [container](https://docs.microsoft.com/azure/cosmos-db/sql/how-to-provision-container-throughput)
+and [database](https://docs.microsoft.com/azure/cosmos-db/sql/how-to-provision-database-throughput) provisioned
+throughput.  By default, spring-data-cosmos will provision throughput for each container created.  If you prefer
+to share throughput between containers, you can enable database provisioned throughput via CosmosConfig.
+
+```java
+@Override
+public CosmosConfig cosmosConfig() {
+    int autoscale = false; 
+    int initialRequestUnits = 400;
+    return CosmosConfig.builder()
+                       .enableDatabaseThroughput(autoscale, initialRequestUnits) 
+                       .build();
+}
+```
+
 ### Define an entity
 - Define a simple entity as item in Azure Cosmos DB.
 

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

@@ -13,6 +13,8 @@ public class CosmosConfig {
 
     private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor;
 
+    private final DatabaseThroughputConfig databaseThroughputConfig;
+
     private final boolean queryMetricsEnabled;
 
     /**
@@ -24,7 +26,22 @@ public class CosmosConfig {
     @ConstructorProperties({"responseDiagnosticsProcessor", "queryMetricsEnabled"})
     public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
                         boolean queryMetricsEnabled) {
+        this(responseDiagnosticsProcessor, null, queryMetricsEnabled);
+    }
+
+    /**
+     * Initialization
+     *
+     * @param responseDiagnosticsProcessor must not be {@literal null}
+     * @param databaseThroughputConfig may be @{literal null}
+     * @param queryMetricsEnabled must not be {@literal null}
+     */
+    @ConstructorProperties({"responseDiagnosticsProcessor", "databaseThroughputConfig", "queryMetricsEnabled"})
+    public CosmosConfig(ResponseDiagnosticsProcessor responseDiagnosticsProcessor,
+                        DatabaseThroughputConfig databaseThroughputConfig,
+                        boolean queryMetricsEnabled) {
         this.responseDiagnosticsProcessor = responseDiagnosticsProcessor;
+        this.databaseThroughputConfig = databaseThroughputConfig;
         this.queryMetricsEnabled = queryMetricsEnabled;
     }
 
@@ -46,6 +63,15 @@ public boolean isQueryMetricsEnabled() {
         return queryMetricsEnabled;
     }
 
+    /**
+     * Gets the database throughput configuration.
+     *
+     * @return DatabaseThroughputConfig, or null if no database throughput is configured
+     */
+    public DatabaseThroughputConfig getDatabaseThroughputConfig() {
+        return databaseThroughputConfig;
+    }
+
     /**
      * Create a CosmosConfigBuilder instance
      *
@@ -60,6 +86,7 @@ public static CosmosConfigBuilder builder() {
      */
     public static class CosmosConfigBuilder {
         private ResponseDiagnosticsProcessor responseDiagnosticsProcessor;
+        private DatabaseThroughputConfig databaseThroughputConfig;
         private boolean queryMetricsEnabled;
         CosmosConfigBuilder() {
         }
@@ -88,19 +115,25 @@ public CosmosConfigBuilder enableQueryMetrics(boolean queryMetricsEnabled) {
             return this;
         }
 
+        public CosmosConfigBuilder enableDatabaseThroughput(boolean autoscale, int requestUnits) {
+            this.databaseThroughputConfig = new DatabaseThroughputConfig(autoscale, requestUnits);
+            return this;
+        }
+
         /**
          * Build a CosmosConfig instance
          *
          * @return CosmosConfig
          */
         public CosmosConfig build() {
-            return new CosmosConfig(this.responseDiagnosticsProcessor, this.queryMetricsEnabled);
+            return new CosmosConfig(this.responseDiagnosticsProcessor, this.databaseThroughputConfig, this.queryMetricsEnabled);
         }
 
         @Override
         public String toString() {
             return "CosmosConfigBuilder{"
                 + "responseDiagnosticsProcessor=" + responseDiagnosticsProcessor
+                + ", databaseThroughputConfig=" + databaseThroughputConfig
                 + ", queryMetricsEnabled=" + queryMetricsEnabled
                 + '}';
         }

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

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.spring.data.cosmos.config;
+
+/**
+ * Throughput config for database creation
+ */
+public class DatabaseThroughputConfig {
+
+    private final boolean autoScale;
+    private final int requestUnits;
+
+    public DatabaseThroughputConfig(boolean autoScale, int requestUnits) {
+        this.autoScale = autoScale;
+        this.requestUnits = requestUnits;
+    }
+
+    public boolean isAutoScale() {
+        return autoScale;
+    }
+
+    public int getRequestUnits() {
+        return requestUnits;
+    }
+
+    @Override
+    public String toString() {
+        return "DatabaseThroughputConfig{"
+            + "autoScale=" + autoScale
+            + ", requestUnits=" + requestUnits
+            + '}';
+    }
+
+}

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

@@ -8,6 +8,7 @@
 import com.azure.cosmos.CosmosAsyncDatabase;
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.CosmosContainerResponse;
+import com.azure.cosmos.models.CosmosDatabaseResponse;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosItemResponse;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
@@ -20,6 +21,7 @@
 import com.azure.spring.data.cosmos.CosmosFactory;
 import com.azure.spring.data.cosmos.common.CosmosUtils;
 import com.azure.spring.data.cosmos.config.CosmosConfig;
+import com.azure.spring.data.cosmos.config.DatabaseThroughputConfig;
 import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter;
 import com.azure.spring.data.cosmos.core.generator.CountQueryGenerator;
 import com.azure.spring.data.cosmos.core.generator.FindQuerySpecGenerator;
@@ -75,6 +77,7 @@ public class CosmosTemplate implements CosmosOperations, ApplicationContextAware
     private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor;
     private final boolean queryMetricsEnabled;
     private final CosmosAsyncClient cosmosAsyncClient;
+    private final DatabaseThroughputConfig databaseThroughputConfig;
 
     private ApplicationContext applicationContext;
 
@@ -126,6 +129,7 @@ public CosmosTemplate(CosmosFactory cosmosFactory,
         this.databaseName = cosmosFactory.getDatabaseName();
         this.responseDiagnosticsProcessor = cosmosConfig.getResponseDiagnosticsProcessor();
         this.queryMetricsEnabled = cosmosConfig.isQueryMetricsEnabled();
+        this.databaseThroughputConfig = cosmosConfig.getDatabaseThroughputConfig();
     }
 
     /**
@@ -458,8 +462,7 @@ public String getContainerName(Class<?> domainType) {
     @Override
     public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformation<?, ?> information) {
 
-        final CosmosContainerResponse response = cosmosAsyncClient
-            .createDatabaseIfNotExists(this.databaseName)
+        final CosmosContainerResponse response = createDatabaseIfNotExists()
             .publishOn(Schedulers.parallel())
             .onErrorResume(throwable ->
                 CosmosExceptionUtils.exceptionHandler("Failed to create database", throwable))
@@ -501,6 +504,19 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat
         return response.getProperties();
     }
 
+    private Mono<CosmosDatabaseResponse> createDatabaseIfNotExists() {
+        if (databaseThroughputConfig == null) {
+            return cosmosAsyncClient
+                .createDatabaseIfNotExists(this.databaseName);
+        } else {
+            ThroughputProperties throughputProperties = databaseThroughputConfig.isAutoScale()
+                ? ThroughputProperties.createAutoscaledThroughput(databaseThroughputConfig.getRequestUnits())
+                : ThroughputProperties.createManualThroughput(databaseThroughputConfig.getRequestUnits());
+            return cosmosAsyncClient
+                .createDatabaseIfNotExists(this.databaseName, throughputProperties);
+        }
+    }
+
     @Override
     public CosmosContainerProperties getContainerProperties(String containerName) {
         final CosmosContainerResponse response = cosmosAsyncClient.getDatabase(this.databaseName)