编码技巧——Mybatis分表插件

﹏ヽ暗。殇╰゛Y 2024-04-24 14:40 2阅读 0赞

原理你们不想看，我就直接放源码了，这里附上我之前的文章的链接：[编码技巧——数据加密（二）Mybatis拦截器\_mybatis 拦截器][Mybatis_mybatis]；

背景是新公司的基础框架不支持分库分表，并且不能使用业界常用的分库分表插件二方包（内网Maven仓库），所以基于Mybatis2.X，基于Mybatis Interceptor自行实现分表插件，后期可作为二方包在公司内推广避免重复造轮子；

下面不废话，直接上你们想要的源码；

### 注解 ###

/**
     * @author Akira
     * @description 分表注解 支持[类和方法]级别 优先方法注解
     * @date 2024-1-17 14:53
     */
    @Target(value = {ElementType.TYPE, ElementType.METHOD})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface Shard {
    
    	// 逻辑表名
    	String logicTbName();
    
    	// 分表属性Name
    	String shardParamName() default "";
    
    	// 对应的分表路由规则类
    	Class<? extends ITableShardRouter> shardRouter();
    
    }

### 分表策略 ###

/**
     * @author Akira
     * @description 分表路由规则
     * @date 2024-1-17 14:36
     */
    public interface ITableShardRouter {
    
    	/**
    	 * 生成分表名
    	 *
    	 * @param logicTableName 逻辑表名
    	 * @param value 分表属性值
    	 */
    	String generateTableName(String logicTableName, Object value);
    
    	/**
    	 * 基础验证
    	 */
    	default void verificationTableNamePrefix(String logicTableName, Object value) {
    		if (StringUtils.isBlank(logicTableName)) {
    			throw new RuntimeException("[tableSharding]logicTableName is null");
    		}
    		if (value == null || StringUtils.isBlank(value.toString())) {
    			throw new RuntimeException("[tableSharding]shardParam value is null");
    		}
    	}
    }

实现类：

/**
     * @author Akira
     * @description
     * @date 2024/1/17 14:40
     */
    @Slf4j
    @Component
    public class TableShardByUidRouter implements ITableShardRouter {
    
    	/**
    	 * 生成分表名
    	 *
    	 * @param logicTableName 逻辑表名
    	 * @param value          分表属性值
    	 */
    	@Override
    	public String generateTableName(String logicTableName, Object value) {
    		try {
    			verificationTableNamePrefix(logicTableName, value);
    		} catch (Exception e) {
    			log.error("[tableSharding]分表执行参数校验失败 [errMessage={}]", e.getMessage());
    			throw new BizException(ResultCodeEnum.BUSINESS_ERROR.getCode(), ResultCodeEnum.BUSINESS_ERROR.getDesc());
    		}
    		String suffix = RouteUtils.getSuffix(value);
    		return logicTableName + "_" + suffix;
    	}
    }

### 拦截器\[核心\] ###

/**
     * @author Akira
     * @description 分表拦截器 按类型自动注入拦截器列表，无需在xml中声明
     * @refer https://blog.csdn.net/minghao0508/article/details/124420442
     * @date 2024/1/17 15:00
     */
    @Slf4j
    @Component
    @Intercepts({
    		@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
    })
    public class TableShardInterceptor implements Interceptor {
    
    	private static final String COUNT_SUFFIX = "_COUNT";
    	private static final ReflectorFactory defaultReflectorFactory = new DefaultReflectorFactory();
    
    	@Override
    	public Object intercept(Invocation invocation) throws Throwable {
    		// mata data
    		MetaObject metaObject = getMetaObject(invocation);
    		BoundSql boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
    		MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
    
    		// 获取Mapper执行方法
    		final String id = mappedStatement.getId();
    		final String mapperClassName = id.substring(0, id.lastIndexOf("."));
    		String methodName = id.substring(id.lastIndexOf(".") + 1);
    
    		// PageHelper分页插件兼容
    		if (methodName.endsWith(COUNT_SUFFIX)) {
    			methodName = methodName.substring(0, methodName.lastIndexOf(COUNT_SUFFIX));
    		}
    
    		// 获取分表注解
    		Shard tableShard = getSuffixShardAnnotation(mapperClassName, methodName);
    
    		// 如果类和方法都没有SuffixShard注解，执行下一个interceptor逻辑
    		if (tableShard == null) {
    			return invocation.proceed();
    		}
    
    		/**
    		 * 分表核心逻辑
    		 */
    		String shardParamName = tableShard.shardParamName();
    		Object parameterObject = boundSql.getParameterObject();
    		// 1.@Param参数列表
    		if (parameterObject instanceof MapperMethod.ParamMap) {
    			MapperMethod.ParamMap parameterMap = (MapperMethod.ParamMap) parameterObject;
    			// 根据字段名获取参数值
    			Object valueObject = parameterMap.get(shardParamName);
    			if (valueObject == null) {
    				log.error("分表属性值字段SuffixShard#shardParamName[{}}]为空", shardParamName);
    				throw new RuntimeException(String.format("分表属性值字段SuffixShard#shardParamName[%s]无匹配", shardParamName));
    			}
    			// 替换sql
    			replaceSql(tableShard, valueObject, metaObject, boundSql);
    		}
    		// 2.单参数实体
    		else {
    			// 如果是基础类型抛出异常
    			if (isBaseType(parameterObject)) {
    				throw new RuntimeException("单参数非法，请使用@Param注解");
    			}
    			// Map类型
    			if (parameterObject instanceof Map) {
    				Map<String, Object> parameterMap = (Map<String, Object>) parameterObject;
    				Object valueObject = parameterMap.get(shardParamName);
    				// 替换sql
    				replaceSql(tableShard, valueObject, metaObject, boundSql);
    			}
    			// 对象类型反射获取值
    			else {
    				Class<?> parameterObjectClass = parameterObject.getClass();
    				Field declaredField = null;
    				try {
    					declaredField = parameterObjectClass.getDeclaredField(shardParamName);
    				} catch (NoSuchFieldException e) {
    					log.error("[tableSharding]未从Mapper参数中获取到分表参数 [errMessage={}]", e.getMessage());
    					throw new RuntimeException("未从Mapper参数中获取到分表参数 请检查mapper方法参数定义");
    				}
    				declaredField.setAccessible(true);
    				Object valueObject = declaredField.get(parameterObject);
    				// 替换sql
    				replaceSql(tableShard, valueObject, metaObject, boundSql);
    			}
    		}
    
    		// 执行逻辑
    		return invocation.proceed();
    	}
    
    	@Override
    	public Object plugin(Object target) {
    		// 当目标类是StatementHandler类型时，才包装目标类，否者直接返回目标本身, 减少目标被代理的次数
    		if (target instanceof StatementHandler) {
    			return Plugin.wrap(target, this);
    		} else {
    			return target;
    		}
    	}
    
    	/**
    	 * 基本数据类型验证
    	 *
    	 * @param object
    	 * @return
    	 */
    	private boolean isBaseType(Object object) {
    		if (object.getClass().isPrimitive()
    				|| object instanceof String
    				|| object instanceof Integer
    				|| object instanceof Double
    				|| object instanceof Float
    				|| object instanceof Long
    				|| object instanceof Boolean
    				|| object instanceof Byte
    				|| object instanceof Short) {
    			return true;
    		} else {
    			return false;
    		}
    	}
    
    	/**
    	 * replaceSql
    	 */
    	private void replaceSql(Shard tableShard, Object value, MetaObject metaObject, BoundSql boundSql) {
    		String logicTableName = tableShard.logicTbName();
    		Class<? extends ITableShardRouter> strategyClazz = tableShard.shardRouter();
    		// 从IOC容器获取策略类
    		ITableShardRouter tableShardStrategy = SpringUtil.getBean(strategyClazz);
    		if (tableShardStrategy == null) {
    			log.error("未找到分表策略 [strategyClazz={}]", strategyClazz.getName());
    		}
    		String shardTableName = Objects.requireNonNull(tableShardStrategy).generateTableName(logicTableName, value);
    		String sql = boundSql.getSql();
    		// 完成表名替换
    		metaObject.setValue("delegate.boundSql.sql", sql.replaceAll(logicTableName, shardTableName));
    		log.warn("[tableSharding]执行Suffix分表逻辑成功 [逻辑表名={} 分表参数={} 分表名={}]", logicTableName, value, shardTableName);
    	}
    
    	/**
    	 * 获取MetaObject对象
    	 */
    	private MetaObject getMetaObject(Invocation invocation) {
    		StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
    		// MetaObject是mybatis里面提供的一个工具类，类似反射的效果
    		MetaObject metaObject = MetaObject.forObject(statementHandler,
    				SystemMetaObject.DEFAULT_OBJECT_FACTORY,
    				SystemMetaObject.DEFAULT_OBJECT_WRAPPER_FACTORY,
    				defaultReflectorFactory
    		);
    		return metaObject;
    	}
    
    	/**
    	 * 获取注解属性
    	 *
    	 * @throws ClassNotFoundException
    	 */
    	private Shard getSuffixShardAnnotation(String mapperClassName, String methodName) throws ClassNotFoundException {
    		Class<?> mapperClass = Class.forName(mapperClassName);
    		Method[] methods = mapperClass.getMethods();
    		for (Method method : methods) {
    			if (method.getName().equalsIgnoreCase(methodName) && method.isAnnotationPresent(Shard.class)) {
    				// 优先使用方法注解
    				return method.getAnnotation(Shard.class);
    			}
    		}
    		// 如果方法没有设置注解，尝试从Mapper接口上面获取注解
    		return mapperClass.getAnnotation(Shard.class);
    	}
    
    }

### 使用 ###

@Shard(logicTbName="t_member", shardParamName="uid", shardRouter=TableShardByUidRouter.class)
    public interface MemberDAO {
    
        List<MemberDO> selectByCondition(@Param("uid") String uid, @Param("qry") MemberQuery qry);
    
    }

参考：[编码技巧——数据加密（二）Mybatis拦截器][Mybatis_mybatis]

[Mybatis_mybatis]: https://blog.csdn.net/minghao0508/article/details/124420442