实现不同数据源动态切换SpringBoot+MyBatis

1,074 阅读2分钟

在实际的业务业务场景中,经常有不同的request请求,需要使用不同的DB数据源。比如此时有请求1需要访问数据库DB1,请求2需要访问数据库DB2,而我们一般的项目中都是固定数据源的,这样的场景就满足不了。这篇文章就是给这种场景尝试提供一种解决方案。

img 闲话少说开肝。

整体架构使用 SpringBoot + Mybatis +Mysql实现。

1、数据源

先给定两个数据源,定义对应的DataSource的Bean。

数据源1:master

/**
    * 主库数据源
    * @return
    */
   @Bean
   @ConfigurationProperties("spring.datasource.master")
   public DataSource master() {
       return DataSourceBuilder.create().build();
  }

数据源2:slave

/**
    * 从库数据源
    * @return
    */
   @Bean
   @ConfigurationProperties("spring.datasource.slave")
   public DataSource slave() {
       return DataSourceBuilder.create().build();
  }

此时需要一个路由数据源,来作为中间层,此时将master和slave数据源传入

/**
    * 路由数据源
    * @param master
    * @param slave
    * @return
    */
   @Bean(name = "DBSource")
   public DataSource dataSourceRoutingDb(@Qualifier("master") DataSource master,
                                         @Qualifier("slave") DataSource slave) {
       Map<Object, Object> targetDataSource = new HashMap<>();
       targetDataSource.put(DBTypeEnum.MASTER, master());
       targetDataSource.put(DBTypeEnum.SLAVE, slave());
       RoutingDataSource dataSourceRoutingDb = new RoutingDataSource();
       dataSourceRoutingDb.setDefaultTargetDataSource(master);
       dataSourceRoutingDb.setTargetDataSources(targetDataSource);

       return dataSourceRoutingDb;
  }

2、数据源上下文设置

public class DBContextHolder {
   private final static ThreadLocal<String> contextHolder = new ThreadLocal<>();

   /**
    * 设置数据源
    * @param dbType
    */
   public static void set(String dbType) {
       contextHolder.set(dbType);
  }

   /**
    * 设置数据源
    * @return
    */
   public static String get() {
       return contextHolder.get();
  }

   /**
    * 清除ThreadLocal中的上下文
    */
   public static void clear() {
       contextHolder.remove();
  }
}

使用ThreadLocal将传入的数据源设置到每个线程的上下文中。

3、数据源路由

这里需要使用到Spring的AbstractRoutingDataSource中的determineCurrentLookupKey方法。

首先看下AbstractRoutingDataSource类结构,继承了AbstractDataSource

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean

既然是AbstractDataSource,当然就是javax.sql.DataSource的子类,于是我们自然地回去看它的getConnection方法:

public Connection getConnection() throws SQLException {  
       return determineTargetDataSource().getConnection();  
}  
public Connection getConnection(String username, String password) throws SQLException {  
   return determineTargetDataSource().getConnection(username, password);  
}

关键就在determineTargetDataSource()里:

/**
    * Retrieve the current target DataSource. Determines the
    * {@link #determineCurrentLookupKey() current lookup key}, performs
    * a lookup in the {@link #setTargetDataSources targetDataSources} map,
    * falls back to the specified
    * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
    * @see #determineCurrentLookupKey()
    */  
   protected DataSource determineTargetDataSource() {  
       Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");  
       Object lookupKey = determineCurrentLookupKey();  
       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 + "]");  
      }  
       return dataSource;  
  }

这里用到了我们需要进行实现的抽象方法determineCurrentLookupKey(),该方法返回需要使用的DataSource的key值,然后根据这个key从resolvedDataSources这个map里取出对应的DataSource,如果找不到,则用默认的resolvedDefaultDataSource。

public void afterPropertiesSet() {  
       if (this.targetDataSources == null) {  
           throw new IllegalArgumentException("Property 'targetDataSources' is required");  
      }  
       this.resolvedDataSources = new HashMap<Object, DataSource>(this.targetDataSources.size());  
       for (Map.Entry entry : this.targetDataSources.entrySet()) {  
           Object lookupKey = resolveSpecifiedLookupKey(entry.getKey());  
           DataSource dataSource = resolveSpecifiedDataSource(entry.getValue());  
           this.resolvedDataSources.put(lookupKey, dataSource);  
      }  
       if (this.defaultTargetDataSource != null) {  
           this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);  
      }  
  }

直接上代码:

取传入的数据源设置

public class RoutingDataSource extends AbstractRoutingDataSource {
   @Override
   protected Object determineCurrentLookupKey() {
       return DBContextHolder.get();
  }
}

4、Mybatis的设置

这里需要给mybatis重新设置SqlSessionFactory数据源,设置事务PlatformTransactionManager数据源

@Configuration
@EnableTransactionManagement
public class MyBatisConfig {

   @Resource(name = "DBSource")
   private DataSource routingDataSource;

   /**
    * SqlSessionFactory的数据源
    * @return
    * @throws Exception
    */
   @Bean
   public SqlSessionFactory sqlSessionFactory() throws Exception {
       SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
       sqlSessionFactoryBean.setDataSource(routingDataSource);
       //可能出错的问题:sqlSessionFactoryBean.setMapperLocations(resolver.getResources(packageSearchPath)),写错成了getResource导致找不到classpath:mapper/*.xml的错误
       sqlSessionFactoryBean.setMapperLocations(new PathMatchingResourcePatternResolver().getResources("classpath:mapper/*.xml"));
       return sqlSessionFactoryBean.getObject();
  }

   /**
    * 设置事务
    * @return
    */
   @Bean
   public PlatformTransactionManager transactionManager() {
       return new DataSourceTransactionManager(routingDataSource);
  }

}

5、使用注解用于数据源切换

设置一个默认的数据源

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD, ElementType.TYPE})
@Documented
public @interface DataSourceSwitcher {

   /**
    * 默认数据源
    * @return
    */
   String value() default DBTypeEnum.MASTER;

   /**
    * 默认清除ThreadLocal中上下文
    * @return
    */
   boolean clear() default true;
}

6、AOP

根据注解传值,将数据源设置到上下文中。

@Aspect
@Component
public class DataSourceContextAop {
   private final static Logger LOGGER = LoggerFactory.getLogger(DataSourceContextAop.class);

   @Around("@annotation(com.lyu.mms.annotation.DataSourceSwitcher)")
   public Object setDataSource(ProceedingJoinPoint point) {
       boolean clear = false;
       try {
           MethodSignature signature = (MethodSignature) point.getSignature();
           Method method = signature.getMethod();
           DataSourceSwitcher dataSourceSwitcher = method.getAnnotation(DataSourceSwitcher.class);
           clear = dataSourceSwitcher.clear();
           DBContextHolder.set(dataSourceSwitcher.value());
           LOGGER.info("切换数据源至->{}", dataSourceSwitcher.value());
           return point.proceed();
      } catch (Throwable throwable) {
           throwable.printStackTrace();
      } finally {
           if (clear) {
               DBContextHolder.clear();
          }
      }
       return null;
  }
}

7、使用

直接在对应的controller上加上相应的注解就愉快的用起来了

/**
    * 保存
    * @param user
    */
   @DataSourceSwitcher(DBTypeEnum.MASTER)
   //@DataSourceSwitcher(DBTypeEnum.SLAVE)
   @Override
   @PutMapping(name = "保存", value = {"/add"})
   public void save(User user) {
       userMapper.insert(user);
       LOGGER.info("新增成功,{}", user);
  }

8、总结

整体借助springboot和mybatis实现数据源切换,关键点使用determineCurrentLookupKey路由数据源,然后为mybatis的SqlSessionFactory和PlatformTransactionManager事务重新设置数据源。

如果需要代码可以联系我获取。

一键三连哦。

我是程序员一二,see you!