Search code examples
javaspringmongodbmulti-tenant

Java Spring Data MongoDB Change Database During Runtime(Multi-tenancy)


I want to have the same structure of collections in multiple databases as follows

Database A:
 CollectionA
 CollectionB

Database B:
 CollectionA
 CollectionB

Lets say i create my documents with the respective repositories Eg.

@Document(collection = collectionA)
@Document(collection = collectionB)

How can i make a configuration to force spring to swap databases for the repository beans between opperations ? I have searched and i have found that i need to create a mongoTemplate bean it dosent seem to work , This method is called only once when the repositories are autowired but not where an insert is made.

Eg.

# Inside a service
@Autowired
DocumentARepository documentARepository;

documentARepository.insert(documentA); # Inserts in Database A Collection A
TenantContext.setTenant("Database B"); # Changes MongoDb Repository Tenant (Database) thread safe class i use for jdbc multitenancy
docucmentARepository.insert(documentB); #Insert in Database B Collection A

Have looked at Multi-tenant mongodb database based with spring data and its an old version of spring data mongodb im using 3.1.3

If anyone can suggest how i can debug what beans are used during the insertion of a document it would be welcome to figure it out on my own.Have no idea how i would go with this to manually figure it out. Thanks in advance


Solution

  • Will solve this by ditching MongoRepositorie interfaces all together and use a MongoTemplate bean instead. Configured like so :

    Mongo Configuration File

    @Configuration
    public class MongoConfiguration {
    
        @Bean
        public MongoClient mongoClient() {
            return MongoClients.create();
        }
    
        @Bean
        public MultiTenantMongoDbFactory multiTenantMongoDatabaseFactory(MongoClient mongoClient) {
            return new MultiTenantMongoDbFactory(mongoClient);
        }
    
        @Bean
        public MongoTemplate mongoTemplate(MultiTenantMongoDbFactory multiTenantMongoDatabaseFactory) {
            return new MongoTemplate(multiTenantMongoDatabaseFactory);
        }
    }
    

    MultiTenantMongoDbFactory File (This is used a the factory that create the mongo templates seems to be used to determine the database for each save although i have not tested for backround batching)

    public class MultiTenantMongoDbFactory implements MongoDatabaseFactory {
    
        private final MongoClient mongoClient;
        private final PersistenceExceptionTranslator exceptionTranslator;
    
        public MultiTenantMongoDbFactory(MongoClient mongoClient) {
            this.mongoClient = mongoClient;
            this.exceptionTranslator = new MongoExceptionTranslator();
        }
    
        public MultiTenantMongoDbFactory(MongoClient mongoClient, ClientSession session) {
            this.mongoClient = mongoClient;
            this.exceptionTranslator = new MongoExceptionTranslator();
        }
    
        @Override
        public MongoDatabase getMongoDatabase() throws DataAccessException {
            return mongoClient.getDatabase(TenantContext.getCurrentTenant());
        }
    
        @Override
        public MongoDatabase getMongoDatabase(String dbName) throws DataAccessException {
            return mongoClient.getDatabase(dbName);
        }
    
        @Override
        public PersistenceExceptionTranslator getExceptionTranslator() {
            return exceptionTranslator;
        }
    
        @Override
        public ClientSession getSession(ClientSessionOptions options) {
            return mongoClient.startSession(options);
        }
    
        @Override
        public MongoDatabaseFactory withSession(ClientSession session) {
            // Create a new MultiTenantMongoDbFactory instance with the same MongoClient and the provided ClientSession.
            return new MultiTenantMongoDbFactory.ClientSessionBoundMongoDbFactory(session, this);
        }
    
        static final private class ClientSessionBoundMongoDbFactory implements MongoDatabaseFactory {
    
            private final ClientSession session;
            private final MongoDatabaseFactory delegate;
    
            public ClientSessionBoundMongoDbFactory(ClientSession session, MongoDatabaseFactory delegate) {
                this.session = session;
                this.delegate = delegate;
            }
    
            /*
             * (non-Javadoc)
             * @see org.springframework.data.mongodb.MongoDbFactory#getMongoDatabase()
             */
            @Override
            public MongoDatabase getMongoDatabase() throws DataAccessException {
                return proxyMongoDatabase(delegate.getMongoDatabase());
            }
    
            /*
             * (non-Javadoc)
             * @see org.springframework.data.mongodb.MongoDbFactory#getMongoDatabase(java.lang.String)
             */
            @Override
            public MongoDatabase getMongoDatabase(String dbName) throws DataAccessException {
                return proxyMongoDatabase(delegate.getMongoDatabase(dbName));
            }
    
            /*
             * (non-Javadoc)
             * @see org.springframework.data.mongodb.MongoDbFactory#getExceptionTranslator()
             */
            @Override
            public PersistenceExceptionTranslator getExceptionTranslator() {
                return delegate.getExceptionTranslator();
            }
    
            /*
             * (non-Javadoc)
             * @see org.springframework.data.mongodb.MongoDbFactory#getSession(com.mongodb.ClientSessionOptions)
             */
            @Override
            public ClientSession getSession(ClientSessionOptions options) {
                return delegate.getSession(options);
            }
    
            /*
             * (non-Javadoc)
             * @see org.springframework.data.mongodb.MongoDbFactory#withSession(com.mongodb.session.ClientSession)
             */
            @Override
            public MongoDatabaseFactory withSession(ClientSession session) {
                return delegate.withSession(session);
            }
    
            /*
             * (non-Javadoc)
             * @see org.springframework.data.mongodb.MongoDbFactory#isTransactionActive()
             */
            @Override
            public boolean isTransactionActive() {
                return session != null && session.hasActiveTransaction();
            }
    
            private MongoDatabase proxyMongoDatabase(MongoDatabase database) {
                return createProxyInstance(session, database, MongoDatabase.class);
            }
    
            private MongoDatabase proxyDatabase(com.mongodb.session.ClientSession session, MongoDatabase database) {
                return createProxyInstance(session, database, MongoDatabase.class);
            }
    
            private MongoCollection<?> proxyCollection(com.mongodb.session.ClientSession session,
                                                       MongoCollection<?> collection) {
                return createProxyInstance(session, collection, MongoCollection.class);
            }
    
            private <T> T createProxyInstance(com.mongodb.session.ClientSession session, T target, Class<T> targetType) {
    
                ProxyFactory factory = new ProxyFactory();
                factory.setTarget(target);
                factory.setInterfaces(targetType);
                factory.setOpaque(true);
    
                factory.addAdvice(new SessionAwareMethodInterceptor<>(session, target, ClientSession.class, MongoDatabase.class,
                        this::proxyDatabase, MongoCollection.class, this::proxyCollection));
    
                return targetType.cast(factory.getProxy(target.getClass().getClassLoader()));
            }
    
            public ClientSession getSession() {
                return this.session;
            }
    
            public MongoDatabaseFactory getDelegate() {
                return this.delegate;
            }
    
            @Override
            public boolean equals(@Nullable Object o) {
                if (this == o)
                    return true;
                if (o == null || getClass() != o.getClass())
                    return false;
    
                MultiTenantMongoDbFactory.ClientSessionBoundMongoDbFactory that = (MultiTenantMongoDbFactory.ClientSessionBoundMongoDbFactory) o;
    
                if (!ObjectUtils.nullSafeEquals(this.session, that.session)) {
                    return false;
                }
                return ObjectUtils.nullSafeEquals(this.delegate, that.delegate);
            }
    
            @Override
            public int hashCode() {
                int result = ObjectUtils.nullSafeHashCode(this.session);
                result = 31 * result + ObjectUtils.nullSafeHashCode(this.delegate);
                return result;
            }
    
            public String toString() {
                return "MongoDatabaseFactorySupport.ClientSessionBoundMongoDbFactory(session=" + this.getSession() + ", delegate="
                        + this.getDelegate() + ")";
            }
        }
    }
    

    TenantContext file what i use to set the tenant context for both jdbc mariadb and mongodb cause i have both

    public class TenantContext {
        private static final ThreadLocal<String> currentTenant = new ThreadLocal<>();
    
        public static String getCurrentTenant() {
            if (currentTenant.get() == null){
                currentTenant.set("DEV_TEST");
            }
            return currentTenant.get();
        }
    
        public static void setCurrentTenant(String tenant) {
            currentTenant.set(tenant);
        }
    
        public static void clear() {
            currentTenant.remove();
        }
    }
    

    Added the full code here if someone wants to copy

    https://github.com/mplein2/spring_data_mongodb_mutlitenancy

    Edit : After making the changes the repositories that are generated are multi tenant also . No need to use the mongo template directly