Java數據源動態切換實現

1 複寫JDBC框架數據路由服務

package com.datasoure.mybatis;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import javax.annotation.PostConstruct;
import javax.sql.DataSource;

import org.jboss.logging.Logger;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.CollectionUtils;

import com.alibaba.druid.pool.DruidDataSource;
import com.common.CommonConstant;
import com.datasoure.holder.DynamicDataSourceContextHolder;
import com.datasoure.redis.RedisService;
import com.utils.JsonUtils;

public class DynamicDataSource extends AbstractRoutingDataSource {
	private static final Logger LOGGER = Logger.getLogger(DynamicDataSource.class);
	private Map<String, DruidDataSource> dataSourceMap;

	@Autowired
	@Qualifier("dataSourceRedisImpl")
	private RedisService redisService;

	@Override
	protected Object determineCurrentLookupKey() {
		String dataSourceType = DynamicDataSourceContextHolder.getDataSourceType();
		LOGGER.info(String.format("Datasource type is [%s]", dataSourceType));
		return dataSourceType;
	}

	@PostConstruct
	public void preUpdate() {
		updateDataSource();
	}

	public Map<String, DruidDataSource> getDataSourceMap() {
		return dataSourceMap;
	}

	public void updateDataSource() {
		Map<String, String> dataSourceConfigMap = redisService.hgetAll(CommonConstant.REDIS_DATASOURCE_CONFIG_KEY);
		if (dataSourceConfigMap == null || dataSourceConfigMap.isEmpty()) {
			LOGGER.info(String.format("DataSource config in redis is empty, redis key is [%s]",
					CommonConstant.REDIS_DATASOURCE_CONFIG_KEY));
			return;
		}
		Map<Object, Object> targetDataSourceMap = new HashMap<>(dataSourceConfigMap.size());
		dataSourceConfigMap.forEach((x, y) -> {
			DataSourceProperty dataSourceProperty = JsonUtils.fromJson(y, DataSourceProperty.class);
			DataSource dataSource = new DruidDataSource();
			BeanUtils.copyProperties(dataSourceProperty, dataSource);
			targetDataSourceMap.put(x, dataSource);
			if (dataSourceProperty.getDefaultDataSource()) {
				this.setDefaultTargetDataSource(dataSource);
			}
		});
		this.setTargetDataSources(targetDataSourceMap);
		this.afterPropertiesSet();
		this.afterDataSourceSet(targetDataSourceMap);
		LOGGER.info(String.format("DataSource update success, map information is [%s]",
				JsonUtils.toJson(dataSourceConfigMap)));
	}

	private void afterDataSourceSet(Map<Object, Object> targetDataSourceMap) {
		if (!CollectionUtils.isEmpty(dataSourceMap)) {
			// Close the existed dataSource
			Iterator<Map.Entry<String, DruidDataSource>> iterator = dataSourceMap.entrySet().iterator();
			while (iterator.hasNext()) {
				Map.Entry<String, DruidDataSource> entry = iterator.next();
				iterator.remove();
				DruidDataSource dataSource = entry.getValue();
				dataSource.close();
				dataSource = null;
			}
		}
		// Backup the new dataSource
		if (dataSourceMap == null) {
			dataSourceMap = new HashMap<>(targetDataSourceMap.size());
		}
		targetDataSourceMap.forEach((x, y) -> {
			dataSourceMap.put((String) x, (DruidDataSource) y);
		});
	}
}

2 數據源上下文配置

package com.datasoure.holder;

import org.jboss.logging.Logger;

import com.common.CommonConstant;
import com.utils.RequestUtil;

public class DynamicDataSourceContextHolder {
	private static final Logger LOGGER = Logger.getLogger(DynamicDataSourceContextHolder.class);
	private static final ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();

	/**
	 * Set the current user's dataSource
	 */
	public static void setDataSourceType() {
		String dataSourceType = RequestUtil.getCompanyInfo() + CommonConstant.DATASOURCE_KEY_SUFFIX;
		LOGGER.info(String.format("User company dataSourceType is [%s]", dataSourceType));
		CONTEXT_HOLDER.set(dataSourceType);
	}

	/**
	 * if project has async method, you could call this method to set the
	 * dataSourceType manually, which could be passed as parms in the async
	 * method
	 * @param dataSourceType
	 */
	public static void setDataSourceType(String dataSourceType) {
		CONTEXT_HOLDER.set(dataSourceType);
	}

	public static String getDataSourceType() {
		return CONTEXT_HOLDER.get();
	}

	public static void clearDataSourceType() {
		CONTEXT_HOLDER.remove();
	}

}

3 動態數據源配置

package com.datasoure.mybatis;

import java.util.HashMap;
import java.util.Map;

import javax.sql.DataSource;

import org.jboss.logging.Logger;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import com.alibaba.druid.pool.DruidDataSource;
import com.common.CommonConstant;

import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;

@Configuration
public class DynamicDataSourceConfig {

	private static final Logger LOGGER = Logger.getLogger(DynamicDataSourceConfig.class);

	@Value("${spring.redis.host}")
	private String host;

	@Value("${spring.redis.port}")
	private int port;

	@Value("${spring.redis.timeout}")
	private String timeout;

	@Value("${spring.redis.jedis.pool.minIdle}")
	private int minIdle;

	@Value("${spring.redis.jedis.pool.maxIdle}")
	private int maxIdle;

	@Value("${spring.redis.jedis.pool.maxTotal}")
	private int maxTotal;

	@Value("${spring.redis.jedis.pool.maxWaitMillis}")
	private String maxWaitMillis;

	@Value("${spring.redis.password}")
	private String password;

	@Bean(name = CommonConstant.ASSET_DATASOURCE_KEY)
	@ConfigurationProperties(prefix = "spring.datasource.mysql")
	public DataSource build() {
		return DataSourceBuilder.create().type(DruidDataSource.class).build();
	}

	@Bean(name = "dynamicDataSource")
	@Primary
	public DynamicDataSource dynamicDataSourceRouting(
			@Qualifier(CommonConstant.ASSET_DATASOURCE_KEY) DataSource masterDataSource) {
		DynamicDataSource dynamicDataSource = new DynamicDataSource();
		Map<Object, Object> targetDataSource = new HashMap<>();
		targetDataSource.put(CommonConstant.ASSET_DATASOURCE_KEY, masterDataSource);
		dynamicDataSource.setTargetDataSources(targetDataSource);
		dynamicDataSource.setDefaultTargetDataSource(masterDataSource);
		return dynamicDataSource;
	}

	@Bean
	public JedisPool redisPoolFactory() {
		JedisPoolConfig jedisPoolConfig = new JedisPoolConfig();
		jedisPoolConfig.setMinIdle(minIdle);
		jedisPoolConfig.setMaxIdle(maxIdle);
		jedisPoolConfig.setMaxTotal(maxTotal);
		jedisPoolConfig.setMaxWaitMillis(Long.parseLong(maxWaitMillis));
		JedisPool jedisPool = new JedisPool(jedisPoolConfig, host, port, Integer.valueOf(timeout), password);
		LOGGER.info("JedisPool created successfully");
		return jedisPool;
	}
}

4 切面攔截處理

package com.datasoure.aspect;

import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import com.datasoure.holder.DynamicDataSourceContextHolder;


@Aspect
@Component
@Order(1)
public class DynamicDataSourceAspect {
	@Pointcut("execution(* com.controller..*.*(..))")
	public void dyDatasourcePointCut() {
	}

	@Before(value = "dyDatasourcePointCut()")
	public void setDateSource() {
		/*
		 * In async method body, the current thread need to
		 * extend father thread's arguments
		 */
		ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
		RequestContextHolder.setRequestAttributes(sra, true);

		DynamicDataSourceContextHolder.setDataSourceType();
	}

	@After(value = "dyDatasourcePointCut()")
	public void clearDateSource() {
		DynamicDataSourceContextHolder.clearDataSourceType();
	}

}

package com.datasoure.aspect;

import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import com.datasoure.holder.DynamicDataSourceContextHolder;


@Aspect
@Component
@Order(1)
public class AsyncServiceDynamicDataSourceAspect {
	@Pointcut("execution(* com.service.threads.impl..*.*(..))")
	public void asyncDyDatasourcePointCut() {
	}

	/*
	 * Before async method be executed, we should reset the
	 * dataSource which belong to the current user's company
	 */
	@Before(value = "asyncDyDatasourcePointCut()")
	public void setDateSource() {
		DynamicDataSourceContextHolder.setDataSourceType();
	}

	@After(value = "asyncDyDatasourcePointCut()")
	public void clearDateSource() {
		DynamicDataSourceContextHolder.clearDataSourceType();
	}

}

5 輔助類

package com.utils;

import javax.servlet.http.HttpServletRequest;

import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import com.common.CommonConstant;

public final class RequestUtil {
	private RequestUtil() {
	};

	public static HttpServletRequest getRequest() {
		ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
		if (attributes == null) {
			return null;
		}
		return attributes.getRequest();
	}

	public static String getUserId() {
		HttpServletRequest request = getRequest();
		if (request == null) {
			return null;
		}
		return request.getHeader(CommonConstant.USER_ID);
	}

	public static String getCompanyInfo() {
		HttpServletRequest request = getRequest();
		if (request == null) {
			return null;
		}
		return request.getHeader(CommonConstant.COMPANY_ID);
	}
}

package com.datasoure;


import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

import org.springframework.context.annotation.Import;
import org.springframework.scheduling.annotation.EnableScheduling;

import com.datasoure.mybatis.DynamicDataSourceConfig;

@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Import(DynamicDataSourceConfig.class)
@Documented
@EnableScheduling
public @interface EnableDynamicDataSource {
}