Implementation of Spring for dynamic registration of multiple data sources

Recently, SaaS is being applied, and the database adopts the single-instance multi-schema architecture (see Reference 1 for details). Each tenant has an independent schema and the whole data source has a shared schema. Therefore, problems of dynamically adding, deleting and switching data sources need to be solved.

After searching a lot of articles on the Internet, most of them are about master-slave data source configuration, or the data source configuration has been determined before the application is launched. It is seldom about how to dynamically load data source without downtime, so I write this article for your reference.

The techniques used

Java8 Spring + SpringMVC + MyBatis Druid connection pool Lombok (The above techniques do not affect the train of thought implementation, just for the convenience of viewing the following code snippet)

Train of thought

When a request comes in, determine the current user belongs to the tenant, switch to the corresponding data source according to the tenant information, and then conduct subsequent business operations.

Code implementation

TenantConfigEntity (Tenant information) 
@EqualsAndHashCode(callSuper = false)
@Data
@FieldDefaults(level = AccessLevel.PRIVATE)
public class TenantConfigEntity {
 /**
  *  The tenant id
  **/
 Integer tenantId;
 /**
  *  The tenant name 
  **/
 String tenantName;
 /**
  *  The tenant name key
  **/
 String tenantKey;
 /**
  *  The database url
  **/
 String dbUrl;
 /**
  *  Database user name 
  **/
 String dbUser;
 /**
  *  Database password 
  **/
 String dbPassword;
 /**
  *  The database public_key
  **/
 String dbPublicKey;
}
DataSourceUtil (Auxiliary tools, not necessary) 
public class DataSourceUtil {
 private static final String DATA_SOURCE_BEAN_KEY_SUFFIX = "_data_source";
 private static final String JDBC_URL_ARGS = "?useUnicode=true&characterEncoding=UTF-8&useOldAliasMetadataBehavior=true&zeroDateTimeBehavior=convertToNull";
 private static final String CONNECTION_PROPERTIES = "config.decrypt=true;config.decrypt.key=";
 /**
  *  Spliced data source spring bean key
  */
 public static String getDataSourceBeanKey(String tenantKey) {
  if (!StringUtils.hasText(tenantKey)) {
   return null;
  }
  return tenantKey + DATA_SOURCE_BEAN_KEY_SUFFIX;
 }
 /**
  *  Pieced together JDBC URL
  */
 public static String getJDBCUrl(String baseUrl) {
  if (!StringUtils.hasText(baseUrl)) {
   return null;
  }
  return baseUrl + JDBC_URL_ARGS;
 }
 /**
  *  Pieced together Druid Connection properties 
  */
 public static String getConnectionProperties(String publicKey) {
  if (!StringUtils.hasText(publicKey)) {
   return null;
  }
  return CONNECTION_PROPERTIES + publicKey;
 }
}

DataSourceContextHolder

Use ThreadLocal to save the data source of the current thread, key name, and implement set, get, clear methods;

public class DataSourceContextHolder {
 private static final ThreadLocal<String> dataSourceKey = new InheritableThreadLocal<>();
 public static void setDataSourceKey(String tenantKey) {
  dataSourceKey.set(tenantKey);
 }
 public static String getDataSourceKey() {
  return dataSourceKey.get();
 }
 public static void clearDataSourceKey() {
  dataSourceKey.remove();
 }
}

DynamicDataSource (Key)

Inherit AbstractRoutingDataSource (it is recommended to read the source code of AbstractRoutingDataSource to understand the process of dynamic switching of data sources) and realize the dynamic selection of data sources;

public class DynamicDataSource extends AbstractRoutingDataSource {
 @Autowired
 private ApplicationContext applicationContext;
 @Lazy
 @Autowired
 private DynamicDataSourceSummoner summoner;
 @Lazy
 @Autowired
 private TenantConfigDAO tenantConfigDAO;
 @Override
 protected String determineCurrentLookupKey() {
  String tenantKey = DataSourceContextHolder.getDataSourceKey();
  return DataSourceUtil.getDataSourceBeanKey(tenantKey);
 }
 @Override
 protected DataSource determineTargetDataSource() {
  String tenantKey = DataSourceContextHolder.getDataSourceKey();
  String beanKey = DataSourceUtil.getDataSourceBeanKey(tenantKey);
  if (!StringUtils.hasText(tenantKey) || applicationContext.containsBean(beanKey)) {
   return super.determineTargetDataSource();
  }
  if (tenantConfigDAO.exist(tenantKey)) {
   summoner.registerDynamicDataSources();
  }
  return super.determineTargetDataSource();
 }
}

DynamicDataSourceSummoner (highlights of highlights)

Load the data source information from the database and dynamically assemble and register spring bean,

@Slf4j
@Component
public class DynamicDataSourceSummoner implements ApplicationListener<ContextRefreshedEvent> {
 //  with spring-data-source.xml Default data source id keep 1 to 
 private static final String DEFAULT_DATA_SOURCE_BEAN_KEY = "defaultDataSource";
 @Autowired
 private ConfigurableApplicationContext applicationContext;
 @Autowired
 private DynamicDataSource dynamicDataSource;
 @Autowired
 private TenantConfigDAO tenantConfigDAO;
 private static boolean loaded = false;
 /**
  * Spring Execute after loading 
  */
 @Override
 public void onApplicationEvent(ContextRefreshedEvent event) {
  //  Prevent duplicate execution 
  if (!loaded) {
   loaded = true;
   try {
    registerDynamicDataSources();
   } catch (Exception e) {
    log.error(" Data source initialization failed , Exception:", e);
   }
  }
 }
 /**
  *  Reads the tenant's from the database DB configuration , Concurrent dynamic injection Spring The container 
  */
 public void registerDynamicDataSources() {
  //  Get all of the tenants DB configuration 
  List<TenantConfigEntity> tenantConfigEntities = tenantConfigDAO.listAll();
  if (CollectionUtils.isEmpty(tenantConfigEntities)) {
   throw new IllegalStateException(" Application initialization failed , Configure the data source first ");
  }
  //  The data source bean Register into the container 
  addDataSourceBeans(tenantConfigEntities);
 }
 /**
  *  According to the DataSource create bean And register it in the container 
  */
 private void addDataSourceBeans(List<TenantConfigEntity> tenantConfigEntities) {
  Map<Object, Object> targetDataSources = Maps.newLinkedHashMap();
  DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) applicationContext.getAutowireCapableBeanFactory();
  for (TenantConfigEntity entity : tenantConfigEntities) {
   String beanKey = DataSourceUtil.getDataSourceBeanKey(entity.getTenantKey());
   //  If the data source is already in spring It's registered , Do not re-register 
   if (applicationContext.containsBean(beanKey)) {
    DruidDataSource existsDataSource = applicationContext.getBean(beanKey, DruidDataSource.class);
    if (isSameDataSource(existsDataSource, entity)) {
     continue;
    }
   }
   //  The assembly bean
   AbstractBeanDefinition beanDefinition = getBeanDefinition(entity, beanKey);
   //  registered bean
   beanFactory.registerBeanDefinition(beanKey, beanDefinition);
   //  In the map , pay attention to 1 It must have been created just now bean object 
   targetDataSources.put(beanKey, applicationContext.getBean(beanKey));
  }
  //  Will create the map object set to  targetDataSources ; 
  dynamicDataSource.setTargetDataSources(targetDataSources);
  //  This operation must be performed before it can be reinitialized AbstractRoutingDataSource  In the  resolvedDataSources And only then will dynamic switching work 
  dynamicDataSource.afterPropertiesSet();
 }
 /**
  *  Assemble data source spring bean
  */
 private AbstractBeanDefinition getBeanDefinition(TenantConfigEntity entity, String beanKey) {
  BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(DruidDataSource.class);
  builder.getBeanDefinition().setAttribute("id", beanKey);
  //  Other configuration inheritance defaultDataSource
  builder.setParentName(DEFAULT_DATA_SOURCE_BEAN_KEY);
  builder.setInitMethodName("init");
  builder.setDestroyMethodName("close");
  builder.addPropertyValue("name", beanKey);
  builder.addPropertyValue("url", DataSourceUtil.getJDBCUrl(entity.getDbUrl()));
  builder.addPropertyValue("username", entity.getDbUser());
  builder.addPropertyValue("password", entity.getDbPassword());
  builder.addPropertyValue("connectionProperties", DataSourceUtil.getConnectionProperties(entity.getDbPublicKey()));
  return builder.getBeanDefinition();
 }
 /**
  *  judge Spring Inside the container DataSource And database DataSource Whether the information 1 to 
  *  note : There is no judgment here public_key, Because the other 3 The basic information can be determined only 1 the 
  */
 private boolean isSameDataSource(DruidDataSource existsDataSource, TenantConfigEntity entity) {
  boolean sameUrl = Objects.equals(existsDataSource.getUrl(), DataSourceUtil.getJDBCUrl(entity.getDbUrl()));
  if (!sameUrl) {
   return false;
  }
  boolean sameUser = Objects.equals(existsDataSource.getUsername(), entity.getDbUser());
  if (!sameUser) {
   return false;
  }
  try {
   String decryptPassword = ConfigTools.decrypt(entity.getDbPublicKey(), entity.getDbPassword());
   return Objects.equals(existsDataSource.getPassword(), decryptPassword);
  } catch (Exception e) {
   log.error(" Data source password validation failed ,Exception:{}", e);
   return false;
  }
 }
}

spring-data-source.xml

<!--  The introduction of jdbc The configuration file  -->
 <context:property-placeholder location="classpath:data.properties" ignore-unresolvable="true"/>
 <!--  public ( The default ) The data source  -->
 <bean id="defaultDataSource" class="com.alibaba.druid.pool.DruidDataSource"
   init-method="init" destroy-method="close">
  <!--  Basic attributes  url , user , password -->
  <property name="url" value="${ds.jdbcUrl}" />
  <property name="username" value="${ds.user}" />
  <property name="password" value="${ds.password}" />
  <!--  Configure the initialization size, minimum, and maximum  -->
  <property name="initialSize" value="5" />
  <property name="minIdle" value="2" />
  <property name="maxActive" value="10" />
  <!--  Configure the time in milliseconds to get the connection wait timeout  -->
  <property name="maxWait" value="1000" />
  <!--  How often does the configuration take place 1 Detect the idle connection that needs to be closed in milliseconds  -->
  <property name="timeBetweenEvictionRunsMillis" value="5000" />
  <!--  configuration 1 The minimum number of milliseconds that a connection can live in the pool  -->
  <property name="minEvictableIdleTimeMillis" value="240000" />
  <property name="validationQuery" value="SELECT 1" />
  <!-- Unit: Seconds, the timeout to check whether the connection is valid -->
  <property name="validationQueryTimeout" value="60" />
  <!-- Recommended configuration is true , does not affect performance, and guarantees security. Detect when applying for connection if idle time is greater than timeBetweenEvictionRunsMillis , the implementation of validationQuery Checks if the connection is valid -->
  <property name="testWhileIdle" value="true" />
  <!-- Execute when applying for connection validationQuery Check to see if the connection is valid. Doing this configuration can degrade performance. -->
  <property name="testOnBorrow" value="true" />
  <!-- Performed when the connection is returned validationQuery Check to see if the connection is valid. Doing this configuration can degrade performance. -->
  <property name="testOnReturn" value="false" />
  <!--Config Filter-->
  <property name="filters" value="config" />
  <property name="connectionProperties" value="config.decrypt=true;config.decrypt.key=${ds.publickey}" />
 </bean>
 <!--  Transaction manager  -->
 <bean id="txManager" class="org.springframework.jdbc.datasource.DataSourceTransactionManager">
  <property name="dataSource" ref="multipleDataSource"/>
 </bean>
 <!-- Multiple data sources -->
 <bean id="multipleDataSource" class="a.b.c.DynamicDataSource">
  <property name="defaultTargetDataSource" ref="defaultDataSource"/>
  <property name="targetDataSources">
   <map>
    <entry key="defaultDataSource" value-ref="defaultDataSource"/>
   </map>
  </property>
 </bean>
 <!--  Annotate the transaction manager  -->
 <!-- Here, order Must be greater than DynamicDataSourceAspectAdvice the order value -->
 <tx:annotation-driven transaction-manager="txManager" order="2"/>
 <!--  create SqlSessionFactory , while specifying the data source  -->
 <bean id="mainSqlSessionFactory" class="org.mybatis.spring.SqlSessionFactoryBean">
  <property name="dataSource" ref="multipleDataSource"/>
 </bean>
 <!-- DAO The package name of the interface, Spring It will automatically find what's under it DAO -->
 <bean id="mainSqlMapper" class="org.mybatis.spring.mapper.MapperScannerConfigurer">
  <property name="sqlSessionFactoryBeanName" value="mainSqlSessionFactory"/>
  <property name="basePackage" value="a.b.c.*.dao"/>
 </bean>
 <bean id="defaultSqlSessionFactory" class="org.mybatis.spring.SqlSessionFactoryBean">
  <property name="dataSource" ref="defaultDataSource"/>
 </bean>
 <bean id="defaultSqlMapper" class="org.mybatis.spring.mapper.MapperScannerConfigurer">
  <property name="sqlSessionFactoryBeanName" value="defaultSqlSessionFactory"/>
  <property name="basePackage" value="a.b.c.base.dal.dao"/>
 </bean>
 <!--  Other configuration omissions  -->

DynamicDataSourceAspectAdvice

Use AOP to automatically switch data sources, for reference only;

@Slf4j
@Aspect
@Component
@Order(1) //  Please note: here order1 To be less than tx:annotation-driven the order , that is, execute first DynamicDataSourceAspectAdvice Section, and then execute the transaction section to obtain the final data source 
@EnableAspectJAutoProxy(proxyTargetClass = true)
public class DynamicDataSourceAspectAdvice {
 @Around("execution(* a.b.c.*.controller.*.*(..))")
 public Object doAround(ProceedingJoinPoint jp) throws Throwable {
  ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
  HttpServletRequest request = sra.getRequest();
  HttpServletResponse response = sra.getResponse();
  String tenantKey = request.getHeader("tenant");
  //  The front end must pass in tenant header,  Otherwise returns 400
  if (!StringUtils.hasText(tenantKey)) {
   WebUtils.toHttp(response).sendError(HttpServletResponse.SC_BAD_REQUEST);
   return null;
  }
  log.info(" The current tenant key:{}", tenantKey);
  DataSourceContextHolder.setDataSourceKey(tenantKey);
  Object result = jp.proceed();
  DataSourceContextHolder.clearDataSourceKey();
  return result;
 }
}

conclusion