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

官方的Generate POJOs.groovy无法定制， 基于官方脚本进行了修改， 支持自定义模板生成代码。

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
 *   PROJECT     project
 *   FILES       files helper
 */

typeMapping = [
        (~/(?i)bigint/)                   : "Long",
        (~/(?i)tinyint/)                  : "Byte",
        (~/(?i)int/)                      : "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 listVmFiles(File dir) {
    return listVmFiles0(dir, new ArrayList<>());
}

List listVmFiles0(File dir, List 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 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,   // 类名
            "comment": table.comment    // 表注释
    ])

    ctx.put("memberName", memberName)
    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)
        }

        fields += [[
                           name   : javaName(col.getName(), false),
                           type   : typeStr,
                           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