GenCodeLocal/src/main/java/gencode/MybatisPlusGenCode.java

package gencode;

import org.apache.velocity.Template;
import org.apache.velocity.VelocityContext;
import org.apache.velocity.app.VelocityEngine;

import java.io.*;
import java.sql.*;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

public class MybatisPlusGenCode {
    //数据驱动
    final static String DRIVER_CLASS = "com.mysql.jdbc.Driver";
    //数据库名
    final static String DATABASE_NAME = "airport_industry_db";
    //数据连接地址
    final static String CONNECTION_URL = "jdbc:mysql://172.16.90.201:3306/" + DATABASE_NAME + "?characterEncoding=UTF-8";
    //数据库登录名
    final static String USER_ID = "root";
    //密码
    final static String PASS_WORD = "123456";
    //数据库表名
    final static String[] TABLE_NAMES = {
            "erp_channel_shop_material",
            "air_goods",};

    /**
     * 不能修改
     */
    //controller模版文件目录
    final static String CONTROLLER_DIR = "src/main/resources/vm/mybatisPlus/java/controller/ModelController.java.vm";
    //service 模版文件目录
    final static String SERVICE_DIR = "src/main/resources/vm/mybatisPlus/java/service/ModelSeevice.java.vm";
    //serviceImpl模版文件目录
    final static String SERVICEIMPL_DIR = "src/main/resources/vm/mybatisPlus/java/service/impl/ModelServiceImpl.java.vm";
    //entity 模板文件目录
    final static String ENTITY_DIR = "src/main/resources/vm/mybatisPlus/java/entity/ModelEntity.java.vm";

    final static String MAPPER_JAVA_DIR = "src/main/resources/vm/mybatisPlus/java/mapper/ModelMapper.java.vm";

    final static String MAPPER_XML_DIR = "src/main/resources/vm/mybatisPlus/java/mapper/ModelMapper.xml.vm";

    //项目目录
    final static String BASE_PACKAGE = "com.hw";
    final static String BASE_PACKAGE_PATH = "src.main.java";
    //基础目录(不允许修改)
    final static String BASE_PACKAGE_NAME = BASE_PACKAGE_PATH + "." + BASE_PACKAGE;

    public static void main(String[] args) throws Exception {
        productionJavaFiles();

    }

    public static List<ColumnVo> getIdType(String tableName) {

        List<ColumnVo> columnVos = new ArrayList<>();

        try {
            //1、注册JDBC驱动
            Class.forName(DRIVER_CLASS);
            //2、获取数据库连接
            Connection connection = DriverManager.getConnection(CONNECTION_URL, USER_ID, PASS_WORD);
            //3、操作数据库
            Statement statement = connection.createStatement();//获取操作数据库的对象

            String sql = "SELECT COLUMN_NAME,IS_NULLABLE,COLUMN_COMMENT,DATA_TYPE,CHARACTER_MAXIMUM_LENGTH,COLUMN_KEY " +
                    "FROM information_schema. COLUMNS WHERE table_schema = '" + DATABASE_NAME + "' AND table_name = '" + tableName + "'";
            ResultSet resultSet = statement.executeQuery(sql);//执行sql，获取结果集
            System.out.println("sql：" + sql);
            while (resultSet.next()) { //遍历结果集，取出数据
                //输出数据
                String key = resultSet.getString("COLUMN_KEY");
                String column_name = resultSet.getString("COLUMN_NAME");
                System.out.println("主键名称：" + column_name);
                String data_type = resultSet.getString("DATA_TYPE");
                String column_comment = resultSet.getString("COLUMN_COMMENT");
                String IS_NULLABLE = resultSet.getString("IS_NULLABLE");

                System.out.println("主键类型：" + data_type);
                MysqlDataType enumType = MysqlDataType.getEnumType(data_type);
                if (enumType != MysqlDataType.ERROR) {
                    ColumnVo columnVo = new ColumnVo(column_name, enumType.getValue(), enumType.getJavaPackage());
                    columnVo.setAttribute(DataUtil.toCase(column_name));
                    columnVo.setAttributeUp(DataUtil.toUp(columnVo.getAttribute()));
                    columnVo.setRemake(column_comment + " 必填:" + IS_NULLABLE);
                    if("PRI".equals(key)){
                        columnVo.setKeyFlag(true);
                    }else{
                        columnVo.setKeyFlag(false);
                    }
                    columnVos.add(columnVo);
                }
            }
            //4、关闭结果集、数据库操作对象、数据库连接
            resultSet.close();
            statement.close();
            connection.close();
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("数据库连接错误!");
        }
        return columnVos;
    }


    private static void productionJavaFiles() throws Exception {

        for (String tableName : TABLE_NAMES) {
            VelocityEngine ve = new VelocityEngine();
            ve.init();
            VelocityContext ctx = new VelocityContext();
            ctx.put("tableName", tableName);
            //获取列数据
            List<ColumnVo> mysqlDataType = getIdType(tableName);
            for (ColumnVo columnVo : mysqlDataType) {
                boolean keyFlag = columnVo.isKeyFlag();
                if(keyFlag){
                    ctx.put("keyPackage", columnVo.getJavaPackage());
                    ctx.put("keyValue", columnVo.getValue());
                    break;
                }

            }
            if (mysqlDataType.size() > 0) {
                Set<String> collect = mysqlDataType.stream().map(ColumnVo::getJavaPackage).collect(Collectors.toSet());
                ctx.put("columns", mysqlDataType);
                ctx.put("javaPackages", collect);
                tableName = DataUtil.toCase(tableName);
                ctx.put("entity", tableName);
                String tableNameNew = DataUtil.toUp(tableName);
                ctx.put("Entity", tableNameNew);

                saveEntity(ve, ctx, tableNameNew);
                saveMapper(ve, ctx, tableNameNew);
                saveMapperXml(ve, ctx, tableNameNew);
                saveService(ve, ctx, tableNameNew);
                saveServiceImpl(ve, ctx, tableNameNew);
                saveController(ve, ctx, tableNameNew);
            }
        }


    }

    private static void saveEntity(VelocityEngine ve, VelocityContext ctx, String entityName) throws IOException {

        createModel(ve, ctx, entityName, "domian", ENTITY_DIR, ".java", "ENTITY_PACKAGE_NAME");
    }

    private static void saveMapper(VelocityEngine ve, VelocityContext ctx, String entityName) throws IOException {

        createModel(ve, ctx, entityName, "mapper", MAPPER_JAVA_DIR, "Mapper.java", "MAPPER_PACKAGE_NAME");
    }
    private static void saveMapperXml(VelocityEngine ve, VelocityContext ctx, String entityName) throws IOException {

        createModel(ve, ctx, entityName, "mapper", MAPPER_XML_DIR, "Mapper.xml", "MAPPER_PACKAGE_NAME");
    }

    private static void saveService(VelocityEngine ve, VelocityContext ctx, String entityName) throws IOException {

        createModel(ve, ctx, entityName, "service", SERVICE_DIR, "Service.java", "service_package_name");
    }

    private static void saveServiceImpl(VelocityEngine ve, VelocityContext ctx, String entityName) throws IOException {

        createModel(ve, ctx, entityName, "service", SERVICEIMPL_DIR, "ServiceImpl.java", "service_impl_package_name");
    }


    private static void saveController(VelocityEngine ve, VelocityContext ctx, String entityName) throws IOException {


        createModel(ve, ctx, entityName, "controller", CONTROLLER_DIR, "Controller.java", "controller_package_name");
    }

    private static void createModel(VelocityEngine ve, VelocityContext ctx, String entityName, String controller, String controllerDir, String s2, String controller_package_name) throws IOException {
        String BaseName = controller;

        Template t = ve.getTemplate(controllerDir);
        //创建文件
        String substring = entityName + s2;
        //输出文件
        String replace = BASE_PACKAGE_NAME.replace('.', '/') + "/" + BaseName + "/";
        String s = replace + substring;

        String str = replace.replace('/', '.') + substring;
        String ENTITY_PACKAGE_NAME = str.replace(BASE_PACKAGE_PATH + ".", "");

        String PACKAGE = BASE_PACKAGE + "." + BaseName;

        ctx.put(controller_package_name, PACKAGE);

        File controllerFile = new File(s);

        saveFile(t, ctx, replace, controllerFile);
    }


    private static void saveFile(Template t, VelocityContext ctx, String filePath, File saveFile) throws IOException {

        File path = new File(filePath);
        //文件，如果不存在，就创建
        if (!path.exists() && !path.isDirectory()) {
            System.out.println("//不存在");
            path.mkdirs();
        } else {
            System.out.println("//目录存在");
        }
        //创建文件输出流
        FileOutputStream outStream = new FileOutputStream(saveFile);
        //因为模板整合的时候，需要提供一个Writer，所以创建一个Writer
        OutputStreamWriter writer = new OutputStreamWriter(outStream);

        //创建一个缓冲流
        BufferedWriter bufferWriter = new BufferedWriter(writer);

        //5 Merge the template and you data toproduce the output
        t.merge(ctx, bufferWriter);

        bufferWriter.flush();//强制刷新

        outStream.close();

        bufferWriter.close();
    }

}