SpringBoot+ThreadLocal+AbstractRoutingDataSource implements dynamic switching of data sources

Hi, everyone, I am the fat boy who grabs my wife’s yogurt.

Recently, when doing business requirements, I need to obtain data from different databases and then write them into the current database, so it involves switching data sources. Originally I wanted to use the dynamic data source SpringBoot starter provided in Mybatis-plus: dynamic-datasource-spring-boot-starter to implement it. After the results were introduced, it was found that it could not be used due to environmental problems in the previous project. Then I studied the data source switching code and decided to use ThreadLocal + AbstractRoutingDataSource to simulate thread data source switching in dynamic-datasource-spring-boot-starter.

1. Introduction

ThreadLocal and AbstractRoutingDataSource are mentioned above. Let’s briefly introduce them.

ThreadLocal: I believe everyone will be familiar with it. Its full name is: thread local variable. Mainly to solve the problem of data inconsistency due to concurrency in multi-threading. ThreadLocal provides a copy of the variable for each thread to ensure that each thread does not access the same object at a certain time. This achieves isolation and increases memory, but greatly reduces the performance consumption during thread synchronization and reduces thread concurrency. Control complexity.

ThreadLocal function: shared in one thread, isolated between different threads

ThreadLocal principle: When ThreadLocal stores a value, it will obtain the current thread instance as the key and store it in the Map in the current thread object.

AbstractRoutingDataSource: Select the current data source according to user-defined rules,

Function: Before executing the query, set the data source to be used, implement the data source of dynamic routing, and execute its abstract method determineCurrentLookupKey() before each database query operation to decide which data source to use.

2. Code implementation

Program environment:

SpringBoot2.4.8

Mybatis-plus3.2.0

Druid1.2.6

lombok1.18.20

commons-lang3 3.10

2.1 Implement ThreadLocal

Create a class to implement ThreadLocal, mainly through the get, set, and remove methods to obtain, set, and delete the data source corresponding to the current thread.

 * @author: jiangjs
 * @description:
 * @date: 2023/7/27 11:21
 **/
public class DataSourceContextHolder {
    
    private static final ThreadLocal<String> DATASOURCE_HOLDER = new ThreadLocal<>();

    
     * Set data source
     * @param dataSourceName data source name
     */
    public static void setDataSource(String dataSourceName){
        DATASOURCE_HOLDER.set(dataSourceName);
    }

    
     * Get the data source of the current thread
     * @return data source name
     */
    public static String getDataSource(){
        return DATASOURCE_HOLDER.get();
    }

    
     * Delete the current data source
     */
    public static void removeDataSource(){
        DATASOURCE_HOLDER.remove();
    }

}

2.2 Implement AbstractRoutingDataSource

Define a dynamic data source class to implement AbstractRoutingDataSource, and associate it with the get method in the ThreadLocal class implemented above through the determineCurrentLookupKey method to achieve dynamic switching of data sources.

 * @author: jiangjs
 * @description: Implement dynamic data sources and route to different data sources according to AbstractRoutingDataSource
 * @date: 2023/7/27 11:18
 **/
public class DynamicDataSource extends AbstractRoutingDataSource {

    public DynamicDataSource(DataSource defaultDataSource,Map<Object, Object> targetDataSources){
        super.setDefaultTargetDataSource(defaultDataSource);
        super.setTargetDataSources(targetDataSources);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSource();
    }
}

In the above code, a construction method of a dynamic data source class is also implemented, mainly to set the default data source and various target data sources saved in Map. The key of the Map is the set data source name, and the value is the corresponding data source (DataSource).

2.3 Configuration database

Configure database information in application.yml:

spring:
  datasource:
    type: com.alibaba.druid.pool.DruidDataSource
    druid:
      master:
        url: jdbc:mysql://xxxxxx:3306/test1?characterEncoding=utf-8 & amp;allowMultiQueries=true & amp;zeroDateTimeBehavior=convertToNull & amp;useSSL=false
        username: root
        password: 123456
        driver-class-name: com.mysql.cj.jdbc.Driver
      slave:
        url: jdbc:mysql://xxxxx:3306/test2?characterEncoding=utf-8 & amp;allowMultiQueries=true & amp;zeroDateTimeBehavior=convertToNull & amp;useSSL=false
        username: root
        password: 123456
        driver-class-name: com.mysql.cj.jdbc.Driver
      initial-size: 15
      min-idle: 15
      max-active: 200
      max-wait: 60000
      time-between-eviction-runs-millis: 60000
      min-evictable-idle-time-millis: 300000
      validation-query: ""
      test-while-idle: true
      test-on-borrow: false
      test-on-return: false
      pool-prepared-statements: false
      connection-properties: false

 * @author: jiangjs
 * @description: Set data source
 * @date: 2023/7/27 11:34
 **/
@Configuration
public class DateSourceConfig {

    @Bean
    @ConfigurationProperties("spring.datasource.druid.master")
    public DataSource masterDataSource(){
        return DruidDataSourceBuilder.create().build();
    }

    @Bean
    @ConfigurationProperties("spring.datasource.druid.slave")
    public DataSource slaveDataSource(){
        return DruidDataSourceBuilder.create().build();
    }

    @Bean(name = "dynamicDataSource")
    @Primary
    public DynamicDataSource createDynamicDataSource(){
        Map<Object,Object> dataSourceMap = new HashMap<>();
        DataSource defaultDataSource = masterDataSource();
        dataSourceMap.put("master",defaultDataSource);
        dataSourceMap.put("slave",slaveDataSource());
        return new DynamicDataSource(defaultDataSource,dataSourceMap);
    }

}

Through the configuration class, the configured database information in the configuration file is converted into a datasource and added to the DynamicDataSource. At the same time, the DynamicDataSource is injected into Spring through @Bean for management. It will be used later when adding dynamic data sources.

2.4 Test

In the master and slave test libraries, add a table test_user respectively, which has only one field user_name.

create table test_user(
  user_name varchar(255) not null comment 'username'
)

Add information to the main library:

insert into test_user (user_name) value ('master');

Add information from the library:

insert into test_user (user_name) value ('slave');

We create a getData method, and the parameter is the name of the data source where the data needs to be queried.

@GetMapping("/getData.do/{datasourceName}")
public String getMasterData(@PathVariable("datasourceName") String datasourceName){
    DataSourceContextHolder.setDataSource(datasourceName);
    TestUser testUser = testUserMapper.selectOne(null);
    DataSourceContextHolder.removeDataSource();
    return testUser.getUserName();
}

You can implement other Mapper and entity classes by yourself.

Results of the:

1. When passing master:

2. When passing slave:

Through the execution results, we can see that passing different data source names, the corresponding databases queried are different, and the returned results are also different.

In the above code, we see DataSourceContextHolder.setDataSource(datasourceName); to set the database that the current thread needs to query, and DataSourceContextHolder.removeDataSource(); to remove the current The data source that the thread has set. Friends who have used Mybatis-plus dynamic data sources should still remember that we use DynamicDataSourceContextHolder.push(String ds); and DynamicDataSourceContextHolder.poll(); when switching data sources. Looking at the source code of these two methods, we will find that the stack is actually used when using ThreadLocal. The advantage of this is that it can use multiple data sources to be nested. I will not show you how to implement them here. For those who are interested, You can take a look at the source code of dynamic data source in Mybatis-plus.

Note: When starting the program, don’t forget to exclude SpringBoot from automatically adding data sources, otherwise circular dependency problems will be reported.

@SpringBootApplication(exclude = DataSourceAutoConfiguration.class)

2.5 Optimization and adjustment

2.5.1 Annotation switching data source

In the above, although dynamic switching of data sources has been implemented, we will find that if multiple businesses are involved in switching data sources, we need to add this piece of code to each implementation class.

Speaking of this, some friends will probably think of using annotations for optimization. Let’s implement it next.

2.5.1.1 Definition annotations

We use the annotation of mybatis dynamic data source switching: DS, the code is as follows:

 * @author: jiangjs
 * @description:
 * @date: 2023/7/27 14:39
 **/
@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
@Inherited
public @interface DS {
    String value() default "master";
}

2.5.1.2 Implement aop
@Aspect
@Component
@Slf4j
public class DSAspect {

    @Pointcut("@annotation(com.jiashn.dynamic_datasource.dynamic.aop.DS)")
    public void dynamicDataSource(){}

    @Around("dynamicDataSource()")
    public Object datasourceAround(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature)point.getSignature();
        Method method = signature.getMethod();
        DS ds = method.getAnnotation(DS.class);
        if (Objects.nonNull(ds)){
            DataSourceContextHolder.setDataSource(ds.value());
        }
        try {
            return point.proceed();
        } finally {
            DataSourceContextHolder.removeDataSource();
        }
    }
}

The code uses @Around, obtains annotation information through ProceedingJoinPoint, gets the annotation to pass the value, and then sets the data source of the current thread. Friends who don’t know about AOP can google or Baidu.

2.5.1.3 Testing

Add two test methods:

@GetMapping("/getMasterData.do")
public String getMasterData(){
    TestUser testUser = testUserMapper.selectOne(null);
    return testUser.getUserName();
}

@GetMapping("/getSlaveData.do")
@DS("slave")
public String getSlaveData(){
    TestUser testUser = testUserMapper.selectOne(null);
    return testUser.getUserName();
}

Since the default value set in @DS is: master, there is no need to add it when calling the main data source.

Results of the:

1. Call the getMasterData.do method:

2. Call the getSlaveData.do method:

Through the execution results, we also switched the data source through @DS, realizing the way of switching data sources through annotations in Mybatis-plus dynamically switching data sources.

2.5.2 Dynamically add data sources

Business scenario: Sometimes our business requires us to add these data sources from database tables that store other data sources, and then switch these data sources according to different situations.

Therefore, we need to transform DynamicDataSource to dynamically load data sources.

2.5.2.1 Data source entity
 * @author: jiangjs
 * @description: data source entity
 * @date: 2023/7/27 15:55
 **/
@Data
@Accessors(chain = true)
public class DataSourceEntity {

    
     * Database address
     */
    private String url;
    
     * database username
     */
    private String userName;
    
     * password
     */
    private String passWord;
    
     * Database driven
     */
    private String driverClassName;
    
     * Database key, which is the key saved in the Map
     */
    private String key;
}

Define the general information of the data source in the entity, and define a key to be used as the key in the Map in DynamicDataSource.

2.5.2.2 Modify DynamicDataSource
 * @author: jiangjs
 * @description: Implement dynamic data sources and route to different data sources according to AbstractRoutingDataSource
 * @date: 2023/7/27 11:18
 **/
@Slf4j
public class DynamicDataSource extends AbstractRoutingDataSource {

    private final Map<Object,Object> targetDataSourceMap;

    public DynamicDataSource(DataSource defaultDataSource,Map<Object, Object> targetDataSources){
        super.setDefaultTargetDataSource(defaultDataSource);
        super.setTargetDataSources(targetDataSources);
        this.targetDataSourceMap = targetDataSources;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSource();
    }

    
     * Add data source information
     * @param dataSources data source entity collection
     * @return return the added result
     */
    public void createDataSource(List<DataSourceEntity> dataSources){
        try {
            if (CollectionUtils.isNotEmpty(dataSources)){
                for (DataSourceEntity ds : dataSources) {
                    
                    Class.forName(ds.getDriverClassName());
                    DriverManager.getConnection(ds.getUrl(),ds.getUserName(),ds.getPassWord());
                    
                    DruidDataSource dataSource = new DruidDataSource();
                    BeanUtils.copyProperties(ds,dataSource);
                    
                    dataSource.setTestOnBorrow(true);
                    
                    
                    dataSource.setTestWhileIdle(true);
                    
                    dataSource.setValidationQuery("select 1 ");
                    dataSource.init();
                    this.targetDataSourceMap.put(ds.getKey(),dataSource);
                }
                super.setTargetDataSources(this.targetDataSourceMap);
                
                super.afterPropertiesSet();
                return Boolean.TRUE;
            }
        }catch (ClassNotFoundException | SQLException e) {
            log.error("---Program error ---:{}", e.getMessage());
        }
        return Boolean.FALSE;
    }

    
     * Verify whether the data source exists
     * @param key key saved by data source
     * @return return result, true: exists, false: does not exist
     */
    public boolean existsDataSource(String key){
        return Objects.nonNull(this.targetDataSourceMap.get(key));
    }
}

In the modified DynamicDataSource, we can add a private final Map targetDataSourceMap. This map will initialize the created Map data source information through the DynamicDataSource construction method when adding the data source configuration file, that is: DateSourceConfig In the createDynamicDataSource() method in the class.

At the same time, we added a createDataSource method to this class to create the data source and add it to the map, and then reassign the target data source through super.setTargetDataSources(this.targetDataSourceMap);

2.5.2.3 Dynamically add data sources

The above code has implemented the method of adding a data source, so let’s simulate adding a data source from a database table, and then we add the data source to the data source Map by calling the method of loading the data source.

Define a database table in the main database to save database information.

create table test_db_info(
    id int auto_increment primary key not null comment 'primary keyId',
    url varchar(255) not null comment 'database URL',
    username varchar(255) not null comment 'username',
    password varchar(255) not null comment 'password',
    driver_class_name varchar(255) not null comment 'database driver'
    name varchar(255) not null comment 'database name'
)

For convenience, we enter the previous slave library into the database and modify the database name.

insert into test_db_info(url, username, password,driver_class_name, name)
value ('jdbc:mysql://xxxxx:3306/test2?characterEncoding=utf-8 & amp;allowMultiQueries=true & amp;zeroDateTimeBehavior=convertToNull & amp;useSSL=false',
       'root','123456','com.mysql.cj.jdbc.Driver','add_slave')

The entities and mappers corresponding to the database table can be added by your friends.

Add a data source when starting SpringBoot:

 * @author: jiangjs
 * @description:
 * @date: 2023/7/27 16:56
 **/
@Component
public class LoadDataSourceRunner implements CommandLineRunner {
    @Resource
    private DynamicDataSource dynamicDataSource;
    @Resource
    private TestDbInfoMapper testDbInfoMapper;
    @Override
    public void run(String... args) throws Exception {
        List<TestDbInfo> testDbInfos = testDbInfoMapper.selectList(null);
        if (CollectionUtils.isNotEmpty(testDbInfos)) {
            List<DataSourceEntity> ds = new ArrayList<>();
            for (TestDbInfo testDbInfo : testDbInfos) {
                DataSourceEntity sourceEntity = new DataSourceEntity();
                BeanUtils.copyProperties(testDbInfo,sourceEntity);
                sourceEntity.setKey(testDbInfo.getName());
                ds.add(sourceEntity);
            }
            dynamicDataSource.createDataSource(ds);
        }
    }
}

After the above SpringBoot startup, the data in the database table has been added to the dynamic data source. We call the previous test method and pass the data source name as a parameter to see the execution results.

2.5.2.4 Testing

Through testing, we found that the database in the database table was dynamically added to the data source, and friends can happily add data sources at will.

Okay, that’s it for today. I hope that my chatting will give you a deeper understanding of the method of dynamically switching data sources.

Work is better than hard work, but play is waste. Let’s take action, friends.

github:github.com/lovejiashn/…[1]

Reference materials

[1]

https://github.com/lovejiashn/dynamic_datasource.git: https://link.juejin.cn/?target=https://github.com/lovejiashn/dynamic_datasource.git

The knowledge points of the article match the official knowledge files, and you can further learn related knowledge. Java Skill TreeHomepageOverview 138192 people are learning the system