occ-toolkit/src/main/java/com/yonyou/occ/codegenerator/CodeGenerator.java

package com.yonyou.occ.codegenerator;

import freemarker.template.Configuration;
import freemarker.template.Template;
import freemarker.template.TemplateException;
import freemarker.template.TemplateExceptionHandler;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.DateFormatUtils;
import org.jdom2.Document;
import org.jdom2.Element;
import org.jdom2.JDOMException;
import org.jdom2.Namespace;
import org.jdom2.filter.Filters;
import org.jdom2.input.SAXBuilder;
import org.jdom2.xpath.XPathExpression;
import org.jdom2.xpath.XPathFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Writer;
import java.util.*;

/**
 * 代码生成器
 *
 * @author wangruiv
 * @date 2017-06-18 11:21:31
 */
public class CodeGenerator {
    private static List<String> EXCLUDE_COLUMNS = null;

    private final Logger log = LoggerFactory.getLogger(CodeGenerator.class);

    private final String oldBasePackage;
    private final String newBasePackage;
    private final String pdmFile;
    private final String tableCode;
    private final String author;
    private final String referenceType;
    private final boolean ifUseCustDocDefRef;
    private final String targetProjectRootPath;
    private final String moduleName;
    private final String templateDir = getClass().getResource("/").getPath() + "generator/";

    public CodeGenerator(String oldBasePackage, String newBasePackage, String pdmFile, String tableCode,
                         String referenceType, boolean ifUseCustDocDefRef, String author, String targetProjectRootPath) {
        this.oldBasePackage = oldBasePackage;
        this.newBasePackage = newBasePackage;
        this.pdmFile = pdmFile;
        this.tableCode = tableCode;
        this.referenceType = referenceType;
        this.ifUseCustDocDefRef = ifUseCustDocDefRef;
        this.author = author;
        this.targetProjectRootPath = targetProjectRootPath;

        moduleName = getModuleNameByTable(tableCode);

        final String[] array = {"ID", "DR", "TS", "CREATOR", "CREATION_TIME", "MODIFIER", "MODIFIED_TIME"};
        EXCLUDE_COLUMNS = Arrays.asList(array);
    }

    public void run() {
        try {
            Table table = getTableFromPdmXml(pdmFile, tableCode, referenceType);

            generateEntity(table);
            generateRepository(table);
            generateDto(table);
            generateMapper(table);
            generateService(table);
            generateController(table);
            if (StringUtils.isNotEmpty(referenceType)) {
                generateRefController(table);
            }
            if (ifUseCustDocDefRef) {
                generateMapperDecorator(table);
            }
            generateMetaDataSql(table);
        } catch (JDOMException e) {
            log.error("解析XML文件时发生错误", e);
        } catch (IOException e) {
            log.error("读取XML文件时发生错误", e);
        } catch (TemplateException e) {
            log.error("解析模板文件时发生错误", e);
        }
    }

    private String getModuleNameByTable(String tableCode) {
        return tableCode.split("_")[0].toLowerCase();
    }

    private Table getTableFromPdmXml(String pdmFile, String tableCode, String referenceType)
            throws JDOMException, IOException {
        // read the XML into a JDOM2 document.
        SAXBuilder jdomBuilder = new SAXBuilder();
        Document jdomDocument = jdomBuilder.build(pdmFile);

        // use the default implementation
        XPathFactory xFactory = XPathFactory.instance();

        // namespaces
        Namespace nsAttribute = jdomDocument.getRootElement().getNamespace("a");
        Namespace nsCollection = jdomDocument.getRootElement().getNamespace("c");
        //Namespace nsObject = jdomDocument.getRootElement().getNamespace("o");

        XPathExpression<Element> expr = xFactory.compile("//a:Code[text()='" + tableCode + "']", Filters.element(),
                null, nsAttribute);
        Element tableElement = expr.evaluateFirst(jdomDocument).getParentElement();
        String tableName = tableElement.getChild("Name", nsAttribute).getText();
        String tableComment;
        Element tableCommentElement = tableElement.getChild("Comment", nsAttribute);
        if (tableCommentElement == null || StringUtils.isBlank(tableComment = tableCommentElement.getText())) {
            tableComment = tableName;
        }
        Table table = new Table(tableName, tableCode, tableComment, referenceType);

        List<Element> columnElements = tableElement.getChild("Columns", nsCollection).getChildren();
        for (Element columnElement : columnElements) {
            String code = columnElement.getChild("Code", nsAttribute).getText();
            if (EXCLUDE_COLUMNS.contains(code.toUpperCase())) {
                continue;
            }
            String name = columnElement.getChild("Name", nsAttribute).getText();
            String comment;
            Element commentElement = columnElement.getChild("Comment", nsAttribute);
            if (commentElement == null || StringUtils.isBlank(comment = commentElement.getText())) {
                comment = name;
            }
            String dataType = columnElement.getChild("DataType", nsAttribute).getText();
            Integer length = null;
            Element lengthElement = columnElement.getChild("Length", nsAttribute);
            if (lengthElement != null) {
                length = Integer.parseInt(lengthElement.getText());
            }
            Integer precision = null;
            Element precisionElement = columnElement.getChild("Precision", nsAttribute);
            if (precisionElement != null) {
                precision = Integer.parseInt(precisionElement.getText());
            }
            Table.Column column = table.new Column(name, code, comment, dataType, length, precision);
            table.getColumns().add(column);
        }

        return table;
    }

    private void generate(Table table, String templateName, String targetSubPath)
            throws IOException, TemplateException {
        Configuration cfg = new Configuration(Configuration.DEFAULT_INCOMPATIBLE_IMPROVEMENTS);
        cfg.setDirectoryForTemplateLoading(new File(templateDir));
        cfg.setTemplateExceptionHandler(TemplateExceptionHandler.IGNORE_HANDLER);
        Template template = cfg.getTemplate(templateName + ".ftl", "UTF-8");

        // 定义模板数据
        Map<String, Object> data = new HashMap<String, Object>();
        data.put("oldBasePackage", oldBasePackage);
        data.put("newBasePackage", newBasePackage);
        data.put("moduleName", moduleName);
        data.put("date", DateFormatUtils.format(new Date(), "yyyy-MM-dd HH:mm:ss"));
        data.put("ifUseCustDocDefRef", ifUseCustDocDefRef);
        data.put("author", author);
        data.put("table", table);

        // 输出文件
        String fileName = table.getClassName();
        // 文件后缀名
        String suffix = ".java";
        // 文件保存目录
        final String outFolderPath;
        if (templateName.equalsIgnoreCase("metaData")) {
            fileName = table.getName().toLowerCase() + "_元数据";
            suffix = ".sql";
            outFolderPath = targetProjectRootPath + targetSubPath;
        } else {
            if (!templateName.equalsIgnoreCase("entity")) {
                fileName += templateName.substring(0, 1).toUpperCase() + templateName.substring(1);
            }
            outFolderPath = targetProjectRootPath + "src/main/java/" + newBasePackage.replace(".", "/") + "/" +
                    moduleName + targetSubPath;
        }

        File outFolder = new File(outFolderPath);
        if (!outFolder.exists() && !outFolder.mkdirs()) {
            System.err.println("文件目录创建失败！");
            return;
        }
        final String outFilePath = outFolderPath + fileName + suffix;
        Writer outFile = new FileWriter(outFilePath);
        template.process(data, outFile);

        System.out.println(templateName + "文件生成成功：" + outFilePath);
    }

    private void generateEntity(Table table) throws IOException, TemplateException {
        generate(table, "entity", "/entity/");
    }

    private void generateRepository(Table table) throws IOException, TemplateException {
        generate(table, "repository", "/repository/");
    }

    private void generateDto(Table table) throws IOException, TemplateException {
        generate(table, "dto", "/service/dto/");
    }

    private void generateMapper(Table table) throws IOException, TemplateException {
        generate(table, "mapper", "/service/mapper/");
    }

    private void generateService(Table table) throws IOException, TemplateException {
        generate(table, "service", "/service/");
    }

    private void generateController(Table table) throws IOException, TemplateException {
        generate(table, "controller", "/web/");
    }

    private void generateRefController(Table table) throws IOException, TemplateException {
        generate(table, "refController", "/web/");
    }

    private void generateMapperDecorator(Table table) throws IOException, TemplateException {
        generate(table, "mapperDecorator", "/service/mapper/");
    }

    private void generateMetaDataSql(Table table) throws IOException, TemplateException {
        generate(table, "metaData", "/sql/");
    }
}