Mybatis实现自定义分表功能

工作笔记

---

前言

由于业务的需要，需要存储每周递增的百万数据，由于数据的特殊性，需要留存往期数据并满足查询要求，在此需求下，月级数据就已经达到了上千万，考虑到功能后期的正常运行，必须要通过一些方式来解决高并发以及大量数据的问题。

一、为什么选择分表？

实际上，解决上述问题有很多种方式：分库分表，读写分离，水平分片，引入缓存等操作。个人选择通过mybatis实现分表，是因为公司使用云存储，考虑到投入产出，以及引入中间件提升系统复杂性等问题。最终和leader沟通后，决定自己通过mybatis实现自定义分表的功能，如果你也有一样的问题,希望这篇文章能给你带来帮助。

二、考虑因素

1. 只有少许业务会涉及分表，分表功能不应该影响到其他业务。
2. 分表规则会有多种，根据业务数据量不同，分表规则也应该有所不通。
3. 应该根据什么原则来进行分表，hash?时间? 如：在我的业务中是通过传入的版本来作为分表依据。
4. 考虑到如果没有对应分表时，应该如何操作。如：insert（创建表）select（查询原始表）

带着上述的思考，那么接下来，我们就直接上代码。

三、具体实现

1.实现分表查询的判断(注解)

代码如下：

@Target({ElementType.METHOD})
@Retention(value = RetentionPolicy.RUNTIME)
public @interface TableShard {

    String WEEK = "week";//周
    String MOUTH = "mouth";//月
    String QUARTER = "quarter";//季
    String YEAR = "year";//年

    //是否开启建表操作
    boolean tableShardOpen() default false;

    //分表策略
    String[] tableShardStrategy();

    //表名
    String[] tableName();

    //对应参数
    String[] columns();

    //默认表名
    String[] defaultName();
}

2.分表策略编写

这里是我的个人实现，具体可以根据业务来变更。
代码如下：
interface:

public interface TableShardStrategy {

    String getTableShardName(String tableName, String version);
}

impl:

@Component("mouth")
public class MouthShard implements TableShardStrategy {


    @Override
    public String getTableShardName(String tableName, String version) {
        //获取DateTime
        DateTime dateTime = DateUtil.parse(version, "yyyyMMdd");
        //获取表名
        return tableName + dateTime.year() +
                dateTime.monthStartFromOne();
    }
}

上述贴了一个按月分表的代码，例如版本2024/01/09，版本则为20231209，原表table,则分表名称为:table202412。

3.选择器以及分表相关操作

通过这里可以选择对应的策略，或者进行创建表或者判断表是否存在的操作
代码如下（示例）：

@Component
public class TableShardSelector {

    @Autowired
    private Map<String, TableShardStrategy> selector;

    public TableShardStrategy select(String type) {
        return selector.get(type);
    }

    /**
     * 查询表是否存在
     * @param tableName 表名
     * @return true or false
     */
    public boolean containsTable(String tableName) {
        DataSource dataSource = SpringContextUtils.getBean(DataSource.class);
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        //验证表是否存在
        List<String> list = jdbcTemplate.queryForList("SHOW TABLES LIKE " + "'" + tableName + "'", String.class);
        return list.isEmpty();
    }

    /**
     * 根据默认表结构生成新表
     * @param defaultTableName 默认表
     * @param tableName 生成表
     */
    public void createTable(String defaultTableName, String tableName) {
        DataSource dataSource = SpringContextUtils.getBean(DataSource.class);
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        //创建表
        jdbcTemplate.execute("CREATE TABLE " + tableName + " LIKE " + defaultTableName);
    }

这里是通过传入默认表的方式创建分表，如：传入了table和table202312,会创建出一个和table结构一致但是名称为table202312的分表，属于取巧的操作，后续表结构是和原始表结构一致，当然，如果有其他的需求也可以变动。

4.分表拦截器

完成了上述的‘前菜’，就应该上核心部分了。分表拦截器的代码：

@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
@Component
@Slf4j
public class ShardInterceptor implements Interceptor {

    @Value("${table.shard.enabled:false}")
    private boolean shardSign;//分表标志 可以从配置开启关闭

    @Autowired
    private TableShardSelector tableShardSelector;//分表选择器

    private static final ReflectorFactory defaultReflectorFactory = new DefaultReflectorFactory();
    private static final String DELEGATE_MAPPER_STATEMENT = "delegate.mappedStatement";
    private static final String DELEGATE_BOUND_SQL = "delegate.boundSql";
    private static final String DELEGATE_BOUND_SQL_SQL = "delegate.boundSql.sql";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        //获取StatementHandler
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        //获取mapper元信息
        MetaObject metaObject = MetaObject.forObject(statementHandler,
                SystemMetaObject.DEFAULT_OBJECT_FACTORY,
                SystemMetaObject.DEFAULT_OBJECT_WRAPPER_FACTORY,
                defaultReflectorFactory);
        //获取MappedStatement
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue(DELEGATE_MAPPER_STATEMENT);
        //获取SQL元数据
        BoundSql boundSql = (BoundSql) metaObject.getValue(DELEGATE_BOUND_SQL);
        //拦截分表逻辑，获取注解
        TableShard annotation = getAnnotation(mappedStatement);
        //当前方法不包含注解，走默认
        if (annotation == null) {
            return invocation.proceed();
        }
        //获取注解参数
        String[] shardStrategys = annotation.tableShardStrategy();
        //获取入参值，根据入参值进行分表
        String[] columnValues = getColumnValue(annotation.columns(), boundSql);
        //获取替换表名称
        String[] tableNames = annotation.tableName();
        //获取默认表名称
        String[] defaultName = annotation.defaultName();
        //新表名称
        String[] newTableNames = new String[tableNames.length];
        //是否开启分表
        if (shardSign) {
            for (int i = 0; i < columnValues.length; i++) {
                //选择生成表策略
                TableShardStrategy select = tableShardSelector.select(shardStrategys[i]);
                //参数为null时走默认表
                newTableNames[i] = columnValues[i] == null ? defaultName[i] : select.getTableShardName(defaultName[i], columnValues[i]);
            }
        } else {
            //未分表走默认表
            newTableNames = defaultName;
        }
        //获取sql
        String sql = boundSql.getSql();
        //是否分表
        boolean tableShardOpen = annotation.tableShardOpen();
        if (tableShardOpen) {  //开启分表 表不存在创建表
            for (int i = 0; i < newTableNames.length; i++) {
                String tab = newTableNames[i];
                if (tableShardSelector.containsTable(tab)) {
                    //保证获取同一锁对象，在执行并行时单线程创建表(多机可以选择分布式锁)
                    synchronized (tab.intern()) {
                        if (tableShardSelector.containsTable(tab)) {
                            //创建表
                            tableShardSelector.createTable(defaultName[i], tab);
                        }
                    }
                }
            }
        } else { //其他查询不存在则用默认表替换
            for (int i = 0; i < newTableNames.length; i++) {
                String tab = newTableNames[i];
                newTableNames[i] = !tableShardSelector.containsTable(tab) ? tab : defaultName[i];
            }
        }
        //替换表名
        for (int i = 0; i < tableNames.length; i++) {
            sql = sql.replace(tableNames[i], newTableNames[i]);
        }
        //拦截并替换sql
        metaObject.setValue(DELEGATE_BOUND_SQL_SQL, sql);
        return invocation.proceed();
    }

    private TableShard getAnnotation(MappedStatement mappedStatement) throws ClassNotFoundException {
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        String methodName = id.substring(id.lastIndexOf(".") + 1);
        //获取pageHelper拦截方法
        if (methodName.endsWith("_COUNT")) {
            methodName = methodName.replace("_COUNT", "");
        }
        Class<?> aClass = Class.forName(className);
        Method[] declaredMethods = aClass.getDeclaredMethods();
        for (Method declaredMethod : declaredMethods) {
            if (declaredMethod.getName().equals(methodName)) {
                return declaredMethod.getAnnotation(TableShard.class);
            }
        }
        return null;
    }

    /**
     * 获取sql参数值
     *
     * @param columns  参数名称
     * @param boundSql sql元数据
     * @return 参数列表
     */
    private String[] getColumnValue(String[] columns, BoundSql boundSql) {
        Object parameterObject = boundSql.getParameterObject();
        if (parameterObject == null) {
            log.info("ShardInterceptor parameterObject is null!!");
            throw new ShardException("分表参数异常！");
        }
        Map map;
        //判断参数类型获取参数值
        if (parameterObject instanceof MapperMethod.ParamMap) {//查询条件是List
            MapperMethod.ParamMap paramMap = (MapperMethod.ParamMap) parameterObject;
            List<Map> list = JSON.parseArray(JSON.toJSONString(paramMap.get("entities")), Map.class);
            map = list.get(0);
        } else if (parameterObject instanceof Map) {//查询条件是Map
            map = (Map) parameterObject;
        } else {//查询条件是其他bean
            map = BeanUtil.beanToMap(parameterObject);
        }
        for (int i = 0; i < columns.length; i++) {
            columns[i] = MapUtils.getString(map, columns[i]);
        }
        return columns;
    }
}

核心是通过mybatis的拦截器，获取Mapper方法的唯一标识，然后拿到对应的注解，通过注解实现分表操作。

5.具体使用

使用分表时只需要通过为对应的Dao加上注解即可。
查询：

    @TableShard(tableShardStrategy = TableShard.WEEK, tableName = "table_new", columns = "version", defaultName = "table")
    List<Data> queryAll(Data data);

上述sql，会通过Week的策略将原sql中表名table_new，根据传入参数中version,替换为新的表名，如果表不存在，则默认使用table。

新增:

    @TableShard(tableShardStrategy = TableShard.WEEK, tableName = "table_new", columns = "version", defaultName = "table", tableShardOpen = true)
    int insertBatch(@Param("entities") List<Data> data);

上述sql，会通过Week的策略将原sql中表名table_new，根据传入参数中version,替换为新的表名，如果表不存在，则会创建该表并插入数据。

涉及多表的查询：

    @TableShard(tableShardStrategy = {TableShard.WEEK, TableShard.WEEK},
            tableName = {"table_new_1", "table_new_2"},
            columns = {"version1", "version2"},
            defaultName = {"table1", "table2"})
    List<Data> queryChange(Data data);

---

总结

内容主要是通过mybatis拦截器对数据实现水平分表和查询操作，过程也学习了mybatis拦截器的使用，希望这篇文章可以帮到你。