Skip to content

Commit

Permalink
Enhancing CosmosTemplate to Support Multi-Tenancy at a DB Level (Azur…
Browse files Browse the repository at this point in the history
…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
trande4884 authored and billwert committed Jan 4, 2023
1 parent 4780c59 commit 9452cfc
Showing 10 changed files with 405 additions and 64 deletions.
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;
}
}
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);
}
}
}
Original file line number Diff line number Diff line change
@@ -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);
}
}
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();
}
}
1 change: 1 addition & 0 deletions sdk/cosmos/azure-spring-data-cosmos/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -3,6 +3,7 @@
### 3.31.0-beta.1 (Unreleased)

#### Features Added
* 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

26 changes: 26 additions & 0 deletions sdk/cosmos/azure-spring-data-cosmos/README.md
Original file line number Diff line number Diff line change
@@ -934,6 +934,32 @@ 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-MultiTenantDBCosmosFactory
public class MultiTenantDBCosmosFactory extends CosmosFactory {

private String tenantId;

/**
* Validate config and initialization
*
* @param cosmosAsyncClient cosmosAsyncClient
* @param databaseName databaseName
*/
public MultiTenantDBCosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName) {
super(cosmosAsyncClient, databaseName);

this.tenantId = databaseName;
}

@Override
public String getDatabaseName() {
return this.getCosmosAsyncClient().getDatabase(this.tenantId).toString();
}
}
```

## 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.
Original file line number Diff line number Diff line change
@@ -20,7 +20,10 @@ public class CosmosFactory {

private final CosmosAsyncClient cosmosAsyncClient;

private final String databaseName;
/**
* Database Name to be used for operations.
*/
protected String databaseName;

private static final String USER_AGENT_SUFFIX =
Constants.USER_AGENT_SUFFIX + PropertyLoader.getProjectVersion();
Loading

0 comments on commit 9452cfc

Please sign in to comment.