Enhancing CosmosTemplate to Support Multi-Tenancy at a DB Level (#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.

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

@@ -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;
+    }
+}

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

@@ -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);
+        }
+    }
+}

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);
     }
 }

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<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();
+    }
+}

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 `CosmosFactory` - See [PR 32516](https://github.com/Azure/azure-sdk-for-java/pull/32516)
 
 #### Breaking Changes
 

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

@@ -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.

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

@@ -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();