自定义Mybatis自动生成代码规则

前言

       大家都清楚mybatis-generate-core 这个工程提供了获取表信息到生成model、dao、xml这三层代码的一个实现，但是这往往有一个痛点，比如需求来了，某个表需要增加字段，肯定需要重新运行mybatis自动生成的脚本，但是会去覆盖之前的代码，如model,dao的java代码，对于xml文件，目前有两种处理方式，一、覆盖，二、追加，本文用的版本是1.3.5版本，默认的是追加方式，本文的目的就是处理xml的一种合并方式，对于java代码的话，我个人认为无论是增加表字段还是其他情况，相对于xml文件都是比较好维护的，这里就不做讨论。

　　对于方式一的话，直接覆盖，肯定会导致之前自定义的sql，直接没了，还需要事先拷贝一份出来，最蛋疼的就是，可能还会在自动生成的代码文件中，增加了一些属性(如主键返回，flushCache属性)，导致后来人员给忽略了，直到某个时刻才爆发出来。所以本文不采用这种方式，而是采用方式2，对于mybatis自定义的合并规则，看下文介绍。本文会对这个合并规则，进行重写，已达到我们的目标。如下

• 在启用自动生成代码后，原有的自定义sql，一律保留，包括,result|sql|select|delete|update|where|insert等标签，只要不是自动生成的
• 自动生成的标签中，手动添加的一些属性，如主键返回useGeneratedKeys="true" keyColumn="id"，刷新一级缓存,flushCache="true"等属性标签也需要保留。

在重写该规则前，肯定是要摸清它的原有流程，下面分为这几个小节进行叙述

一、合并规则原理

二、重写规则

三、简述适用场景

本文采用的数据库是Mysql

一、合并规则原理

先来一段代码，莫慌，这段代码没什么特别，很常见的自动生成代码

package com.qm.mybatis.generate;

import org.mybatis.generator.api.MyBatisGenerator;
import org.mybatis.generator.config.Configuration;
import org.mybatis.generator.config.xml.ConfigurationParser;
import org.mybatis.generator.internal.DefaultShellCallback;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

public class GenerateTest {

    public static void main(String[] args) {
        List<String> warnings = new ArrayList<String>();
        try {
            boolean overwrite = true;
            // 读取配置文件
            InputStream resourceAsStream = GenerateTest.class.getResourceAsStream("/mybatis-generate.xml");
            ConfigurationParser cp = new ConfigurationParser(warnings);
            Configuration config = cp.parseConfiguration(resourceAsStream);
            DefaultShellCallback callback = new DefaultShellCallback(overwrite);
            MyBatisGenerator myBatisGenerator = new MyBatisGenerator(config, callback, warnings);
            myBatisGenerator.generate(null);
        } catch (Exception e) {

            e.printStackTrace();
        }

        warnings.stream().forEach(warn -> {
            System.out.println(warn);
        });
        System.out.println("生成成功！");
    }
}

可见，最终的生成逻辑在MybatisGenerator#generate方法中，

// 最终生成代码的地方
public void generate(ProgressCallback callback, Set<String> contextIds,
            Set<String> fullyQualifiedTableNames, boolean writeFiles) throws SQLException,
            IOException, InterruptedException {

        if (callback == null) {
            callback = new NullProgressCallback();
        }

        generatedJavaFiles.clear();
        generatedXmlFiles.clear();
        ObjectFactory.reset();
        RootClassInfo.reset();

        // calculate the contexts to run
        List<Context> contextsToRun;
        if (contextIds == null || contextIds.size() == 0) {
            contextsToRun = configuration.getContexts();
        } else {
            contextsToRun = new ArrayList<Context>();
            for (Context context : configuration.getContexts()) {
                if (contextIds.contains(context.getId())) {
                    contextsToRun.add(context);
                }
            }
        }

        // setup custom classloader if required
        if (configuration.getClassPathEntries().size() > 0) {
            ClassLoader classLoader = getCustomClassloader(configuration.getClassPathEntries());
            ObjectFactory.addExternalClassLoader(classLoader);
        }

        // now run the introspections...
        int totalSteps = 0;
        for (Context context : contextsToRun) {
            totalSteps += context.getIntrospectionSteps();
        }
        callback.introspectionStarted(totalSteps);

        for (Context context : contextsToRun) {
            context.introspectTables(callback, warnings,
                    fullyQualifiedTableNames);
        }

        // now run the generates
        totalSteps = 0;
        for (Context context : contextsToRun) {
            totalSteps += context.getGenerationSteps();
        }
        callback.generationStarted(totalSteps);

        for (Context context : contextsToRun) {
            context.generateFiles(callback, generatedJavaFiles,
                    generatedXmlFiles, warnings);
        }

        // 前面各种文件都已经生成完毕，在这里进行保存到具体的文件中
        if (writeFiles) {
            callback.saveStarted(generatedXmlFiles.size()
                + generatedJavaFiles.size());

            // 进行xml文件保存(更新)的地方，也是本文的目标
            for (GeneratedXmlFile gxf : generatedXmlFiles) {
                projects.add(gxf.getTargetProject());
                writeGeneratedXmlFile(gxf, callback);
            }

            // 保存java文件，如model,example,dao
            for (GeneratedJavaFile gjf : generatedJavaFiles) {
                projects.add(gjf.getTargetProject());
                writeGeneratedJavaFile(gjf, callback);
            }

            for (String project : projects) {
                shellCallback.refreshProject(project);
            }
        }

        callback.done();
    }

最终的落实地方就在writeGeneratedXmlFile方法内。

private void writeGeneratedXmlFile(GeneratedXmlFile gxf, ProgressCallback callback)
            throws InterruptedException, IOException {
        File targetFile;
        String source;
        try {
            File directory = shellCallback.getDirectory(gxf
                    .getTargetProject(), gxf.getTargetPackage());
            targetFile = new File(directory, gxf.getFileName());
            // 如果为false，基本上就是第一次生成的时候
            if (targetFile.exists()) {
                
                /**
                 * 从这里也可以看出，这个参数决定xml文件的处理方式
                 * 为true时，会执行getMergedSource，透个底，改造也是改造这个方法
                   false，会继续后面两种逻辑。实际生成的内容其实是一样的。这里不做讨论
                 */
                if (gxf.isMergeable()) {
                    source = XmlFileMergerJaxp.getMergedSource(gxf,
                            targetFile);
                } else if (shellCallback.isOverwriteEnabled()) {
                    source = gxf.getFormattedContent();
                    warnings.add(getString("Warning.11", //$NON-NLS-1$
                            targetFile.getAbsolutePath()));
                } else {
                    source = gxf.getFormattedContent();
                    targetFile = getUniqueFileName(directory, gxf
                            .getFileName());
                    warnings.add(getString(
                            "Warning.2", targetFile.getAbsolutePath())); //$NON-NLS-1$
                }
            } else {
                source = gxf.getFormattedContent();
            }

            callback.checkCancel();
            callback.startTask(getString(
                    "Progress.15", targetFile.getName())); //$NON-NLS-1$
            writeFile(targetFile, source, "UTF-8"); //$NON-NLS-1$
        } catch (ShellException e) {
            warnings.add(e.getMessage());
        }
    }

饶了这么多圈，实际上我们要处理的就是重写XmlFileMergerJaxp#getMergedSource方法，或许有的人会提出疑问了，这个类能让你提供扩展吗？让你去继承？然后去改变这个规则，额(⊙o⊙)…，还真没有，这个类实际上就是一个静态方法，那这搞个毛线啊，你即使重写出来了，那你怎么将他插入进去，别告诉你准备重新编译源码。。。。。莫慌，往下看。

说到这，大家可以去了解一下类加载器和其加载的过程，本文不做过多阐述，直接来结论，你要想覆盖一个jar包里的某个方法，你就直接在你项目中，定义这个类(包名和类名需要完全一致)，然后运行的时候，自然会执行你定义的这个类，千万别去想着同样的方法去覆盖jdk自带的类，没用，因为第三方jar包和jdk自带的类的类加载器不是同一个。有兴趣的可以去网上搜索一下。说了这么多，我们就是要这样做。做之前，先了解下这个merge方法的代码。

public static String getMergedSource(InputSource newFile,
            InputSource existingFile, String existingFileName) throws IOException, SAXException,
            ParserConfigurationException, ShellException {

        DocumentBuilderFactory factory = DocumentBuilderFactory
                .newInstance();
        factory.setExpandEntityReferences(false);
        DocumentBuilder builder = factory.newDocumentBuilder();
        builder.setEntityResolver(new NullEntityResolver());

        // 这是xml文件的解析结果，这里就暂且称为旧文件和新文件
        Document existingDocument = builder.parse(existingFile);
        Document newDocument = builder.parse(newFile);

        DocumentType newDocType = newDocument.getDoctype();
        DocumentType existingDocType = existingDocument.getDoctype();

        // 比较两个xml文件是不是同一类型
        if (!newDocType.getName().equals(existingDocType.getName())) {
            throw new ShellException(getString("Warning.12", //$NON-NLS-1$
                    existingFileName));
        }
        
        // 获取根节点
        Element existingRootElement = existingDocument.getDocumentElement();
        Element newRootElement = newDocument.getDocumentElement();

        
        NamedNodeMap attributes = existingRootElement.getAttributes();
        int attributeCount = attributes.getLength();
        for (int i = attributeCount - 1; i >= 0; i--) {
            Node node = attributes.item(i);
            existingRootElement.removeAttribute(node.getNodeName());
        }

        // add attributes from the new root node to the old root node
        attributes = newRootElement.getAttributes();
        attributeCount = attributes.getLength();
        for (int i = 0; i < attributeCount; i++) {
            Node node = attributes.item(i);
            existingRootElement.setAttribute(node.getNodeName(), node
                    .getNodeValue());
        }

        // remove the old generated elements and any
        // white space before the old nodes
        List<Node> nodesToDelete = new ArrayList<Node>();
        NodeList children = existingRootElement.getChildNodes();
        int length = children.getLength();
        for (int i = 0; i < length; i++) {
            Node node = children.item(i);
            if (isGeneratedNode(node)) {
                nodesToDelete.add(node);
            } else if (isWhiteSpace(node)
                    && isGeneratedNode(children.item(i + 1))) {
                nodesToDelete.add(node);
            }
        }

        for (Node node : nodesToDelete) {
            existingRootElement.removeChild(node);
        }

        // add the new generated elements
        children = newRootElement.getChildNodes();
        length = children.getLength();
        Node firstChild = existingRootElement.getFirstChild();
        for (int i = 0; i < length; i++) {
            Node node = children.item(i);
            // don't add the last node if it is only white space
            if (i == length - 1 && isWhiteSpace(node)) {
                break;
            }

            Node newNode = existingDocument.importNode(node, true);
            if (firstChild == null) {
                existingRootElement.appendChild(newNode);
            } else {
                existingRootElement.insertBefore(newNode, firstChild);
            }
        }

        // pretty print the result
        return prettyPrint(existingDocument);
    }

启动的29~43行，的目的是替换mapper节点的namespace，方式重新生成后，namespace有改变

 之后的47~62行，就是删除一些节点，其实按照他这意思就是为了删掉特定的节点，具体实现逻辑在isGeneratedNode方法内，由它决定删不删。

65~81行就是将新文件中的所有节点(非空白节点)全部合并至旧文件中。

private static boolean isGeneratedNode(Node node) {

        boolean rc = false;

        if (node != null && node.getNodeType() == Node.ELEMENT_NODE) {
            Element element = (Element) node;
            String id = element.getAttribute("id"); //$NON-NLS-1$
            if (id != null) {
                for (String prefix : MergeConstants.OLD_XML_ELEMENT_PREFIXES) {
                    if (id.startsWith(prefix)) {
                        rc = true;
                        break;
                    }
                }
            }

            if (rc == false) {
                // check for new node format - if the first non-whitespace node
                // is an XML comment, and the comment includes
                // one of the old element tags,
                // then it is a generated node
                NodeList children = node.getChildNodes();
                int length = children.getLength();
                for (int i = 0; i < length; i++) {
                    Node childNode = children.item(i);
                    if (isWhiteSpace(childNode)) {
                        continue;
                    } else if (childNode.getNodeType() == Node.COMMENT_NODE) {
                        Comment comment = (Comment) childNode;
                        String commentData = comment.getData();
                        for (String tag : MergeConstants.OLD_ELEMENT_TAGS) {
                            if (commentData.contains(tag)) {
                                rc = true;
                                break;
                            }
                        }
                    } else {
                        break;
                    }
                }
            }
        }

        return rc;
    }

逻辑其实也很简单，4~14行的逻辑就是删除id属性值带有一些特定前缀的节点，如果没找到，这删除commentNode节点，看到这，结果就出来了，按照正常情况下，根本不会把之前的就节点给删除掉。还是完完全全的保留。至此，就是我们常说的追加。

二、重写规则

从上述内容中，熟悉了原有的代码合并规则，接下来就是自定义规则了，本文就不放代码了，那样感觉很啰嗦，就直接简述一下实现思路，具体代码会在文末贴出github链接，可自行查看。

一、遍历新文件的所有非空白节点，遍历同时获取到对应旧文件中的节点，这里不考虑旧文件中有删除自动生成的节点情况，若获取到了，则遍历属性，有无增加，若增加，则移植到新文件中对应节点上，同时对该旧文件中的节点进行标记，等遍历完删掉。

二、第一个步骤完成后，然后再将新文件中的所有节点全部移植到旧文件中，最后视情况，需不需要格式化一下xml文件。

具体规则，就是图中红框处的文件

具体效果，大家可自行尝试，这里不贴效果图了。毕竟眼见为实。

注：不保证该规则适用于所有格式的xml文件，这块需要实地尝试。

 三、适用场景

本文这种方式，只适用于代码来生成文件的方式，对于适用maven插件，并不适用，如果需要，这里提供一种无奈方案，就是获取对应源码，替换掉该类，重新编译成jar包，放入到本地仓库里。

代码地址：Mybatis-generate-demo  

四、最后

如果还有其它比较好的方案。欢迎交流。

----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

转载请注明出处