自定义groovy脚本在IDEA中为数据库生成PO实体类

在文章底部有完整的代码实现

前言

我们可能会遇到下列这种问题:

公司的一个小项目,被拆分成了API和后台管理两个服务,但是因为二者共用一个数据库,所以存在着大量相同的数据库实体定义.

因此我们不得不在这两个服务中分别提供一样的实体定义,当我们的表结构发生变更时,我们可能会忘记修改某一个项目中对应的实体,久而久之,我们会发现,某个项目中的某个实体对象和数据库表定义相差甚远.

这时候,我们就可以考虑拆出一个common项目用来存放二者公共的代码,数据实体类作为一个不可变的对象,单独维护在common项目中,并保证common中的实体对象和数据库一一对应.

那么,我们需要手动的去为每个表建立相应的实体类吗?如果表结构发生变化,我们还要手动去更改相应的实体类吗?

答案当然是不,借助于IDEA的数据库插件Database,我们可以通过提供一个groovy脚本来批量生成数据表对应的实体类.

这样,不仅简化了代码编写,同时还能统一实体类的生成逻辑.

关于IDEA中脚本的使用方式参见文档:Generate Java entity classes for tables and views

IDEA中本身为提供了一个名为Generate POJOs.groovy的脚本,该脚本可以完成简单的代码生成工作,但是,在实际业务中,我们的实体类除了属性定义之外,往往还有一些其他的数据要定义在里面,比如:数据库描述相关的注解.

自定义脚本实现

这里,我提供了一个简单的自定义groovy脚本,通过该脚本,你可以完成包括自定义注解在内的代码生成工作.

基本属性定义

脚本的最开始,定义了几个开关用于控制脚本的行为,建议全部开启:

// ============= 常量区域 =================
useLogicDelete = true; /* 启用逻辑删除 */
useLombok = true;/*是否启用lombok*/
useSerializable = true;/*是否启用实体序列化*/

再定义两个Map集合,其中typeMapping负责将数据库类型转换为java类型,javaAliasMap则负责通过java类定义的简短名称获取其全限定名称:

/* 类型转换 ,负责将数据库类型转换为java类型,转换逻辑取自 mybatis */
typeMapping = [
        (~/(?i)real|decimal|numeric/)        : "BigDecimal",
        (~/(?i)bigint/)                      : "Long",
        (~/(?i)tinyint/)                     : "Byte", /* 根据自己的需求,可以考虑转换为Boolean*/
        (~/(?i)int/)                         : "Integer",
        (~/(?i)enum/)                        : "String", /* 枚举统一转换为字符串*/
        (~/(?i)float|double/)                : "Double", /* 根据自己的需求可以考虑其他类型*/
        (~/(?i)datetime|timestamp|date|time/): "Date",
        (~/(?i)/)                            : "String" /*其余的统一转成字符串*/
]


/* java 别名映射 ,负责导包*/
javaAliasMap = [
        "BigDecimal"  : "java.math.BigDecimal",
        "Date"        : "java.util.Date",
        "Getter"      : "lombok.Getter",
        "Setter"      : "lombok.Setter",
        "ToString"    : "lombok.ToString",
        "Serializable": "java.io.Serializable",
        "Table": "cn.jpanda.common.page.annotations.Table",
        "TableID": "cn.jpanda.common.page.annotations.TableID",
        "Column": "cn.jpanda.common.page.annotations.Column",
        "Logic": "cn.jpanda.common.page.annotations.Logic",
]

接着定义了三个生成实体类需要使用的属性,他们负责记录待生成实体类的部分关键信息:

// 代码生成路径,在脚本中通过弹出文件框进行选择,同时会转换为包名称
String packageName = ""
/* 待引入的java类*/
importClass = [];
/* 待使用的类注解 */
usesAnnotations = [];

属性比较简单,上面都有注释,这些属性都是在后面运行时被赋值的.

脚本执行入口

接下来就是脚本的真正入口:

/* 脚本入口 */
/* 选择文件目录 ,并执行生成代码操作 */
FILES.chooseDirectoryAndSave("Choose directory", "Choose where to store generated files") { dir ->
    SELECTION.filter { it instanceof DasTable }.each { generate(it, dir) }
}

该方法会弹出一个文件选择框,供用户选择一个用来存放生成代码的文件夹.

当用户完成选择文件夹后,用户选择的数据库表数据将会依次传入到generate()方法中,用来生成相应的实体对象.

为数据库表生成对应的实体类

通过一个table属性生成相应实体的逻辑比较简单,大概有七步:

def generate(table, dir) {
    // step1: 获取包名
    packageName = getPackageName(dir)
    // step2: 获取当前表名对应的类名称
    def className = javaName(table.getName(), true)
    // step3: 加载lombok注解
    if (useLombok) {
        ["Getter", "Setter", "ToString"].each() {
            convertAliasesAndImportClasses(it)
            usesAnnotations << it;
        };
    }
    // step4: 加载序列化
    if (useSerializable) {
        convertAliasesAndImportClasses("Serializable")
    }
    // step5: 加载自定义类注解
    loadCustomAnnotationToClass(table, className, dir);

    // step6: 处理所有数据列定义
    def properties = processDataColumnDefinition(table);

    // step7: 生成代码
    new File(dir as File, className + ".java").withPrintWriter("utf-8") {
        out ->
            generate(out, className, properties, table)
    }
}

1.根据用户选择的目录生成包名

 // step1: 获取包名
packageName = getPackageName(dir)

static def getPackageName(dir) {
return dir.toString().replaceAll("\\\\", ".").replaceAll("/", ".").replaceAll("^.*src(\\.main\\.java\\.)?", "") + ";"
}

方法的实现并不复杂,默认通过src/main/java目录的子目录层级结构来生成包名称.

2.将表明转换为`java`类型名

 // step2: 获取当前表名对应的类名称
def className = javaName(table.getName(), true)

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]
}

转换规则并不复杂:下划线转驼峰,首字母大写.

3.根据用户开关加载`lombok`注解

if (useLombok) {
    ["Getter", "Setter", "ToString"].each() {
        convertAliasesAndImportClasses(it)
        usesAnnotations << it;
    };
}

convertAliasesAndImportClasses()方法是一个负责完成导包配置的工具方法:

void convertAliasesAndImportClasses(className) {
    def entry = javaAliasMap.find { p, t -> p.equalsIgnoreCase(className) };
    if (entry == null) {
        return;
    }
    def fullName = entry.value
    if (isNotEmpty(fullName)) {
        doImportClass(fullName);
    }
}

这里只简单启用了最常见的三个lombok注解,可以根据自己的需求,调整所需的注解.

4.根据序列化开关选择是否启用序列化

// step4: 加载序列化
if (useSerializable) {
    convertAliasesAndImportClasses("Serializable")
}

5.[扩展点]加载自定义类注解

// step5: 加载自定义类注解
loadCustomAnnotationToClass(table, className, dir);

void loadCustomAnnotationToClass(table, className, dir) {
// 根据自己的需求实现,加载自定义的注解,比如JPA注解

}

loadCustomAnnotationToClass()方法定义在脚本底部,可以根据自己的需要来实现.

6.预处理所有数据列定义

// step6: 处理所有数据列定义
def properties = processDataColumnDefinition(table);

def processDataColumnDefinition(table) {
    // 加载当前数据库主键
    def primaryKey = ""
    def prKey = DasUtil.getPrimaryKey(table);
    if (prKey != null) {
        def keyRef = prKey.getColumnsRef();
        if (keyRef != null) {
            def it = keyRef.iterate();
            while (it.hasNext()) {
                // 默认只处理单主键
                primaryKey = it.next();
            }
            primaryKey = javaName(primaryKey, false);
        }
    }

    // 依次处理每一行数据列
    DasUtil.getColumns(table).reduce([]) { properties, col ->
        // 获取JDBC类型
        def spec = Case.LOWER.apply(col.getDataType().getSpecification())
        // 将JDBC类型映射为JAVA类型
        def typeStr = typeMapping.find { p, t -> p.matcher(spec).find() }.value
        // 判断该类型是否需要导包
        convertAliasesAndImportClasses(typeStr);
        // 将列名称转换为字段名称
        def javaName = javaName(col.getName(), false);
        // 当前列是否为主键
        def isPrimaryKey = javaName.equals(primaryKey);
        // 是否为逻辑删除标记
        def isLogicDeleteFlag = isLogicDelete(javaName);

        def property = [
                name   : javaName,
                colName: col.getName(),
                type   : typeStr,
                comment: col.getComment(),
                isKey  : isPrimaryKey,
                isDel  : isLogicDeleteFlag,
                annos  : []
        ]
        loadCustomAnnotationsToProperty(property);
        properties << property;
    }
}

processDataColumnDefinition()方法看起来实现比较复杂,但是实际上只是负责将数据列定义解析为下列数据结构:

def property = [
        name   : javaName,  /* java属性名*/
        colName: col.getName(), /*jdbc列名*/
        type   : typeStr, /*java类型*/
        comment: col.getComment(),/*注释*/
        isKey  : isPrimaryKey, /*是否为数据表主键*/
        isDel  : isLogicDeleteFlag, /*是否是逻辑删除标志*/
        annos  : [] /*需要添加的属性注解*/
]

property有一个annos属性,他的值是通过loadCustomAnnotationsToProperty()获取的,这是另一个扩展点,用于加载自定义属性注解,比如,我的简单实现是:

void loadCustomAnnotationsToProperty(property) {
    // 根据自己的需求实现,加载自定义属性注解,比如这里,我使用了自定义的注解
    if (property.isDel) {
        convertAliasesAndImportClasses("Logic");
        property.annos << "Logic";
    }
    if (property.isKey) {
        convertAliasesAndImportClasses("TableID");
        property.annos << "TableID";
    }
    convertAliasesAndImportClasses("Column");
    property.annos << "Column";
}

7.生成代码

// step7: 生成代码
new File(dir as File, className + ".java").withPrintWriter("utf-8") {
    out ->
        generate(out, className, properties, table)
}

这一步是使用上面的数据来生成真正的代码,实现很长,但是基本上都是用来处理前面获取到的数据的:

// 生成实体类代码
def generate(out, className, properties, table) {
    // step 1: 输出包名
    out.println "package $packageName"
    out.println ""
    out.println ""
    // step2: 导入类
    importClass.each() {
        out.println "import ${it};"
    }
    out.println ""

    // 补充:生成注释
    generateClassComments(out, table)

    // step3: 生成类注解
    usesAnnotations.each() {
        out.println "@${it}"
    }
    // step4: 生成类名
    out.print "public class $className"
    if (useSerializable) {
        out.println(" implements Serializable {")
        out.println ""
        /* 生成序列化ID*/
        out.println genSerialID()
    } else {
        out.println "{"
    }
    out.println ""
    // step5: 处理每个属性定义
    properties.each() {
        // step5.1: 输出注释
        if (isNotEmpty(it.comment)) {
            out.println "\t/**"
            out.println "\t * ${it.comment}"
            out.println "\t */"
        }
        //step5.2: 输出注解
        it.annos.each() {
            out.println "\t@${it}"
        }
        /*step5.3: 输入字段内容*/
        out.println "\tprivate ${it.type} ${it.name};"
    }

    // step6: 生成getter/setter方法
    if (!useLombok) {
        // 在未启用lombok的情况下生成getter/setter方法
        properties.each() {
            out.println ""
            out.println "  public ${it.type} get${it.name.capitalize()}() {"
            out.println "    return ${it.name};"
            out.println "  }"
            out.println ""
            out.println "  public void set${it.name.capitalize()}(${it.type} ${it.name}) {"
            out.println "    this.${it.name} = ${it.name};"
            out.println "  }"
            out.println ""
        }
    }

    // step7: 结尾
    out.println "}"
}

generate()方法中,值得留意的是:

// 补充:生成注释
generateClassComments(out, table)

这是方法的第三个扩展点,用于加载自定义类注释,我的简单实现是:

void generateClassComments(out, table) {
    /* 生成类注释 */
    ["/**",
     " * ",
     " * ${table.getComment()}",
     " * @author Jpanda [Jpanda@aliyun.com]",
     " * @version 1.0",
     " * @since " + new SimpleDateFormat("yyyy-MM-dd").format(new Date()),
     " */"
    ].each() {
        out.println "${it}"
    }
}

上面就是脚本的基本思路了,脚本比较简单,上面也有详细的注释,基本可以满足使用.

使用脚本

下面是一个简单的代码生成示例.

涉及到的表结构:

create table user
(
    id         int auto_increment primary key,
    iphone     varchar(50)  not null comment '手机号',
    password   varchar(50)  not null comment '密码',
    nikename   varchar(100) null comment '名字',
    group_flag int(1)       null comment '组标识 0 店长 1 组长 2管理员'
)
    comment '用户';

首先,我们把自定脚本放入IDEA的脚本执行目录下:

然后通过Database插件调用脚本:

选择目录:

生成代码:

package cn.jpanda.common.pojos.test;


import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import java.io.Serializable;
import cn.jpanda.common.page.annotations.TableID;
import cn.jpanda.common.page.annotations.Column;

/**
 * 
 * 用户
 * @author Jpanda [Jpanda@aliyun.com]
 * @version 1.0
 * @since 2020-04-22
 */
@Getter
@Setter
@ToString
public class User implements Serializable {

 private static final long serialVersionUID =  2598051958746997471L;

 @TableID
 @Column
 private Integer id;
 /**
  * 手机号
  */
 @Column
 private String iphone;
 /**
  * 密码
  */
 @Column
 private String password;
 /**
  * 名字
  */
 @Column
 private String nikename;
 /**
  * 组标识 0 店长 1 组长 2管理员
  */
 @Column
 private Integer groupFlag;
}

完整代码

下面给出完整的脚本代码,需要根据自己的需求,调节自定义注解相关的数据:

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

import java.text.SimpleDateFormat

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

// ============= 常量区域 =================
useLogicDelete = true; /* 启用逻辑删除 */
useLombok = true;/*是否启用lombok*/
useSerializable = true;/*是否启用实体序列化*/

/* 类型转换 ,负责将数据库类型转换为java类型,转换逻辑取自 mybatis */
typeMapping = [
        (~/(?i)real|decimal|numeric/)        : "BigDecimal",
        (~/(?i)bigint/)                      : "Long",
        (~/(?i)tinyint/)                     : "Byte", /* 根据自己的需求,可以考虑转换为Boolean*/
        (~/(?i)int/)                         : "Integer",
        (~/(?i)enum/)                        : "String", /* 枚举统一转换为字符串*/
        (~/(?i)float|double/)                : "Double", /* 根据自己的需求可以考虑其他类型*/
        (~/(?i)datetime|timestamp|date|time/): "Date",
        (~/(?i)/)                            : "String" /*其余的统一转成字符串*/
]


/* java 别名映射 ,负责导包*/
javaAliasMap = [
        "BigDecimal"  : "java.math.BigDecimal",
        "Date"        : "java.util.Date",
        "Getter"      : "lombok.Getter",
        "Setter"      : "lombok.Setter",
        "ToString"    : "lombok.ToString",
        "Serializable": "java.io.Serializable",
        "Table"       : "cn.jpanda.common.page.annotations.Table",
        "TableID"     : "cn.jpanda.common.page.annotations.TableID",
        "Column"      : "cn.jpanda.common.page.annotations.Column",
        "Logic"       : "cn.jpanda.common.page.annotations.Logic",
]


/* =========  运行时属性定义 ============*/
// 代码生成路径,在脚本中通过弹出文件框进行选择
String packageName = ""
/* 待引入的java类*/
importClass = [];
/* 待使用注解 */
usesAnnotations = [];


/* 脚本入口 */
/* 选择文件目录 ,并执行生成代码操作 */
FILES.chooseDirectoryAndSave("Choose directory", "Choose where to store generated files") { dir ->
    SELECTION.filter { it instanceof DasTable }.each { generate(it, dir) }
}


def generate(table, dir) {
    // step1: 获取包名
    packageName = getPackageName(dir)
    // step2: 获取当前表名对应的类名称
    def className = javaName(table.getName(), true)
    // step3: 加载lombok注解
    if (useLombok) {
        ["Getter", "Setter", "ToString"].each() {
            convertAliasesAndImportClasses(it)
            usesAnnotations << it;
        };
    }
    // step4: 加载序列化
    if (useSerializable) {
        convertAliasesAndImportClasses("Serializable")
    }
    // step5: 加载自定义类注解
    loadCustomAnnotationToClass(table, className, dir);

    // step6: 处理所有数据列定义
    def properties = processDataColumnDefinition(table);

    // step7: 生成代码
    new File(dir as File, className + ".java").withPrintWriter("utf-8") {
        out ->
            generate(out, className, properties, table)
    }
}


// 生成实体类代码
def generate(out, className, properties, table) {
    // step 1: 输出包名
    out.println "package $packageName"
    out.println ""
    out.println ""
    // step2: 导入类
    importClass.each() {
        out.println "import ${it};"
    }
    out.println ""

    // 补充:生成注释
    generateClassComments(out, table)

    // step3: 生成类注解
    usesAnnotations.each() {
        out.println "@${it}"
    }
    // step4: 生成类名
    out.print "public class $className"
    if (useSerializable) {
        out.println(" implements Serializable {")
        out.println ""
        /* 生成序列化ID*/
        out.println genSerialID()
    } else {
        out.println "{"
    }
    out.println ""
    // step5: 处理每个属性定义
    properties.each() {
        // step5.1: 输出注释
        if (isNotEmpty(it.comment)) {
            out.println "\t/**"
            out.println "\t * ${it.comment}"
            out.println "\t */"
        }
        //step5.2: 输出注解
        it.annos.each() {
            out.println "\t@${it}"
        }
        /*step5.3: 输入字段内容*/
        out.println "\tprivate ${it.type} ${it.name};"
    }

    // step6: 生成getter/setter方法
    if (!useLombok) {
        // 在未启用lombok的情况下生成getter/setter方法
        properties.each() {
            out.println ""
            out.println "  public ${it.type} get${it.name.capitalize()}() {"
            out.println "    return ${it.name};"
            out.println "  }"
            out.println ""
            out.println "  public void set${it.name.capitalize()}(${it.type} ${it.name}) {"
            out.println "    this.${it.name} = ${it.name};"
            out.println "  }"
            out.println ""
        }
    }

    // step7: 结尾
    out.println "}"
}

def processDataColumnDefinition(table) {
    // 加载当前数据库主键
    def primaryKey = ""
    def prKey = DasUtil.getPrimaryKey(table);
    if (prKey != null) {
        def keyRef = prKey.getColumnsRef();
        if (keyRef != null) {
            def it = keyRef.iterate();
            while (it.hasNext()) {
                // 默认只处理单主键
                primaryKey = it.next();
            }
            primaryKey = javaName(primaryKey, false);
        }
    }

    // 依次处理每一行数据列
    DasUtil.getColumns(table).reduce([]) { properties, col ->
        // 获取JDBC类型
        def spec = Case.LOWER.apply(col.getDataType().getSpecification())
        // 将JDBC类型映射为JAVA类型
        def typeStr = typeMapping.find { p, t -> p.matcher(spec).find() }.value
        // 判断该类型是否需要导包
        convertAliasesAndImportClasses(typeStr);
        // 将列名称转换为字段名称
        def javaName = javaName(col.getName(), false);
        // 当前列是否为主键
        def isPrimaryKey = javaName.equals(primaryKey);
        // 是否为逻辑删除标记
        def isLogicDeleteFlag = isLogicDelete(javaName);

        def property = [
                name   : javaName,  /* java属性名*/
                colName: col.getName(), /*jdbc列名*/
                type   : typeStr, /*java类型*/
                comment: col.getComment(),/*注释*/
                isKey  : isPrimaryKey, /*是否为数据表主键*/
                isDel  : isLogicDeleteFlag, /*是否是逻辑删除标志*/
                annos  : [] /*需要添加的属性注解*/
        ]
        loadCustomAnnotationsToProperty(property);
        properties << property;
    }
}

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]
}

static def isNotEmpty(content) {
    return content != null && content.toString().trim().length() > 0
}

static String genSerialID() {
    return "\tprivate static final long serialVersionUID =  " + Math.abs(new Random().nextLong()) + "L;"
}
// 生成包名
static def getPackageName(dir) {
    return dir.toString().replaceAll("\\\\", ".").replaceAll("/", ".").replaceAll("^.*src(\\.main\\.java\\.)?", "") + ";"
}

/**
 * 判断一个java属性是否为逻辑删除
 * [根据自己的需求重构该测试]
 * @param property 属性名称
 * @return
 */
boolean isLogicDelete(property) {
    return useLogicDelete && property.equalsIgnoreCase("DeleteFlag") || property.equalsIgnoreCase("DelFlag")
}

/**
 * 导入lombok的类
 */
void importLombokPackage() {
    if (useLombok) {
        ["lombok.Getter", "lombok.Setter", "lombok.ToString"].each() {
            doImportClass(it);
        };
    }
}

/**
 * 导入一个包名称
 * @param className
 */
void doImportClass(className) {
    if (importClass.count { it -> it.equalsIgnoreCase(className) } == 0) {
        importClass << className;
    }
}

void convertAliasesAndImportClasses(className) {
    def entry = javaAliasMap.find { p, t -> p.equalsIgnoreCase(className) };
    if (entry == null) {
        return;
    }
    def fullName = entry.value
    if (isNotEmpty(fullName)) {
        doImportClass(fullName);
    }
}

void loadCustomAnnotationToClass(table, className, dir) {
    // 根据自己的需求实现,加载自定义的注解,比如JPA注解

}

void loadCustomAnnotationsToProperty(property) {
    // 根据自己的需求实现,加载自定义属性注解,比如这里,我使用了自定义的注解
    if (property.isDel) {
        convertAliasesAndImportClasses("Logic");
        property.annos << "Logic";
    }
    if (property.isKey) {
        convertAliasesAndImportClasses("TableID");
        property.annos << "TableID";
    }
    convertAliasesAndImportClasses("Column");
    property.annos << "Column";
}

void generateClassComments(out, table) {
    /* 生成类注释 */
    ["/**",
     " * ",
     " * ${table.getComment()}",
     " * @author Jpanda [Jpanda@aliyun.com]",
     " * @version 1.0",
     " * @since " + new SimpleDateFormat("yyyy-MM-dd").format(new Date()),
     " */"
    ].each() {
        out.println "${it}"
    }
}