From e7e01657b1791c7b0350eb1eedcc91a0a01ddbee Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Thu, 1 Dec 2022 16:09:19 -0500 Subject: [PATCH 01/25] Proof of concept that we can write to two databases from the same session. --- .../core/MultiTenantDBCosmosTemplate.java | 87 +++++++++++++++++++ .../MultiTenantDBCosmosTemplateUnitTest.java | 77 ++++++++++++++++ .../data/cosmos/core/CosmosTemplate.java | 10 ++- 3 files changed, 171 insertions(+), 3 deletions(-) create mode 100644 sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java create mode 100644 sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java new file mode 100644 index 0000000000000..b79eea3782695 --- /dev/null +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.spring.data.cosmos.core; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +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; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.models.UniqueKeyPolicy; +import com.azure.spring.data.cosmos.Constants; +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; +import com.azure.spring.data.cosmos.core.generator.NativeQueryGenerator; +import com.azure.spring.data.cosmos.core.mapping.event.AfterLoadEvent; +import com.azure.spring.data.cosmos.core.mapping.event.CosmosMappingEvent; +import com.azure.spring.data.cosmos.core.query.CosmosPageImpl; +import com.azure.spring.data.cosmos.core.query.CosmosPageRequest; +import com.azure.spring.data.cosmos.core.query.CosmosQuery; +import com.azure.spring.data.cosmos.core.query.CosmosSliceImpl; +import com.azure.spring.data.cosmos.core.query.Criteria; +import com.azure.spring.data.cosmos.core.query.CriteriaType; +import com.azure.spring.data.cosmos.exception.CosmosExceptionUtils; +import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; +import com.fasterxml.jackson.databind.JsonNode; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.BeansException; +import org.springframework.context.ApplicationContext; +import org.springframework.context.ApplicationContextAware; +import org.springframework.data.auditing.IsNewAwareAuditingHandler; +import org.springframework.data.domain.Page; +import org.springframework.data.domain.Pageable; +import org.springframework.data.domain.Slice; +import org.springframework.data.domain.Sort; +import org.springframework.data.repository.query.parser.Part; +import org.springframework.lang.NonNull; +import org.springframework.util.Assert; +import org.springframework.util.ReflectionUtils; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Optional; +import java.util.UUID; +import java.util.stream.Collectors; + +/** + * Template class for cosmos db + */ +public class MultiTenantDBCosmosTemplate extends CosmosTemplate { + + public MultiTenantDBCosmosTemplate(CosmosFactory cosmosFactory, CosmosConfig cosmosConfig, MappingCosmosConverter mappingCosmosConverter, IsNewAwareAuditingHandler cosmosAuditingHandler) { + super(cosmosFactory, cosmosConfig, mappingCosmosConverter, cosmosAuditingHandler); + + super.databaseName = cosmosFactory.getDatabaseName(); + } + + public void setDatabaseName(String dbName, CosmosEntityInformation information) { + super.databaseName = dbName; + + createContainerIfNotExists(information); + } + + public String getDatabaseName() { + return super.databaseName; + } +} diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java new file mode 100644 index 0000000000000..343c32352d3b8 --- /dev/null +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.spring.data.cosmos.core; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.models.PartitionKey; +import com.azure.spring.data.cosmos.CosmosFactory; +import com.azure.spring.data.cosmos.IntegrationTestCollectionManager; +import com.azure.spring.data.cosmos.config.CosmosConfig; +import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; +import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; +import com.azure.spring.data.cosmos.domain.Person; +import com.azure.spring.data.cosmos.repository.TestRepositoryConfig; +import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.domain.EntityScanner; +import org.springframework.context.ApplicationContext; +import org.springframework.data.annotation.Persistent; +import org.springframework.test.context.ContextConfiguration; +import org.springframework.test.context.junit4.SpringJUnit4ClassRunner; + +import static com.azure.spring.data.cosmos.common.TestConstants.ADDRESSES; +import static com.azure.spring.data.cosmos.common.TestConstants.AGE; +import static com.azure.spring.data.cosmos.common.TestConstants.FIRST_NAME; +import static com.azure.spring.data.cosmos.common.TestConstants.HOBBIES; +import static com.azure.spring.data.cosmos.common.TestConstants.ID_1; +import static com.azure.spring.data.cosmos.common.TestConstants.ID_2; +import static com.azure.spring.data.cosmos.common.TestConstants.LAST_NAME; +import static com.azure.spring.data.cosmos.common.TestConstants.PASSPORT_IDS_BY_COUNTRY; +import static org.assertj.core.api.Assertions.assertThat; + +@RunWith(SpringJUnit4ClassRunner.class) +@ContextConfiguration(classes = TestRepositoryConfig.class) +public class MultiTenantDBCosmosTemplateUnitTest { + + private final String testDB1 = "Database1"; + + private final String testDB2 = "Database2"; + + @Autowired + private ApplicationContext applicationContext; + @Autowired + private CosmosConfig cosmosConfig; + @Autowired + private CosmosClientBuilder cosmosClientBuilder; + + @Test + public void testGetDatabaseFunctionality() { + try { + CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); + CosmosFactory cosmosFactory = new CosmosFactory(client, testDB1); + final CosmosMappingContext mappingContext = new CosmosMappingContext(); + mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class)); + final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); + MultiTenantDBCosmosTemplate cosmosTemplate = new MultiTenantDBCosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); + + CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); + + cosmosTemplate.setDatabaseName(testDB1, personInfo); + Person TEST_PERSON = new Person(ID_1, FIRST_NAME, LAST_NAME, HOBBIES, ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY); + assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); + cosmosTemplate.insert(TEST_PERSON, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON))); + + cosmosTemplate.setDatabaseName(testDB2, personInfo); + Person TEST_PERSON2 = new Person(ID_2, FIRST_NAME, LAST_NAME, HOBBIES, ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY); + assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); + cosmosTemplate.insert(TEST_PERSON2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON2))); + } catch (ClassNotFoundException e) { + assertThat("a").isEqualTo("b"); + } + } +} 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 f1b8d79d8a09c..20e5e9252a32a 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 @@ -74,8 +74,8 @@ public class CosmosTemplate implements CosmosOperations, ApplicationContextAware private final MappingCosmosConverter mappingCosmosConverter; private final IsNewAwareAuditingHandler cosmosAuditingHandler; - private final String databaseName; - private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor; + protected String databaseName; + protected final ResponseDiagnosticsProcessor responseDiagnosticsProcessor; private final boolean queryMetricsEnabled; private final int maxDegreeOfParallelism; private final int maxBufferedItemCount; @@ -152,6 +152,10 @@ public CosmosTemplate(CosmosFactory cosmosFactory, this(cosmosFactory, cosmosConfig, mappingCosmosConverter, null); } + public String getDatabaseName() { + return this.databaseName; + } + /** * Sets the application context * @@ -528,7 +532,7 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat return response.getProperties(); } - private Mono createDatabaseIfNotExists() { + protected Mono createDatabaseIfNotExists() { if (databaseThroughputConfig == null) { return cosmosAsyncClient .createDatabaseIfNotExists(this.databaseName); From 7516e97ac461fbb05d8b5103f3d4fcf89814b047 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Mon, 5 Dec 2022 08:49:06 -0500 Subject: [PATCH 02/25] Improving the changes to CosmosTemplate and the test case. --- .../core/MultiTenantDBCosmosTemplate.java | 8 +++- .../MultiTenantDBCosmosTemplateUnitTest.java | 21 ++++++---- .../data/cosmos/core/CosmosTemplate.java | 38 +++++++++---------- 3 files changed, 39 insertions(+), 28 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java index b79eea3782695..cfb0f9dc28bbb 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java @@ -75,10 +75,14 @@ public MultiTenantDBCosmosTemplate(CosmosFactory cosmosFactory, CosmosConfig cos super.databaseName = cosmosFactory.getDatabaseName(); } - public void setDatabaseName(String dbName, CosmosEntityInformation information) { + public void setDatabaseName(String dbName) { super.databaseName = dbName; - createContainerIfNotExists(information); + final Mono response = createDatabaseIfNotExists() + .publishOn(Schedulers.parallel()) + .onErrorResume(throwable -> + CosmosExceptionUtils.exceptionHandler("Failed to create database", throwable, + this.responseDiagnosticsProcessor)); } public String getDatabaseName() { diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java index 343c32352d3b8..6aa6639e45fc1 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java @@ -39,9 +39,14 @@ public class MultiTenantDBCosmosTemplateUnitTest { private final String testDB1 = "Database1"; - private final String testDB2 = "Database2"; + private final Person TEST_PERSON_1 = new Person(ID_1, FIRST_NAME, LAST_NAME, HOBBIES, ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY); + private final Person TEST_PERSON_2 = new Person(ID_2, FIRST_NAME, LAST_NAME, HOBBIES, ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY); + + @ClassRule + public static final IntegrationTestCollectionManager collectionManager = new IntegrationTestCollectionManager(); + @Autowired private ApplicationContext applicationContext; @Autowired @@ -61,15 +66,17 @@ public void testGetDatabaseFunctionality() { CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); - cosmosTemplate.setDatabaseName(testDB1, personInfo); - Person TEST_PERSON = new Person(ID_1, FIRST_NAME, LAST_NAME, HOBBIES, ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY); + cosmosTemplate.setDatabaseName(testDB1); + cosmosTemplate.createContainerIfNotExists(personInfo); + cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); - cosmosTemplate.insert(TEST_PERSON, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON))); + cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); - cosmosTemplate.setDatabaseName(testDB2, personInfo); - Person TEST_PERSON2 = new Person(ID_2, FIRST_NAME, LAST_NAME, HOBBIES, ADDRESSES, AGE, PASSPORT_IDS_BY_COUNTRY); + cosmosTemplate.setDatabaseName(testDB2); + cosmosTemplate.createContainerIfNotExists(personInfo); + cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); - cosmosTemplate.insert(TEST_PERSON2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON2))); + cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); } catch (ClassNotFoundException e) { assertThat("a").isEqualTo("b"); } 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 20e5e9252a32a..d5a01edb3985e 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 @@ -211,14 +211,14 @@ public T insert(String containerName, T objectToSave, PartitionKey partition final JsonNode originalItem = mappingCosmosConverter.writeJsonNode(objectToSave); - LOGGER.debug("execute createItem in database {} container {}", this.databaseName, + LOGGER.debug("execute createItem in database {} container {}", this.getDatabaseName(), containerName); final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); // if the partition key is null, SDK will get the partitionKey from the object final CosmosItemResponse response = cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .createItem(originalItem, partitionKey, options) .publishOn(Schedulers.parallel()) @@ -263,7 +263,7 @@ public T findById(Object id, Class domainType, PartitionKey partitionKey) String idToQuery = CosmosUtils.getStringIDValue(id); final String containerName = getContainerName(domainType); return cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .readItem(idToQuery, partitionKey, JsonNode.class) .publishOn(Schedulers.parallel()) @@ -300,7 +300,7 @@ public T findById(String containerName, Object id, Class domainType) { options.setMaxBufferedItemCount(this.maxBufferedItemCount); options.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); return cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, options, JsonNode.class) .byPage() @@ -359,7 +359,7 @@ public T upsertAndReturnEntity(String containerName, T object) { final JsonNode originalItem = mappingCosmosConverter.writeJsonNode(object); - LOGGER.debug("execute upsert item in database {} container {}", this.databaseName, + LOGGER.debug("execute upsert item in database {} container {}", this.getDatabaseName(), containerName); @SuppressWarnings("unchecked") final Class domainType = (Class) object.getClass(); @@ -368,7 +368,7 @@ public T upsertAndReturnEntity(String containerName, T object) { applyVersioning(domainType, originalItem, options); final CosmosItemResponse cosmosItemResponse = cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .upsertItem(originalItem, options) .publishOn(Schedulers.parallel()) @@ -428,7 +428,7 @@ public Iterable findAll(PartitionKey partitionKey, final Class domainT cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); return cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems("SELECT * FROM r", cosmosQueryRequestOptions, JsonNode.class) .byPage() @@ -462,7 +462,7 @@ public void deleteAll(@NonNull String containerName, @NonNull Class domainTyp @Override public void deleteContainer(@NonNull String containerName) { Assert.hasText(containerName, "containerName should have text."); - cosmosAsyncClient.getDatabase(this.databaseName) + cosmosAsyncClient.getDatabase(this.getDatabaseName()) .getContainer(containerName) .delete() .publishOn(Schedulers.parallel()) @@ -535,19 +535,19 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat protected Mono createDatabaseIfNotExists() { if (databaseThroughputConfig == null) { return cosmosAsyncClient - .createDatabaseIfNotExists(this.databaseName); + .createDatabaseIfNotExists(this.getDatabaseName()); } else { ThroughputProperties throughputProperties = databaseThroughputConfig.isAutoScale() ? ThroughputProperties.createAutoscaledThroughput(databaseThroughputConfig.getRequestUnits()) : ThroughputProperties.createManualThroughput(databaseThroughputConfig.getRequestUnits()); return cosmosAsyncClient - .createDatabaseIfNotExists(this.databaseName, throughputProperties); + .createDatabaseIfNotExists(this.getDatabaseName(), throughputProperties); } } @Override public CosmosContainerProperties getContainerProperties(String containerName) { - final CosmosContainerResponse response = cosmosAsyncClient.getDatabase(this.databaseName) + final CosmosContainerResponse response = cosmosAsyncClient.getDatabase(this.getDatabaseName()) .getContainer(containerName) .read() .block(); @@ -558,7 +558,7 @@ public CosmosContainerProperties getContainerProperties(String containerName) { @Override public CosmosContainerProperties replaceContainerProperties(String containerName, CosmosContainerProperties properties) { - CosmosContainerResponse response = this.cosmosAsyncClient.getDatabase(this.databaseName) + CosmosContainerResponse response = this.cosmosAsyncClient.getDatabase(this.getDatabaseName()) .getContainer(containerName) .replace(properties) .block(); @@ -599,14 +599,14 @@ private void deleteById(String containerName, Object id, PartitionKey partitionK CosmosItemRequestOptions options) { Assert.hasText(containerName, "containerName should not be null, empty or only whitespaces"); String idToDelete = CosmosUtils.getStringIDValue(id); - LOGGER.debug("execute deleteById in database {} container {}", this.databaseName, + LOGGER.debug("execute deleteById in database {} container {}", this.getDatabaseName(), containerName); if (partitionKey == null) { partitionKey = PartitionKey.NONE; } - cosmosAsyncClient.getDatabase(this.databaseName) + cosmosAsyncClient.getDatabase(this.getDatabaseName()) .getContainer(containerName) .deleteItem(idToDelete, partitionKey, options) .publishOn(Schedulers.parallel()) @@ -760,7 +760,7 @@ private Slice sliceQuery(SqlQuerySpec querySpec, }); CosmosAsyncContainer container = - cosmosAsyncClient.getDatabase(this.databaseName).getContainer(containerName); + cosmosAsyncClient.getDatabase(this.getDatabaseName()).getContainer(containerName); Flux> feedResponseFlux; /* @@ -918,7 +918,7 @@ private Long getCountValue(SqlQuerySpec querySpec, String containerName) { private Flux> executeQuery(SqlQuerySpec sqlQuerySpec, String containerName, CosmosQueryRequestOptions options) { - return cosmosAsyncClient.getDatabase(this.databaseName) + return cosmosAsyncClient.getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, options, JsonNode.class) .byPage(); @@ -940,7 +940,7 @@ private Flux findItemsAsFlux(@NonNull CosmosQuery query, }); return cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, cosmosQueryRequestOptions, JsonNode.class) .byPage() @@ -965,7 +965,7 @@ private Flux getJsonNodeFluxFromQuerySpec( cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); return cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, cosmosQueryRequestOptions, JsonNode.class) .byPage() @@ -997,7 +997,7 @@ private T deleteItem(@NonNull JsonNode jsonNode, applyVersioning(domainType, jsonNode, options); return cosmosAsyncClient - .getDatabase(this.databaseName) + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .deleteItem(jsonNode, options) .publishOn(Schedulers.parallel()) From 0125a30b4d6ded80586b74dcf75ededb022fbe97 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Mon, 5 Dec 2022 12:38:30 -0500 Subject: [PATCH 03/25] Moving default setNameAndCreateDatabase() logic into CosmosTemplate. --- .../data/cosmos/core/MultiTenantDBCosmosTemplate.java | 11 +++-------- .../core/MultiTenantDBCosmosTemplateUnitTest.java | 4 ++-- .../azure/spring/data/cosmos/core/CosmosTemplate.java | 10 ++++++++++ 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java index cfb0f9dc28bbb..76c22c1dceb05 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java @@ -75,14 +75,9 @@ public MultiTenantDBCosmosTemplate(CosmosFactory cosmosFactory, CosmosConfig cos super.databaseName = cosmosFactory.getDatabaseName(); } - public void setDatabaseName(String dbName) { - super.databaseName = dbName; - - final Mono response = createDatabaseIfNotExists() - .publishOn(Schedulers.parallel()) - .onErrorResume(throwable -> - CosmosExceptionUtils.exceptionHandler("Failed to create database", throwable, - this.responseDiagnosticsProcessor)); + @Override + public void setNameAndCreateDatabase(String dbName) { + super.setNameAndCreateDatabase(dbName); } public String getDatabaseName() { diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java index 6aa6639e45fc1..1a7dcf243b5d0 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java @@ -66,13 +66,13 @@ public void testGetDatabaseFunctionality() { CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); - cosmosTemplate.setDatabaseName(testDB1); + cosmosTemplate.setNameAndCreateDatabase(testDB1); cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); - cosmosTemplate.setDatabaseName(testDB2); + cosmosTemplate.setNameAndCreateDatabase(testDB2); cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); 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 d5a01edb3985e..81d706c0b8526 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 @@ -156,6 +156,16 @@ public String getDatabaseName() { return this.databaseName; } + public void setNameAndCreateDatabase(String dbName) { + this.databaseName = dbName; + + createDatabaseIfNotExists() + .publishOn(Schedulers.parallel()) + .onErrorResume(throwable -> + CosmosExceptionUtils.exceptionHandler("Failed to create database", throwable, + this.responseDiagnosticsProcessor)); + } + /** * Sets the application context * From 644697b90e8c0d49e7976ac1d0c4a596f2e9e35c Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Mon, 5 Dec 2022 12:50:30 -0500 Subject: [PATCH 04/25] Improving unit test. --- .../MultiTenantDBCosmosTemplateUnitTest.java | 68 +++++++++++++------ 1 file changed, 47 insertions(+), 21 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java index 1a7dcf243b5d0..da0c7901a644a 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java @@ -11,9 +11,11 @@ import com.azure.spring.data.cosmos.config.CosmosConfig; import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; +import com.azure.spring.data.cosmos.domain.Customer; import com.azure.spring.data.cosmos.domain.Person; import com.azure.spring.data.cosmos.repository.TestRepositoryConfig; import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; +import org.junit.Assert; import org.junit.ClassRule; import org.junit.Test; import org.junit.runner.RunWith; @@ -24,6 +26,9 @@ import org.springframework.test.context.ContextConfiguration; import org.springframework.test.context.junit4.SpringJUnit4ClassRunner; +import java.util.ArrayList; +import java.util.List; + import static com.azure.spring.data.cosmos.common.TestConstants.ADDRESSES; import static com.azure.spring.data.cosmos.common.TestConstants.AGE; import static com.azure.spring.data.cosmos.common.TestConstants.FIRST_NAME; @@ -56,29 +61,50 @@ public class MultiTenantDBCosmosTemplateUnitTest { @Test public void testGetDatabaseFunctionality() { + /// Setup + CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); + CosmosFactory cosmosFactory = new CosmosFactory(client, testDB1); + final CosmosMappingContext mappingContext = new CosmosMappingContext(); + try { - CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); - CosmosFactory cosmosFactory = new CosmosFactory(client, testDB1); - final CosmosMappingContext mappingContext = new CosmosMappingContext(); mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class)); - final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); - MultiTenantDBCosmosTemplate cosmosTemplate = new MultiTenantDBCosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); - - CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); - - cosmosTemplate.setNameAndCreateDatabase(testDB1); - cosmosTemplate.createContainerIfNotExists(personInfo); - cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); - cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); - - cosmosTemplate.setNameAndCreateDatabase(testDB2); - cosmosTemplate.createContainerIfNotExists(personInfo); - cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); - cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); - } catch (ClassNotFoundException e) { - assertThat("a").isEqualTo("b"); + } catch (Exception e) { + Assert.fail(); } + + final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); + MultiTenantDBCosmosTemplate cosmosTemplate = new MultiTenantDBCosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); + CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); + + // Create DB1 and add TEST_PERSON_1 to it + cosmosTemplate.setNameAndCreateDatabase(testDB1); + cosmosTemplate.createContainerIfNotExists(personInfo); + cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); + assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); + cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); + + // Create DB2 and add TEST_PERSON_2 to it + cosmosTemplate.setNameAndCreateDatabase(testDB2); + cosmosTemplate.createContainerIfNotExists(personInfo); + cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); + assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); + cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); + + // Check that DB2 has the correct contents + List expectedResultsDB2 = new ArrayList<>(); + expectedResultsDB2.add(TEST_PERSON_2); + Iterable iterableDB2 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); + List resultDB2 = new ArrayList<>(); + iterableDB2.forEach(resultDB2::add); + Assert.assertEquals(expectedResultsDB2, resultDB2); + + // Check that DB1 has the correct contents + cosmosTemplate.setNameAndCreateDatabase(testDB1); + List expectedResultsDB1 = new ArrayList<>(); + expectedResultsDB1.add(TEST_PERSON_1); + Iterable iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); + List resultDB1 = new ArrayList<>(); + iterableDB1.forEach(resultDB1::add); + Assert.assertEquals(expectedResultsDB1, resultDB1); } } From b9c6510180c9e58b5363c49cdfb6b8383f117c4b Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Mon, 5 Dec 2022 12:51:43 -0500 Subject: [PATCH 05/25] Changing function name to be a more accurate description of the functionality. --- .../data/cosmos/core/MultiTenantDBCosmosTemplate.java | 4 ++-- .../cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java | 6 +++--- .../com/azure/spring/data/cosmos/core/CosmosTemplate.java | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java index 76c22c1dceb05..d76b7e86d66cb 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java @@ -76,8 +76,8 @@ public MultiTenantDBCosmosTemplate(CosmosFactory cosmosFactory, CosmosConfig cos } @Override - public void setNameAndCreateDatabase(String dbName) { - super.setNameAndCreateDatabase(dbName); + public void setNameAndCreateDatabaseIfNotExists(String dbName) { + super.setNameAndCreateDatabaseIfNotExists(dbName); } public String getDatabaseName() { diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java index da0c7901a644a..087448c9a0eb4 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java @@ -77,14 +77,14 @@ public void testGetDatabaseFunctionality() { CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); // Create DB1 and add TEST_PERSON_1 to it - cosmosTemplate.setNameAndCreateDatabase(testDB1); + cosmosTemplate.setNameAndCreateDatabaseIfNotExists(testDB1); cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); // Create DB2 and add TEST_PERSON_2 to it - cosmosTemplate.setNameAndCreateDatabase(testDB2); + cosmosTemplate.setNameAndCreateDatabaseIfNotExists(testDB2); cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); @@ -99,7 +99,7 @@ public void testGetDatabaseFunctionality() { Assert.assertEquals(expectedResultsDB2, resultDB2); // Check that DB1 has the correct contents - cosmosTemplate.setNameAndCreateDatabase(testDB1); + cosmosTemplate.setNameAndCreateDatabaseIfNotExists(testDB1); List expectedResultsDB1 = new ArrayList<>(); expectedResultsDB1.add(TEST_PERSON_1); Iterable iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); 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 81d706c0b8526..45d63009b12c1 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 @@ -156,7 +156,7 @@ public String getDatabaseName() { return this.databaseName; } - public void setNameAndCreateDatabase(String dbName) { + public void setNameAndCreateDatabaseIfNotExists(String dbName) { this.databaseName = dbName; createDatabaseIfNotExists() From 4b974b9e9062fade40ba1af922a8db921395cc24 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Thu, 8 Dec 2022 10:07:56 -0500 Subject: [PATCH 06/25] Updating changelog --- sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md b/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md index 6a9bfcfc05a06..d093d82c71b9a 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 3.31.0-beta.1 (Unreleased) #### Features Added +* Added support for multi-tenancy at the Database level via `CosmosTemplate` - See [PR 32516](https://github.com/Azure/azure-sdk-for-java/pull/32516) #### Breaking Changes From ccb755e79f7b8dc4e7996b7e21c60c86adff852f Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Thu, 8 Dec 2022 10:09:26 -0500 Subject: [PATCH 07/25] Removing unused imports. --- .../core/MultiTenantDBCosmosTemplate.java | 56 ------------------- .../MultiTenantDBCosmosTemplateUnitTest.java | 1 - 2 files changed, 57 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java index d76b7e86d66cb..6b7fdefdcdd64 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java @@ -3,66 +3,10 @@ package com.azure.spring.data.cosmos.core; -import com.azure.cosmos.CosmosAsyncClient; -import com.azure.cosmos.CosmosAsyncContainer; -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; -import com.azure.cosmos.models.FeedResponse; -import com.azure.cosmos.models.PartitionKey; -import com.azure.cosmos.models.SqlParameter; -import com.azure.cosmos.models.SqlQuerySpec; -import com.azure.cosmos.models.ThroughputProperties; -import com.azure.cosmos.models.UniqueKeyPolicy; -import com.azure.spring.data.cosmos.Constants; 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; -import com.azure.spring.data.cosmos.core.generator.NativeQueryGenerator; -import com.azure.spring.data.cosmos.core.mapping.event.AfterLoadEvent; -import com.azure.spring.data.cosmos.core.mapping.event.CosmosMappingEvent; -import com.azure.spring.data.cosmos.core.query.CosmosPageImpl; -import com.azure.spring.data.cosmos.core.query.CosmosPageRequest; -import com.azure.spring.data.cosmos.core.query.CosmosQuery; -import com.azure.spring.data.cosmos.core.query.CosmosSliceImpl; -import com.azure.spring.data.cosmos.core.query.Criteria; -import com.azure.spring.data.cosmos.core.query.CriteriaType; -import com.azure.spring.data.cosmos.exception.CosmosExceptionUtils; -import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; -import com.fasterxml.jackson.databind.JsonNode; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.springframework.beans.BeansException; -import org.springframework.context.ApplicationContext; -import org.springframework.context.ApplicationContextAware; import org.springframework.data.auditing.IsNewAwareAuditingHandler; -import org.springframework.data.domain.Page; -import org.springframework.data.domain.Pageable; -import org.springframework.data.domain.Slice; -import org.springframework.data.domain.Sort; -import org.springframework.data.repository.query.parser.Part; -import org.springframework.lang.NonNull; -import org.springframework.util.Assert; -import org.springframework.util.ReflectionUtils; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; - -import java.util.ArrayList; -import java.util.Collections; -import java.util.Iterator; -import java.util.List; -import java.util.Optional; -import java.util.UUID; -import java.util.stream.Collectors; /** * Template class for cosmos db diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java index 087448c9a0eb4..7dcd388799a9d 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java @@ -11,7 +11,6 @@ import com.azure.spring.data.cosmos.config.CosmosConfig; import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; -import com.azure.spring.data.cosmos.domain.Customer; import com.azure.spring.data.cosmos.domain.Person; import com.azure.spring.data.cosmos.repository.TestRepositoryConfig; import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; From f4751abbdf95fec1f35c4d54112fb5994870badd Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Thu, 8 Dec 2022 10:11:52 -0500 Subject: [PATCH 08/25] Code cleanup. --- .../com/azure/spring/data/cosmos/core/CosmosTemplate.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 45d63009b12c1..73d0556dec359 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 @@ -75,7 +75,7 @@ public class CosmosTemplate implements CosmosOperations, ApplicationContextAware private final IsNewAwareAuditingHandler cosmosAuditingHandler; protected String databaseName; - protected final ResponseDiagnosticsProcessor responseDiagnosticsProcessor; + private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor; private final boolean queryMetricsEnabled; private final int maxDegreeOfParallelism; private final int maxBufferedItemCount; @@ -542,7 +542,7 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat return response.getProperties(); } - protected Mono createDatabaseIfNotExists() { + private Mono createDatabaseIfNotExists() { if (databaseThroughputConfig == null) { return cosmosAsyncClient .createDatabaseIfNotExists(this.getDatabaseName()); From 4fdcc4515ecf2a7e62262a4ce3d07682a6a44fea Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Mon, 12 Dec 2022 12:09:35 -0500 Subject: [PATCH 09/25] Refactoring CosmosTemplate to now store the CosmosFactory on the template. With this updated CosmosFactory so that it can be extended to achieve Multi-Tenancy at the database level. The test case was updated also. --- .../core/MultiTenantDBCosmosFactory.java | 34 ++++++++++++ ...> MultiTenantDBCosmosFactoryUnitTest.java} | 11 ++-- .../core/MultiTenantDBCosmosTemplate.java | 30 ----------- .../spring/data/cosmos/CosmosFactory.java | 2 +- .../data/cosmos/core/CosmosTemplate.java | 54 +++++++++---------- 5 files changed, 64 insertions(+), 67 deletions(-) create mode 100644 sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java rename sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/{MultiTenantDBCosmosTemplateUnitTest.java => MultiTenantDBCosmosFactoryUnitTest.java} (91%) delete mode 100644 sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java new file mode 100644 index 0000000000000..f404eb21a57ab --- /dev/null +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.spring.data.cosmos.core; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.spring.data.cosmos.CosmosFactory; +import com.azure.spring.data.cosmos.config.CosmosConfig; +import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; +import org.springframework.data.auditing.IsNewAwareAuditingHandler; + +/** + * Template class for cosmos db + */ +public class MultiTenantDBCosmosFactory extends CosmosFactory { + + public String databaseName; + + /** + * Validate config and initialization + * + * @param cosmosAsyncClient cosmosAsyncClient + * @param databaseName databaseName + */ + public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { + super(cosmosAsyncClient, databaseName); + + this.databaseName = databaseName; + } + + public String getDatabaseName() { + return this.databaseName; + } +} diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java similarity index 91% rename from sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java rename to sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java index 7dcd388799a9d..469cd4d1ee632 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplateUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java @@ -40,7 +40,7 @@ @RunWith(SpringJUnit4ClassRunner.class) @ContextConfiguration(classes = TestRepositoryConfig.class) -public class MultiTenantDBCosmosTemplateUnitTest { +public class MultiTenantDBCosmosFactoryUnitTest { private final String testDB1 = "Database1"; private final String testDB2 = "Database2"; @@ -62,7 +62,7 @@ public class MultiTenantDBCosmosTemplateUnitTest { public void testGetDatabaseFunctionality() { /// Setup CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); - CosmosFactory cosmosFactory = new CosmosFactory(client, testDB1); + MultiTenantDBCosmosFactory cosmosFactory = new MultiTenantDBCosmosFactory(client, testDB1); final CosmosMappingContext mappingContext = new CosmosMappingContext(); try { @@ -72,18 +72,17 @@ public void testGetDatabaseFunctionality() { } final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); - MultiTenantDBCosmosTemplate cosmosTemplate = new MultiTenantDBCosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); + CosmosTemplate cosmosTemplate = new CosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); // Create DB1 and add TEST_PERSON_1 to it - cosmosTemplate.setNameAndCreateDatabaseIfNotExists(testDB1); cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); // Create DB2 and add TEST_PERSON_2 to it - cosmosTemplate.setNameAndCreateDatabaseIfNotExists(testDB2); + cosmosFactory.databaseName = testDB2; cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); @@ -98,7 +97,7 @@ public void testGetDatabaseFunctionality() { Assert.assertEquals(expectedResultsDB2, resultDB2); // Check that DB1 has the correct contents - cosmosTemplate.setNameAndCreateDatabaseIfNotExists(testDB1); + cosmosFactory.databaseName = testDB1; List expectedResultsDB1 = new ArrayList<>(); expectedResultsDB1.add(TEST_PERSON_1); Iterable iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java deleted file mode 100644 index 6b7fdefdcdd64..0000000000000 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosTemplate.java +++ /dev/null @@ -1,30 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.spring.data.cosmos.core; - -import com.azure.spring.data.cosmos.CosmosFactory; -import com.azure.spring.data.cosmos.config.CosmosConfig; -import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; -import org.springframework.data.auditing.IsNewAwareAuditingHandler; - -/** - * Template class for cosmos db - */ -public class MultiTenantDBCosmosTemplate extends CosmosTemplate { - - public MultiTenantDBCosmosTemplate(CosmosFactory cosmosFactory, CosmosConfig cosmosConfig, MappingCosmosConverter mappingCosmosConverter, IsNewAwareAuditingHandler cosmosAuditingHandler) { - super(cosmosFactory, cosmosConfig, mappingCosmosConverter, cosmosAuditingHandler); - - super.databaseName = cosmosFactory.getDatabaseName(); - } - - @Override - public void setNameAndCreateDatabaseIfNotExists(String dbName) { - super.setNameAndCreateDatabaseIfNotExists(dbName); - } - - public String getDatabaseName() { - return super.databaseName; - } -} diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java index 51721280b8471..0262c3a0c8744 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java @@ -20,7 +20,7 @@ public class CosmosFactory { private final CosmosAsyncClient cosmosAsyncClient; - private final String databaseName; + protected String databaseName; private static final String USER_AGENT_SUFFIX = Constants.USER_AGENT_SUFFIX + PropertyLoader.getProjectVersion(); 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 73d0556dec359..6b79f6e60a961 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 @@ -74,13 +74,12 @@ public class CosmosTemplate implements CosmosOperations, ApplicationContextAware private final MappingCosmosConverter mappingCosmosConverter; private final IsNewAwareAuditingHandler cosmosAuditingHandler; - protected String databaseName; + private final CosmosFactory cosmosFactory; private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor; private final boolean queryMetricsEnabled; private final int maxDegreeOfParallelism; private final int maxBufferedItemCount; private final int responseContinuationTokenLimitInKb; - private final CosmosAsyncClient cosmosAsyncClient; private final DatabaseThroughputConfig databaseThroughputConfig; private ApplicationContext applicationContext; @@ -129,8 +128,7 @@ public CosmosTemplate(CosmosFactory cosmosFactory, Assert.notNull(mappingCosmosConverter, "MappingCosmosConverter must not be null!"); this.mappingCosmosConverter = mappingCosmosConverter; this.cosmosAuditingHandler = cosmosAuditingHandler; - this.cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient(); - this.databaseName = cosmosFactory.getDatabaseName(); + this.cosmosFactory = cosmosFactory; this.responseDiagnosticsProcessor = cosmosConfig.getResponseDiagnosticsProcessor(); this.queryMetricsEnabled = cosmosConfig.isQueryMetricsEnabled(); this.maxDegreeOfParallelism = cosmosConfig.getMaxDegreeOfParallelism(); @@ -153,17 +151,11 @@ public CosmosTemplate(CosmosFactory cosmosFactory, } public String getDatabaseName() { - return this.databaseName; + return this.cosmosFactory.getDatabaseName(); } - public void setNameAndCreateDatabaseIfNotExists(String dbName) { - this.databaseName = dbName; - - createDatabaseIfNotExists() - .publishOn(Schedulers.parallel()) - .onErrorResume(throwable -> - CosmosExceptionUtils.exceptionHandler("Failed to create database", throwable, - this.responseDiagnosticsProcessor)); + public CosmosAsyncClient getCosmosAsyncClient() { + return this.cosmosFactory.getCosmosAsyncClient(); } /** @@ -227,7 +219,7 @@ public T insert(String containerName, T objectToSave, PartitionKey partition final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); // if the partition key is null, SDK will get the partitionKey from the object - final CosmosItemResponse response = cosmosAsyncClient + final CosmosItemResponse response = this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .createItem(originalItem, partitionKey, options) @@ -272,7 +264,7 @@ public T findById(Object id, Class domainType, PartitionKey partitionKey) Assert.notNull(partitionKey, "partitionKey should not be null"); String idToQuery = CosmosUtils.getStringIDValue(id); final String containerName = getContainerName(domainType); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .readItem(idToQuery, partitionKey, JsonNode.class) @@ -309,7 +301,7 @@ public T findById(String containerName, Object id, Class domainType) { options.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism); options.setMaxBufferedItemCount(this.maxBufferedItemCount); options.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, options, JsonNode.class) @@ -377,7 +369,7 @@ public T upsertAndReturnEntity(String containerName, T object) { final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); applyVersioning(domainType, originalItem, options); - final CosmosItemResponse cosmosItemResponse = cosmosAsyncClient + final CosmosItemResponse cosmosItemResponse = this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .upsertItem(originalItem, options) @@ -437,7 +429,7 @@ public Iterable findAll(PartitionKey partitionKey, final Class domainT cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount); cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems("SELECT * FROM r", cosmosQueryRequestOptions, JsonNode.class) @@ -472,7 +464,7 @@ public void deleteAll(@NonNull String containerName, @NonNull Class domainTyp @Override public void deleteContainer(@NonNull String containerName) { Assert.hasText(containerName, "containerName should have text."); - cosmosAsyncClient.getDatabase(this.getDatabaseName()) + this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .delete() .publishOn(Schedulers.parallel()) @@ -513,7 +505,7 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat cosmosContainerProperties.setUniqueKeyPolicy(uniqueKeyPolicy); } - CosmosAsyncDatabase cosmosAsyncDatabase = cosmosAsyncClient + CosmosAsyncDatabase cosmosAsyncDatabase = this.getCosmosAsyncClient() .getDatabase(cosmosDatabaseResponse.getProperties().getId()); Mono cosmosContainerResponseMono; @@ -544,20 +536,21 @@ public CosmosContainerProperties createContainerIfNotExists(CosmosEntityInformat private Mono createDatabaseIfNotExists() { if (databaseThroughputConfig == null) { - return cosmosAsyncClient + return this.getCosmosAsyncClient() .createDatabaseIfNotExists(this.getDatabaseName()); } else { ThroughputProperties throughputProperties = databaseThroughputConfig.isAutoScale() ? ThroughputProperties.createAutoscaledThroughput(databaseThroughputConfig.getRequestUnits()) : ThroughputProperties.createManualThroughput(databaseThroughputConfig.getRequestUnits()); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .createDatabaseIfNotExists(this.getDatabaseName(), throughputProperties); } } @Override public CosmosContainerProperties getContainerProperties(String containerName) { - final CosmosContainerResponse response = cosmosAsyncClient.getDatabase(this.getDatabaseName()) + final CosmosContainerResponse response = this.getCosmosAsyncClient() + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .read() .block(); @@ -568,7 +561,8 @@ public CosmosContainerProperties getContainerProperties(String containerName) { @Override public CosmosContainerProperties replaceContainerProperties(String containerName, CosmosContainerProperties properties) { - CosmosContainerResponse response = this.cosmosAsyncClient.getDatabase(this.getDatabaseName()) + CosmosContainerResponse response = this.getCosmosAsyncClient() + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .replace(properties) .block(); @@ -616,7 +610,7 @@ private void deleteById(String containerName, Object id, PartitionKey partitionK partitionKey = PartitionKey.NONE; } - cosmosAsyncClient.getDatabase(this.getDatabaseName()) + this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .deleteItem(idToDelete, partitionKey, options) .publishOn(Schedulers.parallel()) @@ -770,7 +764,7 @@ private Slice sliceQuery(SqlQuerySpec querySpec, }); CosmosAsyncContainer container = - cosmosAsyncClient.getDatabase(this.getDatabaseName()).getContainer(containerName); + this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()).getContainer(containerName); Flux> feedResponseFlux; /* @@ -928,7 +922,7 @@ private Long getCountValue(SqlQuerySpec querySpec, String containerName) { private Flux> executeQuery(SqlQuerySpec sqlQuerySpec, String containerName, CosmosQueryRequestOptions options) { - return cosmosAsyncClient.getDatabase(this.getDatabaseName()) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, options, JsonNode.class) .byPage(); @@ -949,7 +943,7 @@ private Flux findItemsAsFlux(@NonNull CosmosQuery query, cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o)); }); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, cosmosQueryRequestOptions, JsonNode.class) @@ -974,7 +968,7 @@ private Flux getJsonNodeFluxFromQuerySpec( cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount); cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, cosmosQueryRequestOptions, JsonNode.class) @@ -1006,7 +1000,7 @@ private T deleteItem(@NonNull JsonNode jsonNode, final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); applyVersioning(domainType, jsonNode, options); - return cosmosAsyncClient + return this.getCosmosAsyncClient() .getDatabase(this.getDatabaseName()) .getContainer(containerName) .deleteItem(jsonNode, options) From 79338063d1f61383915c59b6d237f4a638169a4c Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 13 Dec 2022 08:30:04 -0500 Subject: [PATCH 10/25] Updating changelog. --- sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md b/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md index d093d82c71b9a..2d039b74398cf 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md @@ -3,7 +3,7 @@ ### 3.31.0-beta.1 (Unreleased) #### Features Added -* Added support for multi-tenancy at the Database level via `CosmosTemplate` - See [PR 32516](https://github.com/Azure/azure-sdk-for-java/pull/32516) +* Added support for multi-tenancy at the Database level via `CosmosFactory` - See [PR 32516](https://github.com/Azure/azure-sdk-for-java/pull/32516) #### Breaking Changes From 8af75dce2ca4fc6605b3e92d8bf47e5f80287f94 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 13 Dec 2022 17:11:19 -0500 Subject: [PATCH 11/25] Making the requested updates in the PR. Adding CosmosFactory to ReactiveCosmosTemplate and adding sample to ReadMe. --- .../core/MultiTenantDBCosmosFactory.java | 2 +- .../MultiTenantDBCosmosFactoryUnitTest.java | 4 +- sdk/cosmos/azure-spring-data-cosmos/README.md | 25 +++++++++ .../data/cosmos/core/CosmosTemplate.java | 4 +- .../cosmos/core/ReactiveCosmosTemplate.java | 56 ++++++++++--------- 5 files changed, 61 insertions(+), 30 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java index f404eb21a57ab..55cd7ee5dd9f5 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java @@ -10,7 +10,7 @@ import org.springframework.data.auditing.IsNewAwareAuditingHandler; /** - * Template class for cosmos db + * Example for extending CosmosFactory for Mutli-Tenancy at the database level */ public class MultiTenantDBCosmosFactory extends CosmosFactory { diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java index 469cd4d1ee632..e0d0d67f31f2c 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java @@ -78,14 +78,14 @@ public void testGetDatabaseFunctionality() { // Create DB1 and add TEST_PERSON_1 to it cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB1); + assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB1); cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); // Create DB2 and add TEST_PERSON_2 to it cosmosFactory.databaseName = testDB2; cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cosmosTemplate.getDatabaseName()).isEqualTo(testDB2); + assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB2); cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); // Check that DB2 has the correct contents diff --git a/sdk/cosmos/azure-spring-data-cosmos/README.md b/sdk/cosmos/azure-spring-data-cosmos/README.md index b76af386b06ff..b3bc94b9a1653 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/README.md +++ b/sdk/cosmos/azure-spring-data-cosmos/README.md @@ -934,6 +934,31 @@ public class MultiDatabaseApplication implements CommandLineRunner { } ``` +### Multi-Tenancy at the Database Level +- Azure-spring-data-cosmos supports multi-tenancy at the database level configuration by extending `CosmosFactory` and overriding the getDatabaseName() function. +```java readme-sample-CosmosFactoryExtension +public class MultiTenantDBCosmosFactory extends CosmosFactory { + + public String databaseName; + + /** + * Validate config and initialization + * + * @param cosmosAsyncClient cosmosAsyncClient + * @param databaseName databaseName + */ + public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { + super(cosmosAsyncClient, databaseName); + + this.databaseName = databaseName; + } + + public String getDatabaseName() { + return this.databaseName; + } +} +``` + ## Beta version package Beta version built from `main` branch are available, you can refer to the [instruction](https://github.com/Azure/azure-sdk-for-java/blob/main/CONTRIBUTING.md#nightly-package-builds) to use beta version packages. 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 6b79f6e60a961..4768cd768e9e4 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 @@ -150,11 +150,11 @@ public CosmosTemplate(CosmosFactory cosmosFactory, this(cosmosFactory, cosmosConfig, mappingCosmosConverter, null); } - public String getDatabaseName() { + private String getDatabaseName() { return this.cosmosFactory.getDatabaseName(); } - public CosmosAsyncClient getCosmosAsyncClient() { + private CosmosAsyncClient getCosmosAsyncClient() { return this.cosmosFactory.getCosmosAsyncClient(); } 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 a533050586953..9becf4409696c 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 @@ -58,14 +58,13 @@ public class ReactiveCosmosTemplate implements ReactiveCosmosOperations, Applica private static final Logger LOGGER = LoggerFactory.getLogger(ReactiveCosmosTemplate.class); + private final CosmosFactory cosmosFactory; private final MappingCosmosConverter mappingCosmosConverter; - private final String databaseName; private final ResponseDiagnosticsProcessor responseDiagnosticsProcessor; private final boolean queryMetricsEnabled; private final int maxDegreeOfParallelism; private final int maxBufferedItemCount; private final int responseContinuationTokenLimitInKb; - private final CosmosAsyncClient cosmosAsyncClient; private final IsNewAwareAuditingHandler cosmosAuditingHandler; private final DatabaseThroughputConfig databaseThroughputConfig; @@ -116,8 +115,7 @@ public ReactiveCosmosTemplate(CosmosFactory cosmosFactory, Assert.notNull(mappingCosmosConverter, "MappingCosmosConverter must not be null!"); this.mappingCosmosConverter = mappingCosmosConverter; - this.cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient(); - this.databaseName = cosmosFactory.getDatabaseName(); + this.cosmosFactory = cosmosFactory; this.responseDiagnosticsProcessor = cosmosConfig.getResponseDiagnosticsProcessor(); this.queryMetricsEnabled = cosmosConfig.isQueryMetricsEnabled(); this.maxDegreeOfParallelism = cosmosConfig.getMaxDegreeOfParallelism(); @@ -140,6 +138,14 @@ public ReactiveCosmosTemplate(CosmosFactory cosmosFactory, this(cosmosFactory, cosmosConfig, mappingCosmosConverter, null); } + private String getDatabaseName() { + return this.cosmosFactory.getDatabaseName(); + } + + private CosmosAsyncClient getCosmosAsyncClient() { + return this.cosmosFactory.getCosmosAsyncClient(); + } + /** * @param applicationContext the application context * @throws BeansException the bean exception @@ -175,7 +181,7 @@ public Mono createContainerIfNotExists(CosmosEntityInfo } CosmosAsyncDatabase database = - cosmosAsyncClient.getDatabase(cosmosDatabaseResponse.getProperties().getId()); + this.getCosmosAsyncClient().getDatabase(cosmosDatabaseResponse.getProperties().getId()); Mono cosmosContainerResponseMono; if (information.getRequestUnit() == null) { @@ -205,20 +211,20 @@ public Mono createContainerIfNotExists(CosmosEntityInfo private Mono createDatabaseIfNotExists() { if (databaseThroughputConfig == null) { - return cosmosAsyncClient - .createDatabaseIfNotExists(this.databaseName); + return this.getCosmosAsyncClient() + .createDatabaseIfNotExists(this.getDatabaseName()); } else { ThroughputProperties throughputProperties = databaseThroughputConfig.isAutoScale() ? ThroughputProperties.createAutoscaledThroughput(databaseThroughputConfig.getRequestUnits()) : ThroughputProperties.createManualThroughput(databaseThroughputConfig.getRequestUnits()); - return cosmosAsyncClient - .createDatabaseIfNotExists(this.databaseName, throughputProperties); + return this.getCosmosAsyncClient() + .createDatabaseIfNotExists(this.getDatabaseName(), throughputProperties); } } @Override public Mono getContainerProperties(String containerName) { - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .read() .map(CosmosContainerResponse::getProperties); @@ -227,7 +233,7 @@ public Mono getContainerProperties(String containerNa @Override public Mono replaceContainerProperties(String containerName, CosmosContainerProperties properties) { - return this.cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .replace(properties) .map(CosmosContainerResponse::getProperties); @@ -273,8 +279,8 @@ public Flux findAll(PartitionKey partitionKey, Class domainType) { cosmosQueryRequestOptions.setMaxBufferedItemCount(this.maxBufferedItemCount); cosmosQueryRequestOptions.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); - return cosmosAsyncClient - .getDatabase(this.databaseName) + return this.getCosmosAsyncClient() + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems("SELECT * FROM r", cosmosQueryRequestOptions, JsonNode.class) .byPage() @@ -325,7 +331,7 @@ public Mono findById(String containerName, Object id, Class domainType options.setMaxBufferedItemCount(this.maxBufferedItemCount); options.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, options, JsonNode.class) .byPage() @@ -360,7 +366,7 @@ public Mono findById(Object id, Class domainType, PartitionKey partiti String idToFind = CosmosUtils.getStringIDValue(id); final String containerName = getContainerName(domainType); - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .readItem(idToFind, partitionKey, JsonNode.class) .publishOn(Schedulers.parallel()) @@ -418,8 +424,8 @@ public Mono insert(String containerName, T objectToSave, final JsonNode originalItem = mappingCosmosConverter.writeJsonNode(objectToSave); final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); // if the partition key is null, SDK will get the partitionKey from the object - return cosmosAsyncClient - .getDatabase(this.databaseName) + return this.getCosmosAsyncClient() + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .createItem(originalItem, partitionKey, options) .publishOn(Schedulers.parallel()) @@ -481,7 +487,7 @@ public Mono upsert(String containerName, T object) { applyVersioning(object.getClass(), originalItem, options); - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .upsertItem(originalItem, options) .publishOn(Schedulers.parallel()) @@ -517,7 +523,7 @@ private Mono deleteById(String containerName, Object id, PartitionKey part partitionKey = PartitionKey.NONE; } - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .deleteItem(idToDelete, partitionKey, cosmosItemRequestOptions) .publishOn(Schedulers.parallel()) @@ -683,7 +689,7 @@ private Flux runQuery(SqlQuerySpec querySpec, Class domainType) { options.setMaxDegreeOfParallelism(this.maxDegreeOfParallelism); options.setMaxBufferedItemCount(this.maxBufferedItemCount); options.setResponseContinuationTokenLimitInKb(this.responseContinuationTokenLimitInKb); - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(querySpec, options, JsonNode.class) .byPage() @@ -721,7 +727,7 @@ private Flux> executeQuery(SqlQuerySpec sqlQuerySpec, String containerName, CosmosQueryRequestOptions options) { - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, options, JsonNode.class) .byPage() @@ -738,7 +744,7 @@ private Flux> executeQuery(SqlQuerySpec sqlQuerySpec, @Override public void deleteContainer(@NonNull String containerName) { Assert.hasText(containerName, "containerName should have text."); - cosmosAsyncClient.getDatabase(this.databaseName) + this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .delete() .doOnNext(cosmosContainerResponse -> @@ -782,8 +788,8 @@ private Flux findItems(@NonNull CosmosQuery query, cosmosQueryRequestOptions.setPartitionKey(new PartitionKey(o)); }); - return cosmosAsyncClient - .getDatabase(this.databaseName) + return this.getCosmosAsyncClient() + .getDatabase(this.getDatabaseName()) .getContainer(containerName) .queryItems(sqlQuerySpec, cosmosQueryRequestOptions, JsonNode.class) .byPage() @@ -804,7 +810,7 @@ private Mono deleteItem(@NonNull JsonNode jsonNode, final CosmosItemRequestOptions options = new CosmosItemRequestOptions(); applyVersioning(domainType, jsonNode, options); - return cosmosAsyncClient.getDatabase(this.databaseName) + return this.getCosmosAsyncClient().getDatabase(this.getDatabaseName()) .getContainer(containerName) .deleteItem(jsonNode, options) .publishOn(Schedulers.parallel()) From cbea46d4936d85ac92f96c93563a3d53db8dee86 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 09:21:15 -0500 Subject: [PATCH 12/25] Making updates for PR comments. --- .../spring/data/cosmos/core/MultiTenantDBCosmosFactory.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java index 55cd7ee5dd9f5..9ddfa5f1c2d73 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java @@ -14,7 +14,7 @@ */ public class MultiTenantDBCosmosFactory extends CosmosFactory { - public String databaseName; + public String manuallySetDatabaseName; /** * Validate config and initialization @@ -28,7 +28,8 @@ public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String da this.databaseName = databaseName; } + @Override public String getDatabaseName() { - return this.databaseName; + return this.manuallySetDatabaseName; } } From 6b1d18bd403549a36b5f1e1934128bb6c29f2f1b Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 10:03:43 -0500 Subject: [PATCH 13/25] Fixing updates to unit test. --- .../spring/data/cosmos/core/MultiTenantDBCosmosFactory.java | 2 +- .../data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java index 9ddfa5f1c2d73..353471d935cd2 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java @@ -25,7 +25,7 @@ public class MultiTenantDBCosmosFactory extends CosmosFactory { public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { super(cosmosAsyncClient, databaseName); - this.databaseName = databaseName; + this.manuallySetDatabaseName = databaseName; } @Override diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java index e0d0d67f31f2c..8d6cc7b2656d0 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java @@ -82,7 +82,7 @@ public void testGetDatabaseFunctionality() { cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); // Create DB2 and add TEST_PERSON_2 to it - cosmosFactory.databaseName = testDB2; + cosmosFactory.manuallySetDatabaseName = testDB2; cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB2); @@ -97,7 +97,7 @@ public void testGetDatabaseFunctionality() { Assert.assertEquals(expectedResultsDB2, resultDB2); // Check that DB1 has the correct contents - cosmosFactory.databaseName = testDB1; + cosmosFactory.manuallySetDatabaseName = testDB1; List expectedResultsDB1 = new ArrayList<>(); expectedResultsDB1.add(TEST_PERSON_1); Iterable iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); From 948a0f2c16f193795620f9651a760f8e856ae6f0 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 10:30:55 -0500 Subject: [PATCH 14/25] Fixing readme --- sdk/cosmos/azure-spring-data-cosmos/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos/README.md b/sdk/cosmos/azure-spring-data-cosmos/README.md index b3bc94b9a1653..9e4e1f32a0ea9 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/README.md +++ b/sdk/cosmos/azure-spring-data-cosmos/README.md @@ -936,7 +936,7 @@ public class MultiDatabaseApplication implements CommandLineRunner { ### Multi-Tenancy at the Database Level - Azure-spring-data-cosmos supports multi-tenancy at the database level configuration by extending `CosmosFactory` and overriding the getDatabaseName() function. -```java readme-sample-CosmosFactoryExtension +```java readme-sample-MultiTenantDBCosmosFactory public class MultiTenantDBCosmosFactory extends CosmosFactory { public String databaseName; From 7b82f126df91e718d671dcbe9af7eb273b8f4cdd Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 10:53:05 -0500 Subject: [PATCH 15/25] Adding file needed for readme. --- .../cosmos/MultiTenantDBCosmosFactory.java | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java new file mode 100644 index 0000000000000..0f6b8a5d398b2 --- /dev/null +++ b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.spring.data.cosmos; + +import com.azure.cosmos.CosmosAsyncClient; + +/** + * Example for extending CosmosFactory for Mutli-Tenancy at the database level + */ +public class MultiTenantDBCosmosFactory extends CosmosFactory { + + public String manuallySetDatabaseName; + + /** + * Validate config and initialization + * + * @param cosmosAsyncClient cosmosAsyncClient + * @param databaseName databaseName + */ + public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { + super(cosmosAsyncClient, databaseName); + + this.manuallySetDatabaseName = databaseName; + } + + @Override + public String getDatabaseName() { + return this.manuallySetDatabaseName; + } +} From b3339d8d3f35f4f37b9cc5bbd35a27d4fac6307c Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 11:12:34 -0500 Subject: [PATCH 16/25] Fixing snippet for readme. --- .../azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java index 0f6b8a5d398b2..e9074e02e973b 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java @@ -8,6 +8,7 @@ /** * Example for extending CosmosFactory for Mutli-Tenancy at the database level */ +// BEGIN: readme-sample-MultiTenantDBCosmosFactory public class MultiTenantDBCosmosFactory extends CosmosFactory { public String manuallySetDatabaseName; @@ -29,3 +30,4 @@ public String getDatabaseName() { return this.manuallySetDatabaseName; } } +// END: readme-sample-MultiTenantDBCosmosFactory From 6e78dab8e71cb6a04e79c8820f69401576da525b Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 12:13:36 -0500 Subject: [PATCH 17/25] Fixing snippet for readme. --- sdk/cosmos/azure-spring-data-cosmos/README.md | 15 ++++++++++++--- .../data/cosmos/MultiTenantDBCosmosFactory.java | 10 +++++++++- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos/README.md b/sdk/cosmos/azure-spring-data-cosmos/README.md index 9e4e1f32a0ea9..9f01db504073a 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/README.md +++ b/sdk/cosmos/azure-spring-data-cosmos/README.md @@ -939,7 +939,7 @@ public class MultiDatabaseApplication implements CommandLineRunner { ```java readme-sample-MultiTenantDBCosmosFactory public class MultiTenantDBCosmosFactory extends CosmosFactory { - public String databaseName; + private String manuallySetDatabaseName; /** * Validate config and initialization @@ -950,11 +950,20 @@ public class MultiTenantDBCosmosFactory extends CosmosFactory { public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { super(cosmosAsyncClient, databaseName); - this.databaseName = databaseName; + this.manuallySetDatabaseName = databaseName; } + @Override public String getDatabaseName() { - return this.databaseName; + return this.manuallySetDatabaseName; + } + + public void setManuallySetDatabaseName(String manuallySetDatabaseName) { + this.manuallySetDatabaseName = manuallySetDatabaseName; + } + + public String getManuallySetDatabaseName() { + return manuallySetDatabaseName; } } ``` diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java index e9074e02e973b..9d46460472df8 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java @@ -11,7 +11,7 @@ // BEGIN: readme-sample-MultiTenantDBCosmosFactory public class MultiTenantDBCosmosFactory extends CosmosFactory { - public String manuallySetDatabaseName; + private String manuallySetDatabaseName; /** * Validate config and initialization @@ -29,5 +29,13 @@ public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String da public String getDatabaseName() { return this.manuallySetDatabaseName; } + + public void setManuallySetDatabaseName(String manuallySetDatabaseName) { + this.manuallySetDatabaseName = manuallySetDatabaseName; + } + + public String getManuallySetDatabaseName() { + return manuallySetDatabaseName; + } } // END: readme-sample-MultiTenantDBCosmosFactory From 40af2432fd150d85ebaaa8cbbf401a4d883271dc Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 15:38:24 -0500 Subject: [PATCH 18/25] Updating readme. --- sdk/cosmos/azure-spring-data-cosmos/README.md | 18 +++++------------- .../cosmos/MultiTenantDBCosmosFactory.java | 14 +++----------- 2 files changed, 8 insertions(+), 24 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos/README.md b/sdk/cosmos/azure-spring-data-cosmos/README.md index 9f01db504073a..e68f6de593fb7 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/README.md +++ b/sdk/cosmos/azure-spring-data-cosmos/README.md @@ -939,8 +939,8 @@ public class MultiDatabaseApplication implements CommandLineRunner { ```java readme-sample-MultiTenantDBCosmosFactory public class MultiTenantDBCosmosFactory extends CosmosFactory { - private String manuallySetDatabaseName; - + private String tenantId; + /** * Validate config and initialization * @@ -949,21 +949,13 @@ public class MultiTenantDBCosmosFactory extends CosmosFactory { */ public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { super(cosmosAsyncClient, databaseName); - - this.manuallySetDatabaseName = databaseName; + + this.tenantId = databaseName; } @Override public String getDatabaseName() { - return this.manuallySetDatabaseName; - } - - public void setManuallySetDatabaseName(String manuallySetDatabaseName) { - this.manuallySetDatabaseName = manuallySetDatabaseName; - } - - public String getManuallySetDatabaseName() { - return manuallySetDatabaseName; + return this.getCosmosAsyncClient().getDatabase(this.tenantId).toString(); } } ``` diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java index 9d46460472df8..fe951fff0b515 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/samples/java/com/azure/spring/data/cosmos/MultiTenantDBCosmosFactory.java @@ -11,7 +11,7 @@ // BEGIN: readme-sample-MultiTenantDBCosmosFactory public class MultiTenantDBCosmosFactory extends CosmosFactory { - private String manuallySetDatabaseName; + private String tenantId; /** * Validate config and initialization @@ -22,20 +22,12 @@ public class MultiTenantDBCosmosFactory extends CosmosFactory { public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) { super(cosmosAsyncClient, databaseName); - this.manuallySetDatabaseName = databaseName; + this.tenantId = databaseName; } @Override public String getDatabaseName() { - return this.manuallySetDatabaseName; - } - - public void setManuallySetDatabaseName(String manuallySetDatabaseName) { - this.manuallySetDatabaseName = manuallySetDatabaseName; - } - - public String getManuallySetDatabaseName() { - return manuallySetDatabaseName; + return this.getCosmosAsyncClient().getDatabase(this.tenantId).toString(); } } // END: readme-sample-MultiTenantDBCosmosFactory From f7b58dc896b0c8acbc208911106be38ba212c38d Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Fri, 16 Dec 2022 16:00:00 -0500 Subject: [PATCH 19/25] Adding javadoc. --- .../main/java/com/azure/spring/data/cosmos/CosmosFactory.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java index 0262c3a0c8744..dfd05f94d4a02 100644 --- a/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/CosmosFactory.java @@ -20,6 +20,9 @@ public class CosmosFactory { private final CosmosAsyncClient cosmosAsyncClient; + /** + * Database Name to be used for operations. + */ protected String databaseName; private static final String USER_AGENT_SUFFIX = From e6b9dfc38c3191d8de820df721f7337868252c65 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 20 Dec 2022 12:08:40 -0500 Subject: [PATCH 20/25] Fixing unit test. --- .../data/cosmos/repository/MultiCosmosTemplateIT.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiCosmosTemplateIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiCosmosTemplateIT.java index 0ae1bb6c320a8..02b00ca05833f 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiCosmosTemplateIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiCosmosTemplateIT.java @@ -4,6 +4,7 @@ import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.models.PartitionKey; +import com.azure.spring.data.cosmos.CosmosFactory; import com.azure.spring.data.cosmos.ReactiveIntegrationTestCollectionManager; import com.azure.spring.data.cosmos.common.TestConstants; import com.azure.spring.data.cosmos.core.ReactiveCosmosTemplate; @@ -80,10 +81,12 @@ public void testSecondaryTemplateWithDiffDatabase() { @Test public void testSingleCosmosClientForMultipleCosmosTemplate() throws IllegalAccessException { - final Field cosmosAsyncClient = FieldUtils.getDeclaredField(ReactiveCosmosTemplate.class, - "cosmosAsyncClient", true); - CosmosAsyncClient client1 = (CosmosAsyncClient) cosmosAsyncClient.get(secondaryReactiveCosmosTemplate); - CosmosAsyncClient client2 = (CosmosAsyncClient) cosmosAsyncClient.get(secondaryDiffDatabaseReactiveCosmosTemplate); + final Field cosmosFactory = FieldUtils.getDeclaredField(ReactiveCosmosTemplate.class, + "cosmosFactory", true); + CosmosFactory cosmosFactory1 = (CosmosFactory) cosmosFactory.get(secondaryReactiveCosmosTemplate); + CosmosAsyncClient client1 = cosmosFactory1.getCosmosAsyncClient(); + CosmosFactory cosmosFactory2 = (CosmosFactory) cosmosFactory.get(secondaryDiffDatabaseReactiveCosmosTemplate); + CosmosAsyncClient client2 = cosmosFactory2.getCosmosAsyncClient(); Assertions.assertThat(client1).isEqualTo(client2); } } From 6674719600ce975407e8460bd39c5219bc9cb437 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 20 Dec 2022 14:00:34 -0500 Subject: [PATCH 21/25] Testing. --- .../data/cosmos/core/MultiTenantDBCosmosFactory.java | 3 --- .../core/MultiTenantDBCosmosFactoryUnitTest.java | 12 ++++++------ 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java index 353471d935cd2..d7212896d6e74 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java @@ -5,9 +5,6 @@ import com.azure.cosmos.CosmosAsyncClient; import com.azure.spring.data.cosmos.CosmosFactory; -import com.azure.spring.data.cosmos.config.CosmosConfig; -import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; -import org.springframework.data.auditing.IsNewAwareAuditingHandler; /** * Example for extending CosmosFactory for Mutli-Tenancy at the database level diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java index 8d6cc7b2656d0..b9543158427fd 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java @@ -62,7 +62,7 @@ public class MultiTenantDBCosmosFactoryUnitTest { public void testGetDatabaseFunctionality() { /// Setup CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); - MultiTenantDBCosmosFactory cosmosFactory = new MultiTenantDBCosmosFactory(client, testDB1); + MultiTenantDBCosmosFactory cf = new MultiTenantDBCosmosFactory(client, testDB1); final CosmosMappingContext mappingContext = new CosmosMappingContext(); try { @@ -72,20 +72,20 @@ public void testGetDatabaseFunctionality() { } final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); - CosmosTemplate cosmosTemplate = new CosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); + CosmosTemplate cosmosTemplate = new CosmosTemplate(cf, cosmosConfig, cosmosConverter, null); CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); // Create DB1 and add TEST_PERSON_1 to it cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB1); + assertThat(cf.getDatabaseName()).isEqualTo(testDB1); cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); // Create DB2 and add TEST_PERSON_2 to it - cosmosFactory.manuallySetDatabaseName = testDB2; + cf.manuallySetDatabaseName = testDB2; cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB2); + assertThat(cf.getDatabaseName()).isEqualTo(testDB2); cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); // Check that DB2 has the correct contents @@ -97,7 +97,7 @@ public void testGetDatabaseFunctionality() { Assert.assertEquals(expectedResultsDB2, resultDB2); // Check that DB1 has the correct contents - cosmosFactory.manuallySetDatabaseName = testDB1; + cf.manuallySetDatabaseName = testDB1; List expectedResultsDB1 = new ArrayList<>(); expectedResultsDB1.add(TEST_PERSON_1); Iterable iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); From 55e793d383d97363588955318f4ae355e1cd1836 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 20 Dec 2022 14:31:33 -0500 Subject: [PATCH 22/25] Testing breaking out setup to be before unit test runs. --- .../MultiTenantDBCosmosFactoryUnitTest.java | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java index b9543158427fd..698c40dad47d6 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java @@ -8,13 +8,17 @@ import com.azure.cosmos.models.PartitionKey; import com.azure.spring.data.cosmos.CosmosFactory; import com.azure.spring.data.cosmos.IntegrationTestCollectionManager; +import com.azure.spring.data.cosmos.common.TestConstants; import com.azure.spring.data.cosmos.config.CosmosConfig; import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; +import com.azure.spring.data.cosmos.domain.AuditableEntity; +import com.azure.spring.data.cosmos.domain.GenIdEntity; import com.azure.spring.data.cosmos.domain.Person; import com.azure.spring.data.cosmos.repository.TestRepositoryConfig; import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; import org.junit.Assert; +import org.junit.Before; import org.junit.ClassRule; import org.junit.Test; import org.junit.runner.RunWith; @@ -58,11 +62,15 @@ public class MultiTenantDBCosmosFactoryUnitTest { @Autowired private CosmosClientBuilder cosmosClientBuilder; - @Test - public void testGetDatabaseFunctionality() { + private MultiTenantDBCosmosFactory cosmosFactory; + private CosmosTemplate cosmosTemplate; + private CosmosEntityInformation personInfo; + + @Before + public void setUp() throws ClassNotFoundException { /// Setup CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); - MultiTenantDBCosmosFactory cf = new MultiTenantDBCosmosFactory(client, testDB1); + cosmosFactory = new MultiTenantDBCosmosFactory(client, testDB1); final CosmosMappingContext mappingContext = new CosmosMappingContext(); try { @@ -72,20 +80,23 @@ public void testGetDatabaseFunctionality() { } final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); - CosmosTemplate cosmosTemplate = new CosmosTemplate(cf, cosmosConfig, cosmosConverter, null); - CosmosEntityInformation personInfo = new CosmosEntityInformation<>(Person.class); + cosmosTemplate = new CosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null); + personInfo = new CosmosEntityInformation<>(Person.class); + } + @Test + public void testGetDatabaseFunctionality() { // Create DB1 and add TEST_PERSON_1 to it cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cf.getDatabaseName()).isEqualTo(testDB1); + assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB1); cosmosTemplate.insert(TEST_PERSON_1, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_1))); // Create DB2 and add TEST_PERSON_2 to it - cf.manuallySetDatabaseName = testDB2; + cosmosFactory.manuallySetDatabaseName = testDB2; cosmosTemplate.createContainerIfNotExists(personInfo); cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); - assertThat(cf.getDatabaseName()).isEqualTo(testDB2); + assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB2); cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); // Check that DB2 has the correct contents @@ -97,7 +108,7 @@ public void testGetDatabaseFunctionality() { Assert.assertEquals(expectedResultsDB2, resultDB2); // Check that DB1 has the correct contents - cf.manuallySetDatabaseName = testDB1; + cosmosFactory.manuallySetDatabaseName = testDB1; List expectedResultsDB1 = new ArrayList<>(); expectedResultsDB1.add(TEST_PERSON_1); Iterable iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); From 52ba22092de0956b60699a5497ac6c27c2a32225 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 20 Dec 2022 15:09:28 -0500 Subject: [PATCH 23/25] Renaming file. --- ...actoryUnitTest.java => MultiTenantDBCosmosFactoryIT.java} | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) rename sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/{MultiTenantDBCosmosFactoryUnitTest.java => MultiTenantDBCosmosFactoryIT.java} (96%) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java similarity index 96% rename from sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java rename to sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java index 698c40dad47d6..39b0084fc7add 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryUnitTest.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java @@ -8,12 +8,9 @@ import com.azure.cosmos.models.PartitionKey; import com.azure.spring.data.cosmos.CosmosFactory; import com.azure.spring.data.cosmos.IntegrationTestCollectionManager; -import com.azure.spring.data.cosmos.common.TestConstants; import com.azure.spring.data.cosmos.config.CosmosConfig; import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; -import com.azure.spring.data.cosmos.domain.AuditableEntity; -import com.azure.spring.data.cosmos.domain.GenIdEntity; import com.azure.spring.data.cosmos.domain.Person; import com.azure.spring.data.cosmos.repository.TestRepositoryConfig; import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; @@ -44,7 +41,7 @@ @RunWith(SpringJUnit4ClassRunner.class) @ContextConfiguration(classes = TestRepositoryConfig.class) -public class MultiTenantDBCosmosFactoryUnitTest { +public class MultiTenantDBCosmosFactoryIT { private final String testDB1 = "Database1"; private final String testDB2 = "Database2"; From 63777fa61cedec21f952f270e25821a67dd9ff22 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 20 Dec 2022 18:34:03 -0500 Subject: [PATCH 24/25] Adding new test config for MultiTenantDB test. --- .../core/MultiTenantDBCosmosFactoryIT.java | 4 +- .../MultiTenantTestRepositoryConfig.java | 96 +++++++++++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiTenantTestRepositoryConfig.java diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java index 39b0084fc7add..ce29146de8f65 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java @@ -12,7 +12,7 @@ import com.azure.spring.data.cosmos.core.convert.MappingCosmosConverter; import com.azure.spring.data.cosmos.core.mapping.CosmosMappingContext; import com.azure.spring.data.cosmos.domain.Person; -import com.azure.spring.data.cosmos.repository.TestRepositoryConfig; +import com.azure.spring.data.cosmos.repository.MultiTenantTestRepositoryConfig; import com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation; import org.junit.Assert; import org.junit.Before; @@ -40,7 +40,7 @@ import static org.assertj.core.api.Assertions.assertThat; @RunWith(SpringJUnit4ClassRunner.class) -@ContextConfiguration(classes = TestRepositoryConfig.class) +@ContextConfiguration(classes = MultiTenantTestRepositoryConfig.class) public class MultiTenantDBCosmosFactoryIT { private final String testDB1 = "Database1"; diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiTenantTestRepositoryConfig.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiTenantTestRepositoryConfig.java new file mode 100644 index 0000000000000..e28533a40997f --- /dev/null +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/repository/MultiTenantTestRepositoryConfig.java @@ -0,0 +1,96 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.spring.data.cosmos.repository; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.spring.data.cosmos.common.ResponseDiagnosticsTestUtils; +import com.azure.spring.data.cosmos.common.TestConstants; +import com.azure.spring.data.cosmos.config.AbstractCosmosConfiguration; +import com.azure.spring.data.cosmos.config.CosmosConfig; +import com.azure.spring.data.cosmos.core.MultiTenantDBCosmosFactory; +import com.azure.spring.data.cosmos.core.mapping.event.SimpleCosmosMappingEventListener; +import com.azure.spring.data.cosmos.repository.config.EnableCosmosRepositories; +import com.azure.spring.data.cosmos.repository.config.EnableReactiveCosmosRepositories; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.PropertySource; +import org.springframework.util.StringUtils; + +import java.util.Arrays; +import java.util.Collection; + +@Configuration +@PropertySource(value = { "classpath:application.properties" }) +@EnableCosmosRepositories +@EnableReactiveCosmosRepositories +public class MultiTenantTestRepositoryConfig extends AbstractCosmosConfiguration { + @Value("${cosmos.uri:}") + private String cosmosDbUri; + + @Value("${cosmos.key:}") + private String cosmosDbKey; + + @Value("${cosmos.database:}") + private String database; + + @Value("${cosmos.queryMetricsEnabled}") + private boolean queryMetricsEnabled; + + @Value("${cosmos.maxDegreeOfParallelism}") + private int maxDegreeOfParallelism; + + @Value("${cosmos.maxBufferedItemCount}") + private int maxBufferedItemCount; + + @Value("${cosmos.responseContinuationTokenLimitInKb}") + private int responseContinuationTokenLimitInKb; + + @Bean + public ResponseDiagnosticsTestUtils responseDiagnosticsTestUtils() { + return new ResponseDiagnosticsTestUtils(); + } + + @Bean + public CosmosClientBuilder cosmosClientBuilder() { + return new CosmosClientBuilder() + .key(cosmosDbKey) + .endpoint(cosmosDbUri) + .contentResponseOnWriteEnabled(true); + } + + @Bean + public MultiTenantDBCosmosFactory cosmosFactory(CosmosAsyncClient cosmosAsyncClient) { + return new MultiTenantDBCosmosFactory(cosmosAsyncClient, getDatabaseName()); + } + + @Bean + @Override + public CosmosConfig cosmosConfig() { + return CosmosConfig.builder() + .enableQueryMetrics(queryMetricsEnabled) + .maxDegreeOfParallelism(maxDegreeOfParallelism) + .maxBufferedItemCount(maxBufferedItemCount) + .responseContinuationTokenLimitInKb(responseContinuationTokenLimitInKb) + .responseDiagnosticsProcessor(responseDiagnosticsTestUtils().getResponseDiagnosticsProcessor()) + .build(); + } + + @Override + protected String getDatabaseName() { + return StringUtils.hasText(this.database) ? this.database : TestConstants.DB_NAME; + } + + @Override + protected Collection getMappingBasePackages() { + final Package mappingBasePackage = getClass().getPackage(); + final String entityPackage = "com.azure.spring.data.cosmos.domain"; + return Arrays.asList(mappingBasePackage.getName(), entityPackage); + } + + @Bean + SimpleCosmosMappingEventListener simpleMappingEventListener() { + return new SimpleCosmosMappingEventListener(); + } +} From 466e5db2cf87dd1a6a130bf4d7f7f1bf2d12b377 Mon Sep 17 00:00:00 2001 From: Trevor Anderson Date: Tue, 20 Dec 2022 21:04:18 -0500 Subject: [PATCH 25/25] Adding cleanup to unit test. --- .../core/MultiTenantDBCosmosFactoryIT.java | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java index ce29146de8f65..dc16252d651a3 100644 --- a/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java +++ b/sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java @@ -4,7 +4,9 @@ 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.models.PartitionKey; import com.azure.spring.data.cosmos.CosmosFactory; import com.azure.spring.data.cosmos.IntegrationTestCollectionManager; @@ -38,6 +40,7 @@ import static com.azure.spring.data.cosmos.common.TestConstants.LAST_NAME; import static com.azure.spring.data.cosmos.common.TestConstants.PASSPORT_IDS_BY_COUNTRY; import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.Assert.assertEquals; @RunWith(SpringJUnit4ClassRunner.class) @ContextConfiguration(classes = MultiTenantTestRepositoryConfig.class) @@ -61,12 +64,13 @@ public class MultiTenantDBCosmosFactoryIT { private MultiTenantDBCosmosFactory cosmosFactory; private CosmosTemplate cosmosTemplate; + private CosmosAsyncClient client; private CosmosEntityInformation personInfo; @Before public void setUp() throws ClassNotFoundException { /// Setup - CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); + client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); cosmosFactory = new MultiTenantDBCosmosFactory(client, testDB1); final CosmosMappingContext mappingContext = new CosmosMappingContext(); @@ -112,5 +116,18 @@ public void testGetDatabaseFunctionality() { List resultDB1 = new ArrayList<>(); iterableDB1.forEach(resultDB1::add); Assert.assertEquals(expectedResultsDB1, resultDB1); + + //Cleanup + deleteDatabaseIfExists(testDB1); + deleteDatabaseIfExists(testDB2); + } + + private void deleteDatabaseIfExists(String dbName) { + CosmosAsyncDatabase database = client.getDatabase(dbName); + try { + database.delete().block(); + } catch (CosmosException ex) { + assertEquals(ex.getStatusCode(), 404); + } } }