利用 ThreadLocal 和 AbstractRoutingDataSource 实现基于注解的动态数据源切换

https://blog.csdn.net/qq_43170312/article/details/124693021

1.数据源配置

DataSource中配置了默认的数据源，如果没有在方法上使用注解，则使用默认数据源

@Configuration
public class DynamicDataSourceConfig {

    @Bean
    @Primary
    public DataSource dynamicDataSource() {
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put(DataSourceConstants.DS_KEY_SLAVE_AUTH, slaveDataSourceAuth());
        dataSourceMap.put(DataSourceConstants.DS_KEY_SLAVE_DATA, slaveDataSourceData());
        //设置动态数据源
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        dynamicDataSource.setTargetDataSources(dataSourceMap);
        // 单点登录系统，设置默认数据源
        dynamicDataSource.setDefaultTargetDataSource(slaveDataSourceAuth());

        return dynamicDataSource;
    }

    /**
     * @return 后勤管理数据源bean
     */
    @Bean("DS_KEY_SLAVE_AUTH")
    @ConfigurationProperties(prefix = "spring.datasource.slave.authserver")
    public DataSource slaveDataSourceAuth() {
        return DataSourceBuilder.create().build();
    }

    /**
     * @return 文件系统数据源bean
     */
    @Bean("DS_KEY_SLAVE_DATA")
    @ConfigurationProperties(prefix = "spring.datasource.slave.data")
    public DataSource slaveDataSourceData() {
        return DataSourceBuilder.create().build();
    }
}

public interface DataSourceConstants {

    /* 从数据库-数据中心 */
    String DS_KEY_SLAVE_DATA = "DS_KEY_SLAVE_DATA";

    /* 从数据库-单点系统 */
    String DS_KEY_SLAVE_AUTH = "DS_KEY_SLAVE_AUTH";
}

2.DynamicDataSource继承AbstractRoutingDataSource

会调用determineCurrentLookupKey方法，从ThreadLocal中获取数据源信息

public class DynamicDataSource extends AbstractRoutingDataSource {
    
    @Override
    protected Object determineCurrentLookupKey() {
        // 当设置数据源 key 到上下文，则从上下文中得到此数据源 key
        return DynamicDataSourceContextHolder.getContextKey();
    }
}

/* 动态数据源数据存储 */
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.DS_KEY_SLAVE_AUTH : key;
    }

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

3.切换数据源的注解

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

4.AOP切面处理

从注解中拿到数据源信息，方法执行前把信息set到ThreadLocal中。线程执行完后

Aspect
@Component
public class DynamicDataSourceAspect {

    @Pointcut("@annotation(com.jnx.smart.base.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 dsAnnotation = targetClass.getAnnotation(DS.class);
        // 先判断类的注解，再判断方法注解
        if (Objects.nonNull(dsAnnotation)) {
            return dsAnnotation;
        } else {
            MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
            return methodSignature.getMethod().getAnnotation(DS.class);
        }
    }
}