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

package com.rthoni.intellij.codefromds.business;

import com.intellij.database.model.DasForeignKey;
import com.intellij.database.model.DataType;
import com.intellij.database.model.MultiRef;
import com.intellij.database.psi.DbDataSource;
import com.intellij.database.util.DasUtil;
import com.intellij.openapi.project.Project;
import com.rthoni.intellij.codefromds.dbo.options.ColumnSelection;
import com.rthoni.intellij.codefromds.dbo.options.GenerateOptions;
import com.rthoni.intellij.codefromds.dbo.options.TableSelection;
import com.rthoni.intellij.codefromds.dbo.options.TypesCastOptions;
import com.rthoni.intellij.codefromds.dbo.template.ColumnDataSourceDbo;
import com.rthoni.intellij.codefromds.dbo.template.DataSourceDbo;
import com.rthoni.intellij.codefromds.dbo.template.ForeignKeyDbo;
import com.rthoni.intellij.codefromds.dbo.template.TableDataSourceDbo;
import org.json.JSONArray;
import org.json.JSONObject;
import org.jtwig.JtwigModel;
import org.jtwig.JtwigTemplate;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Vector;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

/**
 * 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 ForeignKeyDbo convertForeignKey(DasForeignKey fk, DataSourceDbo source)
    {
        ForeignKeyDbo dbo = new ForeignKeyDbo();

        dbo.setName(fk.getName());

        dbo.setSourceForeignKeyName("fk_" + fk.getRefTable().getName());
        dbo.setSourceTable(source.getTables().stream().filter(t -> t.getName().equals(fk.getTable().getName())).findFirst().get());
        MultiRef columnsRef = fk.getColumnsRef();
        dbo.setSourceColumns(dbo.getSourceTable().getColumns().stream().filter(c -> StreamSupport.stream(columnsRef.names().spliterator(), false).anyMatch(cc -> cc.equals(c.getName()))).collect(Collectors.toList()));

        dbo.setTargetForeignKeyName(fk.getTable().getName() + "_fk");
        dbo.setTargetTable(source.getTables().stream().filter(t -> t.getName().equals(fk.getRefTable().getName())).findFirst().get());
        MultiRef refColumns = fk.getRefColumns();
        dbo.setTargetColumns(dbo.getTargetTable().getColumns().stream().filter(c -> StreamSupport.stream(refColumns.names().spliterator(), false).anyMatch(cc -> cc.equals(c.getName()))).collect(Collectors.toList()));

        dbo.getTargetTable().getTargetForeignKeys().add(dbo);

        return dbo;
    }

    public static ColumnDataSourceDbo convertColumn(ColumnSelection columnSelection, TypesCastOptions options)
    {
        ColumnDataSourceDbo dbo = new ColumnDataSourceDbo();
        dbo.setName(columnSelection.getColumn().getName());
        dbo.setPrimary(DasUtil.isPrimary(columnSelection.getColumn()));
        dbo.setSelected(columnSelection.isSelected());

        DataType type = columnSelection.getColumn().getDataType();
        boolean isArray = type.typeName.endsWith("[]");
        String sqlTypeName = isArray ? type.typeName.substring(0, type.typeName.length() - 2) : type.typeName;
        String typeName = type.typeName;
        HashMap<String, HashMap<String, String>> types = options.getTypes();

        if (types.containsKey(sqlTypeName)) {
            HashMap<String, String> subtype = types.get(sqlTypeName);
            if (subtype.containsKey(type.vagueArg)) {
                typeName = subtype.get(type.vagueArg);
            }
            else if (subtype.containsKey("*")) {
                typeName = subtype.get("*");
            }
        }
        if (isArray) {
            typeName = options.getArrayTemplate().replace("%t", typeName);
        }
        dbo.setType(typeName);

        return dbo;
    }

    public static TableDataSourceDbo convertTable(TableSelection tableSelection, TypesCastOptions types)
    {
        TableDataSourceDbo dbo = new TableDataSourceDbo();
        dbo.setName(tableSelection.getTable().getName());
        dbo.setColumns(tableSelection.getColumns().stream().map(c -> convertColumn(c, types)).collect(Collectors.toList()));
        dbo.setPrimaryKeys(dbo.getColumns().stream().filter(ColumnDataSourceDbo::isPrimary).collect(Collectors.toList()));
        dbo.setHasAny(!tableSelection.hasNone());
        dbo.setTargetForeignKeys(new Vector<>());
        return dbo;
    }

    public static DataSourceDbo convertOptions(GenerateOptions options, TypesCastOptions types)
    {
        DataSourceDbo dbo = new DataSourceDbo();

        dbo.setName(options.getSelection().getSource().getName());
        dbo.setTables(options.getSelection().getTables().stream().map(t -> convertTable(t, types)).collect(Collectors.toList()));

        List<TableDataSourceDbo> tables = dbo.getTables();
        List<TableSelection> tableSelections = options.getSelection().getTables();
        for (int i = 0; i < tables.size(); ++i) {
            TableDataSourceDbo table = tables.get(i);
            TableSelection tableSelection = tableSelections.get(i);
            table.setSourceForeignKeys(DasUtil.getForeignKeys(tableSelection.getTable()).toList().stream().map(t -> convertForeignKey(t, dbo)).collect(Collectors.toList()));
        }

        return dbo;
    }

    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 (type != null) {
                typeMap.put("*", type);
            }
            else if (typeObject != null) {
                for (String subtype : typeObject.keySet()) {
                    typeMap.put(subtype, typeObject.getString(subtype));
                }
            }
            map.put(key, typeMap);
        }
        dbo.setTypes(map);

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

        return dbo;
    }

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

        FileOutputStream file = new FileOutputStream(outputFile);

        JtwigTemplate template = JtwigTemplate.inlineTemplate(data);
        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
    {
        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());

        TypesCastOptions types = loadTypesCast(typesCastAbsolutePath);

        DataSourceDbo dbo = convertOptions(options, types);

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