intellij-code-from-ds/src/com/rthoni/intellij/codefromds/business/Generator.java

package com.rthoni.intellij.codefromds.business;

import com.intellij.database.psi.DbDataSource;
import com.intellij.openapi.project.Project;
import com.rthoni.intellij.codefromds.dbo.options.*;
import com.rthoni.intellij.codefromds.dbo.template.*;
import groovy.json.StringEscapeUtils;
import org.json.JSONArray;
import org.json.JSONObject;
import org.jtwig.JtwigModel;
import org.jtwig.JtwigTemplate;
import org.jtwig.environment.EnvironmentConfiguration;
import org.jtwig.environment.EnvironmentConfigurationBuilder;
import org.jtwig.functions.FunctionRequest;
import org.jtwig.functions.SimpleJtwigFunction;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Vector;
import java.util.function.Consumer;

/**
 * Created by robin on 11/15/16.
 */
public abstract class Generator {

    public static void saveOptions(GenerateOptions options) throws Exception
    {
        HashMap<String, Object> map = options.toMap();
        JSONObject obj = new JSONObject(map);
        FileOutputStream file = new FileOutputStream(options.getConfigAbsolutePath());
        file.write(obj.toString(4).getBytes());
        file.close();
    }

    public static GenerateOptions loadOptions(String configPath) throws Exception
    {
        String data = Helper.readFile(configPath);

        JSONObject obj = new JSONObject(data);
        String src = obj.getJSONObject("selection").getString("source");

        DbDataSource dataSource = Helper.findDataSource(src);
        if (dataSource == null) {
            throw new Exception("Data source " + src + " not found");
        }

        GenerateOptions options = new GenerateOptions(dataSource);
        options.setConfigAbsolutePath(configPath);
        options.fromJson(obj);
        return options;
    }

    public static String escapeReservedWords(String id, TypesCastOptions options)
    {
        if (options.getReservedWords().contains(id)) {
            return "_" + id;
        }
        return id;
    }

    public static String escapeString(String str)
    {
        return StringEscapeUtils.escapeJava(str);
    }

    public static void convertSqlType(SqlTypeDbo type, TypesCastOptions options, DataSourceDbo dataSourceDbo)
    {
        if (type == null) {
            return;
        }
        boolean isArray = type.getType().endsWith("[]");
        boolean isSetOf = type.getType().startsWith("setof ");
        String sqlTypeName = isArray ? type.getType().substring(0, type.getType().length() - 2) : type.getType();
        if (isSetOf)
        {
            sqlTypeName = sqlTypeName.substring(6);
        }
        String typeName = null;
        HashMap<String, HashMap<String, String>> types = options.getTypes();

        if (types.containsKey(sqlTypeName)) {
            HashMap<String, String> subtype = types.get(sqlTypeName);
            if (subtype.containsKey(type.getVagueArg())) {
                typeName = subtype.get(type.getVagueArg());
            }
            else if (subtype.containsKey("*")) {
                typeName = subtype.get("*");
            }
        }
        if (typeName == null) {
            TableDataSourceDbo table = dataSourceDbo.findTable(sqlTypeName);
            if (table == null) {
                typeName = types.get("*").get("*");
            }
            else {
                typeName = table.getName();
            }
        }
        if (isArray) {
            typeName = options.getArrayTemplate().replace("%t", typeName);
        }
        if (isSetOf) {
            typeName = options.getSetOfTemplate().replace("%t", typeName);
        }
        type.setLanguageType(typeName);
        type.setLanguageTypeNotNull(options.getNonNullableTypes().contains(typeName));
    }

//    public static boolean isUnionSame(List<ColumnDataSourceDbo> part1, List<ColumnDataSourceDbo> part2, List<ColumnDataSourceDbo> list2)
//    {
//        List<ColumnDataSourceDbo> list1 = new Vector<>(part1);
//        for (ColumnDataSourceDbo col : part2) {
//            if (!list1.contains(col)) {
//                list1.add(col);
//            }
//        }
//        if (list1.size() != list2.size()) {
//            return false;
//        }
//        list1.sort((o1, o2) -> o1.getName().compareTo(o2.getName()));
//        list2.sort((o1, o2) -> o1.getName().compareTo(o2.getName()));
//        return list1.equals(list2);
//    }

    public static DataSourceDbo convertOptions(GenerateOptions options, TypesCastOptions types)
    {
        String dataSourceName = options.getSelection().getSource().getName();
        DataSourceDbo dataSourceDbo = DataSourcesBusiness.getDataSourceDbo(dataSourceName);

        for (TableDataSourceDbo table : dataSourceDbo.getTables())
        {
            TableSelection tableSelection = options.getSelection().getTables().stream()
                    .filter(t -> t.getTable().getName().equals(table.getName()))
                    .findFirst().orElse(null);
            for (ColumnDataSourceDbo column : table.getColumns())
            {
                ColumnSelection columnSelection = tableSelection.getColumns().stream()
                        .filter(c -> c.getColumn().getName().equals(column.getName()))
                        .findFirst().orElse(null);
                convertSqlType(column.getSqlType(), types, dataSourceDbo);
                column.setSelected(columnSelection.isSelected());
            }
        }

        for (StoredProcedureDbo spDbo : dataSourceDbo.getStoredProcedures())
        {
            StoredProcedureSelection sp = options.getSelection().getStoredProcedures().stream()
                    .filter(s -> s.getStoredProcedure().getText().equals(spDbo.getFullName()))
                    .findFirst().orElse(null);
            convertSqlType(spDbo.getSqlType(), types, dataSourceDbo);
            spDbo.setSelected(sp.isSelected());
            for (StoredProcedureArgDbo arg : spDbo.getArguments()) {
                convertSqlType(arg.getSqlType(), types, dataSourceDbo);
            }
        }

        return dataSourceDbo;
    }

    public static TypesCastOptions loadTypesCast(String file) throws IOException
    {
        TypesCastOptions dbo = new TypesCastOptions();
        String data = Helper.readFile(file);

        JSONObject obj = new JSONObject(data);
        JSONObject objTypes = obj.getJSONObject("types");
        HashMap<String, HashMap<String, String>> map = new HashMap<>();

        for (String key : objTypes.keySet()) {
            HashMap<String, String> typeMap = new HashMap<>();
            String type = objTypes.optString(key);
            JSONObject typeObject = objTypes.optJSONObject(key);
            if (typeObject != null) {
                for (String subtype : typeObject.keySet()) {
                    typeMap.put(subtype, typeObject.getString(subtype));
                }
            }
            else if (type != null) {
                typeMap.put("*", type);
            }
            map.put(key, typeMap);
        }
        dbo.setTypes(map);

        List<String> nonNullableTypes = new Vector<>();
        JSONArray array = obj.getJSONArray("non-nullable-types");
        for (int i = 0; i < array.length(); ++i) {
            nonNullableTypes.add(array.getString(i));
        }
        dbo.setNonNullableTypes(nonNullableTypes);

        List<String> reservedWords = new Vector<>();
        array = obj.getJSONArray("reserved-words");
        for (int i = 0; i < array.length(); ++i) {
            reservedWords.add(array.getString(i));
        }
        dbo.setReservedWords(reservedWords);

        dbo.setArrayTemplate(obj.getString("arrayTemplate"));
        dbo.setSetOfTemplate(obj.getString("setOfTemplate"));

        return dbo;
    }

    public static void generateFile(String templatePath, String outputFile, DataSourceDbo dataSource,
                                    TableDataSourceDbo table, TypesCastOptions options) throws IOException
    {
        String data = Helper.readFile(templatePath);

        FileOutputStream file = new FileOutputStream(outputFile);

        final SimpleJtwigFunction escapeReservedWords = new SimpleJtwigFunction() {
            @Override
            public String name() {
                return "escapeReservedWords";
            }

            @Override
            public Object execute(FunctionRequest functionRequest) {
                Object obj = functionRequest.get(0);
                if (obj == null)
                {
                    return null;
                }
                String arg = (String) obj;
                String v = escapeReservedWords(arg, options);
                return v;
            }
        };

        final SimpleJtwigFunction escapeString = new SimpleJtwigFunction() {
            @Override
            public String name() {
                return "escapeString";
            }

            @Override
            public Object execute(FunctionRequest functionRequest) {
                Object obj = functionRequest.get(0);
                if (obj == null)
                {
                    return null;
                }
                String arg = (String) obj;
                String v = escapeString(arg);
                return v;
            }
        };

        final EnvironmentConfiguration configuration = EnvironmentConfigurationBuilder
                .configuration()
                .functions()
                .add(escapeReservedWords)
                .add(escapeString)
                .and()
                .build();

        JtwigTemplate template = JtwigTemplate.inlineTemplate(data, configuration);
        JtwigModel model = JtwigModel.newModel()
                .with("dataSource", dataSource)
                .with("table", table);

        template.render(model, file);

        file.close();
    }

    public static void generate(GenerateOptions options, Project project) throws IOException
    {
        DataSourcesBusiness.getDataSourcesDbo();
        String modelsAbsolutePath = Helper.getAbsolutePath(project, options.getModelsRelativePath());
        String dataSourceTemplateAbsolutePath = Helper.getAbsolutePath(project, options.getDataSourceTemplateRelativePath());
        String modelsTemplateAbsolutePath = Helper.getAbsolutePath(project, options.getModelsTemplateRelativePath());
        String typesCastAbsolutePath = Helper.getAbsolutePath(project, options.getCastFileRelativePath());
        String dataSourceAbsolutePath = Helper.getAbsolutePath(project, options.getDataSourceRelativePath());

        TypesCastOptions types = loadTypesCast(typesCastAbsolutePath);

        DataSourceDbo dbo = convertOptions(options, types);

        generateFile(dataSourceTemplateAbsolutePath, dataSourceAbsolutePath, dbo, null, types);
        for (TableDataSourceDbo table : dbo.getTables()) {
            if (table.hasAny()) {
                String filename = modelsAbsolutePath + File.separator + table.getName() + "." + options.getFilesExtension();
                generateFile(modelsTemplateAbsolutePath, filename, dbo, table, types);
            }
        }
    }
}