Enhancing CosmosTemplate to Support Multi-Tenancy at a DB Level

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

@@ -0,0 +1,35 @@
+// 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;
+
+/**
+ * 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/MultiTenantDBCosmosFactoryUnitTest.java

@@ -0,0 +1,108 @@
+// 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.Assert;
+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;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(classes = TestRepositoryConfig.class)
+public class MultiTenantDBCosmosFactoryUnitTest {
+
+    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;
+
+    @Test
+    public void testGetDatabaseFunctionality() {
+        /// Setup
+        CosmosAsyncClient client = CosmosFactory.createCosmosAsyncClient(cosmosClientBuilder);
+        MultiTenantDBCosmosFactory 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 cosmosTemplate = new CosmosTemplate(cosmosFactory, cosmosConfig, cosmosConverter, null);
+        CosmosEntityInformation<Person, String> 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);
+        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);
+    }
+}

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,40 @@ 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 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;
+    }
+
+    public void setManuallySetDatabaseName(String manuallySetDatabaseName) {
+        this.manuallySetDatabaseName = manuallySetDatabaseName;
+    }
+
+    public String getManuallySetDatabaseName() {
+        return manuallySetDatabaseName;
+    }
+}
+```
+
 ## 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,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();