SpringBoot2.X实现动态数据源

阳光穿透心脏的1/2处 2024-04-18 09:36 84阅读 0赞

## 一、核心原理 ##

动态数据源实现的核心类就是：AbstractRoutingDataSource，在这个类中有五个方法需要特别注意，分别如下：

//设置目标数据源
        public void setTargetDataSources(Map<Object, Object> targetDataSources) {
            this.targetDataSources = targetDataSources;
        }
    
        //设置默认数据源
        public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
            this.defaultTargetDataSource = defaultTargetDataSource;
        }
    
        //选择当前所需数据源
        protected DataSource determineTargetDataSource() {
            Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
            Object lookupKey = this.determineCurrentLookupKey();
            DataSource dataSource = (DataSource)this.resolvedDataSources.get(lookupKey);
            if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
                dataSource = this.resolvedDefaultDataSource;
            }
    
            if (dataSource == null) {
                throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
            } else {
                return dataSource;
            }
        }
    
        @Nullable
        //选择当前数据源所需K
        protected abstract Object determineCurrentLookupKey();
    
        //初始化resolveDataSources和defaultDataSources
        public void afterPropertiesSet() {
            if (this.targetDataSources == null) {
                throw new IllegalArgumentException("Property 'targetDataSources' is required");
            } else {
                this.resolvedDataSources = new HashMap(this.targetDataSources.size());
                this.targetDataSources.forEach((key, value) -> {
                    Object lookupKey = this.resolveSpecifiedLookupKey(key);
                    DataSource dataSource = this.resolveSpecifiedDataSource(value);
                    this.resolvedDataSources.put(lookupKey, dataSource);
                });
                if (this.defaultTargetDataSource != null) {
                    this.resolvedDefaultDataSource = this.resolveSpecifiedDataSource(this.defaultTargetDataSource);
                }
    
            }
        }

从这五个方法可以看出，决定用什么数据源主要由三个条件决定：一、setTargetDataSources方法设置的目标数据源；二、setDefaultTargetDataSource方法设置的默认数据源；三、determineCurrentLookupKey方法返回的key。

## 二、代码实现 ##

### 2.1 数据源配置常量 ###

import com.google.common.collect.Lists;
    
    import java.util.List;
    
    /**
     * @author sunyiran
     * @createTime 2019-09-04 15:37
     */
    public class DataSourceTypeConstant {
        public static final String USERNAME = "username";
        public static final String URL = "url";
        public static final String PASSWORD = "password";
    
        public static final String CONFIG = "config";
    
    
        public static final String ENV_NAME = "spring.datasource.%s.%s";
        public static final String DEFAULT_ENV_NAME = "spring.datasource.%s";
    
    
        public static final String TEST = "test";
    
        public static final String DRIVER = "com.mysql.jdbc.Driver";
    
        public static   List<String> ALL_DATASOURCE = Lists.newArrayList();
    
        static {
            ALL_DATASOURCE.add(TEST);
            ALL_DATASOURCE.add(CONFIG);
    
        }
    
    
    }

### 2.2 数据源上下文 ###

/**
     * 动态数据源上下文
     * @author sunyiran
     * @createTime 2019-09-04 15:43
     */
    public class DynamicDataSourceHolder {
        public static final ThreadLocal<String> DATA_SOURCE_NAME = new ThreadLocal<>();
    
    
        public static void set(String name){
            DATA_SOURCE_NAME.set(name);
        }
    
    
        public static void clear(){
            DATA_SOURCE_NAME.remove();
        }
    
        public static String  get(){
            return DATA_SOURCE_NAME.get();
        }
    }

### 2.3 重写AbstractRoutingDataSource ###

/**
     * @author sunyiran
     * @createTime 2019-09-04 15:42
     */
    public class DynamicDataSource extends AbstractRoutingDataSource  {
    
    
        @Override
        protected Object determineCurrentLookupKey() {
            return DynamicDataSourceHolder.get();
        }
    
    
    }

### 2.4 将重写的数据源放入容器 ###

/**
     * @author sunyiran
     * @createTime 2019-09-04 16:44
     */
    @Configuration
    @Log
    public class DataSourceConfig {
    
        @Resource
        Environment environment;
    
        @Bean
        public DataSource dataSource() {
            DynamicDataSource dynamicDataSource = new DynamicDataSource();
    
            dynamicDataSource.setDefaultTargetDataSource(getDefaultDataSource());
            dynamicDataSource.setTargetDataSources(getDataSourceMap());
    
            return dynamicDataSource;
        }
    
    
        private Object getDefaultDataSource() {
    
            return DataSourceBuilder.create().driverClassName(DataSourceTypeConstant.DRIVER)
                    .password(environment.getProperty(String.format(DataSourceTypeConstant.DEFAULT_ENV_NAME, DataSourceTypeConstant.PASSWORD)))
                    .url(environment.getProperty(String.format(DataSourceTypeConstant.DEFAULT_ENV_NAME, DataSourceTypeConstant.URL)))
                    .username(environment.getProperty(String.format(DataSourceTypeConstant.DEFAULT_ENV_NAME, DataSourceTypeConstant.USERNAME))).build();
        }
    
        private  Map<Object, Object> getDataSourceMap() {
            Map<Object, Object> dataSourceMap = Maps.newHashMap();
    
            List<String> allDatasource = DataSourceTypeConstant.ALL_DATASOURCE;
    
            DataSourceBuilder<?> sourceBuilder = DataSourceBuilder.create().driverClassName(DataSourceTypeConstant.DRIVER);
    
            for (String name : allDatasource) {
                DataSource dataSource = sourceBuilder.password(environment.getProperty(String.format(DataSourceTypeConstant.ENV_NAME, name, DataSourceTypeConstant.PASSWORD)))
                        .url(environment.getProperty(String.format(DataSourceTypeConstant.ENV_NAME, name, DataSourceTypeConstant.URL)))
                        .username(environment.getProperty(String.format(DataSourceTypeConstant.ENV_NAME, name, DataSourceTypeConstant.USERNAME))).build();
                dataSourceMap.put(name, dataSource);
            }
            return dataSourceMap;
        }
    }

### 2.5 自定义注解 ###

/**
     * @author sunyiran
     * @dateTime 2019-09-04 16:53
     */
    @Documented
    @Inherited
    @Target({ElementType.METHOD,ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface DataSourceType {
        String name() ;
    }

### 2.6 注解切面环绕处理 ###

/**
     * @author sunyiran
     * @createTime 2019-09-04 16:58
     */
    @Aspect
    @Component
    @Order(-1)
    @Log4j2
    public class DataSourceTypeHandle {
    
    
        @Around("execution(* com.litb.business.base.service.impl.*.*(..))")
        @Order(1)
        public Object handleType(ProceedingJoinPoint jp)   {
    
            //同时支持类和方法上的注解，以方法上为主
            DataSourceType typeInClass = jp.getTarget().getClass().getAnnotation(DataSourceType.class);
            MethodSignature methodSignature = (MethodSignature) jp.getSignature();
            DataSourceType typeInMethod = methodSignature.getMethod().getDeclaredAnnotation(DataSourceType.class);
    
            String  name;
    
            if (typeInMethod != null) {
                name = typeInMethod.name();
            } else if (typeInClass != null) {
                name = typeInClass.name();
            } else {
                return jp.proceed();
            }
    
            return dealWithTargetDS(jp,name);
        }
    
    
        private Object dealWithTargetDS(ProceedingJoinPoint jp,String name) {
            if (DataSourceTypeConstant.ALL_DATASOURCE.contains(name)) {
                log.info("开始切换数据源，方法{}，数据源{}",jp.getSignature().getName(),name);
                DynamicDataSourceHolder.set(name);
            } else {
                log.error("方法{}，切换非法数据源{}",jp.getSignature().getName(),name);
            }
    
            Object proceed ;
            try {
                proceed = jp.proceed();
            } catch (Throwable throwable) {
                throw new ApiException(BasicResult.ERROR);
            }finally {
    
                DynamicDataSourceHolder.clear();
            }
    
    
         
    
            return proceed;
        }
    }

### 2.7 测试用例 ###

/**
     * @author sunyiran
     * @createTime 2019-09-04 17:27
     */
    @RestController
    @Log4j2
    public class Test {
    
        @Resource
        CityService cityService;
    
        @RequestMapping("/t1")
        @DataSourceType(name = "test")
        public Object test1() {
            log.info("t1----执行------需要test数据源");
            City byId = cityService.findById(1L);
            System.out.println(byId);
            return byId;
        }
    
        @RequestMapping("/t2")
        @DataSourceType(name = "config")
        public Object test2() {
            log.info("t2----执行------需要config数据源");
            City byId = cityService.findById(1L);
            System.out.println(byId);
            return byId;
        }
    
        @RequestMapping("/t3")
        public Object test3() {
            log.info("t3----执行------需要默认数据源");
            City byId = cityService.findById(1L);
            System.out.println(byId);
            return byId;
        }
    }

结果打印：

c.l.b.base.anotation.handle.DataSourceTypeHandle[30] 开始切换数据源，方法test1，数据源test
    com.litb.business.base.controller.test.Test[26] t1----执行------需要test数据源
    City(cityId=1, stateId=0, countryId=0, cityNameEn=resr, cityNameZh=test, cityIsoCode2=1, cityIsoCode3=1, defaultLanguagesId=1, objectsStatusId=1, dateAdded=Thu Sep 09 15:00:00 
    
    
    c.l.b.base.anotation.handle.DataSourceTypeHandle[30] 开始切换数据源，方法test2，数据源config
    com.litb.business.base.controller.test.Test[35] t2----执行------需要config数据源
    City(cityId=1, stateId=1, countryId=223, cityNameEn=ADDISON, cityNameZh=, cityIsoCode2=, cityIsoCode3=, defaultLanguagesId=1, objectsStatusId=1, dateAdded=Thu Sep 09 15:00:00 CST 
    
    com.litb.business.base.controller.test.Test[43] t3----执行------需要默认数据源
    City(cityId=1, stateId=1, countryId=223, cityNameEn=ADDISON, cityNameZh=, cityIsoCode2=, cityIsoCode3=, defaultLanguagesId=1, objectsStatusId=1, dateAdded=Thu Sep 09 15:00:00 CST 1999, lastUpdate=Wed Dec 13 15:50:59 CST 2017)