Read the MyBatis source code: plug-in development

The last article introduced the MyBatis execution SQL query process . After having a certain understanding of the key classes in the source code such as Configuration, Executor, StatementHandler, this article will talk about the plug-in mechanism of MyBatis.

1 Introduction

The MyBatis plug-in is simply understood as an interceptor. It uses a dynamic proxy method to intercept the target method and do some operations before and after.

Based on the plug-in mechanism, you can basically control the various stages of SQL execution, such as the execution stage, the parameter processing stage, the syntax construction stage, and the result set processing stage. The corresponding business logic can be implemented according to the project business.

Methods to support interception:

1. Executor: update, query, commit, rollback and other methods;
2. Parameter handler ParameterHandler: getParameterObject, setParameters and other methods;
3. Result set processor ResultSetHandler: handleResultSets, handleOutputParameters, etc.;
4. SQL syntax builder StatementHandler: prepare, parameterize, batch, update, query and other methods;

2. Plug-in example

The paging plug-in is defined in the MyBatis XML configuration file.

<?xml version="1.0" encoding="UTF-8" ?>
<!DOCTYPE configuration
        PUBLIC "-//mybatis.org//DTD Config 3.0//EN"
        "http://mybatis.org/dtd/mybatis-3-config.dtd">
<configuration>

    <!-- MyBatis XML 配置说明 https://mybatis.org/mybatis-3/zh/configuration.html -->
    <plugins>
        <plugin interceptor="com.sumkor.plugin.PageInterceptor"/>
    </plugins>

</configuration>

The code of the custom paging plug-in PageInterceptor is as follows, which is used to intercept the Executor#query method and modify the SQL statement in the MappedStatement object.

package com.sumkor.plugin;

import com.sumkor.plugin.page.BoundSqlSqlSource;
import com.sumkor.plugin.page.Page;
import com.sumkor.plugin.page.PageUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.Executor;
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.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.Properties;
import java.util.StringJoiner;

/**
 * 拦截 Executor#query 方法
 *
 * @author Sumkor
 * @since 2021/7/26
 */
@Intercepts({
        @Signature(
                type = Executor.class,
                method = "query",
                args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}
        )
})
@Slf4j
public class PageInterceptor implements Interceptor {

    private static final int MAPPED_STATEMENT_INDEX = 0;
    private static final int PARAMETER_INDEX = 1;
    private static final int ROW_BOUNDS_INDEX = 2;

    /**
     * 通过反射工具类 MetaObject 来修改 MappedStatement 对象中的 SQL 语句
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        log.info("------------------PageInterceptor#intercept 开始------------------");
        final Object[] queryArgs = invocation.getArgs();
        final MappedStatement ms = (MappedStatement) queryArgs[MAPPED_STATEMENT_INDEX];
        final Object parameter = queryArgs[PARAMETER_INDEX];

        // 获取分页参数
        Page pagingParam = PageUtil.getPagingParam();
        try {
            if (pagingParam != null) {
                // 构造新的分页查询 SQL 字符串
                final BoundSql boundSql = ms.getBoundSql(parameter);
                String pagingSql = getPagingSql(boundSql.getSql(), pagingParam.getOffset(), pagingParam.getLimit());
                BoundSql newBoundSql = new BoundSql(ms.getConfiguration(), pagingSql, boundSql.getParameterMappings(), boundSql.getParameterObject());

                // 通过反射工具类，重置 MappedStatement 中的 SQL 语句
                // MetaObject metaObject = MetaObject.forObject(ms, new DefaultObjectFactory(), new DefaultObjectWrapperFactory(), new DefaultReflectorFactory());
                MetaObject metaObject = SystemMetaObject.forObject(ms);
                metaObject.setValue("sqlSource", new BoundSqlSqlSource(newBoundSql));

                // 重置 RowBound
                queryArgs[ROW_BOUNDS_INDEX] = new RowBounds(RowBounds.NO_ROW_OFFSET, RowBounds.NO_ROW_LIMIT);
            }
        } catch (Exception e) {
            log.error("PageInterceptor#intercept 异常", e);
        } finally {
            log.info("------------------PageInterceptor#intercept 结束------------------");
            PageUtil.removePagingParam();
        }
        return invocation.proceed();
    }
    
    /**
     * 使得当前插件生效
     */
    @Override
    public Object plugin(Object o) {
        return Plugin.wrap(o, this);
    }

    /**
     * 构造新的 sql： select xxx from xxx where yyy limit offset,limit
     */
    public String getPagingSql(String sql, int offset, int limit) {
        StringBuilder result = new StringBuilder(sql.length() + 100);
        result.append(sql).append(" limit ");

        if (offset > 0) {
            result.append(offset).append(",").append(limit);
        }else{
            result.append(limit);
        }
        return result.toString();
    }
}

Modify the SQL statement in the MappedStatement object through the reflection tool class MetaObject, plus the paging condition of limit m,n

3. Source code analysis

3.1 Parsing the plug-in configuration

Parse the mybatis-config.xml file through SqlSessionFactoryBuilder:

Reader reader = Resources.getResourceAsReader("mybatis-config.xml");
SqlSessionFactory sqlSessionFactory = new SqlSessionFactoryBuilder().build(reader);

Which will parse the plugins tag:

org.apache.ibatis.builder.xml.XMLConfigBuilder#parseConfiguration
org.apache.ibatis.builder.xml.XMLConfigBuilder#pluginElement

  private void pluginElement(XNode parent) throws Exception {
    if (parent != null) {
      for (XNode child : parent.getChildren()) {
        String interceptor = child.getStringAttribute("interceptor");
        Properties properties = child.getChildrenAsProperties();
        Interceptor interceptorInstance = (Interceptor) resolveClass(interceptor).getDeclaredConstructor().newInstance();
        interceptorInstance.setProperties(properties);
        configuration.addInterceptor(interceptorInstance);
      }
    }
  }

Here, the implementation class of Interceptor will be instantiated and registered in the InterceptorChain member variable in the Configuration object.

org.apache.ibatis.session.Configuration#addInterceptor

  protected final InterceptorChain interceptorChain = new InterceptorChain();

  public void addInterceptor(Interceptor interceptor) {
    interceptorChain.addInterceptor(interceptor);
  }

The complete content of the InterceptorChain class is as follows, in which there is a List<Interceptor> collection that stores the plug-in instances registered to the Configuration object.

org.apache.ibatis.plugin.InterceptorChain

public class InterceptorChain {

  private final List<Interceptor> interceptors = new ArrayList<>();

  public Object pluginAll(Object target) {
    for (Interceptor interceptor : interceptors) {
      target = interceptor.plugin(target);
    }
    return target;
  }

  public void addInterceptor(Interceptor interceptor) {
    interceptors.add(interceptor);
  }

  public List<Interceptor> getInterceptors() {
    return Collections.unmodifiableList(interceptors);
  }

}

3.2 Plug-in implementation mechanism

The Executor, ParameterHandler, ResultSetHandler, StatementHandler classes in MyBatis all support plug-in extensions, and the plug-in implementation mechanism is mainly based on dynamic proxy implementation.

The generation of these objects is managed uniformly by the Configuration object.

  public ParameterHandler newParameterHandler(MappedStatement mappedStatement, Object parameterObject, BoundSql boundSql) {
    ParameterHandler parameterHandler = mappedStatement.getLang().createParameterHandler(mappedStatement, parameterObject, boundSql);
    parameterHandler = (ParameterHandler) interceptorChain.pluginAll(parameterHandler);
    return parameterHandler;
  }

  public ResultSetHandler newResultSetHandler(Executor executor, MappedStatement mappedStatement, RowBounds rowBounds, ParameterHandler parameterHandler,
      ResultHandler resultHandler, BoundSql boundSql) {
    ResultSetHandler resultSetHandler = new DefaultResultSetHandler(executor, mappedStatement, parameterHandler, resultHandler, boundSql, rowBounds);
    resultSetHandler = (ResultSetHandler) interceptorChain.pluginAll(resultSetHandler);
    return resultSetHandler;
  }

  public StatementHandler newStatementHandler(Executor executor, MappedStatement mappedStatement, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
    StatementHandler statementHandler = new RoutingStatementHandler(executor, mappedStatement, parameterObject, rowBounds, resultHandler, boundSql);
    statementHandler = (StatementHandler) interceptorChain.pluginAll(statementHandler);
    return statementHandler;
  }

  public Executor newExecutor(Transaction transaction) {
    return newExecutor(transaction, defaultExecutorType);
  }
  
  public Executor newExecutor(Transaction transaction, ExecutorType executorType) {
    executorType = executorType == null ? defaultExecutorType : executorType;
    executorType = executorType == null ? ExecutorType.SIMPLE : executorType;
    Executor executor;
    // 该类型的执行器会批量执行所有更新语句，如果 SELECT 在多个更新中间执行，将在必要时将多条更新语句分隔开来，以方便理解。
    if (ExecutorType.BATCH == executorType) {        
      executor = new BatchExecutor(this, transaction);
    } 
    // 该类型的执行器会复用预处理语句。
    else if (ExecutorType.REUSE == executorType) { 
      executor = new ReuseExecutor(this, transaction);
    } 
    // 该类型的执行器没有特别的行为。它为每个语句的执行创建一个新的预处理语句。
    else {                                         
      executor = new SimpleExecutor(this, transaction);
    }
    if (cacheEnabled) {
      executor = new CachingExecutor(executor);
    }
    // 使用插件来包装 executor
    executor = (Executor) interceptorChain.pluginAll(executor);
    return executor;
  }

3.2.1 Generation of primitive objects

Before exploring the dynamic proxy mechanism of the plug-in, review the generation process of the original object that needs to be proxied:

Executor

Created when opening a SqlSession session:

org.apache.ibatis.session.defaults.DefaultSqlSessionFactory#openSession
org.apache.ibatis.session.defaults.DefaultSqlSessionFactory#openSessionFromDataSource

  private SqlSession openSessionFromDataSource(ExecutorType execType, TransactionIsolationLevel level, boolean autoCommit) {
    Transaction tx = null;
    try {
      final Environment environment = configuration.getEnvironment();
      final TransactionFactory transactionFactory = getTransactionFactoryFromEnvironment(environment);
      tx = transactionFactory.newTransaction(environment.getDataSource(), level, autoCommit); 
      final Executor executor = configuration.newExecutor(tx, execType); // 创建
      return new DefaultSqlSession(configuration, executor, autoCommit);
    } catch (Exception e) {
      closeTransaction(tx); // may have fetched a connection so lets call close()
      throw ExceptionFactory.wrapException("Error opening session.  Cause: " + e, e);
    } finally {
      ErrorContext.instance().reset();
    }
  }

StatementHandler

SqlSession#selectList is created when executing SQL:

org.apache.ibatis.session.defaults.DefaultSqlSession#selectList
org.apache.ibatis.executor.CachingExecutor#query
org.apache.ibatis.executor.BaseExecutor#query
org.apache.ibatis.executor.BaseExecutor#queryFromDatabase
org.apache.ibatis.executor.SimpleExecutor#doQuery

  @Override
  public <E> List<E> doQuery(MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
    Statement stmt = null;
    try {
      Configuration configuration = ms.getConfiguration();
      StatementHandler handler = configuration.newStatementHandler(wrapper, ms, parameter, rowBounds, resultHandler, boundSql); // 创建
      stmt = prepareStatement(handler, ms.getStatementLog());
      return handler.query(stmt, resultHandler);
    } finally {
      closeStatement(stmt);
    }
  }

ParameterHandler、ResultSetHandler

Created in the constructor of StatementHandler:

org.apache.ibatis.session.Configuration#newStatementHandler
org.apache.ibatis.executor.statement.RoutingStatementHandler#RoutingStatementHandler
org.apache.ibatis.executor.statement.PreparedStatementHandler#PreparedStatementHandler
org.apache.ibatis.executor.statement.BaseStatementHandler#BaseStatementHandler

  protected BaseStatementHandler(Executor executor, MappedStatement mappedStatement, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
    this.configuration = mappedStatement.getConfiguration();
    this.executor = executor;
    this.mappedStatement = mappedStatement;
    this.rowBounds = rowBounds;

    this.typeHandlerRegistry = configuration.getTypeHandlerRegistry();
    this.objectFactory = configuration.getObjectFactory();

    if (boundSql == null) { // issue #435, get the key before calculating the statement
      generateKeys(parameterObject);
      boundSql = mappedStatement.getBoundSql(parameterObject);
    }

    this.boundSql = boundSql;

    this.parameterHandler = configuration.newParameterHandler(mappedStatement, parameterObject, boundSql); // 创建 ParameterHandler，支持插件扩展
    this.resultSetHandler = configuration.newResultSetHandler(executor, mappedStatement, rowBounds, parameterHandler, resultHandler, boundSql); // 创建 ResultSetHandler，支持插件扩展
  }

3.2.2 Generation of proxy objects

When creating Executor, ParameterHandler, ResultSetHandler, StatementHandler objects in the Configuration class, the InterceptorChain#pluginAll method is called to extend the plug-in.

org.apache.ibatis.plugin.InterceptorChain#pluginAll

  public Object pluginAll(Object target) {
    for (Interceptor interceptor : interceptors) {
      target = interceptor.plugin(target);
    }
    return target;
  }

The default implementation of the plugin method in the Interceptor interface is as follows:

org.apache.ibatis.plugin.Interceptor#plugin

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

Two input parameters of Plugin#wrap method:

• target represents the primitive objects of Executor, ParameterHandler, ResultSetHandler, StatementHandler.
• Interceptor is a custom plug-in class.

Code flow:

1. Parse the @Intercepts @Signature annotation on the Interceptor plug-in class.
2. According to the annotation, determine whether it is an interception of the target original object, and if it is, a proxy object is generated for the target original object.

org.apache.ibatis.plugin.Plugin#wrap

  public static Object wrap(Object target, Interceptor interceptor) {
    Map<Class<?>, Set<Method>> signatureMap = getSignatureMap(interceptor); // 解析插件类上的 @Intercepts @Signature 注解
    Class<?> type = target.getClass();
    Class<?>[] interfaces = getAllInterfaces(type, signatureMap);
    if (interfaces.length > 0) { // 满足条件的，说明 target 类需要经过插件 interceptor 来拦截，因此为 target 生成代理
      return Proxy.newProxyInstance(
          type.getClassLoader(),
          interfaces,
          new Plugin(target, interceptor, signatureMap)); // 返回代理对象
    }
    return target; // 返回原始对象
  }

3.2.3 Use of proxy objects

Take the above custom paging plugin as an example, the com.sumkor.plugin.PageInterceptor class annotation declares the interception of the Executor#query method, which is to generate a dynamic proxy for the Executor object.

org.apache.ibatis.session.defaults.DefaultSqlSession#selectList

Therefore, the call to the Executor#query method is actually the method of executing the proxy object:

org.apache.ibatis.plugin.Plugin#invoke

public class Plugin implements InvocationHandler {

  private final Object target;
  private final Interceptor interceptor;
  private final Map<Class<?>, Set<Method>> signatureMap;

  private Plugin(Object target, Interceptor interceptor, Map<Class<?>, Set<Method>> signatureMap) {
    this.target = target;
    this.interceptor = interceptor;
    this.signatureMap = signatureMap;
  }
  
  @Override
  public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
    try {
      Set<Method> methods = signatureMap.get(method.getDeclaringClass());
      if (methods != null && methods.contains(method)) { // 如果插件配置类，声明了对指定的方法的拦截，则符合这里的条件
        return interceptor.intercept(new Invocation(target, method, args)); // 将请求转发给插件方法
      }
      return method.invoke(target, args);
    } catch (Exception e) {
      throw ExceptionUtil.unwrapThrowable(e);
    }
  }
}

Further call to the method of the custom plug-in class:

com.sumkor.plugin.PageInterceptor#intercept

4. Summary

The MyBatis plug-in intercepts methods on the four interfaces of Executor, ParameterHandler, ResultSetHandler, and StatementHandler, and uses the JDK dynamic proxy mechanism to create proxy objects for the implementation classes of these interfaces.

When executing the method, first execute the method of the proxy object, so as to execute the interception logic written by yourself.

It is worth noting that the life cycle of these objects such as Executor, ParameterHandler, ResultSetHandler, StatementHandler are within the scope of Session, and a new proxy object will be created every time a SqlSession session is opened.

---

Author: Sumkor
Link: https://segmentfault.com/a/1190000040647196