Spring JPA: How to update 2 different tables in 2 different `DataSource` in the same request?

Question

In our application, we have a common database called central and every customer will have their own database with exactly the same set of tables. Each customer's database might be hosted on our own server or the customer's server based on the requirement of the customer organization.

To handle this multi-tenant requirement, we're extending the AbstractRoutingDataSource from Spring JPA and overriding the determineTargetDataSource() method to create a new DataSource and establish a new connection on the fly based on the incoming customerCode. We also use a simple DatabaseContextHolder class to store the current datasource context in a ThreadLocal variable. Our solution is similar to what is describe in this article.

Let's say in a single request, we'll need to update some data in both the central database and the customer's database as following.

public void createNewEmployeeAccount(EmployeeData employee) {
    DatabaseContextHolder.setDatabaseContext("central");
    // Code to save a user account for logging in to the system in the central database

    DatabaseContextHolder.setDatabaseContext(employee.getCustomerCode());
    // Code to save user details like Name, Designation, etc. in the customer's database
}

This code would only work if determineTargetDataSource() is called every time just before any SQL queries gets executed so that we can switch the DataSource dynamically half way through our method.

However, from this Stackoverflow question, it seems like determineTargetDataSource() is only called once for each HttpRequest when a DataSource is being retrieved for the very first time in that request.

I'd be very grateful if you can give me some insights into when AbstractRoutingDataSource.determineTargetDataSource() actually gets called. Besides, if you've dealt with a similar multi-tenant scenario before, I'd love to hear your opinion on how I should deal with the updating of multiple DataSource in a single request.

See if https://stackoverflow.com/questions/27614301/spring-boot-multiple-datasource can help. Main idea - make use of @Primary annotation — Ian Lim, Dec 13 '18 at 08:02

score 0 · Accepted Answer · answered Dec 13 '18 at 10:32

We found a working solution, which is a mix of static data source settings for our central database and dynamic data source settings for our customer's database.

In essence, we know exactly which table comes from which database. Hence, we were able to separate our @Entity classes into 2 different packages as following.

com.ft.model
   -- central
      -- UserAccount.java
      -- UserAccountRepo.java
   -- customer
      -- UserProfile.java
      -- UserProfileRepo.java

Subsequently, we created two @Configuration classes to set up the data source setting for each package. For our central database, we use static settings as following.

@Configuration
@EnableTransactionManagement
@EnableJpaRepositories(
        entityManagerFactoryRef = "entityManagerFactory",
        transactionManagerRef = "transactionManager",
        basePackages = { "com.ft.model.central" }
)
public class CentralDatabaseConfiguration {
    @Primary
    @Bean(name = "dataSource")
    public DataSource dataSource() {
        return DataSourceBuilder.create(this.getClass().getClassLoader())
                .driverClassName("com.microsoft.sqlserver.jdbc.SQLServerDriver")
                .url("jdbc:sqlserver://localhost;databaseName=central")
                .username("sa")
                .password("mhsatuck")
                .build();
    }

    @Primary
    @Bean(name = "entityManagerFactory")
    public LocalContainerEntityManagerFactoryBean entityManagerFactory(EntityManagerFactoryBuilder builder, @Qualifier("dataSource") DataSource dataSource) {
        return builder
                .dataSource(dataSource)
                .packages("com.ft.model.central")
                .persistenceUnit("central")
                .build();
    }

    @Primary
    @Bean(name = "transactionManager")
    public PlatformTransactionManager transactionManager (@Qualifier("entityManagerFactory") EntityManagerFactory entityManagerFactory) {
        return new JpaTransactionManager(entityManagerFactory);
    }
}

For the @Entity in the customer package, we set up dynamic data source resolver using the following @Configuration.

@Configuration
@EnableTransactionManagement
@EnableJpaRepositories(
        entityManagerFactoryRef = "customerEntityManagerFactory",
        transactionManagerRef = "customerTransactionManager",
        basePackages = { "com.ft.model.customer" }
)
public class CustomerDatabaseConfiguration {
    @Bean(name = "customerDataSource")
    public DataSource dataSource() {
        return new MultitenantDataSourceResolver();
    }

    @Bean(name = "customerEntityManagerFactory")
    public LocalContainerEntityManagerFactoryBean entityManagerFactory(EntityManagerFactoryBuilder builder, @Qualifier("customerDataSource") DataSource dataSource) {
        return builder
                .dataSource(dataSource)
                .packages("com.ft.model.customer")
                .persistenceUnit("customer")
                .build();
    }

    @Bean(name = "customerTransactionManager")
    public PlatformTransactionManager transactionManager(@Qualifier("customerEntityManagerFactory") EntityManagerFactory entityManagerFactory) {
        return new JpaTransactionManager(entityManagerFactory);
    }
}

In the MultitenantDataSourceResolver class, we plan to maintain a Map of the created DataSource using customerCode as key. From each incoming request, we will get the customerCode and inject it into our MultitenantDataSourceResolver to get the correct DataSource within the determineTargetDataSource() method.

public class MultitenantDataSourceResolver extends AbstractRoutingDataSource {
    @Autowired
    private Provider<CustomerWrapper> customerWrapper;

    private static final Map<String, DataSource> dsCache = new HashMap<String, DataSource>();

    @Override
    protected Object determineCurrentLookupKey() {
        try {
            return customerWrapper.get().getCustomerCode();

        } catch (Exception ex) {
            return null;

        }
    }

    @Override
    protected DataSource determineTargetDataSource() {
        String customerCode = (String) this.determineCurrentLookupKey();

        if (customerCode == null)
            return MultitenantDataSourceResolver.getDefaultDataSource();
        else {
            DataSource dataSource = dsCache.get(customerCode);
            if (dataSource == null)
                dataSource = this.buildDataSourceForCustomer();

            return dataSource;
        }
    }

    private synchronized DataSource buildDataSourceForCustomer() {
        CustomerWrapper wrapper = customerWrapper.get();

        if (dsCache.containsKey(wrapper.getCustomerCode()))
            return dsCache.get(wrapper.getCustomerCode() );
        else {
            DataSource dataSource = DataSourceBuilder.create(MultitenantDataSourceResolver.class.getClassLoader())
                    .driverClassName("com.microsoft.sqlserver.jdbc.SQLServerDriver")
                    .url(wrapper.getJdbcUrl())
                    .username(wrapper.getDbUsername())
                    .password(wrapper.getDbPassword())
                    .build();

            dsCache.put(wrapper.getCustomerCode(), dataSource);

            return dataSource;
        }
    }

    private static DataSource getDefaultDataSource() {
        return DataSourceBuilder.create(CustomerDatabaseConfiguration.class.getClassLoader())
                .driverClassName("com.microsoft.sqlserver.jdbc.SQLServerDriver")
                .url("jdbc:sqlserver://localhost;databaseName=central")
                .username("sa")
                .password("mhsatuck")
                .build();
    }
}

The CustomerWrapper is a @RequestScope object whose values will be populated on each request by the @Controller. We use java.inject.Provider to inject it into our MultitenantDataSourceResolver.

Lastly, even though, logically, we will never save anything using the default DataSource because all requests will always contain a customerCode, at startup time, there is no customerCode available. Hence, we still need to provide a valid default DataSource. Otherwise, the application will not be able to start.

If you have any comments or a better solution, please let me know.

Spring JPA: How to update 2 different tables in 2 different `DataSource` in the same request?

1 Answers1