io.github.bootystar.mybatisplus.generator.engine.EnhanceVelocityTemplateEngine Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of booty-mybatis-plus-generator Show documentation
The newest version!
package io.github.bootystar.mybatisplus.generator.engine;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.builder.CustomFile;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.AbstractTemplateEngine;
import com.baomidou.mybatisplus.generator.engine.VelocityTemplateEngine;
import io.github.bootystar.mybatisplus.generator.config.IConfig;

import java.io.File;
import java.util.List;
import java.util.Map;
import java.util.Optional;

/**
 * 自定义Velocity引擎
 * @author booty
 *
 */
public class EnhanceVelocityTemplateEngine extends VelocityTemplateEngine {

    private IConfig customConfig;


    public EnhanceVelocityTemplateEngine() {
    }


    public EnhanceVelocityTemplateEngine(IConfig customConfig) {
        this.customConfig = customConfig;
    }

    @Override
    protected void outputCustomFile(List customFiles, TableInfo tableInfo, Map objectMap) {
        String entityName = tableInfo.getEntityName();
        String parentPath = getPathInfo(OutputFile.parent);
        Boolean fileOverride = customConfig.getFileOverride();
        customFiles.forEach(file -> {
            String filePath = StringUtils.isNotBlank(file.getFilePath()) ? file.getFilePath() : parentPath;
            if (StringUtils.isNotBlank(file.getPackageName())) {
                filePath = filePath + File.separator + file.getPackageName();
            }
            String fileName = filePath + File.separator + entityName + file.getFileName();

//            outputFile(new File(fileName), objectMap, file.getTemplatePath(), file.isFileOverride());
            outputFile(new File(fileName), objectMap, file.getTemplatePath(), fileOverride);
        });
    }

    @Override
    public Map getObjectMap(ConfigBuilder config, TableInfo tableInfo) {
        Map objectMap = super.getObjectMap(config, tableInfo);
        if (customConfig!=null){
            Map customData = customConfig.renderData(tableInfo);
            if (customData!=null){
                objectMap.putAll(customData);
            }
        }
        objectMap.put("basePackage",config.getPackageConfig().getParent());

        return objectMap ;
    }


    @Override
    public AbstractTemplateEngine batchOutput() {
        try {
            ConfigBuilder config = this.getConfigBuilder();
            List tableInfoList = config.getTableInfoList();
            tableInfoList.forEach(tableInfo -> {
                Map objectMap = this.getObjectMap(config, tableInfo);
                Optional.ofNullable(config.getInjectionConfig()).ifPresent(t -> {
                    // 添加自定义属性
                    t.beforeOutputFile(tableInfo, objectMap);
                    // 输出自定义文件
                    outputCustomFile(t.getCustomFiles(), tableInfo, objectMap);
                });
                // entity
                outputEntity(tableInfo, objectMap);
                // mapper and xml
                outputMapper(tableInfo, objectMap);
                // service
                outputService(tableInfo, objectMap);
                // controller
                outputController(tableInfo, objectMap);

                Optional.ofNullable(customConfig).ifPresent(t -> {
                    // 输出自定义文件
                    outputCustomFile(t.getCustomFiles(), tableInfo, objectMap);
                });

            });
        } catch (Exception e) {
            throw new RuntimeException("无法创建文件，请检查配置信息！", e);
        }
        return this;
    }


}