forked from Azure/azure-sdk-for-java
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Enhancing CosmosTemplate to Support Multi-Tenancy at a DB Level (Azur…
…e#32516) * Proof of concept that we can write to two databases from the same session. * Improving the changes to CosmosTemplate and the test case. * Moving default setNameAndCreateDatabase() logic into CosmosTemplate. * Improving unit test. * Changing function name to be a more accurate description of the functionality. * Updating changelog * Removing unused imports. * Code cleanup. * 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. * Updating changelog. * Making the requested updates in the PR. Adding CosmosFactory to ReactiveCosmosTemplate and adding sample to ReadMe. * Making updates for PR comments. * Fixing updates to unit test. * Fixing readme * Adding file needed for readme. * Fixing snippet for readme. * Fixing snippet for readme. * Updating readme. * Adding javadoc. * Fixing unit test. * Testing. * Testing breaking out setup to be before unit test runs. * Renaming file. * Adding new test config for MultiTenantDB test. * Adding cleanup to unit test.
- Loading branch information
1 parent
4780c59
commit 9452cfc
Showing
10 changed files
with
405 additions
and
64 deletions.
There are no files selected for viewing
32 changes: 32 additions & 0 deletions
32
...smos-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
// 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; | ||
|
||
/** | ||
* 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; | ||
} | ||
} |
133 changes: 133 additions & 0 deletions
133
...os-test/src/test/java/com/azure/spring/data/cosmos/core/MultiTenantDBCosmosFactoryIT.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,133 @@ | ||
// 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.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; | ||
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.MultiTenantTestRepositoryConfig; | ||
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; | ||
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 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; | ||
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; | ||
import static org.junit.Assert.assertEquals; | ||
|
||
@RunWith(SpringJUnit4ClassRunner.class) | ||
@ContextConfiguration(classes = MultiTenantTestRepositoryConfig.class) | ||
public class MultiTenantDBCosmosFactoryIT { | ||
|
||
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 | ||
private CosmosConfig cosmosConfig; | ||
@Autowired | ||
private CosmosClientBuilder cosmosClientBuilder; | ||
|
||
private MultiTenantDBCosmosFactory cosmosFactory; | ||
private CosmosTemplate cosmosTemplate; | ||
private CosmosAsyncClient client; | ||
private CosmosEntityInformation<Person, String> personInfo; | ||
|
||
@Before | ||
public void setUp() throws ClassNotFoundException { | ||
/// Setup | ||
client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder); | ||
cosmosFactory = new MultiTenantDBCosmosFactory(client, testDB1); | ||
final CosmosMappingContext mappingContext = new CosmosMappingContext(); | ||
|
||
try { | ||
mappingContext.setInitialEntitySet(new EntityScanner(this.applicationContext).scan(Persistent.class)); | ||
} catch (Exception e) { | ||
Assert.fail(); | ||
} | ||
|
||
final MappingCosmosConverter cosmosConverter = new MappingCosmosConverter(mappingContext, null); | ||
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(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.manuallySetDatabaseName = testDB2; | ||
cosmosTemplate.createContainerIfNotExists(personInfo); | ||
cosmosTemplate.deleteAll(personInfo.getContainerName(), Person.class); | ||
assertThat(cosmosFactory.getDatabaseName()).isEqualTo(testDB2); | ||
cosmosTemplate.insert(TEST_PERSON_2, new PartitionKey(personInfo.getPartitionKeyFieldValue(TEST_PERSON_2))); | ||
|
||
// Check that DB2 has the correct contents | ||
List<Person> expectedResultsDB2 = new ArrayList<>(); | ||
expectedResultsDB2.add(TEST_PERSON_2); | ||
Iterable<Person> iterableDB2 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); | ||
List<Person> resultDB2 = new ArrayList<>(); | ||
iterableDB2.forEach(resultDB2::add); | ||
Assert.assertEquals(expectedResultsDB2, resultDB2); | ||
|
||
// Check that DB1 has the correct contents | ||
cosmosFactory.manuallySetDatabaseName = testDB1; | ||
List<Person> expectedResultsDB1 = new ArrayList<>(); | ||
expectedResultsDB1.add(TEST_PERSON_1); | ||
Iterable<Person> iterableDB1 = cosmosTemplate.findAll(personInfo.getContainerName(), Person.class); | ||
List<Person> 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); | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
96 changes: 96 additions & 0 deletions
96
...rc/test/java/com/azure/spring/data/cosmos/repository/MultiTenantTestRepositoryConfig.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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<String> 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(); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.