When I use AbstractRoutingDataSource, some scenarios need to switch the data source in the transaction. Of course, there are ways to avoid this problem from the code level, but it is not so worry-free to use.

Later, I was reminded by the automatic injection scheme of HttpServletRequest. I used the following method, which not only perfectly solved the problem of switching data sources in the transaction, but also solved the transaction consistency of multiple data sources to some extent.

My idea is to delay obtaining the real Connection object. When calling AbstractRoutingDataSource.getConnection(), I get a proxy object. Only when the specific method is actually executed, the real data source is obtained.

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.StringUtils;

import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;


public class RoutingDataSource extends AbstractRoutingDataSource {

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

    private static String defaultKey;

    public RoutingDataSource(String defaultDataSource, Map<Object, Object> targetDataSource) {
        defaultKey = defaultDataSource;
        super.setTargetDataSources(targetDataSource);
    }


    /**
     * Rewrite the method to get the dynamic proxy object
     * @return
     * @throws SQLException
     */
    @Override
    public Connection getConnection() throws SQLException {
        return (Connection) Proxy.newProxyInstance(this.getClass().getClassLoader(),
                new Class[]{Connection.class},
                new RoutingProxyConnectionInvocationHandler(this));
    }

    /**
     * Get the real Connection object
     * @return
     * @throws SQLException
     */
    public Connection getRealConnection() throws SQLException {
        return super.getConnection();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return getDataSourceKey();
    }

    public static String getDataSourceKey() {
        String key = CONTEXT_HOLDER.get();
        if (StringUtils.hasLength(key)) {
            key = defaultKey;
        }
        return key;
    }

    public static void setDataSourceKey(String dataSource) {
        CONTEXT_HOLDER.set(dataSource);
    }

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

}

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;


public class RoutingProxyConnectionInvocationHandler implements InvocationHandler {

    private static final String METHOD_CLOSE = "close";
    private static final String METHOD_SET_AUTOCOMMIT = "setAutoCommit";
    private static final String METHOD_COMMIT = "commit";
    private static final String METHOD_ROLLBACK = "rollback";

    private RoutingDataSource dataSource;

    private Map<Object, Connection> connectionMap = new HashMap<>();

    public RoutingProxyConnectionInvocationHandler(RoutingDataSource dataSource) {
        this.dataSource = dataSource;
    }

    /**
     * Whether to submit automatically
     */
    private boolean autoCommit = true;

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        //Transaction-related method processing
        {
            if (METHOD_CLOSE.equals(method.getName())) {
                doClose();
                return null;
            }

            if (METHOD_SET_AUTOCOMMIT.equals(method.getName())) {
                doSetAutoCommit(args);
                return null;
            }

            if (METHOD_COMMIT.equals(method.getName())) {
                doCommit();
                return null;
            }

            if (METHOD_ROLLBACK.equals(method.getName())) {
                doRollback();
                return null;
            }
        }

        //Obtain the real connection according to the route and execute the formal method
        Object lookupKey = dataSource.determineCurrentLookupKey();

        Connection realConnection = connectionMap.get(lookupKey);
        // If it already exists, you don't need to get it from the data source
        if (realConnection == null) {
            realConnection = dataSource.getRealConnection();

            // Set whether to submit automatically
            if (!autoCommit) {
                realConnection.setAutoCommit(false);
            }
            connectionMap.put(lookupKey, realConnection);
        }


        return method.invoke(realConnection, args);
    }

    /**
     * close all connections
     * @throws SQLException
     */
    private void doClose() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.close();
        }
    }

    /**
     * Set up autocommit
     * @param args
     * @throws SQLException
     */
    private void doSetAutoCommit(Object[] args) throws SQLException {
        this.autoCommit = (Boolean) args[0];
        for (Connection connection : connectionMap.values()) {
            connection.setAutoCommit(autoCommit);
        }
    }

    /**
     * submit
     * @throws SQLException
     */
    private void doCommit() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.commit();
        }
    }

    /**
     *  rollBack
     * @throws SQLException
     */
    private void doRollback() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.rollback();
        }
    }
}

Comment From: jackycjw

If the official thinks the plan is feasible, I hope you can provide support

Comment From: sbrannen

My idea is to delay obtaining the real Connection object. When calling AbstractRoutingDataSource.getConnection(), I get a proxy object. Only when the specific method is actually executed, the real data source is obtained.

Are you aware of org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy?

Comment From: jackycjw

Thank you for your reply

My business scenario is more complicated, and I need to perform insert operations on two databases in one transaction.

I tried org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy this scheme. It is still not possible to switch data sources within a transaction. Maybe it's the wrong way I'm using it. Below is my configuration class code.

@Configuration
public class DatasourceConfiguration {

    @Bean
    @ConfigurationProperties(prefix = "spring.datasource.db1")
    public DataSource db1(){
        return new HikariDataSource();
    }

    @Bean
    @ConfigurationProperties(prefix = "spring.datasource.db2")
    public DataSource db2(){
        return new HikariDataSource();
    }

    @Bean
    @Primary
    public DataSource dataSource(){
        Map<Object, Object> targetDataSource = Map.of("db1",db1(), "db2", db2());
        RoutingDatasource routingDatasource = new RoutingDatasource("db1", targetDataSource);
        return new LazyConnectionDataSourceProxy(routingDatasource);
    }

    public static class RoutingDatasource extends AbstractRoutingDataSource{

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

        private static String defaultKey;

        public RoutingDatasource(String defaultDataSource, Map<Object, Object> targetDataSource) {
            defaultKey = defaultDataSource;
            super.setTargetDataSources(targetDataSource);
        }

        @Override
        protected Object determineCurrentLookupKey() {
            return getDataSourceKey();
        }

        public static String getDataSourceKey() {
            String key = CONTEXT_HOLDER.get();
            if (StringUtils.isEmpty(key)) {
                key = defaultKey;
            }
            return key;
        }

        public static void setDataSourceKey(String dataSource) {
            CONTEXT_HOLDER.set(dataSource);
        }

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

}

Comment From: sbrannen

LazyConnectionDataSourceProxy already supports lazy Connection retrieval, and we do not have any plans to introduce a concrete RoutingDataSource implementation with lazy Connection retrieval and thread-local DataSource tracking.

I need to perform insert operations on two databases in one transaction.

If you need to support that, you should investigate support for XA transactions with JTA or simulate XA transactions.

In that case, you may find the Distributed transactions in Spring, with and without XA article helpful.

In light of the above, I am closing this issue.

Comment From: jackycjw

Thanks for your help.