Spring整合Sharding实现查询走指定数据库

一、项目背景

• 数据库采用MySQL一主多从，使用Sharding实现主从读写分离。由于大数据报表，MySQL查询很慢，所以报表的查询走列式数据库。现在需要指定某些查询走列数据库，某些走主库，其他查询实现MySQL从库负载均衡。

二、实现思路

1. Spring整合Sharding实现主从读写分离。
2. 在Mapper上添加自定义注解，注解用于标识哪些方法走哪些库
3. 自定义Mybatis拦截器拦截所有查询请求，获取请求的Mapper以及具体方法，获取Mapper上的注解根据方法名称找指定的数据库，将数据库标识设置到线程副本中
4. 自定义Sharding负载均衡策略，根据线程副本获取指定数据库标识，将该数据库返回，从而实现查询走指定数据库
5. 配合nacos配置中心，实现动态开启关闭使用列数据库查询。

三、Spring整合Sharding，配置主从数据库

• 网上一搜一大把，这里不再赘述

四、自定义注解

• MasterAnotation 标注的方法，表示查询走主库

@Target({ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface MasterAnotation {
	
	String[] methods() ;

}

• ColumnarAnotation 标注的方法，表示查询列数据库

@Target({ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface ColumnarAnotation {
	
	String[] methods() ;

}

五、添加MyBatis拦截器，拦截所有查询

• ShardingInterceptor 用于拦截所有查询请求
• 拦截器中注入的ShardingUtil对象，用于获取具体Mapper以及方法，根据注解将数据库标识设置到线程副本中，查询结束清空线程副本。

@Intercepts({
		@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class,RowBounds.class, ResultHandler.class }),
		@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class,RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class }),
 })
public class ShardingInterceptor implements Interceptor {
	
	private ShardingUtil shardingUtil;
	
	public ShardingInterceptor(ShardingUtil shardingUtil){
		this.shardingUtil=shardingUtil;
	}
	
	
	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		Object[] args = invocation.getArgs();
		MappedStatement ms=(MappedStatement) args[0];
		//查询之前做一些处理
		shardingUtil.beforeStarRocksProcess(ms);
		Object obj = invocation.proceed();
		//查询之后做一些处理
		shardingUtil.afterStarRocksProcess();
		return obj;
	}
	
	@Override
	public Object plugin(Object target) {
		if (target instanceof Executor) {
			return Plugin.wrap(target, this);
		}
		return target;
	}

	@Override
	public void setProperties(Properties properties) {
		// TODO Auto-generated method stub

	}

}

六、创建线程副本管理器

public class ShardingManager {
	
	
	/**
	 * 存储当前线程的数据库信息
	 */
	private static final ThreadLocal<String> SLAVE_INFO = new ThreadLocal<>();
	
	public static String getSlave(){
		return SLAVE_INFO.get();
	}
	
	public static void setSlave(String slaveName){
		SLAVE_INFO.set(slaveName);
	}
	
	public static void clear(){
		SLAVE_INFO.remove();
	}

}

七、自定义负载均衡策略

• 自定义负载均衡策略，根据线程副本获取当前线程的数据库标识，查询则返回该数据库

@Getter
@Setter
public  class ShardingMasterSlaveLoadBalanceAlgorithm implements MasterSlaveLoadBalanceAlgorithm {

	private static final ConcurrentHashMap<String, AtomicInteger> COUNTS = new ConcurrentHashMap<>();

	private Properties properties = new Properties();

	@Override
	public String getType() {
		return "CUSTOM";
	}

	@Override
	public String getDataSource(final String name, final String masterDataSourceName,final List<String> slaveDataSourceNames) {
		//获取线程副本
		String slaveName=Optional.ofNullable(ShardingManager.getSlave()).orElse(DBConstant.SLAVE);
		//判断是否走主数据库
		if(slaveName.equals(masterDataSourceName)){
			return masterDataSourceName;
		}
		//判断是否走列数据库
		if(slaveName.equals(DBConstant.COLUMNAR)){
			return DBConstant.COLUMNAR;
		}
		//获取从库
		List<String> realSlaveDataSourceNames=new ArrayList<String>();
		slaveDataSourceNames.forEach(item->{
			if(!item.equals(DBConstant.COLUMNAR)){
				realSlaveDataSourceNames.add(item);
			}
		});
		//不存在从库，直接走主库
		if(ObjectUtils.isEmpty(realSlaveDataSourceNames)){
			return masterDataSourceName;
		}
		//随机获取从库
		AtomicInteger count = COUNTS.containsKey(name) ? COUNTS.get(name) : new AtomicInteger(0);
		COUNTS.putIfAbsent(name, count);
		count.compareAndSet(realSlaveDataSourceNames.size(), 0);
		return realSlaveDataSourceNames.get(Math.abs(count.getAndIncrement()) % realSlaveDataSourceNames.size());
	}

}

八、处理拦截的查询

• ShardingUtil 用于处理拦截的查询，获取Mapper方法指定的数据库，将数据库标识存入线程副本。
• 同时引入enableColumnar属性用于动态启动关闭列数据库，当列数据库出现问题或者列数据库需重新全量同步MySQL数据时，通过该属性动态关闭使用列数据库。

@Slf4j
@Component
public class ShardingUtil {
	
	/**
 	 * 动态刷新配置，可在nacos中动态切换,列数据库出现问题，切换该参数为false直接走MySQL数据库
 	 */
 	@Value("${enable.columnar:true}")
    private boolean enableColumnar;

	
	public  void beforeStarRocksProcess(MappedStatement mappedStatement) throws Exception {
		//清空线程副本，防止遗留副本产生影响
		ShardingManager.clear();
		//走主库
		if(isMaster(mappedStatement)){
			//不使用HinManager强制主库路由，在事务中，强制路由后，后续查询都走主库
			ShardingManager.setSlave(DBConstant.MASTER);
			return ;
		}
		//走MySQL数据库
		if(!enableStarRocks){
			ShardingManager.setSlave(DBConstant.SLAVE);
			return ;
		}
		//走列数据库
		if(isColumnar(mappedStatement)){
			ShardingManager.setSlave(DBConstant.COLUMNAR);
		}
	
	}


	public  void afterStarRocksProcess() throws Exception {
		ShardingManager.clear();
	}
	

	private  boolean isColumnar(MappedStatement mappedStatement)throws Exception{
		return isAnnotation(mappedStatement, classType->classType.isAnnotationPresent(StarRocksAnotation.class),
				classType->classType.getAnnotation(StarRocksAnotation.class).methods());
	}
	

	private  boolean isMaster(MappedStatement mappedStatement)throws Exception{
		return isAnnotation(mappedStatement, classType->classType.isAnnotationPresent(MasterAnotation.class),
				classType-> classType.getAnnotation(MasterAnotation.class).methods());
	}
	
	
	private boolean isAnnotation(MappedStatement mappedStatement,Predicate<Class<?>> isAnnotationPresent,Function<Class<?>, String[]> methodFunction) throws Exception{
		String id = mappedStatement.getId();
		String classMapper = id.substring(0, id.lastIndexOf("."));
		Class<?> classType = Class.forName(classMapper);
		if (isAnnotationPresent.test(classType)) {
			//获取mapper的方法
			String method = id.replace(String.format("%s.", classMapper), "");
			String[] methodArr=methodFunction.apply(classType);
			List<String> methods = new ArrayList<>(Arrays.asList(methodArr)); 
			methods.add("selectList_COUNT");
			return methods.contains(method);
		}
		return false;
	}
	
	public void setColumnarStatus(String value){
		if(!ObjectUtils.isEmpty(value)){
			this.enableStarRocks=value.equals("false")?false:true;
			if(enableStarRocks){
				log.info("=====>启用Columnar");
			}else{
				log.info("=====>禁用Columnar");
			}
		}
	}
	
}

九、注入拦截器

• 将ShardingInterceptor拦截器通过Spring SPI 等技术注入，同时将ShardingUtil 注入到该拦截器中。方法多样，这里就不再赘述。

至此撒花结束！！