mybatis攔截器Interceptor實現分表

mybatis攔截器Interceptor實現對分表

整體思路：通過mybatis攔截器攔截sql語句，對其中的表名進行修改，再執行

github：https://github.com/cjx913/cbatis

具體代碼如下：

SubTableInterceptor

package com.cjx913.cbatis.core;

import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.HashSet;
import java.util.Properties;
import java.util.Set;

@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}),
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
})
public class SubTableInterceptor implements Interceptor {

    private SubTableHandler subTableHandler;
    private Set <String> handleSqlIds;

    public Object intercept(Invocation invocation) throws Throwable {
        Executor executor = (Executor) invocation.getTarget();

        Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameter = args[1];

        String sqlId = mappedStatement.getId();
        //是否處理
        if (handleSqlIds == null || handleSqlIds.isEmpty() || !handleSqlIds.contains(sqlId)) {
            return invocation.proceed();
        }


        if (mappedStatement.getSqlCommandType().compareTo(SqlCommandType.SELECT) == 0) {
            RowBounds rowBounds = (RowBounds) args[2];
            ResultHandler resultHandler = (ResultHandler) args[3];

            CacheKey cacheKey;
            BoundSql boundSql;
            if (args.length == 4) {
                boundSql = mappedStatement.getBoundSql(parameter);
                cacheKey = executor.createCacheKey(mappedStatement, parameter, rowBounds, boundSql);
            } else {
                boundSql = (BoundSql) args[5];
                cacheKey = (CacheKey) args[4];
            }

            subTableHandler.boundSqlHandler(boundSql);

            return executor.query(mappedStatement, parameter, rowBounds, resultHandler, cacheKey, boundSql);
        } else {
            MappedStatement newMappedStatement = subTableHandler.mappedStatementHandler(mappedStatement, parameter);
            return executor.update(newMappedStatement, parameter);
        }
    }

    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    public void setProperties(Properties properties) {
        String tableNameHandlerClassName = properties.getProperty("table-name-handler").trim();
        if (tableNameHandlerClassName == null || "".equalsIgnoreCase(tableNameHandlerClassName)) {
            throw new CbatisException("無效的名錶處理類！請在插件添加<property name=\"table-name-tableNameHandler\" value=\"表明處理類完全限定名\"/>\n");
        }
        subTableHandler = new SubTableHandler(tableNameHandlerClassName);

        String[] sqlIds = properties.getProperty("handle-sql-ids").split(",");
        if (sqlIds != null && sqlIds.length > 0) {
            handleSqlIds = new HashSet <>();
            for (String s : sqlIds) {
                handleSqlIds.add(s.trim());
            }
        }
    }

    public Set <String> getHandleSqlIds() {
        return handleSqlIds;
    }

    public void setHandleSqlIds(Set <String> handleSqlIds) {
        this.handleSqlIds = handleSqlIds;
    }
}

SubTableHandler

package com.cjx913.cbatis.core;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.ibatis.builder.StaticSqlSource;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;

import java.lang.reflect.Field;
import java.util.*;

/**
 * 分表處理器
 */
public class SubTableHandler {
    private AbstractTableNameHandler tableNameHandler;

    public SubTableHandler(String tableNameHandlerClassName) {
        try {
            tableNameHandler = (AbstractTableNameHandler) Class.forName(tableNameHandlerClassName).newInstance();
        } catch (InstantiationException | IllegalAccessException | ClassNotFoundException e) {
            throw new CbatisException("分表處理器錯誤->表名處理類錯誤！請確認" + tableNameHandlerClassName + "是否存在", e);
        }
    }

    public MappedStatement mappedStatementHandler(MappedStatement mappedStatement, Object parameter) {
        SqlSource sqlSource = mappedStatement.getSqlSource();

        BoundSql boundSql = sqlSource.getBoundSql(parameter);
        boundSqlHandler(boundSql);

        StaticSqlSource newSqlSource = new StaticSqlSource(mappedStatement.getConfiguration(), boundSql.getSql(), boundSql.getParameterMappings());

        MappedStatement newMappedStatement = new MappedStatement.Builder(mappedStatement.getConfiguration(), mappedStatement.getId(), newSqlSource, mappedStatement.getSqlCommandType())
                .resource(mappedStatement.getResource())
                .parameterMap(mappedStatement.getParameterMap())
                .flushCacheRequired(mappedStatement.isFlushCacheRequired())
                .keyGenerator(mappedStatement.getKeyGenerator())
                .keyProperty(CbatisUtil.split(mappedStatement.getKeyProperties()))
                .keyColumn(CbatisUtil.split(mappedStatement.getKeyColumns()))
                .timeout(mappedStatement.getTimeout())
                .cache(mappedStatement.getCache())
                .useCache(mappedStatement.isUseCache())
                .databaseId(mappedStatement.getDatabaseId())
                .fetchSize(mappedStatement.getFetchSize())
                .lang(mappedStatement.getLang())
                .resultMaps(mappedStatement.getResultMaps())
                .resultOrdered(mappedStatement.isResultOrdered())
                .build();
        return newMappedStatement;
    }


    public String boundSqlHandler(BoundSql boundSql) {
        try {
            String sql = boundSql.getSql();
            Object parameter = boundSql.getParameterObject();

            sql = sqlHandler(sql);

            Set <String> tableNames = getTableNames(sql);
            if (!tableNames.isEmpty()) {
                Map <String, String> newTableNames = getNewTableNames(tableNames, parameter);
                if (!newTableNames.isEmpty()) {
                    sql = replaceTableName(sql, newTableNames);
                    //通過反射修改sql語句
                    Field field = boundSql.getClass().getDeclaredField("sql");
                    field.setAccessible(true);
                    field.set(boundSql, sql);
                }
            }
            return sql;
        } catch (Exception e) {
            throw new CbatisException("Sql修改錯誤！", e);
        }
    }

    /**
     * sql處理，去掉換行和多餘空格
     *
     * @param sql
     * @return
     */
    private String sqlHandler(String sql) {
        String[] strings = sql.trim().replace("\n", " ").split("\\s+");
        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < strings.length; i++) {
            if (strings[i].endsWith(",")) {
                stringBuilder.append(strings[i]);
            } else {
                stringBuilder.append(strings[i] + " ");
            }
        }
        sql = stringBuilder.toString().trim().toUpperCase();
        return sql;
    }

    /**
     * @param sql
     * @return 獲取sql所有的表名
     */
    private Set <String> getTableNames(String sql) {
        Statement statement = null;
        try {
            statement = CCJSqlParserUtil.parse(sql);
        } catch (JSQLParserException e) {
            throw new CbatisException("解析sql語句錯誤！sql:" + sql, e);
        }
        TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
        List <String> tableList = tablesNamesFinder.getTableList(statement);
        Set <String> tableNames = new HashSet <>();
        for (String tableName : tableList) {
            //獲取去掉“`”的表名
            if (tableName.startsWith("`") && tableName.endsWith("`")) {
                tableNames.add(tableName.substring(1, tableName.length()-1));
            }else {
                tableNames.add(tableName);
            }
        }

        return tableNames;
    }

    /**
     * @param tableNames
     * @param parameter
     * @return 原表名與新表名的Map
     */
    private Map <String, String> getNewTableNames(Set <String> tableNames, Object parameter) {
        Map <String, String> newTableNames = new HashMap <>();

        for (String tableName : tableNames) {
            //獲取新表名（新表名都帶“`”）
            String newTableName = "`"+tableNameHandler.tableNameHandler(tableName, parameter)+"`";
            newTableNames.put(tableName, newTableName);
        }
        return newTableNames;
    }

    /**
     * 表名替換
     *
     * @param sql
     * @param tableNames
     * @return
     */
    private String replaceTableName(String sql, Map <String, String> tableNames) {
        //去掉sq的“`”
        sql = sql.replace("`", "");
        Set <Map.Entry <String, String>> entrySet = tableNames.entrySet();
        for (Map.Entry <String, String> entry : entrySet) {
            if (!entry.getKey().equalsIgnoreCase(entry.getValue())) {
                sql = sql.replace(" " + entry.getKey() + " ", " " + entry.getValue() + " ")
                        .replace(" " + entry.getKey() + ",", " " + entry.getValue() + ",")
                        .replace("," + entry.getKey() + " ", "," + entry.getValue() + " ")
                        .replace("," + entry.getKey() + ",", "," + entry.getValue() + ",")
                        .replace(" " + entry.getKey() + "(", " " + entry.getValue() + "(");
            }
        }
        return sql.toUpperCase();
    }
}

AbstractTableNameHandler

package com.cjx913.cbatis.core;

import org.apache.ibatis.binding.MapperMethod;

import java.beans.IntrospectionException;
import java.beans.PropertyDescriptor;
import java.lang.annotation.Annotation;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.*;

/**
 * 表名處理器
 */
public abstract class AbstractTableNameHandler {
    private Map <String, TableNameHandlerMethod> tableMethodMap = new HashMap <>();

    public AbstractTableNameHandler() {
        initTableMethodMap();
    }

    public final String tableNameHandler(String tableName, Object parameter) {

        TableNameHandlerMethod tableNameHandlerMethod = tableMethodMap.get(tableName.toUpperCase());
        if (tableNameHandlerMethod == null) {
            return tableName.trim().toUpperCase();
        }
        Method method = tableNameHandlerMethod.getMethod();

        //獲取方法參數值，可以爲null
        Object[] params = paramHandler(method, parameter);
//        if (params == null) {
//            throw new CbatisException("調用分表處理方法錯誤！沒有指定參數\n表名:" + tableName + ",方法:" + method.getName());
//        }

        Object newTableName = null;
        try {
            newTableName = method.invoke(tableNameHandlerMethod.getInstanse(), params);
        } catch (IllegalAccessException | InvocationTargetException e) {
            throw new CbatisException("調用分表處理方法錯誤！表名:" + tableName + ",方法:" + method.getName(), e);
        }
        if (newTableName instanceof String) {
            return ((String) newTableName).trim().toUpperCase();
        } else {
            throw new CbatisException("調用分表處理方法錯誤！->返回非字符串。表名:" + tableName);
        }
    }

    /**
     * 獲取參數值
     *
     * @param method
     * @param parameter
     * @return
     */
    private Object[] paramHandler(Method method, Object parameter) {

        MapperMethod.ParamMap <Object> paramMap = null;
        if (parameter instanceof MapperMethod.ParamMap) {
            paramMap = (MapperMethod.ParamMap <Object>) parameter;
        }

        List <String> params = new ArrayList <>();
        Annotation parameterAnnotations[][] = method.getParameterAnnotations();
        for (int i = 0; i < parameterAnnotations.length; i++) {
            for (Annotation annotation : parameterAnnotations[i]) {
                if (Param.class.equals(annotation.annotationType())) {
                    params.add(((Param) annotation).value());
                }
            }
        }
        //獲取參數值
        List <Object> parameters = new ArrayList <>();
        try {
            if (paramMap != null) {
                for (String param : params) {
                    Object p = paramMap.get(param);
                    parameters.add(p);
                }
            } else {
                for (String param : params) {
                    PropertyDescriptor pd = new PropertyDescriptor(param, parameter.getClass());
                    Method getMethod = pd.getReadMethod();
                    Object p = getMethod.invoke(parameter);
                    parameters.add(p);
                }
            }
        } catch (IllegalAccessException | InvocationTargetException | IntrospectionException e) {
            throw new CbatisException("獲取參數值錯誤", e);
        }

//        if (parameters.isEmpty()) {
//            return null;
//        } else {
//            return parameters.toArray();
//        }
        return parameters.toArray();
    }

    /**
     * @return 表名作爲鍵，處理該表的方法作爲值。<br/>
     * table name as key ,tableNameHandler method as value.
     */
    private void initTableMethodMap() {
        try {
            for (TableNameHandlerMethod tableNameHandlerMethod : setTableMethodSet()) {
                this.tableMethodMap.put(tableNameHandlerMethod.getMethod().getName().toUpperCase(), tableNameHandlerMethod);
            }
        } catch (NoSuchMethodException e) {
            throw new CbatisException("調用參數處理方法錯誤！", e);
        }
    }

    /**
     * 表名處理方法
     *
     * @return
     * @throws NoSuchMethodException
     */
    public abstract Set <TableNameHandlerMethod> setTableMethodSet() throws NoSuchMethodException;

    protected TableNameHandlerMethod getThisTableNameHandlerMethod(Method method) {
        return new TableNameHandlerMethod(this, method);
    }


    public class TableNameHandlerMethod {
        private Object instanse;
        private Method method;

        /**
         * 限制在AbstractTableNameHandler子類寫表名處理方法
         *
         * @param method
         */
        public TableNameHandlerMethod(Method method) {
            instanse = AbstractTableNameHandler.this;
            this.method = method;
        }

        private TableNameHandlerMethod(Object instanse, Method method) {
            this.instanse = instanse;
            this.method = method;
        }

        public Object getInstanse() {
            return instanse;
        }

        public void setInstanse(Object instanse) {
            this.instanse = instanse;
        }

        public Method getMethod() {
            return method;
        }

        public void setMethod(Method method) {
            this.method = method;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;

            TableNameHandlerMethod that = (TableNameHandlerMethod) o;

            return method.getName().equalsIgnoreCase(that.method.getName());
        }

        @Override
        public int hashCode() {
            return method.getName().toUpperCase().hashCode();
        }
    }
}

Param

package com.cjx913.cbatis.core;

import java.lang.annotation.*;

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.PARAMETER)
public @interface Param {
    String value();
}

CbatisException

package com.cjx913.cbatis.core;

public class CbatisException extends RuntimeException {
    public CbatisException() {
    }

    public CbatisException(String message) {
        super(message);
    }

    public CbatisException(String message, Throwable cause) {
        super(message, cause);
    }

    public CbatisException(Throwable cause) {
        super(cause);
    }

    public CbatisException(String message, Throwable cause, boolean enableSuppression, boolean writableStackTrace) {
        super(message, cause, enableSuppression, writableStackTrace);
    }
}

CbatisUtil

package com.cjx913.cbatis.core;

public class CbatisUtil {
    public static String split(String[] strings){
        if (strings==null){
            return null;
        }
        StringBuilder stringBuilder =new StringBuilder();
        for(String str:strings){
            stringBuilder.append(str+",");
        }
        return stringBuilder.deleteCharAt(stringBuilder.length() - 1).toString();
    }
}

在mybatis的配置文件添加插件

<plugins>
        <plugin interceptor="com.cjx913.cbatis.core.SubTableInterceptor">
            <property name="table-name-handler" value="com.cjx913.cbatis.SimpleTableNameHandler"/>
            <property name="handle-sql-ids" value="
                com.cjx913.cbatis.mapper.UserMapper.insertUser,
                com.cjx913.cbatis.mapper.UserMapper.userBaseMapper,
                com.cjx913.cbatis.mapper.OrderMapper.insertOrder"/>
        </plugin>
    </plugins>

 

分表策略事例SimpleTableNameHandler

package com.cjx913.cbatis;

import com.cjx913.cbatis.core.AbstractTableNameHandler;
import com.cjx913.cbatis.core.Param;

import java.util.HashSet;
import java.util.Set;

public class SimpleTableNameHandler extends AbstractTableNameHandler {


    public String user(@Param("username") String username) {

        if (username == null || "".equalsIgnoreCase(username.trim())) {
            return "user";
        }
        if (username.hashCode() % 2 == 0) {
            return "user_0";
        } else {
            return "user_1";
        }
    }

    public String order(@Param("userId") Integer userId, @Param("id") Integer id) {

        if (userId == null) {
            return "order";
        }
        if (id == null) id = 0;
        int i = (userId.hashCode() + id.hashCode()) % 3;
        if (i == 0) {
            return "order_1";
        } else if (i == 1) {
            return "order_2";
        } else {
            return "order_3";
        }
    }

    @Override
    public Set <TableNameHandlerMethod> setTableMethodSet() throws NoSuchMethodException {
        Set <TableNameHandlerMethod> tableMethodMap = new HashSet <>();
        tableMethodMap.add(new TableNameHandlerMethod(this.getClass().getMethod("user", String.class)));
        tableMethodMap.add(getThisTableNameHandlerMethod(this.getClass().getMethod("order", Integer.class, Integer.class)));
        return tableMethodMap;
    }
}

在 public Set <TableNameHandlerMethod> setTableMethodSet() throws NoSuchMethodException 方法添加表名處理邏輯的方法，方法名爲表名。如order表可通過order()方法改成order_1，order_2，order_3。方法參數要使用@Param註解（是com.cjx913.cbatis.core.Param）value是從mybatis的mapper的方法獲取參數值的屬性名，可以獲取到mapper的方法參數值