MyBatis generator - 简单使用（4）- 自定义生成（修改源码）

不念不忘少年蓝@ 2023-01-17 07:43 242阅读 0赞

### 一、官方相关文档 ###

*  github：https://github.com/mybatis/generator
 *  文档：http://mybatis.org/generator/index.html

### 二、自定义生成 ###

在上一篇博客（[MyBatis generator - 简单使用（3）- 自定义注释][MyBatis generator - _3_-]）中自定义了注释生成器，以及博客（[MyBatis generator - 简单使用（2）- 配置详解][MyBatis generator - _2_-]）详细介绍了MBG常用的一些标签的配置，但是还是不能满足一些个性化需求，或者根本不支持通过配置修改，所以这篇博客示例下直接修改源码后生成符合需求的模板代码，不局限于注释。

从github下载源码就可开搞了，我下载的版本是`1.4.0`，要求JDK1.8。

和上一篇博客一样，需求如下：

*  实体类：
    
     *  注释：生成版权、类注释、参数注释
     *  注解：类上添加注解并导包
     *  不生成getter、setter方法
     *  数据库的BIT（Boolean）、SMALLINT（Short）、TINYINT（Byte）映射成Java的Integer
     *  不生成WithBlobs的实体类
 *  Dao：
    
     *  注释：生成版权、类注释，方法不生成注释
     *  方法名修改：selectByPrimaryKey -> findById，updateByPrimaryKey-> update，deleteByPrimaryKey -> deleteById
     *  不生成WithBlobs的方法
 *  XML：
    
     *  不生成注释
     *  不生成WithBlobs的标签

注：

*  实体类由`org.mybatis.generator.codegen.mybatis3.model.BaseRecordGenerator`生成
 *  Dao由`org.mybatis.generator.codegen.mybatis3.javamapper.JavaMapperGenerator`生成
 *  注释默认由`org.mybatis.generator.internal.DefaultCommentGenerator`生成

#### 2.1 自定义注释 ####

为了获取到表的备注信息，在配置MBG时，需要在MySQL连接添加`characterEncoding=utf8&useInformationSchema=true`，否则获取不要表的REMARK信息；若只添加`useInformationSchema=true`，获取到的备注信息可能会乱码。

示例：

（1）版权信息

修改DefaultCommentGenerator\#addJavaFileComment()方法，定义类的版权信息，支持 实体类、Dao

public void addJavaFileComment(CompilationUnit compilationUnit) { 
        if (suppressAllComments) { 
            return;
        }
    
        // 版权信息
        compilationUnit.addFileCommentLine("/*");
        compilationUnit.addFileCommentLine(" * Copyright © MOMO All Rights Reserved.");
        compilationUnit.addFileCommentLine(" */");
    }

（2）实体类注释

修改`DefaultCommentGenerator#addModelClassComment()`方法，自定义符合公司规范的注释模板。

/** * 实体类的注释 */
    public void addModelClassComment(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) { 
        if (suppressAllComments || !addRemarkComments) { 
            return;
        }
        // 实体类注释
        // BaseRecordGenerator类没有调用addClassComment()方法，调用了addModelClassComment()，所以在这里添加实体类的注释
        topLevelClass.addJavaDocLine("/**");
        // introspectedTable.getRemarks()：获取表的备注
        topLevelClass.addJavaDocLine(" * " + introspectedTable.getRemarks() + "（表：" + introspectedTable.getFullyQualifiedTable() + "）");
        topLevelClass.addJavaDocLine(" *");
        topLevelClass.addJavaDocLine(" * @author MQG");
        topLevelClass.addJavaDocLine(" * @date " + getDateString());
        topLevelClass.addJavaDocLine(" */");
    }

（3）实体类字段注释

修改DefaultCommentGenerator\#addFieldComment()方法，定义实体类的字段注释

/** * 实体类字段注释 */
    @Override
    public void addFieldComment(Field field, IntrospectedTable introspectedTable, IntrospectedColumn introspectedColumn) { 
        if (suppressAllComments) { 
            return;
        }
    	// 由BaseRecordGenerator调用JavaBeansUtil.getJavaBeansField()方法 -> JavaBeansUtil.addGeneratedJavaDoc()方法
    	//最终调用 DefaultCommentGenerator#addFieldComment()方法
        field.addJavaDocLine("/**");
        field.addJavaDocLine(" * " + introspectedColumn.getRemarks());
        field.addJavaDocLine(" */");
    }

（4）Dao的类注释

`org.mybatis.generator.codegen.mybatis3.javamapper.JavaMapperGenerator`负责Dao的生成，但是并没有调用生成类注释，所以修改`JavaMapperGenerator#getCompilationUnits()`方法，添加生成类注释的入口：

/** * MyBatis Mapper文件生成器 */
    public class JavaMapperGenerator extends AbstractJavaClientGenerator { 
    	@Override
        public List<CompilationUnit> getCompilationUnits() { 
            // 略
    		Interface interfaze = new Interface(type);
            interfaze.setVisibility(JavaVisibility.PUBLIC);
            commentGenerator.addJavaFileComment(interfaze);
            
            // 新增：添加类注释
            commentGenerator.addClassComment(interfaze, introspectedTable);
            
            // 略
    	}
    }

修改`DefaultCommentGenerator#addClassComment()`方法，定义实体类的字段注释

/** * Dao类注释 */
    @Override
    public void addClassComment(InnerInterface innerInterface, IntrospectedTable introspectedTable) { 
    	if (suppressAllComments) { 
           return;
        }
    	// Dao注释
        innerInterface.addJavaDocLine("/**");
        // introspectedTable.getRemarks()：获取表的备注
        innerInterface.addJavaDocLine(" * " + introspectedTable.getRemarks() + "Dao");
        innerInterface.addJavaDocLine(" *");
        innerInterface.addJavaDocLine(" * @author MQG");
        innerInterface.addJavaDocLine(" * @date " + getDateString());
        innerInterface.addJavaDocLine(" */");
    }

（5）Dao的方法不生成注释

在生成Dao方法的各个类中，都会调用`context.getCommentGenerator().addGeneralMethodComment();`生成方法的注释，例如：

*  `org.mybatis.generator.codegen.mybatis3.javamapper.elements.InsertMethodGenerator#addInterfaceElements()`：生成insert()方法
 *  `org.mybatis.generator.codegen.mybatis3.javamapper.elements.InsertSelectiveMethodGenerator#addInterfaceElements()`：生成insertSelective()方法
 *  …

这里我们直接去掉注释内容：

/** * 添加普通方法注释 * * <p>适用于Dao等的方法注释 */
    @Override
    public void addGeneralMethodComment(Method method, IntrospectedTable introspectedTable) { 
    }

（6）XML不生成注释

和Dao方法的注释类似，都是调用`context.getCommentGenerator().addComment(XmlElement);`方法生成注释，例如：

*  `org.mybatis.generator.codegen.mybatis3.xmlmapper.elements.InsertElementGenerator#addElements()`：生成insert查询语句
 *  `org.mybatis.generator.codegen.mybatis3.xmlmapper.elements.InsertSelectiveElementGenerator#addElements()`：生成insertSelective查询语句
 *  …

这里直接去掉注释内容：

/** * 添加XML元素的注释 * * <p>适用于Dao等的方法注释 */
    @Override
    public void addComment(XmlElement xmlElement) { 
    }

#### 2.2 实体类注解 ####

为实体类自动添加@Data、@Builder等注解。

由于`DefaultCommentGenerator`实现的`org.mybatis.generator.api.CommentGenerator`接口只有以下三个添加注解的方法：

*  addClassAnnotation：为类添加@Generated注解，由`DynamicSqlSupportClassGenerator`在调用该方法
 *  addFieldAnnotation：为字段添加@Generated注解
 *  addGeneralMethodAnnotation：为方法添加@Generated注解

看了下`addClassAnnotation()`方法中的内容以及调用的类，不太适合为实体类添加自定义的注解，所以我们自定义个为实体类添加注解的方法：

（1）`org.mybatis.generator.api.CommentGenerator`接口新增方法：

/** * 为实体类添加注解（自定义的方法） * * @param topLevelClass the top level class */
    default void addModelClassAnnotation(TopLevelClass topLevelClass) { 
    }

（2）`DefaultCommentGenerator`重写该方法，添加自定义的注解：

/** * 实体类注解（自定义的方法） */
    @Override
    public void addModelClassAnnotation(TopLevelClass topLevelClass) { 
        // import ModelProperty、Data、Builer、 NoArgsConstructor、AllArgsConstructor
        Map<String, String> importTypeMap = new HashMap<>(16);
        importTypeMap.put("@Data", "lombok.Data");
        importTypeMap.put("@Builder", "lombok.Builder");
        importTypeMap.put("@NoArgsConstructor", "lombok.NoArgsConstructor");
        importTypeMap.put("@AllArgsConstructor", "lombok.AllArgsConstructor");
    
        for (Map.Entry<String, String> importType : importTypeMap.entrySet()) { 
            // 添加import
            FullyQualifiedJavaType fullyQualifiedJavaType = new FullyQualifiedJavaType(importType.getValue());
            topLevelClass.addImportedType(fullyQualifiedJavaType);
            // 添加注解
            topLevelClass.addAnnotation(importType.getKey());
        }
    }

（3）修改`BaseRecordGenerator#getCompilationUnits()`方法，调用该方法：

public List<CompilationUnit> getCompilationUnits() { 
    	// 略
    	
        // 类文件注释（版权）
        commentGenerator.addJavaFileComment(topLevelClass);
        // 实体类注释
        commentGenerator.addModelClassComment(topLevelClass, introspectedTable);
    
        // 新增：添加实体类注解
        commentGenerator.addModelClassAnnotation(topLevelClass);
    	// 略
    }

#### 2.3 实体类不生成getter、setter ####

在`BaseRecordGenerator#getCompilationUnits()`方法中有生成getter、setter方法的逻辑，注掉即可：

@Override
    public List<CompilationUnit> getCompilationUnits() { 
        // 略
    
        String rootClass = getRootClass();
        for (IntrospectedColumn introspectedColumn : introspectedColumns) { 
            if (RootClassInfo.getInstance(rootClass, warnings).containsProperty(introspectedColumn)) { 
                continue;
            }
    
            // 生成实体类字段，包括注释
            Field field = getJavaBeansField(introspectedColumn, context, introspectedTable);
            if (plugins.modelFieldGenerated(field, topLevelClass, introspectedColumn, introspectedTable,  Plugin.ModelClassType.BASE_RECORD)) { 
                topLevelClass.addField(field);
                topLevelClass.addImportedType(field.getType());
            }
    
            // 去除getter、setter
    // Method method = getJavaBeansGetter(introspectedColumn, context, introspectedTable);
    // if (plugins.modelGetterMethodGenerated(method, topLevelClass, introspectedColumn, introspectedTable, Plugin.ModelClassType.BASE_RECORD)) { 
    // topLevelClass.addMethod(method);
    // }
    //
    // if (!introspectedTable.isImmutable()) { 
    // method = getJavaBeansSetter(introspectedColumn, context, introspectedTable);
    // if (plugins.modelSetterMethodGenerated(method, topLevelClass, introspectedColumn, introspectedTable, Plugin.ModelClassType.BASE_RECORD)) { 
    // topLevelClass.addMethod(method);
    // }
    // }
        }
    
        // 略
    }

#### 2.4 实体类字段类型 ####

MBG中数值类型的BIT、TINYINT、SMALLINT映射成的Java类型分别为Boolean、Byte、Short，我们要把它修改成Integer，对应生成实体类的字段类型就可以更改。

MBG中是写死了MySQL数据类型和Java类的映射关系，对应的映射关系在`org.mybatis.generator.internal.types.JavaTypeResolverDefaultImpl`的typeMap中，在实例化该类的时候，构造方法中会初始化typeMap，修改如下：

public JavaTypeResolverDefaultImpl() { 
        super();
        properties = new Properties();
        typeMap = new HashMap<>();
        
        // 略
        // BIT
    // typeMap.put(Types.BIT, new JdbcTypeInformation("BIT", new FullyQualifiedJavaType(Boolean.class.getName())));
        typeMap.put(Types.BIT, new JdbcTypeInformation("BIT", new FullyQualifiedJavaType(Integer.class.getName())));
    
    	// SMALLINT
    // typeMap.put(Types.SMALLINT, new JdbcTypeInformation("SMALLINT", new FullyQualifiedJavaType(Short.class.getName())));
        typeMap.put(Types.SMALLINT, new JdbcTypeInformation("SMALLINT", new FullyQualifiedJavaType(Integer.class.getName())));
    	
    	// TINYINT
    // typeMap.put(Types.TINYINT, new JdbcTypeInformation("TINYINT", new FullyQualifiedJavaType(Byte.class.getName())));
        typeMap.put(Types.TINYINT, new JdbcTypeInformation("TINYINT", new FullyQualifiedJavaType(Integer.class.getName())));
    	// 略
    }

#### 2.5 Dao方法名称 ####

`org.mybatis.generator.api.IntrospectedTable#calculateXmlAttributes()`方法中保存了生成Dao方法名称（XML中id属性名称也一样）的映射关系，找到需要修改的方法名修改即可：

protected void calculateXmlAttributes() { 
       // ...
    
        // deleteByPrimaryKey -> deleteById
    // setDeleteByPrimaryKeyStatementId("deleteByPrimaryKey");
        setDeleteByPrimaryKeyStatementId("deleteById");
    
        // selectByPrimaryKey -> findById
    // setSelectByPrimaryKeyStatementId("selectByPrimaryKey");
        setSelectByPrimaryKeyStatementId("findById");
    
        // updateByPrimaryKey -> update
    // setUpdateByPrimaryKeyStatementId("updateByPrimaryKey");
        setUpdateByPrimaryKeyStatementId("update");
    
        //...
    }

#### 2.6 不生成带有BLOBS的实体类、Dao中的方法（包括XML） ####

MBG支持通过设置`<table>`标签的`modelType="flat"`设置只生成一个实体类，即使有表中有text之类的列，但是在Dao和XML中仍然会生成带有BLOBS的方法。

所以，最简单的方式，在MBG解析数据库表中列时，即使是`text`、`mediumtext`等类型的列，都解析为普通的类型。

在`org.mybatis.generator.api.IntrospectedTable#addColumn()`方法中，将所有列都添加至`baseColumns`普通列中：

/** * 添加列（区分blob或基本列） * * @param introspectedColumn */
    public void addColumn(IntrospectedColumn introspectedColumn) { 
    // if (introspectedColumn.isBLOBColumn()) { 
    // blobColumns.add(introspectedColumn);
    // } else { 
    // baseColumns.add(introspectedColumn);
    // }
    	// 不区分blob，都添加至baseColums
        baseColumns.add(introspectedColumn);
    
        introspectedColumn.setIntrospectedTable(this);
    }

#### 2.7 生成的文件 ####

（1）实体类  
![在这里插入图片描述][20210628171639841.png]  
（2）Dao  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L215dHRfMTA1NjY_size_16_color_FFFFFF_t_70]  
（3）XML

此时只会有`Base_Column_List`（data是text类型）：  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L215dHRfMTA1NjY_size_16_color_FFFFFF_t_70 1]

总结：

*  通过修改源码就可以满足自定义的生成需求
 *  在公司可以打成jar供其他小伙伴使用，再进一步做个页面，在页面上勾勾选选就可以生成符合需求的样板代码那就更舒服了。。

[MyBatis generator - _3_-]: https://blog.csdn.net/mytt_10566/article/details/115984935
[MyBatis generator - _2_-]: https://blog.csdn.net/mytt_10566/article/details/115877723
[20210628171639841.png]: /images/20221021/be622cde9d334cdfafdeaa949384e260.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L215dHRfMTA1NjY_size_16_color_FFFFFF_t_70]: /images/20221021/248b597b91804ceb87a2043e68998ea1.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L215dHRfMTA1NjY_size_16_color_FFFFFF_t_70 1]: /images/20221021/72e656af9c6e4236a1419a7d7bf11f49.png