自定义注解实现多数据源动态切换

背景

前几天在做数据表分库逻辑，需要前端传入dbName实现自由选择数据源的分库逻辑。本打算集成MyBatis-Plus提供的dynamic-datasource工具包，但是数据源配置都写死在了yml里，而公司的项目是基于微服务的，配置类都统一访问远程配置中心获取，于是只能利用AOP在service层做拦截，自己造轮子实现分库逻辑，过程如下：

1、注册多数据源：

@Configuration
@ImportResource("classpath:mybatis/spring-mybatis.xml")
public class DataSourceConfig {

    @Bean(name = "DataSource1")
    @Qualifier("DataSource1")
    @Primary
    public DataSource DataSource1() {
        BasicDataSource dataSource = new BasicDataSource();
        dataSource.setDriverClassName("com.mysql.jdbc.Driver");
        dataSource.setUrl("jdbc:mysql://localhost:3306/datasource1？characterEncoding=UTF8&useSSL=false");
        dataSource.setUsername("root");
        dataSource.setPassword("");
        return dataSource;
    }

    @Bean(name = "DataSource2")
    @Qualifier("DataSource2")
    public DataSource DataSource2() {
        BasicDataSource dataSource = new BasicDataSource();
        dataSource.setDriverClassName("com.mysql.jdbc.Driver");
        dataSource.setUrl("jdbc:mysql://localhost:3306/datasource2？characterEncoding=UTF8&useSSL=false");
        dataSource.setUsername("root");
        dataSource.setPassword("");
        return dataSource;
    }

	// hashmap存放多数据源
    @Bean(name = "dynamicDataSource")
    public DataSource dynamicDataSource(@Autowired @Qualifier("DataSource1") DataSource primary, @Autowired @Qualifier("DataSource2") DataSource ds2) {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        // 默认数据源
        dynamicDataSource.setDefaultTargetDataSource(primary);
        // 配置多数据源
        Map<Object, Object> dsMap = new HashMap<Object, Object>(2);
        dsMap.put("DataSource1", primary);
        dsMap.put("DataSource2", ds2);
        dynamicDataSource.setTargetDataSources(dsMap);
        return dynamicDataSource;
    }

    @Bean
    public PlatformTransactionManager txManager(DataSource dynamicDataSource) {
        return new DataSourceTransactionManager(dynamicDataSource);
    }
}

2、在spring配置文件中配置这两个数据源

核心配置项如下：

<!-- dynamicDataSource -->
<bean id="dataSource" class="com.douyu.ocean.dataworks.vienna.core.config.DynamicDataSource">
    <!-- 为targetDataSources注入两个数据源 -->
    <property name="targetDataSources">
        <map key-type="java.lang.String">
            <entry key="ds1" value-ref="DataSource1"/>
            <entry key="ds2" value-ref="DataSource2"/>
        </map>
    </property>
    <!-- 为指定数据源RoutingDataSource注入默认的数据源-->
    <property name="defaultTargetDataSource" ref="DataSource1"/>
</bean>

<!-- sql session factory bean -->
<bean id="sqlSessionFactory" class="org.mybatis.spring.SqlSessionFactoryBean" primary="true">
    <property name="dataSource" ref="dataSource"/>
    <property name="mapperLocations" value="classpath*:mybatis/mapper/*Mapper.xml"/>
    <property name="configLocation" value="classpath:mybatis/mybatis-conf.xml"/>
    <!--分页插件-->
    <property name="plugins">
        <array>
            <ref bean="pageInterceptor" />
        </array>
    </property>
</bean>

<!-- mapper scanner-->
<bean class="org.mybatis.spring.mapper.MapperScannerConfigurer" primary="true">
    <property name="sqlSessionFactoryBeanName" value="sqlSessionFactory"/>
    <property name="basePackage" value="com.douyu.ocean.dataworks.vienna.core.dao.mapper"/>
</bean>

<!--配置事务-->
<bean id="transactionManager" class="org.springframework.jdbc.datasource.DataSourceTransactionManager"
      primary="true">
    <property name="dataSource" ref="dataSource"/>
</bean>

<!--使用事务注解的方式-->
<tx:annotation-driven transaction-manager="transactionManager"/>

3、自定义一个数据源类，继承 org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource 并重写 determineCurrentLookupKey()

@Component
public class DynamicDataSource extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDB();
    }
}

4、实现数据源持有类DataSourceContextHolder

@Component
public class DataSourceContextHolder {
    private static final ThreadLocal<String> contextHolder = new ThreadLocal<>();
    // 设置数据源名
    public static void setDB(String dbType) {
        contextHolder.set(dbType);
    }
    // 获取数据源名
    public static String getDB() {
        return contextHolder.get();
    }
    // 清除数据源名
    public static void clearDB() {
        contextHolder.remove();
    }
}

5、自定义注解

import java.lang.annotation.*;

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.TYPE, ElementType.METHOD})
@Documented
public @interface DS {
    String value() default "";
}

6、利用AOP实现动态数据源切换

@Aspect
@Order(-1)
@Component
public class DynamicDataSourceAspect {
    private Logger logger = LoggerFactory.getLogger(this.getClass());

    //切点
    @Pointcut("execution(* com.douyu.ocean.dataworks.vienna.core.service.impl..*(..))")
    public void aspect() { }

    @Before("aspect()")
    private void before(JoinPoint point) {
        Object target = point.getTarget();
        String method = point.getSignature().getName();
        Class<?> classz = target.getClass();// 获取目标类
        Class<?>[] parameterTypes = ((MethodSignature) point.getSignature())
                .getMethod().getParameterTypes();
        try {
            Method m = classz.getMethod(method, parameterTypes);
            if (m != null && m.isAnnotationPresent(DS.class)) {
                DS data = m.getAnnotation(DS.class);
                logger.info("method :{},datasource:{}",m.getName() ,data.value());
                DataSourceContextHolder.setDB(data.value()); // 数据源放到当前线程中
            }
        } catch (Exception e) {
            logger.error("get datasource error ", e);
            DataSourceContextHolder.setDB("DataSource1");
        }
    }

    @AfterReturning("aspect()")
    public void after(JoinPoint point) {
        DataSourceContextHolder.clearDB();
    }
}

7、若方法体上没有注解，则使用默认数据源，如果有以下注解，则使用指定的数据源

@Override
@DS("ds2")
public List<User> selectDataSource() {
    return userMapper.selectUsers();
}

8、此时分库逻辑已实现，但是无法通过前端传入dbName实现动态切换数据源，阅读dynamic-datasource-spring-boot-starter源码，学习一下别人的实现方式

/**
 * Copyright © 2018 organization baomidou
 * <pre>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * <pre/>
 */
package com.baomidou.dynamic.datasource.processor;

import org.aopalliance.intercept.MethodInvocation;
import org.springframework.context.expression.MethodBasedEvaluationContext;
import org.springframework.core.DefaultParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.ParserContext;
import org.springframework.expression.spel.standard.SpelExpressionParser;

import java.lang.reflect.Method;

/**
 * @author TaoYu
 * @since 2.5.0
 */
public class DsSpelExpressionProcessor extends DsProcessor {

    /**
     * 参数发现器
     */
    private static final ParameterNameDiscoverer NAME_DISCOVERER = new DefaultParameterNameDiscoverer();
    /**
     * Express语法解析器
     */
    private static final ExpressionParser PARSER = new SpelExpressionParser();
    /**
     * 解析上下文的模板
     * 对于默认不设置的情况下,从参数中取值的方式 #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 false;
        }

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

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

    @Override
    public boolean matches(String key) {
        return true;
    }

    @Override
    public String doDetermineDatasource(MethodInvocation invocation, String key) {
        Method method = invocation.getMethod();
        Object[] arguments = invocation.getArguments();
        EvaluationContext context = new MethodBasedEvaluationContext(null, method, arguments, NAME_DISCOVERER);
        final Object value = PARSER.parseExpression(key, parserContext).getValue(context);
        return value == null ? null : value.toString();
    }

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

9、在此可以看到MyBatis-Plus提供的方式是通过SPEL表达式+Express语法解析器实现的，于是在AOP层加入SPEL表达式解析，拦截前缀为 # 的注解值

@Aspect
@Order(-1)
@Component
public class DynamicDataSourceAspect {
    private static final String DYNAMIC_PREFIX = "#";
    private Logger logger = LoggerFactory.getLogger(this.getClass());

    //切点
    @Pointcut("execution(* com.douyu.ocean.dataworks.vienna.core.service.impl..*(..))")
    public void aspect() { }

    @Before("aspect()")
    private void before(JoinPoint point) {
        Object target = point.getTarget();
        String method = point.getSignature().getName();
        Class<?> classz = target.getClass();// 获取目标类
        Class<?>[] parameterTypes = ((MethodSignature) point.getSignature())
                .getMethod().getParameterTypes();
        try {
            Method m = classz.getMethod(method, parameterTypes);
            if (m != null && m.isAnnotationPresent(DS.class)) {
                DS data = m.getAnnotation(DS.class);
                logger.info("method :{},datasource:{}",m.getName() ,data.value());
                String value = data.value();

                Object[] args = point.getArgs();
                //获取被拦截方法参数名列表(使用Spring支持类库)
                LocalVariableTableParameterNameDiscoverer localVariableTable = new LocalVariableTableParameterNameDiscoverer();
                String[] paraNameArr = localVariableTable.getParameterNames(m);
                //使用SPEL进行key的解析
                ExpressionParser parser = new SpelExpressionParser();
                //SPEL上下文
                StandardEvaluationContext context = new StandardEvaluationContext();
                //把方法参数放入SPEL上下文中
                for(int i=0;i<paraNameArr.length;i++) {
                    context.setVariable(paraNameArr[i], args[i]);
                }
                // 使用变量方式传入业务动态数据
                if(value.startsWith(DYNAMIC_PREFIX )) {
                    value = parser.parseExpression(value).getValue(context, String.class);
                }
                System.out.println("加工后的value" + value);
                DataSourceContextHolder.setDB(value);// 数据源放到当前线程中
            }
        } catch (Exception e) {
            logger.error("get datasource error ", e);
            DataSourceContextHolder.setDB("masterDataSource");
        }
    }

    @AfterReturning("aspect()")
    public void after(JoinPoint point) {
        DataSourceContextHolder.clearDB();
    }
}

10、使用SPEL表达式绑定注解动态变量参数值

@Override
@DS("#dbName")
public List<User> selectDataSource(String dbName) {
    return userMapper.selectUsers();
}

11、完成自定义注解实现多数据源动态切换

注解为@DS("#dbName")，前端传入dbName为dy_tmp，后台成功识别解析