AbstractRoutingDataSource doesn't change connection

Question

I use AbstractRoutingDataSource to change data source dynamically and ThreadLocal to set up currentLookupKey. It works nice when I use only one data source per http request. I use JpaRepository

@Component
@Primary
public class RoutingDataSource extends AbstractRoutingDataSource {

    @Autowired
    private DatabaseMap databaseMap;

    @Override
    public void afterPropertiesSet() {
        setTargetDataSources(databaseMap.getSourcesMap());
        setDefaultTargetDataSource(databaseMap.getSourcesMap().get("DEFAULT"));
        super.afterPropertiesSet();
    }

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

}

public class DatabaseContextHolder {

    private static final ThreadLocal<String> contextHolder = new ThreadLocal<String>();

    public static void setDatabaseType(String string) {
        contextHolder.set(string);
    }

    public static String getDatabaseType() {
        return (String) contextHolder.get();
    }

    public static void clearDatabaseType() {
        contextHolder.remove();
    }
}

When I try to get data in my REST controller I get data only from one database.

Some code in my REST controller

DatabaseContextHolder.setDatabaseType("db1");
//here I get data from db1 as expected
//I use JpaRepository
DatabaseContextHolder.clearDatabaseType();
DatabaseContextHolder.setDatabaseType("db2");
//here I should get data from db2 but get from db1

I tried to debug and it looks like Spring obtains data source only once in http request.

This method is called only once.

@Override
public Connection getConnection() throws SQLException {
    return determineTargetDataSource().getConnection();
}

Is there any way to force Spring to change data source.

Answer 1

I encountered the same problem and investigated it further by debugging my Hibernate implementation.

When Hibernate needs a connection, it calls the getConnection method and stores it in the Hibernate session. It then reuses that connection whenever needed.

To solve this, you can get the session before setting another databaseType in DatabaseContextHolder and close the connection using the disconnect method in your REST controller:

// You will need the EntityManager
@Autowired
private EntityManager entityManager;

[..]

DatabaseContextHolder.setDatabaseType("db1");
//here I get data from db1 as expected
//I use JpaRepository
DatabaseContextHolder.clearDatabaseType();
//Add this line
entityManager.unwrap(Session.class).disconnect();
DatabaseContextHolder.setDatabaseType("db2");
//here I should get data from db2 but get from db1

Answer 2

In my case, I followed the suggestion from jccampanero and it work: separating two methods in the service class, each method has a database connection with a different data source, switching datasource by AbstractRoutingDataSource.

I think the key point here is the database configuration where we give the AbstractRoutingDataSource to EntityManagerFactory and TransactionManager instead of DataSource as normal.

public class RoutingDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return ReadOnlyContext.isReadOnly() ? DSType.READ_ONLY : DSType.WRITE;
    }
}

and the DatabaseConfiguration:

@Bean
public RoutingDataSource actualDataSource(
    @Qualifier("dataSource") DataSource readWriteDataSource,
    @Qualifier("dataSourceReadOnly") DataSource readOnlyDataSource
) {
    Map<Object, Object> targetDataSources = new HashMap<>();
    targetDataSources.put(DSType.READ_ONLY, readOnlyDataSource);
    targetDataSources.put(DSType.WRITE, readWriteDataSource);

    RoutingDataSource routingDataSource = new RoutingDataSource();
    routingDataSource.setTargetDataSources(targetDataSources);
    routingDataSource.setDefaultTargetDataSource(readOnlyDataSource);

    return routingDataSource;
}

@Bean(name = "entityManagerFactory")
public LocalContainerEntityManagerFactoryBean entityManagerFactory(
    EntityManagerFactoryBuilder builder,
    RoutingDataSource actualDataSource
) {
    return builder.dataSource(actualDataSource).packages("the.domain.package").persistenceUnit("persistenUnitName").build();
}

@Bean(name = "transactionManager")
public PlatformTransactionManager transactionManager(RoutingDataSource actualDataSource) {
    return new DataSourceTransactionManager(actualDataSource);
}

With the above configuration, JPA Repositories will use the entityManagerFactory bean to obtain DB connections (in each service method, Spring will call the method determineCurrentLookupKey in RoutingDataSource to get the datasource we specified before).

Answer 3

Set spring.jpa.open-in-view to false.

Answer 4

I had the same issue, none of the above solution could fix it.. but making my Service method final (in my REST Controller)

public final Response

Answer 5

Your problem could be related with transaction delimitation.

When you define a @Transactional annotation in your code, Spring will create on your behalf all the stuff necessary to begin and end, and commiting or rollback if required, a transaction.

As you can see in the doBegin method in the source code of the DataSourceTransactionManager class - the same applies to other transaction managers - , Spring obtains a Connection when the transaction is initialized - this is why the method getConnection is invoked only once - , and it will reuse that connection across all the underlying operations against the database within that transaction (it makes sense for ACID preservation).

So, if you need to connect to several data sources within the same request processing, you can define different methods in you service code, every one annotated with a @Transactional annotation, and change the underlying data source as you require before invoke them:

DatabaseContextHolder.setDatabaseType("db1");
// Invoke a service method annotated with @Transactional
// It can use the underlying JpaRepositories that you need
DatabaseContextHolder.clearDatabaseType();
DatabaseContextHolder.setDatabaseType("db2");
// Invoke again another (or the same, what you need) service method
// annotated with @Transactional, You should get data from db2 this time

Answer 6

• My suspicion here is you have a method annotated with @Transactional annotation. Before calling that transactional method, you first specify one datasource key and you call the transactional method.Inside the transactional method, you first call repository and it works as expected with datasource look up key you set. However then you set different key, inside the transactional method, and call another repository and it still uses the key you set first time.

• DataSource will be chosen by the framework when the transaction starts so if you are using @Transactional annotation, whatever switching you do inside the method is useless. Because the datasource would have been chosen by proxy created for @Transactional annotation. Best option is to have the branching logic in non transactional service or use TransactionTemplate instead of @Transactional

• For example, make sure YourRestController does not have class level @Transactional as well as no @Transactional annotation in this yourRestControllerMethod, you will keep them to your service.

     @RestController
     public class YourRestController {

      @Autowired
      TransactionalService transactional

      public void yourRestControllerMethod(){
        //set the datasource look up key to A1 and then
        transactional.methodA1();
        //change datasource look up key to A2 and then
        transactional.methodA2();
      }

     }

    @Service
    public class TransactionalService {

       @Transactional
       public void methodA1(){

       }
       
       @Transactional
       public void methodA2() {

       }

    }