SpringBoot中实现多数据源切换并保证事务一致

项目中通常会遇到在不同数据库进行数据操作的场景，比如：在A库插入数据后往B、C库中再插入数据。针对此种情况可以使用AbstractRoutingDataSource抽象类对数据源进行统一管理。AbstractRoutingDataSource在这里不再进行详细介绍，它的原理也很简单，看源码完全可以看懂，大致原理图放在下面。这里主要介绍使用以及如何保证事务统一。

本文使用到的技术

AbstractRoutingDataSource、Transaction、SpringManagedTransactionFactory、AOP、动态代理

实现AbstractRoutingDataSource相关功能

以下主要实现多数据源的相关配置

配置文件定义多数据源

spring.datasource.druid.default.driver-class-name=com.mysql.jdbc.Driver
spring.datasource.druid.default.url=jdbc:mysql://10.10.10.10:3306/test?useUnicode=true&characterEncoding=utf-8&useAffectedRows=true&allowMultiQueries=true&useSSL=false
spring.datasource.druid.default.username=root
spring.datasource.druid.default.password=123456

dynamic.datasource.first.driver-class-name=com.mysql.jdbc.Driver
dynamic.datasource.first.url=jdbc:mysql://127.0.0.1:3306/mytest?autoReconnect=true&useSSL=false&failOverReadOnly=false
dynamic.datasource.first.username=root
dynamic.datasource.first.password=654321

定义数据源配置实体类

@Data
public class DataSourceProperty {
    private String driverClassName;
    private String url;
    private String username;
    private String password;

}

定义数据源配置接收类

/**
* 该类主要接收配置文件中配置的多个数据源
*/
@Data
@ConfigurationProperties(prefix = "dynamic")
public class DynamicDataSourceProperty {

    private Map<String, DataSourceProperty> datasource = new LinkedHashMap<>();
}

定义初始化数据源工厂类

public class DynamicDataSourceFactory {
    protected static Logger logger = LoggerFactory.getLogger(DynamicDataSourceFactory.class);

    public static DruidDataSource buildDruidDataSource(DataSourceProperty properties) {

        DruidDataSource druidDataSource = new DruidDataSource();

        druidDataSource.setDriverClassName(properties.getDriverClassName());
        druidDataSource.setUrl(properties.getUrl());
        druidDataSource.setUsername(properties.getUsername());
        druidDataSource.setPassword(properties.getPassword());

        try {
            druidDataSource.init();
        } catch (SQLException e) {
            logger.error("DynamicDataSourceFactory is error:" + e);
        }
        return druidDataSource;
    }
}

实现AbstractRoutingDataSource抽象类

public class DynamicDataSource extends AbstractRoutingDataSource {

    private final Logger logger = LoggerFactory.getLogger(DynamicDataSource.class);

    @Override
    protected Object determineCurrentLookupKey() {
        String dataSourceName = DynamicContextUtils.peek();
        if (dataSourceName != null) {
            logger.info("DynamicDataSource," + dataSourceName);
        }
        return dataSourceName;
    }
}

实现多数据源配置类

@Configuration
@EnableConfigurationProperties(DynamicDataSourceProperty.class)
public class DynamicDataSourceConfig {

    private Logger logger = LoggerFactory.getLogger(DynamicDataSourceConfig.class);

    @Autowired
    private DynamicDataSourceProperty properties;

    /**
     * 默认数据源
     *
     * @return
     */
    @Bean
    @ConfigurationProperties(prefix = "spring.datasource.druid.default")
    public DataSourceProperty dataSourceProperty() {
        return new DataSourceProperty();
    }

    /**
     * @param dataSourceProperty 数据库连接属性配置
     * @return 数据源
     */
    @Bean
    public DynamicDataSource dynamicDataSource(DataSourceProperty dataSourceProperty) {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        //设置动态数据源
        dynamicDataSource.setTargetDataSources(getDynamicDataSource());

        //默认数据源
        logger.info("defaultDataSource," + JSON.toJSONString(dataSourceProperty));
        DruidDataSource defaultDataSource = DynamicDataSourceFactory.buildDruidDataSource(dataSourceProperty);
        dynamicDataSource.setDefaultTargetDataSource(defaultDataSource);
        return dynamicDataSource;
    }

    private Map<Object, Object> getDynamicDataSource() {
        //多个数据源
        Map<String, DataSourceProperty> dataSourcePropertyMap = properties.getDatasource();

        //数据源的数量
        Map<Object, Object> targetDataSources = new ConcurrentHashMap<>(dataSourcePropertyMap.size());
        //初始化设置的多个数据源
        dataSourcePropertyMap.forEach((k, v) -> {
            DruidDataSource druidDataSource = DynamicDataSourceFactory.buildDruidDataSource(v);
            logger.info("dynamic," + k + ",property," + JSON.toJSONString(v));
            targetDataSources.put(k, druidDataSource);
        });

        return targetDataSources;
    }
}

注解实现并配置AOP

在进行多数据源操作时，为了方便可以自定义@DataSource注解，并在AOP切面类中实现对注解的拦截，通过ThreadLoacal工具安全操作拦截到的注解上的值，并通过该工具数据源名称的传递。

自定义注解

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
public @interface DataSource {
    String value() default "";
}

实现切面类拦截注解

@Aspect
@Component
@Order(Ordered.HIGHEST_PRECEDENCE)
public class DataSourceAspect {

    protected Logger logger = LoggerFactory.getLogger(DataSourceAspect.class);

    @Pointcut("@annotation(com.test.work.dynamicDataSource.DataSource)" +
            "|| @within(com.test.work.dynamicDataSource.DataSource)")
    public  void dataSourcePointCut(){

    }

    @Around("dataSourcePointCut()")
    public  Object around(@NotNull ProceedingJoinPoint point) throws Throwable{
        MethodSignature signature = (MethodSignature) point.getSignature();
        Class targetClass = point.getTarget().getClass();
        Method method = signature.getMethod();

        DataSource targetDataSource = (DataSource) targetClass.getAnnotation(DataSource.class);
        DataSource methodDataSource = method.getAnnotation(DataSource.class);
        if(Objects.nonNull(targetDataSource) || Objects.nonNull(methodDataSource)){

            String value = Objects.nonNull(methodDataSource) ? methodDataSource.value() : targetDataSource.value();
            DynamicContextUtils.push(value);
            logger.info("set datasource is {}", value);
        }

        try{
            return point.proceed();
        }finally {
            DynamicContextUtils.poll();
            logger.info("clean datasource");
        }
    }
}

实现ThreadLocal工具类

public class DynamicContextUtils {

    private static final ThreadLocal<Deque<String>> CONTEXT
            = ThreadLocal.withInitial(ArrayDeque::new);

    /**
     * 获得当前线程数据源
     *
     * @return 数据源名称
     */
    public static String peek() {
        return CONTEXT.get().peek();
    }

    /**
     * 设置当前线程数据源
     *
     * @param dataSource 数据源名称
     */
    public static void push(String dataSource) {
        CONTEXT.get().push(dataSource);
    }

    /**
     * 清空当前线程数据源
     */
    public static void poll() {
        Deque<String> deque = CONTEXT.get();
        deque.poll();
        if (deque.isEmpty()) {
            CONTEXT.remove();
        }
    }

}

总结1

以上功能即可实现在多个数据源之间进行切换执行，只要在方法或者类上使用@DataSource注解即可。但是几个明显的问题：

使用AbstractRoutingDataSource抽象类无法实现动态增加数据源（本文也没有解决）；
在同一个事务中不同数据源切换会报错，原因是AOP只拦截方法或者类上的注解，对于在A方法中使用B方法，且A、B方法分别操作不同的数据库，这种情况只会使用拦截并解析到的那个数据源。再深一点的原因是事务开始时，会把数据库连接等信息缓存下来，后续事务操作只会从缓存的信息中获取，在执行方法时，只缓存了A对应的数据源，就只能使用这个了。
无法保证全局事务保证问题。

针对以上2和3的问题，通过查阅资料，找到了两种解决方案，实现Transaction接口和动态代理。

实现Transaction接口方案（常用）

实现Transaction接口

/**
 * @author Ash
 * @date 2024/6/27 15:59
 * @description: 多数据源事务类
 */

public class MultiDataSourceTransaction implements Transaction {
    private static final Log LOGGER = LogFactory.getLog(MultiDataSourceTransaction.class);

    private final DataSource dataSource;

    private Connection mainConnection;

    private String mainDatabaseIdentification;

    private ConcurrentMap<String, Connection> otherConnectionMap;


    private boolean isConnectionTransactional;

    private boolean autoCommit;


    public MultiDataSourceTransaction(DataSource dataSource) {
        notNull(dataSource, "No DataSource specified");
        this.dataSource = dataSource;
        otherConnectionMap = new ConcurrentHashMap<>();
        mainDatabaseIdentification = DynamicContextUtils.peek();
    }


    /**
     * {@inheritDoc}
     */
    @Override
    public Connection getConnection() throws SQLException {
        String databaseIdentification = DynamicContextUtils.peek();
        if (databaseIdentification.equals(mainDatabaseIdentification)) {
            if (mainConnection != null) {
                return mainConnection;
            } else {
                openMainConnection();
                mainDatabaseIdentification = databaseIdentification;
                return mainConnection;
            }
        } else {
            if (!otherConnectionMap.containsKey(databaseIdentification)) {
                try {
                    Connection conn = dataSource.getConnection();
                    otherConnectionMap.put(databaseIdentification, conn);
                } catch (SQLException ex) {
                    throw new CannotGetJdbcConnectionException("Could not get JDBC Connection", ex);
                }
            }
            return otherConnectionMap.get(databaseIdentification);
        }

    }


    private void openMainConnection() throws SQLException {
        this.mainConnection = DataSourceUtils.getConnection(this.dataSource);
        this.autoCommit = this.mainConnection.getAutoCommit();
        this.isConnectionTransactional = DataSourceUtils.isConnectionTransactional(this.mainConnection, this.dataSource);

        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug(
                    "JDBC Connection ["
                            + this.mainConnection
                            + "] will"
                            + (this.isConnectionTransactional ? " " : " not ")
                            + "be managed by Spring");
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void commit() throws SQLException {
        if (this.mainConnection != null && !this.isConnectionTransactional && !this.autoCommit) {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("Committing JDBC Connection [" + this.mainConnection + "]");
            }
            this.mainConnection.commit();
            for (Connection connection : otherConnectionMap.values()) {
                connection.commit();
            }
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void rollback() throws SQLException {
        if (this.mainConnection != null && !this.isConnectionTransactional && !this.autoCommit) {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("Rolling back JDBC Connection [" + this.mainConnection + "]");
            }
            this.mainConnection.rollback();
            for (Connection connection : otherConnectionMap.values()) {
                connection.rollback();
            }
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void close() throws SQLException {
        DataSourceUtils.releaseConnection(this.mainConnection, this.dataSource);
        for (Connection connection : otherConnectionMap.values()) {
            DataSourceUtils.releaseConnection(connection, this.dataSource);
        }
    }

    @Override
    public Integer getTimeout() throws SQLException {
        return null;
    }
}

实现SpringManagedTransactionFactory接口

/**
 * @author Ash
 * @date 2024/6/27 16:12
 * @description: 多数据源工厂类
 */
public class MultiDataSourceTransactionFactory extends SpringManagedTransactionFactory {
    @Override
    public Transaction newTransaction(DataSource dataSource, TransactionIsolationLevel level, boolean autoCommit) {
        DynamicContextUtils.push("default");
        return new MultiDataSourceTransaction(dataSource);
    }
}

DynamicDataSourceConfig类中增加以下内容

@Bean("sqlSessionFactory")
    public SqlSessionFactory sqlSessionFactory() throws Exception {
        SqlSessionFactoryBean sqlSessionFactory = new SqlSessionFactoryBean();
        sqlSessionFactory.setDataSource(dynamicDataSource(dataSourceProperty()));
        // 使用自定义的多数据源事务工厂，如采用JdbcTemplate方式可不配置
        sqlSessionFactory.setTransactionFactory(new MultiDataSourceTransactionFactory());
        //添加XML目录
        ResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
        sqlSessionFactory.setMapperLocations(resolver.getResources("classpath*:mapper/*.xml"));
        return sqlSessionFactory.getObject();
    }

动态代理

这里复习一下动态代理的使用：

java动态代理机制中有两个重要的类和接口InvocationHandler（接口）和Proxy（类），这一个类Proxy和接口InvocationHandler是我们实现动态代理的核心。

InvocationHandler接口是Proxy代理实例的调用处理程序实现的一个接口，每一个Proxy代理实例都有一个关联的调用处理程序；在代理实例调用方法时，方法调用被编码分派到调用处理程序的invoke方法。

Proxy类就是用来创建一个代理对象的类，它提供了很多方法，但是我们最常用的是newProxyInstance方法。

/**
 * loader：一个classloader对象，定义了由哪个classloader对象对生成的代理类进行加载
 * interfaces：一个interface对象数组，表示我们将要给我们的代理对象提供一组什么样的接口，如果我们提供了这样一个接口对象数组，那么也就是声明了代理类实现了这些接口，代理类就可以调用接口中声明的所有方法。
 * h：一个InvocationHandler对象，表示的是当动态代理对象调用方法的时候会关联到哪一个InvocationHandler对象上，并最终由其调用。
 */
public static Object newProxyInstance(ClassLoader loader, 
                                            Class<?>[] interfaces, 
                                            InvocationHandler h)

代码举例：

// 定义一个接口
public interface People {

    public String work();
}


// 定义接口实现类
public class Teacher implements People{

    @Override
    public String work() {
        System.out.println("我是你老师！！！");
    }
}

创建代理类的调用处理方法：

/**
 * @author Ash
 * @date 2024/6/27 11:36
 * @description: 代理对象调用方法
 */
public class TeacherHandle implements InvocationHandler {
	// 真正调用的类
    private Object object;

    public TeacherHandle(Object object) {
        this.object = object;
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        System.out.println("你能🐎我吗？");
        Object invoke = method.invoke(object, args);
        System.out.println("你不讲武德！！！");
        return invoke;
    }
}

测试：

@Test
    public void testTeacher() {
        People teacher = new Teacher();

        TeacherHandle teacherHandle = new TeacherHandle(teacher);

        People proxy = (People) Proxy.newProxyInstance(teacherHandle.getClass().getClassLoader(), teacher.getClass().getInterfaces(), teacherHandle);
        proxy.work();
    }

为什么使用动态代理？

上面说了，AbstractRoutingDataSource 是用一个总数据源代理若干个子数据源，模仿其思想，使用一个总Connection，代理若干个子Connection。

从路由数据源中获取Connection对象，并不直接返回真正的Connection，而是返回一个Connection的动态代理对象，真正的Connection对象的获取，延迟到Connection的方法执行上，也就是在方法执行时再判断需要使用哪个数据源的链接。在动态代理的切面上，根据路由信息，获取真正的Connection对象将获取的真正的Connection对象保存起来，在后续的commit/rollback/close方法上，对所有的真实连接进行commit/rollback/close。

修改DynamicDataSource类

public class DynamicDataSource extends AbstractRoutingDataSource {

    private final Logger logger = LoggerFactory.getLogger(DynamicDataSource.class);

    @Override
    protected Object determineCurrentLookupKey() {
        String dataSourceName = DynamicContextUtils.peek();
        if (dataSourceName != null) {
            logger.info("DynamicDataSource," + dataSourceName);
        }
        return dataSourceName;
    }

    //使用动态代理返回一个代理对象
    @Override
    public Connection getConnection() throws SQLException {
        return (Connection) Proxy.newProxyInstance(this.getClass().getClassLoader(), new Class[]{Connection.class},
                new DynamicDataSourceHandler(this));
    }
	//获取方法执行时的真实连接
    public Connection getRealConnection() throws SQLException {
        return super.getConnection();
    }
}

实现代理对象的处理方法

/**
 * @author Ash
 * @date 2024/6/27 11:54
 * @description: 代理类操作方法
 */
public class DynamicDataSourceHandler 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 DynamicDataSource dataSource;

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

    public DynamicDataSourceHandler(DynamicDataSource dynamicDataSource) {
        this.dataSource = dynamicDataSource;
    }

    private boolean autoCommit = false;

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        //特殊方法处理
        {
            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;
            }
        }

        //根据路由获取真正的连接，去执行正式的方法
        Object lookupKey = dataSource.determineCurrentLookupKey();

        Connection realConnection = connectionMap.get(lookupKey);
        //如果已经存在，则无需再从数据源中获取
        if (realConnection == null) {
            realConnection = dataSource.getRealConnection();

            //设置是否自动提交
            if (!autoCommit) {
                realConnection.setAutoCommit(false);
            }
            connectionMap.put(lookupKey, realConnection);
        }

        //执行
        return method.invoke(realConnection, args);
    }

    /**
     * 关闭所有连接
     * @throws SQLException
     */
    private void doClose() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.close();
        }

    }

    /**
     * 设置自动提交
     * @param args
     * @throws SQLException
     */
    private void doSetAutoCommit(Object[] args) throws SQLException {
        this.autoCommit = (Boolean) args[0];
        for (Connection connection : connectionMap.values()) {
            connection.setAutoCommit(autoCommit);
        }
    }

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

    /**
     * 统一回滚
     * @throws SQLException
     */
    protected void doRollback() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.rollback();
        }
    }
}

总结2

该方案和实现Transaction接口的方案基本思想是一致的，都是用总代理代替子连接，然后将执行方法时的真实连接保存，最后进行提交或者回滚时一起进行，从而保证事务的一致性。但是使用动态代理的思想很巧妙，感觉很新颖，还是值得学来扎实基础的。

SpringBoot中实现多数据源切换并保证事务一致

SpringBoot中实现多数据源切换并保证事务一致

本文使用到的技术

实现AbstractRoutingDataSource相关功能

配置文件定义多数据源

定义数据源配置实体类

定义数据源配置接收类

定义初始化数据源工厂类

实现AbstractRoutingDataSource抽象类

实现多数据源配置类

注解实现并配置AOP

自定义注解

实现切面类拦截注解

实现ThreadLocal工具类

总结1

实现Transaction接口方案（常用）

实现Transaction接口

实现SpringManagedTransactionFactory接口

DynamicDataSourceConfig类中增加以下内容

动态代理

为什么使用动态代理？

修改DynamicDataSource类

实现代理对象的处理方法

总结2

参考文章