Multi-tenancy: Managing multiple datasources with Spring Data JPA

To implement multi-tenancy with Spring Boot we can use AbstractRoutingDataSource as base DataSource class for all ‘tenant databases‘.

It has one abstract method determineCurrentLookupKey that we have to override. It tells the AbstractRoutingDataSource which of the tenant datasource it has to provide at the moment to work with. Because it works in the multi-threading environment, the information of the chosen tenant should be stored in ThreadLocal variable.

The AbstractRoutingDataSource stores the info of the tenant datasources in its private Map<Object, Object> targetDataSources. The key of this map is a tenant identifier (for example the String type) and the value – the tenant datasource. To put our tenant datasources to this map we have to use its setter setTargetDataSources.

The AbstractRoutingDataSource will not work without ‘default’ datasource which we have to set with method setDefaultTargetDataSource(Object defaultTargetDataSource).

After we set the tenant datasources and the default one, we have to invoke method afterPropertiesSet() to tell the AbstractRoutingDataSource to update its state.

So our ‘MultiTenantManager’ class can be like this:

@Configuration
public class MultiTenantManager {

    private final ThreadLocal<String> currentTenant = new ThreadLocal<>();
    private final Map<Object, Object> tenantDataSources = new ConcurrentHashMap<>();
    private final DataSourceProperties properties;

    private AbstractRoutingDataSource multiTenantDataSource;

    public MultiTenantManager(DataSourceProperties properties) {
        this.properties = properties;
    }

    @Bean
    public DataSource dataSource() {
        multiTenantDataSource = new AbstractRoutingDataSource() {
            @Override
            protected Object determineCurrentLookupKey() {
                return currentTenant.get();
            }
        };
        multiTenantDataSource.setTargetDataSources(tenantDataSources);
        multiTenantDataSource.setDefaultTargetDataSource(defaultDataSource());
        multiTenantDataSource.afterPropertiesSet();
        return multiTenantDataSource;
    }

    public void addTenant(String tenantId, String url, String username, String password) throws SQLException {

        DataSource dataSource = DataSourceBuilder.create()
                .driverClassName(properties.getDriverClassName())
                .url(url)
                .username(username)
                .password(password)
                .build();

        // Check that new connection is 'live'. If not - throw exception
        try(Connection c = dataSource.getConnection()) {
            tenantDataSources.put(tenantId, dataSource);
            multiTenantDataSource.afterPropertiesSet();
        }
    }

    public void setCurrentTenant(String tenantId) {
        currentTenant.set(tenantId);
    }

    private DriverManagerDataSource defaultDataSource() {
        DriverManagerDataSource defaultDataSource = new DriverManagerDataSource();
        defaultDataSource.setDriverClassName("org.h2.Driver");
        defaultDataSource.setUrl("jdbc:h2:mem:default");
        defaultDataSource.setUsername("default");
        defaultDataSource.setPassword("default");
        return defaultDataSource;
    }
}

Brief explanation:

• map tenantDataSources it’s our local tenant datasource storage which we put to the setTargetDataSources setter;

• DataSourceProperties properties is used to get Database Driver Class name of tenant database from the spring.datasource.driverClassName of the ‘application.properties’ (for example, org.postgresql.Driver);

• method addTenant is used to add a new tenant and its datasource to our local tenant datasource storage. We can do this on the fly – thanks to the method afterPropertiesSet();

• method setCurrentTenant(String tenantId) is used to ‘switch’ onto datasource of the given tenant. We can use this method, for example, in the REST controller when handling a request to work with database. The request should contain the ‘tenantId’, for example in the X-TenantId header, that we can retrieve and put to this method;

• defaultDataSource() is build with in-memory H2 Database to avoid using the default database on the working SQL server.

Note: you must set spring.jpa.hibernate.ddl-auto parameter to none to disable the Hibernate make changes in the database schema. You have to create a schema of tenant databases beforehand.

A full example of this class and more you can find in my repo.

UPDATED

This branch demonstrates an example of using the dedicated database to store tenant DB properties instead of property files (see the question of @MarcoGustavo below).