product-promotion/poyee-base/src/main/java/com/poyee/util/TableInitUtil.java

package com.poyee.util;

import com.baomidou.mybatisplus.annotation.IdType;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.FastAutoGenerator;
import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.poyee.base.controller.BaseController;
import com.poyee.base.mapper.IBaseMapper;
import com.poyee.common.db.EnhanceOtherTemplateEngine;
import com.poyee.common.db.TableInitCustomerUtil;
import com.poyee.common.dto.DbConfigModal;
import com.poyee.common.dto.TableColumnConfigModal;
import com.poyee.common.dto.TableConfigModal;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.yaml.snakeyaml.Yaml;

import java.io.InputStream;
import java.sql.*;
import java.util.*;

/**
 * table工具类
 */
@Slf4j
public class TableInitUtil {

    static final String tmp_path = "/templates/";

    static final String project_path = "C:\\mine\\poyee\\poyee-generator\\src\\main\\";

    public static void main(String[] args) {
        new TableInitUtil().checkTableConfig(null);
    }

    public static List<DbConfigModal> checkTableConfig(DbConfigModal req) {
        List<DbConfigModal> list = new ArrayList<>();
        String sourcePath = "spring.datasource.dynamic.datasource";
        //获取配置文件中数据源信息
        String dataSource = Objects.isNull(req) ? null : req.getDataSource();
        Yaml yaml = new Yaml();
        InputStream inputStream = TableInitUtil.class.getClassLoader().getResourceAsStream("application-dev.yml");
        Map<String, Object> config = yaml.load(inputStream);
        try {
            config = getSources(config, sourcePath);
        }catch (Exception e){
            log.error("获取数据源信息失败:{}",e.getMessage());
            throw new RuntimeException("获取数据源信息失败");
        }
        if(StringUtils.isBlank(dataSource)){
            //如果为空则获取所有
            Optional.ofNullable((Map<String, Object>) config)
                    .ifPresent(map -> map.forEach((k, v) -> {
                        if (v instanceof Map) {
                            Map<String, String> mmap = (Map<String, String>) v;
                            DbConfigModal req1 = new DbConfigModal();
                            req1.setDbUrl(mmap.get("url"));
                            req1.setDbUsername(mmap.get("username"));
                            req1.setDbPwd(String.valueOf(mmap.get("password")));
                            req1.setDriverClassName(mmap.get("driver-class-name"));
                            req1.setDataSource(k);
                            list.add(req1);
                        }
                    }));
        }else{
            Map<String,String> dataSourceConfig = (Map<String, String>) config.get(dataSource);
            log.info("数据源信息:{}",dataSourceConfig);
            req.setDbUrl(dataSourceConfig.get("url"));
            req.setDbUsername(dataSourceConfig.get("username"));
            req.setDbPwd(String.valueOf(dataSourceConfig.get("password")));
            req.setDriverClassName(dataSourceConfig.get("driver-class-name"));
            list.add(req);
        }
        return list;
    }

    static Map<String, Object> getSources(Map<String, Object> config,String sourcePath) throws Exception{
        if(StringUtils.isBlank(sourcePath)){
            return config;
        }
        if (sourcePath.indexOf(".") == 0) {
            sourcePath = sourcePath.substring(1);
        }
        String key = sourcePath.split("\\.")[0];
        sourcePath = sourcePath.replaceFirst(key, "");
        config = (Map<String, Object>) config.get(key);
        return getSources(config,sourcePath);
    }

    public static List<TableConfigModal> checkTableInfos(@NotNull DbConfigModal config) {
        Connection conn = null;
        List<TableConfigModal> tableConfigs = new ArrayList<>();
        try {
            Class.forName(config.getDriverClassName()).newInstance();
            conn = DriverManager.getConnection(config.getDbUrl(), config.getDbUsername(), config.getDbPwd());
            DatabaseMetaData metaData = conn.getMetaData();
            ResultSet tables = metaData.getTables(null, null, null, new String[]{"TABLE"});
            while (tables.next()) {
                String tableName = tables.getString("TABLE_NAME");
                String tableComment = tables.getString("REMARKS");
                tableConfigs.add(new TableConfigModal(tableName, tableComment, null));
            }
            return tableConfigs;
        } catch (SQLException | ClassNotFoundException | IllegalAccessException | InstantiationException e) {
            throw new RuntimeException(e);
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }

    public static List<TableConfigModal> checkTableColumns(@NotNull DbConfigModal config) {
        Connection conn = null;
        List<TableConfigModal> tableConfigs = new ArrayList<>();
        try {
            Class.forName(config.getDriverClassName()).newInstance();
            conn = DriverManager.getConnection(config.getDbUrl(), config.getDbUsername(), config.getDbPwd());
            DatabaseMetaData metaData = conn.getMetaData();
            ResultSet tables = metaData.getTables(null, null, config.getTableName(), new String[]{"TABLE"});
            while (tables.next()) {
                String tableName = tables.getString("TABLE_NAME");
                String tableComment = tables.getString("REMARKS");
                log.info("表名：{} , 注释：{}", tableName, tableComment);
                ResultSet udTs = metaData.getUDTs(null, null, tableName, null);
                while (udTs.next()) {
                    String udTsREMARKS = udTs.getString("REMARKS");
                    log.info("udTsREMARKS：{} ", udTsREMARKS);
                }
                List<TableColumnConfigModal> columns = new ArrayList<>();
                ResultSet columnsResult = metaData.getColumns(null, null, tableName, null);
                while (columnsResult.next()) {
                    String columnName = columnsResult.getString("COLUMN_NAME");
                    Integer ordinalPosition = columnsResult.getInt("ORDINAL_POSITION");
                    String columnDef = columnsResult.getString("COLUMN_DEF");
                    String isNullable = columnsResult.getString("IS_NULLABLE");
                    String dataType = columnsResult.getString("DATA_TYPE");
                    Integer charOctetLength = columnsResult.getInt("CHAR_OCTET_LENGTH");
                    Integer numericPrecision = columnsResult.getInt("NUM_PREC_RADIX");
                    Integer numericScale = columnsResult.getInt("DECIMAL_DIGITS");
                    Integer datetimePrecision = columnsResult.getInt("COLUMN_SIZE");
                    String columnType = columnsResult.getString("TYPE_NAME");
                    String columnComment = columnsResult.getString("REMARKS");
                    TableColumnConfigModal column = new TableColumnConfigModal();
                    column.setColumnName(columnName);
                    column.setColumnType(columnType);
                    column.setColumnComment(columnComment);
                    column.setColumnDefault(columnDef);
                    column.setIsNullable(isNullable);
                    column.setCharacterOctetLength(charOctetLength);
                    column.setDataType(dataType);
                    column.setDatetimePrecision(datetimePrecision);
                    column.setNumericPrecision(numericPrecision);
                    column.setNumericScale(numericScale);
                    column.setColumnType(columnType);
                    column.setOrdinalPosition(ordinalPosition);
                    columns.add(column);
                }
                tableConfigs.add(new TableConfigModal(tableName, tableComment, columns));
            }
            return tableConfigs;
        } catch (SQLException | ClassNotFoundException | IllegalAccessException | InstantiationException e) {
            throw new RuntimeException(e);
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }

    /**
     * 生成工具
     *
     * @param config
     */
    public static void generation(TableConfigModal config) {
        FastAutoGenerator.create(config.getDbUrl(), config.getDbUsername(), config.getDbPwd())
                         .globalConfig(builder -> {
                             builder.author(config.getAuthor()) // 设置作者
                                    .fileOverride() // 覆盖已生成文件
                                    .outputDir(config.getJavaProjectPath() + "java"); // 指定输出目录
                             if (config.isOpenSwagger()) { // 开启 swagger 模式
                                 builder.enableSwagger();
                             }
                         })
                         .packageConfig(builder -> {
                             Map<OutputFile, String> outputFileStringMap = new HashMap<>();
                             outputFileStringMap.put(OutputFile.xml, config.getJavaProjectPath() + "resources\\mapper"); // 设置mapperXml生成路径
                             outputFileStringMap.put(OutputFile.other, config.getVueProjectPath() + "views");// 设置pages生成路径
                             builder.parent(config.getJavaPackage()) // 设置父包名
                                    .moduleName(config.getJavaModule())//模块名
                                    .other("dto")
                                    .pathInfo(outputFileStringMap)
                                    ;
                             System.out.println(builder);
                         })
                         .injectionConfig(builder -> {
                             builder.customFile(checkOtherTemplates(config))
                                    .beforeOutputFile((tableInfo, stringObjectMap) -> {
                                        List<TableField> fields = tableInfo.getFields();
                                        fields.forEach(field -> {
                                            String propertyName = field.getPropertyName();
                                            if (StringUtils.isNotBlank(propertyName)) {
                                                if (!Objects.isNull(config.getCustomMap())) {
                                                    Map<String, Object> customeMap =
                                                            Optional.ofNullable(config.getCustomMap()
                                                                                      .get(propertyName))
                                                                    .orElse(new HashMap<>());
                                                    field.setCustomMap(customeMap);
                                                } else {
                                                    field.setCustomMap(TableInitCustomerUtil.checkCustomMap(field));
                                                }
                                            }
                                        });
                                    });
                         })
                         .templateConfig(builder -> {
                             builder.disable()
                                    .entity(config.getTemplates()
                                                  .contains("entity") ? tmp_path + "entity.java.vm" : "")
                                    .controller(config.getTemplates()
                                                      .contains("controller") ? tmp_path + "controller.java.vm" : "")
                                    .service(config.getTemplates()
                                                   .contains("service") ? tmp_path + "service.java.vm" : "")
                                    .serviceImpl(config.getTemplates()
                                                       .contains("serviceImpl") ? tmp_path + "serviceImpl.java.vm" : "")
                                    .mapper(config.getTemplates()
                                                  .contains("mapper") ? tmp_path + "mapper.java.vm" : "")
                                    .xml(config.getTemplates()
                                               .contains("mapperXml") ? tmp_path + "mapper.xml.vm" : "")
                                    .build();
                         })
                         .strategyConfig(builder -> {
                             builder.addInclude(config.getTableName()) // 设置需要生成的表名
                                    .addTablePrefix("t_", "c_"); // 设置过滤表前缀
                             if (config.getTemplates()
                                       .contains("entity")) {
                                 builder.entityBuilder()//配置entity
                                        .enableLombok()
                                        .idType(IdType.AUTO)
                                        .enableTableFieldAnnotation()
                                        .fileOverride()
                                        .build();
                             }
                             if (config.getTemplates()
                                       .contains("mapper")) {
                                 builder.mapperBuilder()//配置mapper
                                        .superClass(IBaseMapper.class)
                                        .enableBaseResultMap()
                                        .enableBaseColumnList()
                                        .enableMapperAnnotation()
                                        .fileOverride()
                                        .build();
                             }
                             if (config.getTemplates()
                                       .contains("service")) {
                                 builder.serviceBuilder()//配置service
                                        .superServiceClass("com.poyee")
                                        .superServiceImplClass("com.poyee")
                                        .fileOverride()
                                        .build();
                             }
                             if (config.getTemplates()
                                       .contains("controller")) {
                                 builder.controllerBuilder()
                                        .superClass(BaseController.class)
                                        .enableRestStyle()
                                        .fileOverride()
                                        .build();
                             }
                         })
                         .templateEngine(new EnhanceOtherTemplateEngine())
                         .execute();
    }

    /**
     * @param config
     * @return
     */
    @NotNull
    private static Map<String, String> checkOtherTemplates(TableConfigModal config) {
        Map<String, String> map = new HashMap<>();
        config.getTemplates()
              .forEach(t -> {
                  switch (t) {
                      case "Dto":
                          map.put("Dto.java", tmp_path + "dto.java.vm");
                          break;
                      case "index":
                          map.put("index.vue", tmp_path + "index.vue.vm");
                          break;
                      case "index2":
                          map.put("index.vue", tmp_path + "index2.vue.vm");
                          break;
                      case "hold":
                          map.put("hold.vue", tmp_path + "hold.vue.vm");
                          break;
                  }
              });
        return map;
    }

    /**
     * @param tableName
     */
    public static void generation(String tableName) {
        FastAutoGenerator.create("jdbc:mysql://192.168.207.128:3306/semm", "semm", "123456")
//        FastAutoGenerator.create("jdbc:postgresql://localhost:5432/hobby_stocks", "postgres", "123456")
                         .globalConfig(builder -> {
                             builder.author("lsz") // 设置作者
                                    .enableSwagger() // 开启 swagger 模式
                                    .fileOverride() // 覆盖已生成文件
                                    .outputDir(project_path + "java"); // 指定输出目录
                         })
                         .packageConfig(builder -> {
                             Map<OutputFile, String> outputFileStringMap = new HashMap<>();
                             outputFileStringMap.put(OutputFile.xml, project_path + "resources\\mapper"); // 设置mapperXml生成路径
                             outputFileStringMap.put(OutputFile.other, project_path + "resources\\pages");// 设置pages生成路径
                             builder.parent("com.poyee") // 设置父包名
                                    .other("dto")
                                    .pathInfo(outputFileStringMap);
                         })
                         .injectionConfig(builder -> {
                             Map<String, String> map = new HashMap<>();
                             map.put("Dto.java", tmp_path + "dto.java.vm");
                             map.put("index.vue", tmp_path + "index.vue.vm");
                             map.put("hold.vue", tmp_path + "hold.vue.vm");
                             builder.customFile(map)
                                    .beforeOutputFile((tableInfo, stringObjectMap) -> {
                                        List<TableField> fields = tableInfo.getFields();
                                        fields.forEach(field -> field.setCustomMap(TableInitCustomerUtil.checkCustomMap(field)));
                                    });
                         })
                         .templateConfig(builder -> {
                             builder.disable()
                                    .entity(tmp_path + "entity.java.vm")
                                    .controller(tmp_path + "controller.java.vm")
                                    .service(tmp_path + "service.java.vm")
                                    .serviceImpl(tmp_path + "serviceImpl.java.vm")
                                    .mapper(tmp_path + "mapper.java.vm")
                                    .xml(tmp_path + "mapper.xml.vm")
//                            .entity(path+"dto.java.vm")
                                    .build();
                         })
                         .strategyConfig(builder -> { //策略配置
                             builder.addInclude(tableName) // 设置需要生成的表名
                                    .addTablePrefix("t_", "c_"); // 设置过滤表前缀
                             builder.entityBuilder()//配置entity
                                    .enableLombok()
                                    .idType(IdType.AUTO)
                                    .enableTableFieldAnnotation()
                                    .fileOverride()
                                    .build();
                             builder.mapperBuilder()//配置mapper
                                    .enableBaseResultMap()
                                    .enableBaseColumnList()
                                    .enableMapperAnnotation()
                                    .fileOverride()
                                    .build();
                             builder.serviceBuilder()//配置service
                                    .fileOverride()
                                    .build();
                             builder.controllerBuilder()
                                    .enableRestStyle()
                                    .fileOverride()
                                    .build();
                         })
                         .templateEngine(new EnhanceOtherTemplateEngine())
                         .execute();
    }


}