feat:增加适配atomikos数据源,支持同一个事务下多次切换数据源 (#481)

* feat:增加适配atomikos数据源,支持同一个事务下多次切换数据源

---------

Co-authored-by: jiazhifeng <wb.jiazhifeng01@mesg.corp.netease.com>

core/pom.xml

@@ -49,6 +49,11 @@
             <artifactId>druid-spring-boot-starter</artifactId>
             <optional>true</optional>
         </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-jta-atomikos</artifactId>
+            <optional>true</optional>
+        </dependency>
         <dependency>
             <groupId>com.github.chris2018998</groupId>
             <artifactId>beecp</artifactId>

core/src/main/java/com/baomidou/dynamic/datasource/creator/AtomikosDataSourceCreator.java

@@ -0,0 +1,65 @@
+package com.baomidou.dynamic.datasource.creator;
+
+import com.baomidou.dynamic.datasource.enums.XADataSourceEnum;
+import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DataSourceProperty;
+import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.atomikos.AtomikosConfig;
+import com.baomidou.dynamic.datasource.toolkit.ConfigMergeCreator;
+import com.baomidou.mybatisplus.annotation.DbType;
+import com.baomidou.mybatisplus.extension.toolkit.JdbcUtils;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.boot.jta.atomikos.AtomikosDataSourceBean;
+
+import javax.sql.DataSource;
+import java.util.Properties;
+
+import static com.baomidou.dynamic.datasource.support.DdConstants.ATOMIKOS_DATASOURCE;
+
+/**
+ * Atomikos数据源配置
+ *
+ * @author <a href="mailto:312290710@qq.com">jiazhifeng</a>
+ * @date 2023/03/02 10:20
+ */
+public class AtomikosDataSourceCreator extends AbstractDataSourceCreator implements DataSourceCreator, InitializingBean {
+    private static final ConfigMergeCreator<AtomikosConfig, AtomikosConfig> MERGE_CREATOR = new ConfigMergeCreator<>("AtomikosConfig", AtomikosConfig.class, AtomikosConfig.class);
+    private AtomikosConfig atomikosConfig;
+
+    @Override
+    public DataSource doCreateDataSource(DataSourceProperty dataSourceProperty) {
+        AtomikosConfig config = MERGE_CREATOR.create(atomikosConfig, dataSourceProperty.getAtomikos());
+        AtomikosDataSourceBean xaDataSource = new AtomikosDataSourceBean();
+
+        DbType dbType = JdbcUtils.getDbType(dataSourceProperty.getUrl());
+        xaDataSource.setXaDataSourceClassName(XADataSourceEnum.getByDbType(dbType));
+
+        Properties xaProperties = new Properties();
+        xaProperties.setProperty("url", dataSourceProperty.getUrl());
+        xaProperties.setProperty("user", dataSourceProperty.getUsername());
+        xaProperties.setProperty("password", dataSourceProperty.getPassword());
+        xaDataSource.setXaProperties(xaProperties);
+
+        xaDataSource.setUniqueResourceName(dataSourceProperty.getPoolName());
+        xaDataSource.setMinPoolSize(config.getMinPoolSize());
+        xaDataSource.setMaxPoolSize(config.getMaxPoolSize());
+        xaDataSource.setBorrowConnectionTimeout(config.getBorrowConnectionTimeout());
+        xaDataSource.setReapTimeout(config.getReapTimeout());
+        xaDataSource.setMaxIdleTime(config.getMaxIdleTime());
+        xaDataSource.setTestQuery(config.getTestQuery());
+        xaDataSource.setMaintenanceInterval(config.getMaintenanceInterval());
+        xaDataSource.setDefaultIsolationLevel(config.getDefaultIsolationLevel());
+        xaDataSource.setMaxLifetime(config.getMaxLifetime());
+        return xaDataSource;
+    }
+
+    @Override
+    public boolean support(DataSourceProperty dataSourceProperty) {
+        Class<? extends DataSource> type = dataSourceProperty.getType();
+        DbType dbType = JdbcUtils.getDbType(dataSourceProperty.getUrl());
+        return (type == null || ATOMIKOS_DATASOURCE.equals(type.getName())) && XADataSourceEnum.contains(dbType);
+    }
+
+    @Override
+    public void afterPropertiesSet() throws Exception {
+        atomikosConfig = properties.getAtomikos();
+    }
+}

core/src/main/java/com/baomidou/dynamic/datasource/enums/XADataSourceEnum.java

@@ -0,0 +1,45 @@
+package com.baomidou.dynamic.datasource.enums;
+
+import com.baomidou.mybatisplus.annotation.DbType;
+import lombok.Getter;
+
+/**
+ * 目前支持的XA数据源
+ *
+ * @author <a href="mailto:312290710@qq.com">jiazhifeng</a>
+ * @date 2023/03/02 23:05
+ */
+@Getter
+public enum XADataSourceEnum {
+    ORACLE(DbType.ORACLE,"oracle.jdbc.xa.client.OracleXADataSource"),
+    MYSQL(DbType.MYSQL, "com.mysql.cj.jdbc.MysqlXADataSource"),
+    POSTGRE_SQL(DbType.POSTGRE_SQL, "org.postgresql.xa.PGXADataSource"),
+    H2(DbType.H2, "org.h2.jdbcx.JdbcDataSource"),
+    ;
+
+    private final DbType dbType;
+    private final String xaDataSourceClassName;
+
+    XADataSourceEnum(DbType dbType, String xaDataSourceClassName) {
+        this.dbType = dbType;
+        this.xaDataSourceClassName = xaDataSourceClassName;
+    }
+
+    public static boolean contains(DbType dbType){
+        for (XADataSourceEnum item : values()) {
+            if (item.getDbType() == dbType) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    public static String getByDbType(DbType dbType){
+        for (XADataSourceEnum item : values()) {
+            if (item.getDbType() == dbType) {
+                return item.getXaDataSourceClassName();
+            }
+        }
+        return null;
+    }
+}

core/src/main/java/com/baomidou/dynamic/datasource/spring/boot/autoconfigure/DataSourceProperty.java

@@ -15,6 +15,7 @@
  */
 package com.baomidou.dynamic.datasource.spring.boot.autoconfigure;
 
+import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.atomikos.AtomikosConfig;
 import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.beecp.BeeCpConfig;
 import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.dbcp2.Dbcp2Config;
 import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.druid.DruidConfig;
@@ -100,6 +101,11 @@ public class DataSourceProperty {
      */
     @NestedConfigurationProperty
     private Dbcp2Config dbcp2 = new Dbcp2Config();
+    /**
+     * atomikos参数配置
+     */
+    @NestedConfigurationProperty
+    private AtomikosConfig atomikos = new AtomikosConfig();
 
     /**
      * 解密公匙(如果未设置默认使用全局的)

core/src/main/java/com/baomidou/dynamic/datasource/spring/boot/autoconfigure/DynamicDataSourceCreatorAutoConfiguration.java

@@ -17,9 +17,12 @@ package com.baomidou.dynamic.datasource.spring.boot.autoconfigure;
 
 import cn.beecp.BeeDataSource;
 import com.alibaba.druid.pool.DruidDataSource;
+import com.atomikos.jdbc.AtomikosDataSourceBean;
 import com.baomidou.dynamic.datasource.creator.*;
+import com.baomidou.dynamic.datasource.tx.AtomikosTransactionFactory;
 import com.zaxxer.hikari.HikariDataSource;
 import org.apache.commons.dbcp2.BasicDataSource;
+import org.apache.ibatis.transaction.TransactionFactory;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.context.annotation.Bean;
@@ -40,7 +43,8 @@ public class DynamicDataSourceCreatorAutoConfiguration {
     public static final int HIKARI_ORDER = 3000;
     public static final int BEECP_ORDER = 4000;
     public static final int DBCP2_ORDER = 5000;
-    public static final int DEFAULT_ORDER = 6000;
+    public static final int ATOMIKOS_ORDER = 6000;
+    public static final int DEFAULT_ORDER = 7000;
 
     @Primary
     @Bean
@@ -118,4 +122,23 @@ public class DynamicDataSourceCreatorAutoConfiguration {
         }
     }
 
+    /**
+     * 存在Atomikos数据源时, 加入创建器
+     */
+    @ConditionalOnClass({AtomikosDataSourceBean.class})
+    @Configuration
+    static class AtomikosDataSourceCreatorConfiguration {
+
+        @Bean
+        @Order(ATOMIKOS_ORDER)
+        public AtomikosDataSourceCreator atomikosDataSourceCreator() {
+            return new AtomikosDataSourceCreator();
+        }
+
+        @Bean
+        public TransactionFactory atomikosTransactionFactory() {
+            return new AtomikosTransactionFactory();
+        }
+    }
+
 }

core/src/main/java/com/baomidou/dynamic/datasource/spring/boot/autoconfigure/DynamicDataSourceProperties.java

@@ -16,6 +16,7 @@
 package com.baomidou.dynamic.datasource.spring.boot.autoconfigure;
 
 import com.baomidou.dynamic.datasource.enums.SeataMode;
+import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.atomikos.AtomikosConfig;
 import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.beecp.BeeCpConfig;
 import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.dbcp2.Dbcp2Config;
 import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.druid.DruidConfig;
@@ -104,6 +105,11 @@ public class DynamicDataSourceProperties {
      */
     @NestedConfigurationProperty
     private Dbcp2Config dbcp2 = new Dbcp2Config();
+    /**
+     * atomikos全局参数配置
+     */
+    @NestedConfigurationProperty
+    private AtomikosConfig atomikos = new AtomikosConfig();
 
     /**
      * aop with default ds annotation

core/src/main/java/com/baomidou/dynamic/datasource/spring/boot/autoconfigure/atomikos/AtomikosConfig.java

@@ -0,0 +1,49 @@
+package com.baomidou.dynamic.datasource.spring.boot.autoconfigure.atomikos;
+
+import lombok.Data;
+
+/**
+ * Atomikos 配置
+ *
+ * @author <a href="mailto:312290710@qq.com">jiazhifeng</a>
+ * @date 2023/03/02 10:20
+ */
+@Data
+public class AtomikosConfig {
+    /**
+     * 设置最小池大小。连接的数量不会低于该值。池将在初始化期间打开此数量的连接。可选，默认为 1
+     */
+    private int minPoolSize = 1;
+    /**
+     * 设置最大池大小。池连接的数量不会超过这个值。可选，默认为 10。
+     */
+    private int maxPoolSize = 10;
+    /**
+     * 设置池在池为空时等待连接在池中可用的最长时间（以秒为单位）。默认为 30
+     */
+    private int borrowConnectionTimeout = 30;
+    /**
+     * 设置连接池在声明连接之前允许使用连接的时间量（以秒为单位）。默认值为 0（无超时）
+     */
+    private int reapTimeout = 0;
+    /**
+     * 设置未使用的多余连接应保留在池中的最大秒数。选修的。注意：超额连接是在 minPoolSize 限制之上创建的连接。默认值为 60 秒
+     */
+    private int maxIdleTime = 60;
+    /**
+     * 设置用于在返回连接之前验证连接的 SQL 查询或语句
+     */
+    private String testQuery = "SELECT 1";
+    /**
+     * 设置池维护线程的维护间隔。以秒为单位的时间间隔。如果未设置或不是正数，则将使用池的默认值（60 秒）。
+     */
+    private int maintenanceInterval = 60;
+    /**
+     * 设置此数据源返回的连接的默认隔离级别。负值将被忽略并导致特定于供应商的 JDBC 驱动程序或 DBMS 内部默认值。
+     */
+    private int defaultIsolationLevel = -1;
+    /**
+     * 设置连接在自动销毁之前保留在池中的最大秒数。可选，默认为 0（无限制）
+     */
+    private int maxLifetime = 60;
+}

core/src/main/java/com/baomidou/dynamic/datasource/support/DdConstants.java

@@ -48,4 +48,8 @@ public interface DdConstants {
      * DBCP2数据源
      */
     String DBCP2_DATASOURCE = "org.apache.commons.dbcp2.BasicDataSource";
+    /**
+     * Atomikos数据源
+     */
+    String ATOMIKOS_DATASOURCE = "com.atomikos.jdbc.AtomikosDataSourceBean";
 }

core/src/main/java/com/baomidou/dynamic/datasource/tx/AtomikosTransactionFactory.java

@@ -0,0 +1,29 @@
+package com.baomidou.dynamic.datasource.tx;
+
+import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
+import org.apache.ibatis.session.TransactionIsolationLevel;
+import org.apache.ibatis.transaction.Transaction;
+import org.mybatis.spring.transaction.SpringManagedTransaction;
+import org.mybatis.spring.transaction.SpringManagedTransactionFactory;
+
+import javax.sql.DataSource;
+
+/**
+ * Atomikos事务适配-多数据源切换
+ *
+ * @author <a href="mailto:312290710@qq.com">jiazhifeng</a>
+ * @date 2023/03/02 10:20
+ */
+public class AtomikosTransactionFactory extends SpringManagedTransactionFactory {
+
+    @Override
+    public Transaction newTransaction(DataSource dataSource, TransactionIsolationLevel level, boolean autoCommit) {
+        DataSource determineDataSource = dataSource;
+
+        // e.g:ConnectionHolder conHolder = (ConnectionHolder) TransactionSynchronizationManager.getResource(dataSource);
+        if (dataSource instanceof DynamicRoutingDataSource) {
+            determineDataSource = ((DynamicRoutingDataSource)dataSource).determineDataSource();
+        }
+        return new SpringManagedTransaction(determineDataSource);
+    }
+}

pom.xml

@@ -50,6 +50,7 @@
         <mybatis.plus.version>3.5.2</mybatis.plus.version>
         <hikaricp.version>2.4.13</hikaricp.version>
         <druid.version>1.2.14</druid.version>
+        <atomikos.version>2.6.9</atomikos.version>
         <beeCp.version>3.2.9</beeCp.version>
         <commons-dbcp2.version>2.8.0</commons-dbcp2.version>
         <p6spy.version>3.9.1</p6spy.version>
@@ -114,6 +115,11 @@
                 <artifactId>druid-spring-boot-starter</artifactId>
                 <version>${druid.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.springframework.boot</groupId>
+                <artifactId>spring-boot-starter-jta-atomikos</artifactId>
+                <version>${atomikos.version}</version>
+            </dependency>
             <dependency>
                 <groupId>com.github.chris2018998</groupId>
                 <artifactId>beecp</artifactId>