spring boot 使用AbstractRoutingDataSource 动态数据源切换

环境基于spring boot 3.0.6 && mybatis-spring:3.0.2

想要实现动态数据源切换，方案很多，很多中间件比如ShardingSphere-jdbc，mycat2 等都可实现，这里使用继承AbstractRoutingDataSource类,重写其 determineCurrentLookupKey()方法来实现， crud基础代码就不贴出来了，小伙子们可以自行发挥。

typescript 复制代码

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceHook.getDataSource();
    }
}

DynamicDataSourceHook 实现如下，利用threadlocal特性解决多线程间通信问题


public class DynamicDataSourceHook {

    private static ThreadLocal<String> threadLocal = new ThreadLocal<>();

    public static String getDataSource() {
        return threadLocal.get();
    }

    public static void setDataSource(String dataSource) {
        threadLocal.remove();
        threadLocal.set(dataSource);
    }

    public static void clearDataSource() {
        threadLocal.remove();
    }
}

测试一下：

可以看到结果是不一样的，当然这里需要去对应的数据库C一下数据。相应的也需要注册对应的datasource bean。如下：当然这里可以随便写下，直接赋值能通就行，应该后面准备优化掉，总不能我加一个数据源就得改代码重新发布吧。

思路：可以参考spring 众多的xxxAutoConfiguration ,利用配置文件自动注册

less 复制代码

@Bean
@ConditionalOnProperty(prefix = "ylli.datasource", value = "enabled", havingValue = "true")
public DataSource masterDataSource() {
    DynamicDataSourceProperties.DataSourceProperties dataSourceProperties = dynamicDataSourceProperties.getDataSource("master");

    return DynamicDataSourceBuilder.custom()
            .driverClassName(dataSourceProperties.getDriverClassName())
            .url(dataSourceProperties.getUrl())
            .username(dataSourceProperties.getUsername())
            .password(dataSourceProperties.getPassword())
            .type(DEFAULT_TYPE)
            .build();
}

@Bean
@ConditionalOnProperty(prefix = "ylli.datasource", value = "enabled", havingValue = "true")
public DataSource slaveDataSource() {
    DynamicDataSourceProperties.DataSourceProperties dataSourceProperties = dynamicDataSourceProperties.getDataSource("slave");
    return DynamicDataSourceBuilder.from(new DynamicDataSourceBuilder(
                    dataSourceProperties.getDriverClassName(),
                    dataSourceProperties.getUrl(),
                    dataSourceProperties.getUsername(),
                    dataSourceProperties.getPassword(),
                    DEFAULT_TYPE))
            .build();
}

优化1: 使用注解来替换 DynamicDataSourceHook.setDataSource(datasource);

话不多说，直接贴代码

less 复制代码

@Target({ElementType.TYPE, ElementType.METHOD})
@Retention(java.lang.annotation.RetentionPolicy.RUNTIME)
@Documented
//遗传，只能作用于类上，被此注解标识的类，其子类也会继承此注解
@Inherited
@Conditional(DataSource.DataSourceCondition.class)
public @interface DataSource {

    String value();

    public class DataSourceCondition implements Condition {
        @Override
        public boolean matches(ConditionContext context, AnnotatedTypeMetadata metadata) {
            String datasourceEnabled = context.getEnvironment().getProperty("ylli.datasource.enabled");
            return StringUtils.hasText(datasourceEnabled) && Boolean.parseBoolean(datasourceEnabled);
        }
    }
}

利用condition来加一个开关控制注解是否生效。

less 复制代码

@Aspect
@Component
public class DataSourceAspect {

    @Pointcut("@annotation(com.ylli.api.base.annotation.DataSource) || @within(com.ylli.api.base.annotation.DataSource)")
    public void guard() {
    }

    @Around("guard()")
    public Object doGuard(ProceedingJoinPoint point) throws Throwable {
        Method method = ((MethodSignature) point.getSignature()).getMethod();
        //先获取方法上的注解
        DataSource dataSource = method.getAnnotation(DataSource.class);
        if (dataSource == null) {
            //得到类上的访问注解
            dataSource = point.getTarget().getClass().getAnnotation(DataSource.class);
        }

        if (Objects.nonNull(dataSource)) {
            DynamicDataSourceHook.setDataSource(dataSource.value());
        }
        try {
            return point.proceed();
        } finally {
            DynamicDataSourceHook.clearDataSource();
        }
    }
}

测试一下效果

可以看到当我们什么都不加的时候会查询默认数据库，那么当我们加上注解呢。

可以看到返回结果是不同的。

优化二：使用配置来动态注册bean.

这里设计一个新配置格式作为spring.datasource.xxx 的补充，而不是取代或兼容，因为这样会增加复杂度。

typescript 复制代码

@ConfigurationProperties(prefix = "ylli.datasource")
public class DynamicDataSourceProperties {

    private boolean enabled;

    //ylli.datasource.multi.name.url;
    //eg.
    //ylli.datasource.multi.master.username=root
    //ylli.datasource.multi.master.url=jdbc:mysql://localhost:3306
    private LinkedHashMap<String, DataSourceProperties> multi;


    public boolean isEnabled() {
        return enabled;
    }

    public void setEnabled(boolean enabled) {
        this.enabled = enabled;
    }

    public LinkedHashMap<String, DataSourceProperties> getMulti() {
        return multi;
    }

    public void setMulti(LinkedHashMap<String, DataSourceProperties> multi) {
        this.multi = multi;
    }

    public DataSourceProperties getDataSource(String name) {
        return multi.get(name);
    }

    @Data
    public static class DataSourceProperties {
        private String url;
        private String username;
        private String password;
        private String driverClassName;
        //默认使用 HikariDataSource.class;
        private Class<? extends DataSource> type;
    }
}

定义一个属性类来接收参数，有其他需求可以在DataSourceProperties 中新增后实现。

自动装配类如下：注册默认数据源 & 注册动态数据源

ini 复制代码

@Configuration
@EnableConfigurationProperties(DynamicDataSourceProperties.class)
@ConditionalOnProperty(prefix = "ylli.datasource", value = "enabled", havingValue = "true")
public class DynamicDataSourceAutoConfig implements ApplicationContextAware {

    private static String DEFAULT_DRIVER_CLASS_NAME;

    private static String DEFAULT_URL;

    private static String DEFAULT_USERNAME;

    private static String DEFAULT_PASSWORD;

    private static Class<? extends DataSource> DEFAULT_TYPE = HikariDataSource.class;

    @Autowired
    DynamicDataSourceProperties properties;
    ApplicationContext applicationContext;

    public DynamicDataSourceAutoConfig(DynamicDataSourceProperties properties) {
        this.properties = properties;
    }

    //default datasource
    @Bean
    //@ConditionalOnMissingBean(DataSource.class)
    @Conditional(Condition0.class)
    public DataSource defaultDataSource() {
        return DynamicDataSourceBuilder.ofDefaults();
    }

    @Bean
    @ConditionalOnBean(name = "defaultDataSource")
    @Primary
    public DynamicDataSource dynamicDataSource() {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        //默认数据源
        dynamicDataSource.setDefaultTargetDataSource(defaultDataSource());

        //数据源库
        dynamicDataSource.setTargetDataSources(register());
        return dynamicDataSource;
    }

    /*
     * 注册数据源
     * 设计为不取代spring.datasource,而是希望作为DLC的形式，对额外的数据源进行补充，降低复杂度。
     */
    @Lazy
    Map<Object, Object> register() {
        DefaultListableBeanFactory defaultListableBeanFactory = (DefaultListableBeanFactory) applicationContext.getAutowireCapableBeanFactory();
        Map<Object, Object> targetDataSources = new HashMap<>();

        if (properties.isEnabled()) {

            properties.getMulti().entrySet().forEach(entry -> {
                BeanDefinitionBuilder beanDefinitionBuilder = BeanDefinitionBuilder
                        //若是其他类型可以替换
                        .genericBeanDefinition(HikariDataSource.class);
//
//                //  如有其他属性，可以在这里添加 && DataSourceProperties 中需定义 & builder 支持
//                if (entry.getValue().getType() != null) {
//                    beanDefinitionBuilder.addPropertyValue("type", entry.getValue().getType());
//                }
                defaultListableBeanFactory.registerBeanDefinition(entry.getKey(), beanDefinitionBuilder.getBeanDefinition());
                HikariDataSource dataSource = (HikariDataSource) applicationContext.getBean(entry.getKey());
                dataSource.setJdbcUrl(entry.getValue().getUrl());
                dataSource.setUsername(entry.getValue().getUsername());
                dataSource.setPassword(entry.getValue().getPassword());
                dataSource.setDriverClassName(entry.getValue().getDriverClassName());

                targetDataSources.put(entry.getKey(), dataSource);
            });

        }
        return targetDataSources;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }


    public static class Condition0 implements Condition {

        @Override
        public boolean matches(ConditionContext context, AnnotatedTypeMetadata metadata) {
            Environment environment = context.getEnvironment();
            if (StringUtils.hasText(DEFAULT_URL = environment.getProperty("spring.datasource.url"))
                    && StringUtils.hasText(DEFAULT_USERNAME = environment.getProperty("spring.datasource.username"))
                    && StringUtils.hasText(DEFAULT_PASSWORD = environment.getProperty("spring.datasource.password"))
                    && StringUtils.hasText(DEFAULT_DRIVER_CLASS_NAME = environment.getProperty("spring.datasource.driverClassName"))
            ) {
                return true;
            }
            return false;
        }
    }

    static class DynamicDataSourceBuilder {
        private String driverClassName;

        private String url;

        private String username;

        private String password;

        private Class<? extends DataSource> type;

        /**
         * 设置初始化信息.
         */
        public DynamicDataSourceBuilder() {
            this.driverClassName = DEFAULT_DRIVER_CLASS_NAME;
            this.url = DEFAULT_URL;
            this.username = DEFAULT_USERNAME;
            this.password = DEFAULT_PASSWORD;
            this.type = DEFAULT_TYPE;
        }

        public DynamicDataSourceBuilder(String driverClassName,
                                        String url,
                                        String username,
                                        String password,
                                        Class<? extends DataSource> type) {
            this.driverClassName = driverClassName;
            this.url = url;
            this.username = username;
            this.password = password;
            this.type = type;
        }

        public DynamicDataSourceBuilder(DynamicDataSourceBuilder builder) {
            this.driverClassName = builder.driverClassName;
            this.url = builder.url;
            this.username = builder.username;
            this.password = builder.password;
            this.type = builder.type;
        }

        //初始化会赋默认值
        public static DynamicDataSourceBuilder custom() {
            return new DynamicDataSourceBuilder();
        }

        public static DynamicDataSourceBuilder from(DynamicDataSourceBuilder builder) {
            return new DynamicDataSourceBuilder(builder);
        }

        public static DataSource ofDefaults() {
            return new DynamicDataSourceBuilder().build();
        }

        public DynamicDataSourceBuilder driverClassName(String driverClassName) {
            this.driverClassName = driverClassName;
            return this;
        }

        public DynamicDataSourceBuilder url(String url) {
            this.url = url;
            return this;
        }

        public DynamicDataSourceBuilder username(String username) {
            this.username = username;
            return this;
        }

        public DynamicDataSourceBuilder password(String password) {
            this.password = password;
            return this;
        }

        public DynamicDataSourceBuilder type(Class<? extends DataSource> type) {
            this.type = type;
            return this;
        }

        /**
         * 返回数据源连接
         */
        public DataSource build() {
            DataSource dataSource = DataSourceBuilder.create()
                    .driverClassName(driverClassName)
                    .url(url)
                    .username(username)
                    .password(password)
                    .type(type)
                    .build();
            return dataSource;
        }
    }

}

配置格式如下

ini 复制代码

spring.datasource.url=jdbc:mysql://127.0.0.1:3306/example?useUnicode=true&useJDBCCompliantTimezoneShift=true&serverTimezone=Asia/Shanghai
spring.datasource.username=root
spring.datasource.password=123456
spring.datasource.driverClassName=com.mysql.cj.jdbc.Driver
spring.datasource.hikari.connection-init-sql=SELECT 1
# --------------------------------------------------------------------------
ylli.datasource.enabled=false
ylli.datasource.multi.slave.url=jdbc:mysql://127.0.0.1:3307/example?useUnicode=true&useJDBCCompliantTimezoneShift=true&serverTimezone=Asia/Shanghai
ylli.datasource.multi.slave.username=root
ylli.datasource.multi.slave.password=123456
ylli.datasource.multi.slave.driverClassName=com.mysql.cj.jdbc.Driver
# --------------------------------------------------------------------------

如此一个简单的自动配置便好了，在不影响原生功能的同时，可以相对达到开箱即用，当然还有很多属性未支持，比如配置池，可以自行添加。

DataSourceConfig 可以有，但没必要。有需求可以加

spring boot 使用AbstractRoutingDataSource 动态数据源切换

优化1: 使用注解来 替换 DynamicDataSourceHook.setDataSource(datasource);

优化二：使用配置来动态注册bean.

优化1: 使用注解来替换 DynamicDataSourceHook.setDataSource(datasource);