mybatisPlus自动生成代码

1.创建FileCreateStragey类

package com.rtsm.zhjs.background.generator.mybatis;

import com.baomidou.mybatisplus.generator.config.IFileCreate;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.rules.FileType;
import com.rtsm.zhjs.background.util.XmlUtils;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.*;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * @author xin-bing
 * @date 3/2/2019
 */
class FileCreateStrategy implements IFileCreate {
    private boolean createCtr;
    private boolean createServ;
    private Map<String, Document> oldXmlDoc = new HashMap<>();
    public FileCreateStrategy(boolean createCtr, boolean createServ) {
        this.createCtr = createCtr;
        this.createServ = createServ;
    }

    @Override
    public boolean isCreate(ConfigBuilder configBuilder, FileType fileType, String filePath) {
        boolean fileExists = new File(filePath).exists();
        switch (fileType) {
            case CONTROLLER:
                return !fileExists && createCtr; //如果文件存在，则不创建
            case SERVICE:
                return !fileExists && createServ; //如果文件存在，则不创建
            case SERVICE_IMPL:
                return !fileExists && createServ; //如果文件存在，则不创建
            case MAPPER:
                return !fileExists; //如果文件不存在，则创建
            case ENTITY: //直接创建覆盖entity
                return true;
            case XML:
                return handleXMLCreate(fileExists, configBuilder, filePath); //合并XML
            case OTHER:
                if(filePath.endsWith(".xml")) {
                  return  handleXMLCreate(fileExists, configBuilder, filePath);
                } else {
                    return !fileExists;
                }
        }
        return true;
    }

    private boolean handleXMLCreate(boolean exist, ConfigBuilder configBuilder, String filePath) {
        if(!exist) {
            return true;
        }
        // 记录老的文件内容
        // 等代码生成器工作完成后，再进行合并操作
        try {
            Document doc = XmlUtils.newDocumentBuilder().parse(filePath);
            System.out.println(filePath + "==========record old xml string==============================");
            System.out.println(readXmlFile(filePath));
            oldXmlDoc.put(filePath, doc);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return true;
    }

    public void doMergeXmlFiles() {
        for(Map.Entry<String, Document> entry : oldXmlDoc.entrySet()) {
            doMergeXmlFile(entry.getKey(), entry.getValue());
        }
    }

    /**
     * 合并Mapper.xml文件
     * time: 3/2/2019 20:43
     * @author xin-bing
     * @version 1.0
     * @param filePath 要合并的文件路径
     * @param oldDoc 老的xmlDoc
     */
    private void doMergeXmlFile(String filePath, Document oldDoc) {
        String finallyXML = null;
        try (StringWriter writer = new StringWriter()) {
            Element oldRoot = oldDoc.getDocumentElement();
            Document newDoc = XmlUtils.newDocumentBuilder().parse(filePath);
            Map<String, Node> oldNodes = getMapperNodes(oldDoc);
            Map<String, Node> newNodes = getMapperNodes(newDoc);
            for(Map.Entry<String, Node> newNodeEntry : newNodes.entrySet()) {
                if(oldNodes.containsKey(newNodeEntry.getKey())) { //原有的mapper中已经存在，替换
                    Node newNode = oldDoc.importNode(newNodeEntry.getValue(), true);
                    oldRoot.replaceChild(newNode, oldNodes.get(newNodeEntry.getKey()));
                } else { //如果老的中不存在，属于新增的部分
                    oldRoot.appendChild(newNodeEntry.getValue());
                }
            }
            TransformerFactory tf = TransformerFactory.newInstance();
            Transformer transformer = tf.newTransformer();
            DOMSource source = new DOMSource(oldDoc);
            transformer.setOutputProperty(OutputKeys.ENCODING, "UTF-8");
            transformer.setOutputProperty(OutputKeys.INDENT, "yes");
            StreamResult result = new StreamResult(writer);
            transformer.transform(source, result);
            String output = writer.getBuffer().toString(); // .replaceAll("\n|\r", "");
            String[] arr = output.split("\n", 2); // 将字符串分为两部分
            // 处理读取之后 documentType dtd等信息没有的问题
            // 没有找的有效的解决的办法，直接操作字符串来处理
            StringBuilder sb = new StringBuilder(arr[0])
                        .append("\n")
                        .append("<!DOCTYPE mapper PUBLIC \"-//mybatis.org//DTD Mapper 3.0//EN\" \"http://mybatis.org/dtd/mybatis-3-mapper.dtd\">")
                        .append("\n")
                        .append(arr[1]);
            //将String写入文件中
            finallyXML = sb.toString();
        } catch (Exception e) {
            throw new RuntimeException("doMergeXmlFile error,file path:" + filePath, e);
        }
        try(FileOutputStream fos = new FileOutputStream(filePath);
            OutputStreamWriter osw = new OutputStreamWriter(fos, "utf-8");
            BufferedWriter bw = new BufferedWriter(osw)) {
            bw.write(finallyXML);
        } catch (Exception e) {
            System.out.println("--------------------" + filePath + " new XML String ------------------------");
            System.out.println(finallyXML);
            throw new RuntimeException("write to " + filePath + " failed", e);
        }
    }
    private Map<String, Node> getMapperNodes(Document doc) {
        Map<String, Node> map = new LinkedHashMap<>();
        Element root = doc.getDocumentElement();
        NodeList childNodes = root.getChildNodes();
        for(int i=0;i<childNodes.getLength();i++) {
            Node item = childNodes.item(i);
            if(item.getNodeType() == 1) { // node
                map.put(item.getAttributes().getNamedItem("id").toString(), item);
                System.out.println(item.getAttributes().getNamedItem("id"));
            }
        }
        return map;
    }

    /**
     * 读取xml内容
     * time: 2/3/2019 20:43
     * @author xin-bing
     * @version 1.0
     * @param filePath
     * @return
     */
    private String readXmlFile(String filePath) {
        try(FileInputStream fis = new FileInputStream(filePath);
            InputStreamReader isr = new InputStreamReader(fis, "utf-8");
            BufferedReader br = new BufferedReader(isr)){
            StringBuilder sb = new StringBuilder();
            br.lines().forEach(line -> sb.append(line).append("\n"));
            return sb.toString();
        }catch (IOException e) {
            throw new RuntimeException("read " + filePath + " error", e);
        }
    }
}

2.创建 MybatisPlusEnhancedGen 类

package com.rtsm.zhjs.background.generator.mybatis;

import com.baomidou.mybatisplus.annotation.DbType;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.generator.AutoGenerator;
import com.baomidou.mybatisplus.generator.InjectionConfig;
import com.baomidou.mybatisplus.generator.config.*;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.rules.DateType;
import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
import com.baomidou.mybatisplus.generator.engine.VelocityTemplateEngine;

import java.io.File;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 基于mybatis-plus的代码生成器，进行了一点点的定制化
 * 本生成器会采取如下的策略：
 * entity: 不管之前是否存在，都直接生成
 * controller: 如果存在该controller，那么不会再生成，否则根据运行代码生成器时的输入来执行
 * service/serviceIml：同上
 * mapper：暂时同上
 * mapper.xml：会对新生成的和老的文件进行合并，所以如果数据库发生变更，可以直接用代码生成器来生成覆盖
 *
 * @author xin-bing
 * @date 3/2/2019
 */
public class MybatisPlusEnhancedGen {
    private static final String jdbcUrl = "jdbc:mysql:// localhost:3306/aa?useUnicode=true&useSSL=false&serverTimezone=GMT%2B8&characterEncoding=utf8";
    private static final String driverName = "com.mysql.cj.jdbc.Driver";
    private static final String userName = "";
    private static final String password = "";
    private static final String basePackage = "com.rtsm.zhjs.background.modules";
    public static void main(String[] args) {
        // 代码生成器
        AutoGenerator mpg = new AutoGenerator();
        // 全局配置
        GlobalConfig gc = new GlobalConfig();
        String projectPath = System.getProperty("user.dir");
        gc.setOutputDir(projectPath + "/src/main/java")
                .setAuthor(System.getProperty("user.name"))
                .setBaseResultMap(true)
                .setBaseColumnList(true)
                //.setEntityName("%sEntity")
                .setServiceName("%sService")
                .setOpen(false) //是否打开文件
                .setDateType(DateType.TIME_PACK);
        mpg.setGlobalConfig(gc);
        // 收集输入
        String input = scanner("模块名");
        String moduleName = input;
        input = scanner("表名，（如有多个，以逗号分隔）");
        Set<String> tables = Arrays.stream(input.split("[,，]")).map(String::trim).filter(s->!s.isEmpty()).collect(Collectors.toSet());
        String[] tableNames = new String[tables.size()];
        tables.toArray(tableNames);
        input = scanner("是否生成controller（0否，1是）");
        boolean contollerFlag = Integer.parseInt(input) > 0;
        input = scanner("是否生成service（0否，1是）");
        boolean serviceFlag = Integer.parseInt(input) > 0;
        // 数据源配置
        DataSourceConfig dsc = new DataSourceConfig();
        dsc.setUrl(jdbcUrl)
                .setDbType(DbType.MYSQL)
                .setDriverName(driverName)
                .setUsername(userName)
                .setPassword(password);
        mpg.setDataSource(dsc);
        // 包配置
        PackageConfig pc = new PackageConfig();
        pc.setModuleName(moduleName);
        pc.setParent(basePackage);
        mpg.setPackageInfo(pc);

        // 自定义配置
        InjectionConfig cfg = new InjectionConfig() {
            @Override
            public void initMap() {
                Map<String, Object> map = new HashMap<>();
                map.put("currDateTime", LocalDateTime.now().format(DateTimeFormatter.BASIC_ISO_DATE));
                this.setMap(map);
            }
        };
        // 如果模板引擎是 freemarker
        String xmlTemplatePath = "/mybatis-plus-temp/mapper.xml.vm";
        String mapperTemplatePath = "/mybatis-plus-temp/mapper.java";
        String controllerTemplatePath = "/mybatis-plus-temp/controller.java";
        String serviceTemplatePath = "/mybatis-plus-temp/service.java";
        String serviceImplTempatePath = "/mybatis-plus-temp/serviceImpl.java";
        String entityTemplatePath = "/mybatis-plus-temp/entity.java";
        // 自定义输出配置
        List<FileOutConfig> focList = new ArrayList<>();
        // 自定义配置会被优先输出
        focList.add(new FileOutConfig(xmlTemplatePath) {
            @Override
            public String outputFile(TableInfo tableInfo) {
                // 自定义输出文件名 ， 如果你 Entity 设置了前后缀、此处注意 xml 的名称会跟着发生变化！！
                String fileDir = projectPath + "/src/main/resources/mapper/" + moduleName + "/";
                String filePath =  fileDir + tableInfo.getEntityName() + "Mapper" + StringPool.DOT_XML;
                File file = new File(filePath);
                if(!file.getParentFile().exists()) {
                    file.getParentFile().mkdirs();
                }
                return filePath;
            }
        });
        cfg.setFileOutConfigList(focList);
        FileCreateStrategy fileCreateStrategy = new FileCreateStrategy(contollerFlag, serviceFlag);
        cfg.setFileCreate(fileCreateStrategy);
        mpg.setCfg(cfg);
        // 配置模板
        TemplateConfig templateConfig = new TemplateConfig();
        // 配置自定义输出模板
        //指定自定义模板路径，注意不要带上.ftl/.vm, 会根据使用的模板引擎自动识别
        templateConfig
                .setXml(null) // 因为上面已经指定了xml模板的FileOutConfig，这里再次将xml的模板添加，会造成重复生成，
                .setMapper(mapperTemplatePath)
                .setController(controllerTemplatePath)
                .setService(serviceTemplatePath)
                .setServiceImpl(serviceImplTempatePath)
                .setEntity(entityTemplatePath);
        mpg.setTemplate(templateConfig);
        // 策略配置
        StrategyConfig strategy = new StrategyConfig();
        strategy.setCapitalMode(true)
                .setSkipView(false)
                .setNaming(NamingStrategy.underline_to_camel) //数据库表映射到实体的命名策略
                .setColumnNaming(NamingStrategy.underline_to_camel) //数据库表字段映射到实体的命名策略, 未指定按照 naming 执行
                .setEntityLombokModel(true)
                .setRestControllerStyle(true)
                .setInclude(tableNames)
                .setControllerMappingHyphenStyle(true)
                .setTablePrefix(pc.getModuleName() + "_");
        mpg.setStrategy(strategy);
        mpg.setTemplateEngine(new VelocityTemplateEngine());
        mpg.execute();
        fileCreateStrategy.doMergeXmlFiles(); //合并XML
    }
    private static String scanner(String tip) {
        Scanner scanner = new Scanner(System.in);
        StringBuilder help = new StringBuilder();
        help.append("请输入").append(tip).append("：");
        System.out.println(help.toString());
        String ipt = scanner.next().trim();
        if(!org.apache.commons.lang3.StringUtils.isBlank(ipt)) {
            return ipt;
        } else {
            throw new RuntimeException("请输入正确的" + tip);
        }
    }
}