idea使用groovy脚本自动生成数据库实体

迈不过友情╰ 2023-06-25 03:21 15阅读 0赞

import com.intellij.database.model.DasTable
    import com.intellij.database.util.Case
    import com.intellij.database.util.DasUtil
    
    import java.time.LocalDate
    import java.time.LocalDateTime
    import java.util.Date
    
    /* * Available context bindings: * SELECTION Iterable<DasObject> * PROJECT project * FILES files helper */
    
    // 此处指定包路径，路径需要自行维护;
    packageName = "com.demo.api.tesla.entity;"
    // 此处指定对应的类型映射，可按需修改，目前tinyint如果要映射到自定义枚举类型，只能手动修改
    typeMapping = [
            (~/(?i)bigint/)                   : "Long",
            (~/(?i)int/)                      : "Integer",
            (~/(?i)tinyint/)                  : "Boolean",
            (~/(?i)float|double|decimal|real/): "BigDecimal",
            (~/(?i)time|datetime|timestamp/)  : "Date",
            (~/(?i)date/)                     : "LocalDate",
            (~/(?i)/)                         : "String"
    ]
    
    // 上面用到类和它的导入路径的之间的映射
    importMap = [
            "BigDecimal"   : "java.math.BigDecimal",
            "LocalDate"    : "java.time.LocalDate",
            "LocalDateTime": "java.time.LocalDateTime",
            "Date": "java.util.Date",
    ]
    
    // 导入路径列表，下面引用的时候会去重，也可以直接声明成一个 HashSet
    importList = []
    
    // 基类中已经有的属性，就无需在新生成的类中声明了
    userTimeList = [
            "example"
    ]
    
    // 弹出选择文件的对话框
    FILES.chooseDirectoryAndSave("Choose directory", "Choose where to store generated files") {  dir ->
        SELECTION.filter {  it instanceof DasTable }.each {  generate(it, dir) }
    }
    
    def generate(table, dir) { 
        def className = javaName(table.getName(), true) + "DO"
        def writer = new PrintWriter(new OutputStreamWriter(new FileOutputStream(new File(dir, className + ".java")), "utf-8"))
        def fields = calcFields(table)
        writer.withPrintWriter {  out -> generate(out, className, fields, table) }
    }
    
    // 从这里开始，拼实体类的具体逻辑代码
    def generate(out, className, fields, table) { 
        out.println "package $packageName"
        out.println ""
        out.println ""
        // 引入所需的包
        out.println "import lombok.Data;"
        out.println "import com.demo.basic.annotation.Column;"
        out.println "import com.demo.basic.annotation.Table;"
        // 去重后导入列表
        importList.unique().each() {  pkg ->
            out.println "import " + pkg + ";"
        }
        out.println ""
        // 添加类注释
        out.println "/**"
        // 如果添加了表注释，会加到类注释上
        if (isNotEmpty(table.getComment())) { 
            out.println " * " + table.getComment()
        }
        out.println " *"
        out.println " * @author wulongtao"
        out.println " * @since " + LocalDate.now()
        out.println " */"
        // 添加类注解
        out.println "@Data"
        out.println "@Table(\"${table.getName()}\")"
        out.println "public class $className {"
        out.println ""
        boolean isId = true
        // 遍历字段，按下面的规则生成
        fields.each() { 
            // 输出注释
            if (isNotEmpty(it.comment)) { 
                out.println "\t/**"
                out.println "\t * ${it.comment}"
                out.println "\t */"
            }
            // 这边默认第一个字段为主键，实际情况大多数如此，遇到特殊情况可能需要手动修改
            if (isId) { 
                out.println "\t@Column(isPK = true, name = \"${it.realName}\")"
                isId = false
            } else { 
                out.println "\t@Column(name = \"${it.realName}\")"
            }
            out.println "\tprivate ${it.type} ${it.name};"
            out.println ""
        }
        out.println "}"
    }
    
    def calcFields(table) { 
        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
            // 如果是BaseTimeUserIdPO中已有的字段，则不再添加
            if (!userTimeList.contains(col.getName())) { 
                // 如果是需要导入的类型，则往list里面添加
                if (importMap.containsKey(typeStr)) { 
                    importList.add(importMap.get(typeStr))
                }
                fields += [[
                                   realName: col.getName(),
                                   name    : javaName(col.getName(), false),
                                   type    : typeStr,
                                   dbType  : col.getDataType(),
                                   comment : col.getComment()
                           ]]
            } else { 
                fields += []
            }
        }
    }
    
    def isNotEmpty(content) { 
        return content != null && content.toString().trim().length() > 0
    }
    
    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]
    }