基于注解的动态数据源实现

需求

有些项目不只访问一个数据库,可能需要访问多个数据库,那么就会有一个问题,怎么进行数据源的切换.

动态数据源

解决这个需求的一个常见解决方案是使用动态数据源.下面将按部就班的来介绍一下如何实现基于注解的动态数据源.完整的代码请参考https://github.com/CodeShowZz/data-source/tree/master/dynamic-data-source.

第一步:配置数据源

将项目中需要使用的数据源放到一个配置文件中,比如叫做jdbc.properties,在我的例子中,我有两个数据源,一个是learning库,另外一个是test库.

数据库配置文件:

spring.datasource.test.driver-class-name=com.mysql.jdbc.Driver
spring.datasource.test.jdbc-url=jdbc:mysql://localhost:3306/test?useSSL=false&serverTimezone=GMT%2B8&characterEncoding=UTF-8
spring.datasource.test.username=root
spring.datasource.test.password=123456

spring.datasource.learning.driver-class-name=com.mysql.jdbc.Driver
spring.datasource.learning.jdbc-url=jdbc:mysql://localhost:3306/learning?useSSL=false&serverTimezone=GMT%2B8&characterEncoding=UTF-8
spring.datasource.learning.username=root
spring.datasource.learning.password=123456

数据源常量类:

public class DataSourceConstants {

    public static final String DB_LEARNING = "learning";

    public static final String DB_TEST= "test";
}

动态数据源类:

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getContextKey();
    }
}

这里使用了一个DynamicDataSourceContextHolder类,将在下面进行讲解.

数据源配置:

@EnableAutoConfiguration(exclude = {DataSourceAutoConfiguration.class})
@Configuration
@PropertySource("classpath:jdbc.properties")
@MapperScan(basePackages = "com.dynamic.datasource.dao")
public class DynamicDataSourceConfig {

    @Bean(DataSourceConstants.DB_LEARNING)
    @ConfigurationProperties(prefix = "spring.datasource.learning")
    public DataSource learningDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean(DataSourceConstants.DB_TEST)
    @ConfigurationProperties(prefix = "spring.datasource.test")
    public DataSource testDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean
    @Primary
    public DataSource dynamicDataSource() {
        Map<Object, Object> dataSourceMap = new HashMap(2);
        dataSourceMap.put(DataSourceConstants.DB_LEARNING, learningDataSource());
        dataSourceMap.put(DataSourceConstants.DB_TEST, testDataSource());
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        dynamicDataSource.setTargetDataSources(dataSourceMap);
        dynamicDataSource.setDefaultTargetDataSource(testDataSource());
        return dynamicDataSource;
    }
}

在这里讲一下具体的原理,首先我们定义了两个数据源,然后在dynamicDataSource方法中定义了一个Map,将两个数据源以(名称,数据源)的形式放入.接着调用setTargetDataSources将Map设置进去,并通过setDefaultTargetDataSource设置了默认数据源.在每次执行sql语句时,将通过DynamicDataSource类实现的determineCurrentLookupKey方法返回的key从Map中找到对应的数据源,如果没有找到,将使用默认数据源.

了解了这个原理,那么改变determineCurrentLookupKey方法返回的key就可以实现数据源的切换,那如何改造这个方法使得可以动态切换数据源呢?通常来说,会将它放在ThreadLocal中.

第二步:引入ThreadLocal

定义ThreadLocal对象:

public class DynamicDataSourceContextHolder {

    /**
     * 动态数据源名称上下文
     */
    private static final ThreadLocal<String> DATASOURCE_CONTEXT_KEY_HOLDER = new ThreadLocal<>();
    /**
     * 设置/切换数据源
     */
    public static void setContextKey(String key){
        DATASOURCE_CONTEXT_KEY_HOLDER.set(key);
    }
    /**
     * 获取数据源名称
     */
    public static String getContextKey(){
        String key = DATASOURCE_CONTEXT_KEY_HOLDER.get();
        return key == null? DataSourceConstants.DB_TEST:key;
    }

    /**
     * 删除当前数据源
     */
    public static void removeContextKey(){
        DATASOURCE_CONTEXT_KEY_HOLDER.remove();
    }
}

很清晰可以看到上面通过ThreadLocal来动态的修改数据源对应的key值,以此来决定某次数据库操作使用的是哪个数据源.至此,一个简单的动态数据源实现就搞定了,接下来可以测试一下.

第三步:测试

@Test
public void testDynamicDataSource() {
    Student student = studentDao.queryById(1);
    System.out.println(student);
    DynamicDataSourceContextHolder.setContextKey(DataSourceConstants.DB_LEARNING);
    System.out.println(userDao.selectById(1));
    DynamicDataSourceContextHolder.removeContextKey();
    DynamicDataSourceContextHolder.setContextKey(DataSourceConstants.DB_TEST);
    System.out.println(studentDao.queryById(1));
    DynamicDataSourceContextHolder.removeContextKey();
}

这样,就可以实现动态数据源了,但是可以很清楚的看到,我们需要在做数据库操作时设置ThreadLocal的值,使用后还要清除值,如果能够尽可能消除这种样板代码就更好了.我们可以引入AOP,并自定义注解来做这件事.

第四步:引入AOP

注解:

@Target({ElementType.METHOD,ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
public @interface DS {
    /**
     * 数据源名称
     */
    String value() default DataSourceConstants.DB_TEST;
}

AOP:

@Aspect
@Component
public class DynamicDataSourceAspect {

    @Pointcut("@annotation(com.dynamic.datasource.annotation.DS)")
    public void dataSourcePointCut() {

    }

    @Around("dataSourcePointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        String dsKey = getDSAnnotation(joinPoint).value();
        DynamicDataSourceContextHolder.setContextKey(dsKey);
        try{
            return joinPoint.proceed();
        }finally {
            DynamicDataSourceContextHolder.removeContextKey();
        }
    }

    /**
     * 根据类或方法获取数据源注解指定的值
     */
    private DS getDSAnnotation(ProceedingJoinPoint joinPoint) {
        Class<?> targetClass = joinPoint.getTarget().getClass();
        DS classAnnotation = targetClass.getAnnotation(DS.class);
        if (classAnnotation != null) {
            return classAnnotation;
        }
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        return methodSignature.getMethod().getAnnotation(DS.class);
    }

}

在Dao层接口的类或方法上添加注解:

@Mapper
public interface StudentDao {
    @DS(DataSourceConstants.DB_TEST)
    Student queryById(Integer id);
}

@Mapper
@DS(DataSourceConstants.DB_LEARNING)
public interface UserDao {
    User selectById(Integer id);
}

第五步:再次测试

@Test
   public void testDynamicDataSourceUseAnnotation() {
       Student student = studentDao.queryById(1);
       System.out.println(student);
       System.out.println(userDao.selectById(1));
       System.out.println(studentDao.queryById(1));
   }

这样基于注解的动态数据源就实现完成了.