Kim Rudolph

Spring DataSource Routing

A default setup of spring-data applications is a package with Entity definitions and one with Repository definitions. There is a solution for using those same entities against multiple databases without starting a second application.

Use cases are

The following example shows a solution for a single tool with access to development, testing and production data. It enables data display migration without the need to export and import data to another format. That might not be the best choice for most applications as it opens the door to accidently modifiying data in the wrong environment.

The example is based on Spring Boot. Selected dependencies at http://start.spring.io/ are JPA and Cache.

The source code with tests is available at https://github.com/krudolph/spring-datasource-routing.

Defining multiple Environments

As mentioned, there are three defined environments.

public enum DatabaseEnvironment {
  DEVELOPMENT, TESTING, PRODUCTION
}

The current DatabaseEnvironment is held in a ThreadLocal based context.

public class DatabaseContextHolder {

    private static final ThreadLocal<DatabaseEnvironment> CONTEXT =
        new ThreadLocal<>();

    public static void set(DatabaseEnvironment databaseEnvironment) {
        CONTEXT.set(databaseEnvironment);
    }

    public static DatabaseEnvironment getEnvironment() {
        return CONTEXT.get();
    }

    public static void clear() {
        CONTEXT.remove();
    }

}

Switching the context is achieved by setting the DatabaseEnvironment.

DatabaseContextHolder.set(DatabaseEnvironment.TESTING);

Configuring the DataSource Routing

Spring Boot provides a lot of AutoConfiguration magic, including DataSource setup. That needs to be excluded.

@SpringBootApplication(exclude = { 
    DataSourceAutoConfiguration.class,
    DataSourceTransactionManagerAutoConfiguration.class,
    HibernateJpaAutoConfiguration.class })
public class RoutingApplication {

    public static void main(String[] args) {
        SpringApplication.run(RoutingApplication.class, args);
    }
}

Using AbstractRoutingDataSource enables the mapping between the actual DataSources and the DatabaseEnvironment.

public class DataSourceRouter extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return DatabaseContextHolder.getEnvironment();
    }
}

Disabling the AutoConfiguration enables a custom configuration of the DataSource initialization.

@Configuration
@EnableJpaRepositories(basePackageClasses = CustomerRepository.class, 
                       entityManagerFactoryRef = "customerEntityManager", 
                       transactionManagerRef = "customerTransactionManager")
@EnableTransactionManagement
public class DatasourceConfiguration {

...

DataSources for each environment are defined by prefixed properties. Mandatory properties are jdbcUrl, username and password. See the application-test.properties file for a full example.

 
@Bean
@ConfigurationProperties(prefix = "app.customer.development.datasource")
public DataSource customerDevelopmentDataSource() {
    return DataSourceBuilder.create().build();
}

@Bean
@ConfigurationProperties(prefix = "app.customer.testing.datasource")
public DataSource customerTestingDataSource() {
    return DataSourceBuilder.create().build();
}

@Bean
@ConfigurationProperties(prefix = "app.customer.production.datasource")
public DataSource customerProductionDataSource() {
    return DataSourceBuilder.create().build();
}

A Map filled with those DataSources is used by the DataSourceRouter to select the DataSource that belongs to the current environment.

@Bean
@Primary
public DataSource customerDataSource() {
    DataSourceRouter router = new DataSourceRouter();

    final HashMap<Object, Object> map = new HashMap<>(3);
    map.put(DatabaseEnvironment.DEVELOPMENT, customerDevelopmentDataSource());
    map.put(DatabaseEnvironment.TESTING, customerTestingDataSource());
    map.put(DatabaseEnvironment.PRODUCTION, customerProductionDataSource());
    router.setTargetDataSources(map);
    return router;
}

The last part of the configuration is the default JPA/Hibernate boilerplate code. The JPA prefix enables shared configurations like app.customer.jpa.properties.hibernate.hbm2ddl.auto=create.

@Autowired(required = false)
private PersistenceUnitManager persistenceUnitManager;

@Bean
@Primary
@ConfigurationProperties("app.customer.jpa")
public JpaProperties customerJpaProperties() {
    return new JpaProperties();
}

@Bean
@Primary
public LocalContainerEntityManagerFactoryBean customerEntityManager(
    final JpaProperties customerJpaProperties) {

    EntityManagerFactoryBuilder builder =
        createEntityManagerFactoryBuilder(customerJpaProperties);

    return builder.dataSource(customerDataSource()).packages(Customer.class)
        .persistenceUnit("customerEntityManager").build();
}

@Bean
@Primary
public JpaTransactionManager customerTransactionManager(
    @Qualifier("customerEntityManager") final EntityManagerFactory factory) {
    return new JpaTransactionManager(factory);
}

private EntityManagerFactoryBuilder createEntityManagerFactoryBuilder(
    JpaProperties customerJpaProperties) {
    JpaVendorAdapter jpaVendorAdapter =
        createJpaVendorAdapter(customerJpaProperties);
    return new EntityManagerFactoryBuilder(jpaVendorAdapter,
        customerJpaProperties.getProperties(), this.persistenceUnitManager);
}

private JpaVendorAdapter createJpaVendorAdapter(
    JpaProperties jpaProperties) {
    AbstractJpaVendorAdapter adapter = new HibernateJpaVendorAdapter();
    adapter.setShowSql(jpaProperties.isShowSql());
    adapter.setDatabase(jpaProperties.getDatabase());
    adapter.setDatabasePlatform(jpaProperties.getDatabasePlatform());
    adapter.setGenerateDdl(jpaProperties.isGenerateDdl());
    return adapter;
}

What about Caching?

Manually configuring cache key attributes (@Cacheable( key = "...")) will cause bugs, because they are not environment aware. The solution is using only @Cacheable without key attributes and adding a global KeyGenerator.

@Configuration
@EnableCaching
public class CachingConfiguration extends CachingConfigurerSupport {

    @Override
    public KeyGenerator keyGenerator() {
        return new EnvironmentAwareCacheKeyGenerator();
    }

}

That custom KeyGenerator sets the current database environment as a prefix.

public class EnvironmentAwareCacheKeyGenerator implements KeyGenerator {

    @Override
    public Object generate(Object target, Method method, Object... params) {

        String key = DatabaseContextHolder.getEnvironment().name() + "-" + (
            method == null ? "" : method.getName() + "-") + StringUtils
            .collectionToDelimitedString(Arrays.asList(params), "-");

        return key;
    }

}

Adding more DataSources with multiple Environments

It is possible to add another DataSourceConfiguration if there is the need to connect to other databases. The prefix would be something like app.contracts.. Entity and Repository classes have to be in a separate package. An additional EntityManager will be available and needs to be sprecified with a unit name, e.g. @PersistenceContext(unitName = "contractEntityManager") if injected. Be aware that the already existing DataSource needs to be definied as @Primary.