使用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 {

    /**
     * 为什么要用链表存储(准确的是栈)
     * <pre>
     * 为了支持嵌套切换，如ABC三个service都是不同的数据源
     * 其中A的某个业务要调B的方法，B的方法需要调用C的方法。一级一级调用切换，形成了链。
     * 传统的只设置当前线程的方式不能满足此业务需求，必须使用栈，后进先出。
     * </pre>
     */
    private static final ThreadLocal<Deque<String>> LOOKUP_KEY_HOLDER = new NamedThreadLocal<Deque<String>>("dynamic-datasource") {
        @Override
        protected Deque<String> initialValue() {
            return new ArrayDeque<>();
        }
    };

    private DynamicDataSourceContextHolder() {
    }

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

    /**
     * 设置当前线程数据源
     * <p>
     * 如非必要不要手动调用，调用后确保最终清除
     * </p>
     *
     * @param ds 数据源名称
     */
    public static String push(String ds) {
        String dataSourceStr = StringUtils.isEmpty(ds) ? "" : ds;
        LOOKUP_KEY_HOLDER.get().push(dataSourceStr);
        return dataSourceStr;
    }

    /**
     * 清空当前线程数据源
     * <p>
     * 如果当前线程是连续切换数据源 只会移除掉当前线程的数据源名称
     * </p>
     */
    public static void poll() {
        Deque<String> deque = LOOKUP_KEY_HOLDER.get();
        deque.poll();
        if (deque.isEmpty()) {
            LOOKUP_KEY_HOLDER.remove();
        }
    }

    /**
     * 强制清空本地线程
     * <p>
     * 防止内存泄漏，如手动调用了push可调用此方法确保清除
     * </p>
     */
    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<User> {
}

Paginate.java

public class Paginate<T> extends Page {

    private String dataSourceId = "master";

    protected long size = 10;

    protected long current = 1;

    public String getDataSourceId() {
        return dataSourceId;
    }

    /**
     * 默认为master
     * @param dataSourceId
     */
    public void setDataSourceId(String dataSourceId) {
        this.dataSourceId = StringUtils.hasLength(dataSourceId) ? dataSourceId : this.dataSourceId;
    }

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

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

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

    @Override
    public Page<T> 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<User> pageUser(Paginate<User> 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<User> lambdaQueryWrapper = new QueryWrapper().lambda();
        IPage<User> userIPage = userMapper.selectPage(page, lambdaQueryWrapper);
        return userIPage.getRecords();
    }

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

IndexController.java

@RestController
public class IndexController {

    @Autowired
    UserService userService;

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

    @RequestMapping("api/pageUser")
    public List<User> pageUser(@RequestBody Paginate<User> 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"
}

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