mybatis分表攔截器實現

MyBatis提供了一種插件(plugin)的功能，雖然叫做插件，但其實這是攔截器功能。那么攔截器攔截MyBatis中的哪些內容呢？

我們進入官網看一看：

MyBatis 允許你在已映射語句執行過程中的某一點進行攔截調用。默認情況下，MyBatis 允許使用插件來攔截的方法調用包括：

1. Executor (update, query, flushStatements, commit, rollback, getTransaction, close, isClosed)
2. ParameterHandler (getParameterObject, setParameters)
3. ResultSetHandler (handleResultSets, handleOutputParameters)
4. StatementHandler (prepare, parameterize, batch, update, query)

我們看到了可以攔截Executor接口的部分方法，比如update，query，commit，rollback等方法，還有其他接口的一些方法等。

總體概括為：

1. 攔截執行器的方法
2. 攔截參數的處理
3. 攔截結果集的處理
4. 攔截Sql語法構建的處理

對於攔截器Mybatis為我們提供了一個Interceptor接口，通過實現該接口就可以定義我們自己的攔截器。我們先來看一下這個接口的定義：

package org.apache.ibatis.plugin;
 
import java.util.Properties;
 
public interface Interceptor {
 
  Object intercept(Invocation invocation) throws Throwable;
 
  Object plugin(Object target);
 
  void setProperties(Properties properties);
 
}

我們可以看到在該接口中一共定義有三個方法，intercept、plugin和setProperties。plugin方法是攔截器用於封裝目標對象的，通過該方法我們可以返回目標對象本身，也可以返回一個它的代理。當返回的是代理的時候我們可以對其中的方法進行攔截來調用intercept方法，當然也可以調用其他方法。setProperties方法是用於在Mybatis配置文件中指定一些屬性的。

分表代碼實現如下：

首先定義一個分表的注解

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

@Retention(RetentionPolicy.RUNTIME)
@Target({ ElementType.TYPE })
public @interface TableSharde {
    /**
     * 待分表的表名
     * @return
     */
    String tableName();

    /**
     * 分表策略
     * @return
     */
    Class<?> strategy();

    /**
     * 分表條件
     * @return
     */
    String[] shardeBy();
}

分表攔截器實現

import com.twsz.annotation.TableSharde;
import com.twsz.exception.BusinessException;
import com.twsz.tablesharde.TableStrategy;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.ReflectorFactory;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.factory.ObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.apache.ibatis.reflection.wrapper.ObjectWrapperFactory;
import org.springframework.util.StringUtils;

import java.sql.Connection;
import java.util.Properties;

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

    private static Log log = LogFactory.getLog(TableShardeInterceptor.class);
    private static final ObjectFactory DEFAULT_OBJECT_FACTORY = new DefaultObjectFactory();
    private static final ObjectWrapperFactory DEFAULT_OBJECT_WRAPPER_FACTORY = new DefaultObjectWrapperFactory();
    private static final ReflectorFactory REFLECTOR_FACTORY = new DefaultReflectorFactory();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        MetaObject metaStatementHandler = MetaObject.forObject(statementHandler, DEFAULT_OBJECT_FACTORY, DEFAULT_OBJECT_WRAPPER_FACTORY, REFLECTOR_FACTORY);
        BoundSql boundSql = (BoundSql) metaStatementHandler.getValue("delegate.boundSql");//獲取sql語句
        String originSql = boundSql.getSql();
        //log.info("boundSql:" + originSql);
        if (!StringUtils.isEmpty(originSql)) {
            MappedStatement mappedStatement = (MappedStatement) metaStatementHandler.getValue("delegate.mappedStatement");
            String id = mappedStatement.getId();
            String className = id.substring(0, id.lastIndexOf("."));
            Class<?> clazz = Class.forName(className);
            TableSharde tableSharde = clazz.getAnnotation(TableSharde.class);
            String newSql = "";
            if (tableSharde != null) {
                String tableName = tableSharde.tableName();
                String[] shardeBy = tableSharde.shardeBy();
                Class<?> strategyClazz = tableSharde.strategy();
                TableStrategy tableStrategy = (TableStrategy)strategyClazz.newInstance();
                newSql = tableStrategy.doSharde(metaStatementHandler, tableName, shardeBy);
                //log.info("newSql:" + newSql);
                metaStatementHandler.setValue("delegate.boundSql.sql", newSql);
            }

        }
        else {
            log.error("TableShardeInterceptor error,sql is empty");
            throw new BusinessException("TableShardeInterceptor error,sql is empty");
        }

        //Object parameterObject = metaStatementHandler.getValue("delegate.boundSql.parameterObject");//獲取參數

        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        // 當目標類是StatementHandler類型時，才包裝目標類，否者直接返回目標本身,減少目標被代理的次數
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    @Override
    public void setProperties(Properties properties) {

    }

分表策略實現，我這里是根據用戶id實現分表的策略，大家也可以實現自己的策略

import com.twsz.tablesharde.TableStrategy;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.reflection.MetaObject;

import java.lang.reflect.Field;
import java.util.Map;
import java.util.Set;

public class ShardeByUserCodeStrategy implements TableStrategy {
    @Override
    public String doSharde(MetaObject metaStatementHandler, String tableName, String[] shardeBy) throws Exception{
        BoundSql boundSql = (BoundSql) metaStatementHandler.getValue("delegate.boundSql");//獲取sql語句
        String originSql = boundSql.getSql();
        boundSql.getParameterMappings();
        String userCode = shardeBy[0];
        Object parameterObject = metaStatementHandler.getValue("delegate.boundSql.parameterObject");//獲取參數
        if (parameterObject instanceof String) {
            originSql = originSql.replaceAll(tableName, tableName + "_" + parameterObject);
        }
        else if (parameterObject instanceof Map) {
            Map<String, Object> map = (Map<String, Object>)parameterObject;
            Set<String> set = map.keySet();
            String value = "";
            for (String key: set) {
                if (key.equals(userCode)) {
                    value = map.get(userCode).toString();
                    break;
                }
            }
            originSql = originSql.replaceAll(tableName, tableName + "_" + value);
        }
        else {
            Class<?> clazz = parameterObject.getClass();
            String value = "";
            Field[] fields = clazz.getDeclaredFields();
            for (Field field: fields) {
                field.setAccessible(true);
                String fieldName = field.getName();
                if (fieldName.equals(userCode)) {
                    value = field.get(parameterObject).toString();
                    break;
                }
            }
            originSql = originSql.replaceAll(tableName, tableName + "_" + value);
        }
        return originSql;
    }

最后需要在數據庫會話工廠中配置插件

@Bean(name = "masterSqlSessionFactory")
    @Primary
    SqlSessionFactory masterSqlSessionFactory(@Qualifier("masterDataSource") DataSource dataSource) throws Exception {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        sqlSessionFactoryBean.setDataSource(dataSource);
        sqlSessionFactoryBean.setMapperLocations(new PathMatchingResourcePatternResolver().getResources("classpath:mapper/*.xml"));
        sqlSessionFactoryBean.setConfigLocation(new PathMatchingResourcePatternResolver().getResource("classpath:mybatis-config.xml"));
        sqlSessionFactoryBean.setPlugins(new Interceptor[] {new TableShardeInterceptor()});
        return sqlSessionFactoryBean.getObject();
    }