使用mybatisplus实现用户级的动态数据源切换

昨天在B站看了基于SpringBoot和MyBatis-Plus多数据源分析的视频，评论区有小伙伴想实现页面级的数据源切换，经过初步的分析就有了下面的想法。

通过分析源码可知道，MyBatis-Plus提供了dynamic-datasource-spring-boot-starter 以支持多数据源的需求，在使用方式上，运用注解可灵活便捷地实现指定数据源进行表操作。能否实现用户级的多数据源切换，用户在前端页面做出选择，后端匹配对应的数据源? 很简单，下面来详细谈谈。

dynamic-datasource 提供了一个重要的类DynamicRoutingDataSource ，这个类是实现动态数据源的核心，其中：

@Override
public DataSource determineDataSource() {
    String dsKey = DynamicDataSourceContextHolder.peek();
    return getDataSource(dsKey);
}

这里的DynamicDataSourceContextHolder.peek() 操作的是ThreadLocal变量，源码如下：

public final class DynamicDataSourceContextHolder {

    
    private static final ThreadLocal> LOOKUP_KEY_HOLDER = new NamedThreadLocal>("dynamic-datasource") {
        @Override
        protected Deque initialValue() {
            return new ArrayDeque<>();
        }
    };

    private DynamicDataSourceContextHolder() {
    }

    
    public static String peek() {
        return LOOKUP_KEY_HOLDER.get().peek();
    }

    
    public static String push(String ds) {
        String dataSourceStr = StringUtils.isEmpty(ds) ? "" : ds;
        LOOKUP_KEY_HOLDER.get().push(dataSourceStr);
        return dataSourceStr;
    }

    
    public static void poll() {
        Deque deque = LOOKUP_KEY_HOLDER.get();
        deque.poll();
        if (deque.isEmpty()) {
            LOOKUP_KEY_HOLDER.remove();
        }
    }

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

所以LOOKUP_KEY_HOLDER是重中之重，若我们在执行业务逻辑之前，改变该变量的值，即可达到我们的目的。并且官方源码中也提供了相应的静态方法可供使用。

基础文件清单

User.java

@Data
public class User {
    private Long id;
    private String name;
    private Integer age;
    private String email;
}

UserMapper.java

@Mapper
public interface UserMapper extends baseMapper {
}

Paginate.java

public class Paginate extends Page {

    private String dataSourceId = "master";

    protected long size = 10;

    protected long current = 1;

    public String getDataSourceId() {
        return dataSourceId;
    }

    
    public void setDataSourceId(String dataSourceId) {
        this.dataSourceId = StringUtils.hasLength(dataSourceId) ? dataSourceId : this.dataSourceId;
    }

    @Override
    public long getSize() {
        return size;
    }

    @Override
    public Page setSize(long size) {
        this.size = size;
        return this;
    }

    @Override
    public long getCurrent() {
        return current;
    }

    @Override
    public Page setCurrent(long current) {
        this.current = current;
        return this;
    }
}

UserService.java

@Service
public class UserService {

    private UserMapper userMapper;

    private DynamicRoutingDataSource routingDataSource;

    public UserService() {}

    @Autowired
    private UserService(UserMapper userMapper, DynamicRoutingDataSource routingDataSource) {
        this.userMapper = userMapper;
        this.routingDataSource = routingDataSource;
    }

    public List pageUser(Paginate page) {
        // 核心代码
        Assert.assertTrue("不支持的数据源id",
                listDataSource().stream().anyMatch(
                        dataSourceId-> page.getDataSourceId().equalsIgnoreCase(dataSourceId)));

        if(!page.getDataSourceId().equalsIgnoreCase(DynamicDataSourceContextHolder.peek())) {
            DynamicDataSourceContextHolder.poll();
            DynamicDataSourceContextHolder.push(page.getDataSourceId());
        }

        LambdaQueryWrapper lambdaQueryWrapper = new QueryWrapper().lambda();
        IPage userIPage = userMapper.selectPage(page, lambdaQueryWrapper);
        return userIPage.getRecords();
    }

    public List listDataSource() {
        return new ArrayList<>(routingDataSource.getDataSources().keySet());
    }
}

IndexController.java

@RestController
public class IndexController {

    @Autowired
    UserService userService;

    @RequestMapping("api/listDataSource")
    public List listDataSource() {
        return userService.listDataSource();
    }

    @RequestMapping("api/pageUser")
    public List pageUser(@RequestBody Paginate page) {
        return userService.pageUser(page);
    }
}

application.yaml

spring:
    application:
        name: demo
    datasource:
        dynamic:
            datasource:
                master:
                    driver-class-name: com.p6spy.engine.spy.P6SpyDriver
                    url: jdbc:p6spy:mysql://127.0.0.1:3306/spring
                    username: root
                    password: root
                slave:
                    driver-class-name: com.p6spy.engine.spy.P6SpyDriver
                    url: jdbc:p6spy:mysql://127.0.0.1:3306/slave
                    username: root
                    password: root

    sql:
        init:
            data-locations: classpath:db/mysql/data.sql
            schema-locations: classpath:db/mysql/schema.sql
            mode: NEVER
    output:
        ansi:
            enabled: always

mybatis-plus:
    mapper-locations: classpath:mappers/*.xml
    configuration:
        log-impl: org.apache.ibatis.logging.stdout.StdOutImpl
server:
    port: 80

###
GET http://localhost:8080/api/listDataSource
Content-Type: application/json

###
GET http://localhost:8080/api/pageUser
Content-Type: application/json

{
  "size": 10,
  "current": 1,
  "dataSourceId": ""
}

###
GET http://localhost:8080/api/pageUser
Content-Type: application/json

{
  "size": 10,
  "current": 1,
  "dataSourceId": "slave"
}

经过测试，此方式可以实现用户级动态切换数据源，并且线程安全。
核心逻辑还可以抽取重构，这里仅提供功能测试。