仿写@DS 多数据源动态切换

分享干货的你

发布于 2021-05-11 10:34:01

1.7K0

发布于 2021-05-11 10:34:01

文章被收录于专栏：分享干货的你

最近公司在做项目，用到了多数据源，我在网上找了好多的开源项目。

其中 https://github.com/baomidou/dynamic-datasource-spring-boot-starter 这个挺不错的。内容丰富，完善。

我把代码clone 下来了自己仿写了一个乞丐版的。但是主要功能还是有的。

1，先按照标准的yml 格式定义两个配置类

@Component
@ConfigurationProperties(prefix = DynamicDataSourceProperties.PREFIX)
public class DynamicDataSourceProperties {

    public static final String PREFIX = "spring.datasource.dynamic";
    /**
     * 必须设置默认的库,默认master
     */
    private String primary = "master";

    /**
     * 每一个数据源
     */
    private Map<String, DataSourceProperty> datasource = new LinkedHashMap<>();

    private Map<Object, Object> dataSourceMap = new LinkedHashMap<>();

// get  set  省略 

     // 初始化构造数据源集合
    @PostConstruct
    public void initDataSource(){
        for(Map.Entry<String, DataSourceProperty> key:this.datasource.entrySet()) {
            DataSourceProperty property = key.getValue();
            DruidDataSource druidDataSource = new DruidDataSource();
            druidDataSource.setDriverClassName(property.getDriverClassName());
            druidDataSource.setUrl(property.getUrl());
            druidDataSource.setUsername(property.getUsername());
            druidDataSource.setPassword(property.getPassword());
            dataSourceMap.put(key.getKey(),druidDataSource);
        }
    }
 

   // 通过路由切换，动态绑定数据源
    @Bean
    public DataSource dynamicDataSource() {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        
        // 设置主数据源
        dynamicDataSource.setTargetDataSources(getDataSourceMap());
        
        // 设置数据源集合
        dynamicDataSource.setDefaultTargetDataSource(getDataSourceMap().get(primary));
        return dynamicDataSource;
    }
}

public class DataSourceProperty {

    /**
     * JDBC driver
     */
    private String driverClassName;
    /**
     * JDBC url 地址
     */
    private String url;
    /**
     * JDBC 用户名
     */
    private String username;
    /**
     * JDBC 密码
     */
    private String password;
    
    // get set  省略
    }

在创建路由切换的类：

public class DynamicDataSource extends AbstractRoutingDataSource {

    //默认支持两种， dataSource 和String  , key 就是上面代码set 进去Map 的key 
    // 发现当前数据源 
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.peek();
    }
}

数据源设置类：

/**
 * 核心基于ThreadLocal的切换数据源工具类
 *
 * @author TaoYu Kanyuxia
 * @since 1.0.0
 */
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();
    }

自定义注解，可以作用在类和方法上面，数据源嵌套，先进后出，使用最深层的注解数据源。

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

AOP 实现

@Component
@Aspect
@Order(-1) // 优先级最高
public class AspectDs {

  
    @Pointcut("@within(xxxxxx.DS)||@annotation(xxxxxx.DS)")
    public void ds() {

    }

    @Before(value = "ds()")
    public void dataSource(JoinPoint joinPoint) {
        Class cls = joinPoint.getSignature().getDeclaringType();
        // 如果类上面有就使用 类上面的
        DS ds = AnnotationUtils.findAnnotation(cls, DS.class);
        if(ds != null && !"".equals(ds)) {
            DynamicDataSourceContextHolder.push(ds.value());
            System.out.println( cls.getSimpleName() + "类上面的 动态 数据源 :"+ds.value());
        }
        MethodSignature signature = (MethodSignature)joinPoint.getSignature();
        Method method = signature.getMethod();
        DS annotation = AnnotationUtils.findAnnotation(method, DS.class);
        if(annotation != null && !"".equals(annotation)){
            DynamicDataSourceContextHolder.push(ds.value());
            System.out.println(method.getName()+ "方法上面的动态数据源 :"+ds.value());
        }
    }
}

application.yml 配置

spring:
  datasource:
    dynamic:
      primary: master #设置默认的数据源或者数据源组,默认值即为master
      datasource:
        master:
          url: jdbc:mysql://xxxxxxxxxxx:3306/db1
          username: root
          password: xxxxxxxxx
          driver-class-name: com.mysql.jdbc.Driver # 3.2.0开始支持SPI可省略此配置
        slave_1:
          url: jdbc:mysql://xxxxxxxxx:3306/db2
          username: root
          password: xxxxxxxxxxxxxxxx
          driver-class-name: com.mysql.jdbc.Driver

测试：