Mybatis-Plus多租户，配合@DS动态切换数据源

背景是这样，@DS这个注解会标记具体使用哪个数据源，但是我们这个场景是多个租户，多个数据库，每个租户需要查询的数据源是不一样的，但是执行方式是一样的，那么仅仅用@DS("master")这种方式是不满足的，因为对于@DS里面的内容，需要根据租户来判断具体查询哪个数据源。接下来开始表演，spring spel配合@DS注解来动态修改数据源

1、首先是依赖跑不掉，mybatis-plus动态数据源

<!--mybatis-plus-->
<dependency>
  <groupId>com.baomidou</groupId>
  <artifactId>mybatis-plus-boot-starter</artifactId>
  <version>3.4.3.4</version>
</dependency>
 
<!--dynamic-datasource-->
<dependency>
  <groupId>com.baomidou</groupId>
  <artifactId>dynamic-datasource-spring-boot-starter</artifactId>
  <version>3.4.1</version>
</dependency>

2、自定义默认处理器DsProcessor

这个处理器是干什么的呢，是在执行@DS注解之前，进行解析@DS内容，我的场景如下。对于master数据源，我租户1的数据源是master1,租户2的数据源是master2。那么我需要进行替换@DS内容，将@DS("#master")在执行之前变成@DS("master1"),那么这个处理器就起作用了。

@Slf4j
public class DsExpressionProcessor extends DsProcessor {

    private StarRocksMapProperties rocksMapProperties;

    public void setRocksMapProperties(StarRocksMapProperties properties) {
        this.rocksMapProperties = properties;
    }

    /**
     * 参数发现器
     */
    private static final ParameterNameDiscoverer NAME_DISCOVERER = new DefaultParameterNameDiscoverer();
    /**
     * Express语法解析器
     */
    private static final ExpressionParser PARSER = new SpelExpressionParser();

    private static final String[] STAR_ROCKS_PREFIX_STR = {"#master1", "#slave1"};

    private String prefixStr;

    /**
     * 解析上下文的模板
     * 对于默认不设置的情况下,从参数中取值的方式 #param1
     * 设置指定模板 ParserContext.TEMPLATE_EXPRESSION 后的取值方式: #{#param1}
     * issues: https://github.com/baomidou/dynamic-datasource-spring-boot-starter/issues/199
     */
    private ParserContext parserContext = new ParserContext() {

        @Override
        public boolean isTemplate() {
            return true;
        }

        @Override
        public String getExpressionPrefix() {
            return null;
        }

        @Override
        public String getExpressionSuffix() {
            return null;
        }
    };

    private BeanResolver beanResolver;

    public DsSrExpressionProcessor() {
    }

    @Override
    public boolean matches(String key) {
        this.prefixStr = Arrays.stream(STAR_ROCKS_PREFIX_STR).filter(key::startsWith).findAny().orElse(null);
        boolean res = StrUtil.isNotBlank(this.prefixStr);
        if (res) {
            // 自定义 ParserContext
            initParserContext(key);
        }
        return res;
    }

    @Override
    public String doDetermineDatasource(MethodInvocation invocation, String key) {
        Map<String, String> dsConfigMap = initDsConfigMap();
        if (Objects.isNull(dsConfigMap)) {
            return null;
        }
        //先从 security 上下文获取 tenantId
        String tenantId = SecurityUtils.getTenantId();
        // ThreadLocalCache 获取 tenantId
        if(StringUtils.isBlank(tenantId)){
            tenantId = ThreadLocalCache.getTenantId();
        }

        if (key.length() <= this.prefixStr.length()) {
            return dsConfigMap.getOrDefault(tenantId, null);
        }
        try {
            // 参数获取 tenantId
            Method method = invocation.getMethod();
            Object[] arguments = invocation.getArguments();
            StandardEvaluationContext context = new MethodBasedEvaluationContext((Object) null, method, arguments, NAME_DISCOVERER);
            context.setBeanResolver(this.beanResolver);
            Object value = PARSER.parseExpression(key, this.parserContext).getValue(context);
            return dsConfigMap.getOrDefault(String.valueOf(value), null);
        } catch (Exception e) {
            // ThreadLocalCache 获取 tenantId
            return dsConfigMap.getOrDefault(ThreadLocalCache.getTenantId(), null);
        }
    }

    /**
     * 自定义ParserContext
     * @param key
     */
    private void initParserContext(String key) {

        this.parserContext = new ParserContext() {

            @Override
            public boolean isTemplate() {
                return true;
            }

            @Override
            public String getExpressionPrefix() {
                if (key.length() > prefixStr.length()) {
                    return prefixStr + "{";
                }
                return prefixStr;

            }

            @Override
            public String getExpressionSuffix() {
                if (key.length() > prefixStr.length()) {
                    return "}";
                }
                return null;
            }
        };
    }


    private Map<String,String> initDsConfigMap(){

        if (Objects.isNull(this.rocksMapProperties)) {
            return null;
        }
        Map<String,String> dsConfigMap = new HashMap<>();
        JSONObject jsonObject = JSONUtil.parseObj(this.rocksMapProperties);
        for (Map.Entry<String, Object> entry : jsonObject.entrySet()) {
            if(CharSequenceUtil.equalsIgnoreCase(this.prefixStr.replace("#","")+"Map",entry.getKey())){
                dsConfigMap = (Map<String,String>)entry.getValue();
                break;
            }
        }
        return dsConfigMap;
    }

    public void setParserContext(ParserContext parserContext) {
        this.parserContext = parserContext;
    }

    public void setBeanResolver(BeanResolver beanResolver) {
        this.beanResolver = beanResolver;
    }

}

这段是什么意思呢，就是在配置文件中，将对应租户的数据源配置好，然后在上下文中获取到租户号，根据租户号获取他的数据源，然后替换其中的@DS变量。

3、然后再把处理器注册上

@Bean
    @ConditionalOnMissingBean
    public DsProcessor dsProcessor() {
        DsHeaderProcessor headerProcessor = new DsHeaderProcessor();
        DsSessionProcessor sessionProcessor = new DsSessionProcessor();
        DsSpelExpressionProcessor spelExpressionProcessor = new DsSpelExpressionProcessor();
        headerProcessor.setNextProcessor(sessionProcessor);
        sessionProcessor.setNextProcessor(spelExpressionProcessor);
        return headerProcessor;
    }

4、spring事务会导致数据源切换失效，请使用@DS的事务