IDEA 根据模板生成代码Generate POJOs.groovy

官方的Generate POJOs.groovy无法定制， 基于官方脚本进行了修改， 支持自定义Volocity模板生成代码， 无需额外增加任何依赖。

1.放置脚本

Generate Template POJOs.groovy

import com.intellij.database.model.DasTable
import com.intellij.database.model.ObjectKind
import com.intellij.database.util.Case
import com.intellij.database.util.DasUtil

import java.io.StringWriter
import java.util.HashSet

import org.apache.velocity.VelocityContext
import org.apache.velocity.Template
import org.apache.velocity.app.VelocityEngine
import org.apache.velocity.app.Velocity
import org.apache.velocity.exception.ResourceNotFoundException
import org.apache.velocity.exception.ParseErrorException
import org.apache.velocity.exception.MethodInvocationException
import org.apache.velocity.runtime.RuntimeConstants

/*
 * Available context bindings:
 *   SELECTION   Iterable<DasObject>
 *   PROJECT     project
 *   FILES       files helper
 */

typeMapping = [
        (~/(?i)bit/)                   : "Boolean",
        (~/(?i)bigint/)                   : "Long",
        (~/(?i)bigint unsigned/)          : "Long",
        (~/(?i)tinyint/)                  : "Byte",
        (~/(?i)int/)                      : "Integer",
        (~/(?i)int unsigned/)                      : "Integer",
        (~/(?i)mediumint/)                      : "Integer",
        (~/(?i)float|double|decimal|real/): "java.math.BigDecimal",
        (~/(?i)datetime|timestamp/)       : "java.util.Date",
        (~/(?i)date/)                     : "java.util.Date",
        (~/(?i)time/)                     : "java.sql.Time",
        (~/(?i)/)                         : "String"
]

FILES.chooseDirectoryAndSave("Choose directory", "Choose where to store generated files") { dir ->
    SELECTION.filter { it instanceof DasTable && it.getKind() == ObjectKind.TABLE }.each { generate(it, dir) }
}


VelocityEngine velocityEngine() {
    // 初始化模版引擎
    VelocityEngine ve = new VelocityEngine()
    // 这两个属性可以从RuntimeConstants常量中找到， 引用常量有些版本报错， 就直接写死了
    ve.setProperty("runtime.log", PROJECT.getBaseDir().path + "/dbTools.log") // 对应RuntimeConstants.RUNTIME_LOG
    ve.setProperty("file.resource.loader.path", PROJECT.getBaseDir().path)
    // 对应RuntimeConstants.FILE_RESOURCE_LOADER_PATH
    ve.init()
    ve
}

List<File> listVmFiles(File dir) {
    return listVmFiles0(dir, new ArrayList<>());
}

List<File> listVmFiles0(File dir, List<File> fileList) {
    File[] files = dir.listFiles();
    if(files == null) {
        return fileList;
    }

    for(File file: files) {
        if(file.isDirectory()) {
            listVmFiles0(file, fileList);
        } else {
            String fileName = file.getName();
            if(fileName.startsWith("Generate_POJO_") && fileName.endsWith(".vm")) {
                fileList.add(file);
            }
        }
    }
    return fileList;
}

def generate(table, dir) {
    // 转驼峰， 如t_user转换为TUser
    def className = javaName(table.getName(), true)

    // 读取所有模板文件
    List<File> templateFileList = listVmFiles(dir)
    if(templateFileList == null || templateFileList.isEmpty()) {
        return
    }

    String templateEncoding = "UTF-8"
    VelocityEngine ve = velocityEngine()

    // 根据模板循环生成代码
    for(File templateFile: templateFileList) {
        def templatePath = templateFile.getAbsolutePath().substring(PROJECT.getBaseDir().path.length() + 1)
        Template template = ve.getTemplate(templatePath, templateEncoding)

        VelocityContext ctx = new VelocityContext()
        // 设置变量
        setContextProperty(ctx, table, className, dir)

        // 输出
        StringWriter sw = new StringWriter()
        template.merge(ctx, sw)
        String generateContent = sw.toString()

        // 默认生成的文件名
        String fileName = className + ".java";

        // 读取属性定义
        def lines = generateContent.split("\n")

        // 存储正文(去除属性定义)
        sw = new StringWriter()

        boolean isDefinedReading = true
        for(int i = 0; i < lines.length; i ++){
            def line = lines[i]
            if(isDefinedReading) {
                // 先读取属性定义
                if(line.startsWith("#fileName=")) {
                    fileName = line.substring("#filename=".length()).trim()
                    continue
                }

                isDefinedReading = false
            }
            // 再读取正文
            sw.append(line)
        }

//        new File(new File("D:\\"), fileName).append(fileName)

        new File(templateFile.getParentFile(),  fileName).withPrintWriter("UTF-8") { out -> out.print sw }
    }
}

def setContextProperty(ctx, table, className, dir) {
    // 首字母小写
    def memberName = String.valueOf(Character.toLowerCase(className.charAt(0))) + className.substring(1)

    // 将类信息放入模板变量
    ctx.put("class", [
            "name"   : className,   // 类名， 首字母大写
            "member"   : memberName,   // 类名, 首字母小写
            "comment": table.comment    // 表注释
    ])

    ctx.put("table", [
            "name": table.getName()
    ])

    def cmbFields = calcFields(table)
    // 将字段信息放入模板变量
    ctx.put("imports", cmbFields.imports)
    ctx.put("fields", cmbFields.fields)
}

def calcFields(table) {
    def imports = [] as HashSet
    def fields = DasUtil.getColumns(table).reduce([]) { fields, col ->
        def spec = Case.LOWER.apply(col.getDataType().getSpecification())
        def typeStr = typeMapping.find { p, t -> p.matcher(spec).find() }.value

        if (typeStr.contains(".")) {
            imports << typeStr
            typeStr = typeStr.substring(typeStr.lastIndexOf(".") + 1)
        }

        def jdbcType = Case.UPPER.apply(col.getDataType().typeName)

        switch (jdbcType) {
            case "TINYINT": {
                if (col.getDataType().size == 1) {
                    // 一个长度的TINYINT, 用Boolean类型
                    typeStr = "Boolean"
                    jdbcType = "BIT"
                }
                break
            }
            case "INT":
            case "INT UNSIGNED":
            case "MEDIUMINT": {
                typeStr = "Integer"
                jdbcType = "INTEGER"
                break
            }
            case "TEXT": {
                typeStr = "String"
                jdbcType = "VARCHAR"
                break
            }
        }

        fields += [[
                           name   : javaName(col.getName(), false),
                           type   : typeStr,
                           colName: col.getName(),
//                           jdbcType: Case.LOWER.apply(col.getDataType().getSpecification()),
                           jdbcType: jdbcType == "INT" ? "INTEGER" : jdbcType,
                           sizeUnit: col.getDataType().sizeUnit,
                           size: col.getDataType().size,
                           comment: col.comment,
                           annos  : ""]]

    }
    ["fields": fields, "imports": imports]
}

def javaName(str, capitalize) {
    def s = com.intellij.psi.codeStyle.NameUtil.splitNameIntoWords(str)
            .collect { Case.LOWER.apply(it).capitalize() }
            .join("")
            .replaceAll(/[^\p{javaJavaIdentifierPart}[_]]/, "_")
    capitalize || s.length() == 1 ? s : Case.LOWER.apply(s[0]) + s[1..-1]
}

放到schema目录:

2. 定义模板

在POJO同目录下创建Volocity模板文件, 定义生成模板

文件模板命名规则： Generate_POJO_{自定义}.vm, 必须按指定前缀命名
示例模板Generate_POJO_ReqDTO.vm

#fileName=${class.name}ReqDTO.java
package com.woo.service.dto.req;

import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import com.woo.dao.entity.${class.name};

#foreach($imp in $imports)
import $imp;
#end

/**
* $class.comment
*/
@ToString
@Getter
@Setter
public class ${class.name}ReqDTO extends ${class.name} {

}

其中第一行#fileName=${class.name}ReqDTO.java表示生成后的文件名, 可以是任意文件，如.java、.xml等

3. 生成POJO

选中需要生成POJO的表， 支持多选

4. 指定模块路径

在上一步弹框中， 选择模块路径

可以指定根路径， 会自动查找下级路径的模板； 也可以指定到子路径， 灵活控制生成的代码范围

脚本将自动搜索指定目录下所有的模板定义Generate_POJO_{自定义}.vm, 并在模板定义目录下生成对应的POJO:

我选中了两张表， 所以生成了两个POJO

模板中可使用的变量:

### 类名(表明转驼峰，首字母大写)
${class.name}

### 表注释
${class.comment}

### 表字段需要import的类， 如java.math.BigDecimal
${imports}
### 示例
#foreach($imp in $imports)
import $imp;
#end

### 字段名(列名转驼峰，首字母小写)
${fields}
### 示例
#foreach($field in $fields)
/** ${field.comment} */
private ${field.type} ${field.name};
#end

# 表名
${table.name}

模板示例

最新示例和模板代码参考GitHub项目:

https://github.com/lile1991/IDEA-Database-Tools

生成实体模板: Generate_POJO_Entity.vm

#fileName=${class.name}.java
package com.xx.dao.entity;

import com.baomidou.mybatisplus.annotation.*;
import com.baomidou.mybatisplus.extension.activerecord.Model;
import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
#foreach($imp in $imports)
import $imp;
#end

/**
* $class.comment
*/
@ToString
@Getter
@Setter
@TableName("$table.name")
public class ${class.name} extends Model<${class.name}> {
#foreach($field in $fields)
    /** ${field.comment} */
#if($foreach.index == 0)
    @TableId(value = "${field.colName}", type = IdType.INPUT)
#elseif(${field.name} == 'createTime' || ${field.name} == 'creator')
    @TableField(value = "${field.colName}", fill = FieldFill.INSERT)
#elseif(${field.name} == 'updateTime' || ${field.name} == 'updater')
    @TableField(value = "${field.colName}", fill = FieldFill.INSERT_UPDATE)
#else
    @TableField(value = "${field.colName}")
#end
    private ${field.type} ${field.name};
#end
}

生成DTO模板: Generate_POJO_DTO.vm

#fileName=${class.name}DTO.java
package com.xx.dto;

import io.swagger.v3.oas.annotations.media.Schema;
import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
#foreach($imp in $imports)
import $imp;
#end

/**
* $class.comment
*/
@ToString
@Getter
@Setter
@Schema(description = "$class.comment")
public class ${class.name}DTO {
#foreach($field in $fields)
    @Schema(description = "${field.comment}")
    private ${field.type} ${field.name};
#end
}

生成Mapper接口模板: Generate_POJO_Mapper.vm

#fileName=${class.name}Mapper.java
package com.xx.dao;

import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.xx.dao.entity.${class.name};

/**
* $class.comment
*/
public interface ${class.name}Mapper extends BaseMapper<${class.name}> {

}

生成Service模板: Generate_POJO_Service.vm

#fileName=${class.name}Service.java
package com.xx.service;

import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.xx.dao.${class.name}Mapper;
import com.xx.dao.entity.${class.name};
import org.springframework.stereotype.Service;

/**
* $class.comment
*/
@Service
public class ${class.name}Service extends ServiceImpl<${class.name}Mapper, ${class.name}> {
}

生成Mapper.xml模板: Generate_POJO_MapperXml.vm

#fileName=${class.name}Mapper.xml
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN" "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
<mapper namespace="com.xx.dao.${class.name}Mapper">
    <resultMap id="BaseResultMap" type="com.xxx.dao.entity.${class.name}">
        #foreach($field in $fields)
            #if($foreach.index == 0)
            <id column="${field.colName}" jdbcType="${field.jdbcType}" property="${field.name}" />
            #else
            <result column="${field.colName}" jdbcType="${field.jdbcType}" property="${field.name}" />
            #end
        #end
    </resultMap>
    <sql id="Base_Column_List">
        #foreach($field in $fields)
            `${field.colName}`#if(${foreach.index} < ${fields.size()} - 1),
        #end#end

    </sql>
</mapper>