From 6b96d75e91dd9d9f78d92228fc432d5fa1a26c2e Mon Sep 17 00:00:00 2001 From: Blackbaud-MikeLueders Date: Mon, 4 Oct 2021 11:13:28 -0500 Subject: [PATCH] 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 --- .../data/cosmos/core/CosmosTemplateIT.java | 50 +++++++++++++++---- .../cosmos/core/ReactiveCosmosTemplateIT.java | 49 ++++++++++++++---- sdk/cosmos/azure-spring-data-cosmos/README.md | 18 +++++++ .../data/cosmos/config/CosmosConfig.java | 35 ++++++++++++- .../config/DatabaseThroughputConfig.java | 34 +++++++++++++ .../data/cosmos/core/CosmosTemplate.java | 20 +++++++- .../cosmos/core/ReactiveCosmosTemplate.java | 20 +++++++- 7 files changed, 202 insertions(+), 24 deletions(-) create mode 100644 sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/DatabaseThroughputConfig.java diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java index ff82e533c30ac..a06887c999e06 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplateIT.java +++ b/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 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); + } + } + } diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplateIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplateIT.java index ede2e69118dcf..ae4676eb498c1 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplateIT.java +++ b/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 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); + } + } + } diff --git a/sdk/cosmos/azure-spring-data-cosmos/README.md b/sdk/cosmos/azure-spring-data-cosmos/README.md index 09c71fe5a6fe0..5ec4994de3ae0 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/README.md +++ b/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. diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/CosmosConfig.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/CosmosConfig.java index 97d61aa158239..156d28623469f 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/CosmosConfig.java +++ b/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 + '}'; } diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/DatabaseThroughputConfig.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/config/DatabaseThroughputConfig.java new file mode 100644 index 0000000000000..c593588cf8132 --- /dev/null +++ b/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 + + '}'; + } + +} diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java index 946b802f23133..655c7ed189375 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java +++ b/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 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) diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java index 0e0d8ede36a2a..44566c01130b5 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplate.java @@ -7,6 +7,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.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedResponse; @@ -18,6 +19,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; @@ -61,6 +63,7 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica private final boolean queryMetricsEnabled; private final CosmosAsyncClient cosmosAsyncClient; private final IsNewAwareAuditingHandler cosmosAuditingHandler; + private final DatabaseThroughputConfig databaseThroughputConfig; private ApplicationContext applicationContext; @@ -114,6 +117,7 @@ public ReactiveCosmosTemplate(CosmosFactory cosmosFactory, this.responseDiagnosticsProcessor = cosmosConfig.getResponseDiagnosticsProcessor(); this.queryMetricsEnabled = cosmosConfig.isQueryMetricsEnabled(); this.cosmosAuditingHandler = cosmosAuditingHandler; + this.databaseThroughputConfig = cosmosConfig.getDatabaseThroughputConfig(); } /** @@ -146,8 +150,7 @@ public void setApplicationContext(@NonNull ApplicationContext applicationContext @Override public Mono createContainerIfNotExists(CosmosEntityInformation information) { - return cosmosAsyncClient - .createDatabaseIfNotExists(this.databaseName) + return createDatabaseIfNotExists() .publishOn(Schedulers.parallel()) .onErrorResume(throwable -> CosmosExceptionUtils.exceptionHandler("Failed to create database", throwable)) @@ -188,6 +191,19 @@ public Mono createContainerIfNotExists(CosmosEntityInfo } + private Mono 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 Mono getContainerProperties(String containerName) { return cosmosAsyncClient.getDatabase(this.databaseName)