基于注解的Mybatis Generator插件

在使用Mybatis时,苦于在Mapper和XML之间切换带来的烦恼,好在有注解的方式来解决这个问题.但在使用Mybatis Generator时发现其生成的方法有限,所以便有了理由研究这方面的东西.先来赶下效果:

首先引入 maven 的依赖:

<dependency>
   <groupId>org.mybatis.generator</groupId>
   <artifactId>mybatis-generator-core</artifactId>
   <version>1.3.5</version>
</dependency>

插件类 MapperPlugin:

import org.mybatis.generator.api.GeneratedJavaFile;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.java.*;

import java.util.List;

import static org.mybatis.generator.internal.util.StringUtility.stringHasValue;

/**
 * @author WangXingYu
 * @date 2018-01-31
 */
public class MapperPlugin extends PluginAdapter {

    private String daoTargetPackage;

    @Override
    public boolean validate(List<String> warnings) {

        boolean valid1 = stringHasValue(properties.getProperty("targetProject"));

        daoTargetPackage = properties.getProperty("targetPackage");
        boolean valid2 = stringHasValue(daoTargetPackage);

        return valid1 && valid2;
    }

    @Override
    public List<GeneratedJavaFile> contextGenerateAdditionalJavaFiles(IntrospectedTable introspectedTable) {
        //获取所有生成的文件
        List<GeneratedJavaFile> generatedJavaFiles = introspectedTable.getGeneratedJavaFiles();
        for (GeneratedJavaFile javaFile : generatedJavaFiles) {
            //获取编辑项(存放了生成文件时所需要的信息)
            CompilationUnit unit = javaFile.getCompilationUnit();
            FullyQualifiedJavaType baseModelJavaType = unit.getType();
            String shortName = baseModelJavaType.getShortName();
            //判断生成的文件是否是Mapper接口
            if (shortName.endsWith("Mapper")) {
                if (stringHasValue(daoTargetPackage)) {
                    if (unit instanceof Interface) {
                        /*为文件添加扩展的方法*/
                        //添加import项
                        unit.addImportedType(new FullyQualifiedJavaType("java.util.List"));
                        //新建一个方法
                        Method method = new Method();
                        //设置访问权限
                        method.setVisibility(JavaVisibility.PUBLIC);
                        //设置方法名称
                        method.setName("selectAll");
                        //设置返回类型

                        FullyQualifiedJavaType returnFullyQualifiedJavaType = new FullyQualifiedJavaType(String.format("List<%s>", introspectedTable.getBaseRecordType()));
                        method.setReturnType(returnFullyQualifiedJavaType);
                        //设置注解
                        AnnotatedSelectAllMethodGengrator annotatedSelectAllMethodGengrator = new AnnotatedSelectAllMethodGengrator(introspectedTable, false, false);
                        annotatedSelectAllMethodGengrator.addMapperAnnotations(((Interface) unit), method);
                        //把方法添加到编辑项
                        ((Interface) unit).addMethod(method);
                    }
                }
            }
        }
        return generatedJavaFiles;
    }
}

Mybatis Generator的插件需要继承PluginAdapter类,并且重写validate().contextGenerateAdditionalJavaFiles()是用来添加生成文件的,Mybatis Generator会生成此方法返回的List<GeneratedJavaFile>,具体操作已经在上面有注释了.重点是设置注解这里,方法的注解就是在这里生成的.

SelectAllMethodGengrator 类

import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.dom.java.*;
import org.mybatis.generator.codegen.mybatis3.javamapper.elements.AbstractJavaMapperMethodGenerator;

import java.util.List;
import java.util.Set;
import java.util.TreeSet;

/**
 * @author WangXingYu
 * @date 2018-02-02
 */
public class SelectAllMethodGengrator extends AbstractJavaMapperMethodGenerator {
    private boolean isSimple;

    SelectAllMethodGengrator(boolean isSimple) {
        super();
        this.isSimple = isSimple;
    }

    @Override
    public void addInterfaceElements(Interface interfaze) {
        Set<FullyQualifiedJavaType> importedTypes = new TreeSet<>();
        Method method = new Method();
        method.setVisibility(JavaVisibility.PUBLIC);

        FullyQualifiedJavaType returnType = introspectedTable.getRules().calculateAllFieldsClass();
        method.setReturnType(returnType);
        importedTypes.add(returnType);

        method.setName(introspectedTable.getSelectByPrimaryKeyStatementId());

        if (!isSimple && introspectedTable.getRules().generatePrimaryKeyClass()) {
            FullyQualifiedJavaType type = new FullyQualifiedJavaType(introspectedTable.getPrimaryKeyType());
            importedTypes.add(type);
            method.addParameter(new Parameter(type, "key"));
        } else {
            // no primary key class - fields are in the base class
            // if more than one PK field, then we need to annotate the
            // parameters
            // for MyBatis3
            List<IntrospectedColumn> introspectedColumns = introspectedTable.getPrimaryKeyColumns();
            boolean annotate = introspectedColumns.size() > 1;
            if (annotate) {
                importedTypes.add(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Param"));
            }
            StringBuilder sb = new StringBuilder();
            for (IntrospectedColumn introspectedColumn : introspectedColumns) {
                FullyQualifiedJavaType type = introspectedColumn.getFullyQualifiedJavaType();
                importedTypes.add(type);
                Parameter parameter = new Parameter(type, introspectedColumn.getJavaProperty());
                if (annotate) {
                    sb.setLength(0);
                    sb.append("@Param(\"");
                    sb.append(introspectedColumn.getJavaProperty());
                    sb.append("\")");
                    parameter.addAnnotation(sb.toString());
                }
                method.addParameter(parameter);
            }
        }

        addMapperAnnotations(interfaze, method);

        context.getCommentGenerator().addGeneralMethodComment(method, introspectedTable);

        if (context.getPlugins().clientSelectByPrimaryKeyMethodGenerated(method, interfaze, introspectedTable)) {
            interfaze.addImportedTypes(importedTypes);
            interfaze.addMethod(method);
        }
    }

    public void addMapperAnnotations(Interface interfaze, Method method) {
    }
}

AnnotatedSelectAllMethodGengrator

import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.api.dom.java.Interface;
import org.mybatis.generator.api.dom.java.Method;

import java.util.Iterator;

import static org.mybatis.generator.api.dom.OutputUtilities.javaIndent;
import static org.mybatis.generator.codegen.mybatis3.MyBatis3FormattingUtilities.getSelectListPhrase;
import static org.mybatis.generator.internal.util.StringUtility.escapeStringForJava;

/**
 * @author WangXingYu
 * @date 2018-02-02
 */
public class AnnotatedSelectAllMethodGengrator extends SelectAllMethodGengrator {

    private boolean useResultMapIfAvailable;

    AnnotatedSelectAllMethodGengrator(IntrospectedTable introspectedTable, boolean useResultMapIfAvailable, boolean isSimple) {
        super(isSimple);
        this.useResultMapIfAvailable = useResultMapIfAvailable;
        this.introspectedTable = introspectedTable;
    }

    @Override
    public void addMapperAnnotations(Interface interfaze, Method method) {

        interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Select"));

        StringBuilder sb = new StringBuilder();
        method.addAnnotation("@Select({");
        javaIndent(sb, 1);
        sb.append("\"select\",");
        method.addAnnotation(sb.toString());

        Iterator<IntrospectedColumn> iter = introspectedTable.getAllColumns().iterator();
        sb.setLength(0);
        javaIndent(sb, 1);
        sb.append('"');
        boolean hasColumns = false;
        while (iter.hasNext()) {
            sb.append(escapeStringForJava(getSelectListPhrase(iter.next())));
            hasColumns = true;

            if (iter.hasNext()) {
                sb.append(", ");
            }

            if (sb.length() > 80) {
                sb.append("\",");
                method.addAnnotation(sb.toString());

                sb.setLength(0);
                javaIndent(sb, 1);
                sb.append('"');
                hasColumns = false;
            }
        }

        if (hasColumns) {
            sb.append("\",");
            method.addAnnotation(sb.toString());
        }

        sb.setLength(0);
        javaIndent(sb, 1);
        sb.append("\"from ");
        sb.append(escapeStringForJava(introspectedTable.getAliasedFullyQualifiedTableNameAtRuntime()));
        sb.append("\"");
        method.addAnnotation(sb.toString());
        method.addAnnotation("})");

        if (useResultMapIfAvailable) {
            if (introspectedTable.getRules().generateBaseResultMap() || introspectedTable.getRules().generateResultMapWithBLOBs()) {
                addResultMapAnnotation(interfaze, method);
            } else {
                addAnnotatedResults(interfaze, method);
            }
        } else {
            addAnnotatedResults(interfaze, method);
        }
    }

    private void addResultMapAnnotation(Interface interfaze, Method method) {
        interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.annotations.ResultMap"));

        String annotation = String.format("@ResultMap(\"%s\")", introspectedTable.getBaseRecordType().substring(introspectedTable.getBaseRecordType().lastIndexOf(".") + 1).toLowerCase());
        method.addAnnotation(annotation);
    }

    private void addAnnotatedResults(Interface interfaze, Method method) {
        interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.type.JdbcType"));

        if (introspectedTable.isConstructorBased()) {
            interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Arg"));
            interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.annotations.ConstructorArgs"));
            method.addAnnotation("@ConstructorArgs({");
        } else {
            interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Result"));
            interfaze.addImportedType(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Results"));
            method.addAnnotation(String.format("@Results(id = \"%s\", value= {", introspectedTable.getBaseRecordType().substring(introspectedTable.getBaseRecordType().lastIndexOf(".") + 1).toLowerCase()));
        }
        StringBuilder sb = new StringBuilder();

        Iterator<IntrospectedColumn> iterPk = introspectedTable.getPrimaryKeyColumns().iterator();
        Iterator<IntrospectedColumn> iterNonPk = introspectedTable.getNonPrimaryKeyColumns().iterator();
        while (iterPk.hasNext()) {
            IntrospectedColumn introspectedColumn = iterPk.next();
            sb.setLength(0);
            javaIndent(sb, 1);
            sb.append(getResultAnnotation(interfaze, introspectedColumn, true, introspectedTable.isConstructorBased()));

            if (iterPk.hasNext() || iterNonPk.hasNext()) {
                sb.append(',');
            }
            method.addAnnotation(sb.toString());
        }

        while (iterNonPk.hasNext()) {
            IntrospectedColumn introspectedColumn = iterNonPk.next();
            sb.setLength(0);
            javaIndent(sb, 1);
            sb.append(getResultAnnotation(interfaze, introspectedColumn, false, introspectedTable.isConstructorBased()));

            if (iterNonPk.hasNext()) {
                sb.append(',');
            }
            method.addAnnotation(sb.toString());
        }
        method.addAnnotation("})");
    }
}

在 Mybatis Genrator 的配置文件中添加自定义的插件:

<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE generatorConfiguration  PUBLIC "-//mybatis.org//DTD MyBatis Generator Configuration 1.0//EN"
        "http://mybatis.org/dtd/mybatis-generator-config_1_0.dtd">
<generatorConfiguration>
    <context id="MysqlContext" targetRuntime="MyBatis3" defaultModelType="flat">

        <!-- 生成的model实现Serializable借口 -->
        <plugin type="org.mybatis.generator.plugins.SerializablePlugin">
            <property name="suppressJavaInterface" value="false"/>
        </plugin>

        <!-- 自动生成equals方法和hashcode方法 -->
        <plugin type="org.mybatis.generator.plugins.EqualsHashCodePlugin"/>

        <!-- 该插件给实体类添加toString() -->
        <plugin type="org.mybatis.generator.plugins.ToStringPlugin"/>

        <!-- 自定义方法 -->
        <plugin type="插件名">
            <property name="targetProject" value="src/main/java"/>
            <property name="targetPackage" value="mapper所在的包名"/>
        </plugin>

        <!--关闭注释 -->
        <commentGenerator>
            <property name="suppressAllComments" value="true"/>
        </commentGenerator>

        <!--数据库连接信息 -->
        <jdbcConnection driverClass="com.mysql.jdbc.Driver" connectionURL="****" userId="****" password="****">
        </jdbcConnection>
        <!--生成的entity 包路径 -->
        <javaModelGenerator targetPackage="实体所在的包名" targetProject="src/main/java">
            <property name="enableSubPackages" value="ture"/>
            <property name="trimStrings" value="true"/>
            <property name="rootClass" value="实体基类"/>
        </javaModelGenerator>

        <!-- type="ANNOTATEDMAPPER",生成基于注解的Mapper接口
                type="MIXEDMAPPER",生成最小的基于注解的Mapper接口
                type="XMLMAPPER",生成SQLMap XML文件和独立的Mapper接口 -->
        <javaClientGenerator type="ANNOTATEDMAPPER" targetPackage="mapper所在的包名" targetProject="src/main/java">
            <property name="enableSubPackages" value="ture"/>
        </javaClientGenerator>
        <!--对应数据库表名 -->
        <table tableName="sec_user" domainObjectName="User" enableCountByExample="false" enableUpdateByExample="false" enableDeleteByExample="false" enableSelectByExample="false" selectByExampleQueryId="false">
            <generatedKey column="id" sqlStatement="MySql" identity="true"/>
        </table>
    </context>
</generatorConfiguration>

如果用 maven 的方式运行的话需要把插件打包添加到依赖中, 感觉很麻烦, 而且不便于调试. 所以我用代码的方式运行, 这个看个人偏好, 都可以的. 下面附上我自己的代码:

import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mybatis.generator.api.MyBatisGenerator;
import org.mybatis.generator.config.Configuration;
import org.mybatis.generator.config.xml.ConfigurationParser;
import org.mybatis.generator.internal.DefaultShellCallback;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;
import org.springframework.util.ResourceUtils;

import java.io.File;
import java.io.FileNotFoundException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author WangXingYu
 * @date 2018-01-31
 */
@RunWith(SpringRunner.class)
@SpringBootTest
public class Generator {

    private File configFile;

    @Before
    public void before() throws FileNotFoundException {
        configFile = ResourceUtils.getFile("classpath:mybatis-generator.xml");
    }

    @Test
    public void generate() throws Exception {
        List<String> warnings = new ArrayList<>();
        ConfigurationParser cp = new ConfigurationParser(warnings);
        Configuration config = cp.parseConfiguration(configFile);
        DefaultShellCallback callback = new DefaultShellCallback(true);
        MyBatisGenerator myBatisGenerator = new MyBatisGenerator(config, callback, warnings);
        myBatisGenerator.generate(null);
    }
}

这里只添加一个 selectAll() 的简单例子, 后面再慢慢研究