intellij-code-from-ds/src/com/rthoni/intellij/codefromds/DataAccess/DataSourcesDataAccess.java

package com.rthoni.intellij.codefromds.DataAccess;

import com.intellij.database.model.DasArgument;
import com.intellij.database.psi.*;
import com.intellij.database.util.DasUtil;
import com.rthoni.intellij.codefromds.business.Helper;
import com.rthoni.intellij.codefromds.dbo.template.*;

import java.util.List;
import java.util.stream.Collectors;

/**
 * Created by robin on 4/13/17.
 */
public class DataSourcesDataAccess {

    public static List<DataSourceDbo> getDataSourcesDbo()
    {
        return Helper.getDataSources().stream()
                .map(DataSourcesDataAccess::dataSourceToDbo)
                .collect(Collectors.toList());
    }

    public static DataSourceDbo getDataSourceDbo(String name)
    {
        return Helper.getDataSources().stream()
                .filter(dataSource -> dataSource.getName().equals(name))
                .map(DataSourcesDataAccess::dataSourceToDbo)
                .findFirst().orElse(null);
    }

    private static DataSourceDbo dataSourceToDbo(DbDataSource dataSource)
    {
        DataSourceDbo dataSourceDbo = new DataSourceDbo();
        dataSourceDbo.setName(dataSource.getName());

        final DataHolder<TableDataSourceDbo> tableDataSourceDbo = new DataHolder<>();
        final DataHolder<ColumnDataSourceDbo> columnDataSourceDbo = new DataHolder<>();

        dataSource.getModel().traverser().forEach(dasObject ->
        {
            if (dasObject instanceof DbTable) {
                DbTable table = (DbTable) dasObject;
                tableDataSourceDbo.data = new TableDataSourceDbo();
                tableDataSourceDbo.data.setName(table.getName());
                dataSourceDbo.addTable(tableDataSourceDbo.data);
            }
            else if (dasObject instanceof DbColumn) {
                DbColumn column = (DbColumn) dasObject;
                columnDataSourceDbo.data = new ColumnDataSourceDbo();
                columnDataSourceDbo.data.setName(column.getName());
                columnDataSourceDbo.data.setNotNull(column.isNotNull());
                SqlTypeDbo type = new SqlTypeDbo();
                type.setName(column.getDataType().typeName);
                type.guessArrayAndSetOf();
                type.setVagueArg(column.getDataType().vagueArg);
                columnDataSourceDbo.data.setType(type);
                columnDataSourceDbo.data.setDefaultValue(column.getDefault());
                columnDataSourceDbo.data.setPrimary(DasUtil.isPrimary(column));
                tableDataSourceDbo.data.addColumn(columnDataSourceDbo.data);
            }
        });

        dataSource.getModel().traverser().forEach(dasObject ->
        {
            if (dasObject instanceof DbTable) {
                DbTable table = (DbTable) dasObject;
                tableDataSourceDbo.data = dataSourceDbo.findTable(table.getName());
                DasUtil.getForeignKeys(table).forEach(fk ->
                {
                    if (fk.getRefTable() == null) {
                        return;
                    }

                    TableDataSourceDbo srcTable = dataSourceDbo.findTable(fk.getTable().getName());
                    TableDataSourceDbo targetTable = dataSourceDbo.findTable(fk.getRefTable().getName());

                    ForeignKeyDbo fkDbo = new ForeignKeyDbo();
                    fkDbo.setName(fk.getName());
                    fkDbo.setSourceTable(srcTable);
                    fkDbo.setSourceForeignKeyName(fkDbo.getName());//TODO find if it already exists
                    fkDbo.setTargetTable(targetTable);
                    fkDbo.setTargetForeignKeyName(fkDbo.getName());//TODO find if it already exists

                    fk.getColumnsRef().names().forEach(s -> fkDbo.addSourceColumn(srcTable.findColumn(s)));
                    fk.getRefColumns().names().forEach(s -> fkDbo.addTargetColumn(targetTable.findColumn(s)));

                    srcTable.addSourceForeignKey(fkDbo);
                    targetTable.addTargetForeignKey(fkDbo);

                });
            }
            else if (dasObject instanceof DbColumn) {
                DbColumn column = (DbColumn) dasObject;
                columnDataSourceDbo.data = tableDataSourceDbo.data.findColumn(column.getName());
            }
            else if (dasObject instanceof DbIndex) {
                DbIndex index = (DbIndex) dasObject;
            }
            else if (dasObject instanceof DbConstraint) {
                DbConstraint constraint = (DbConstraint) dasObject;
            }
            else if (dasObject instanceof DbRoutine) {
                DbRoutine sp = (DbRoutine) dasObject;
                StoredProcedureDbo spDbo = new StoredProcedureDbo();
                spDbo.setSelected(false);
                spDbo.setName(sp.getName());
                spDbo.setFullName(sp.getText());
                dataSourceDbo.addStoredProcedure(spDbo);
                DataHolder<Integer> position = new DataHolder<>();
                position.data = 0;
                TableDataSourceDbo returnType = new TableDataSourceDbo();
                returnType.setName(spDbo.getName());
                sp.getArguments().forEach(o ->
                {
                    SqlTypeDbo sqlTypeDbo = new SqlTypeDbo();
                    sqlTypeDbo.setName(o.getDataType().typeName);
                    sqlTypeDbo.guessArrayAndSetOf();
                    sqlTypeDbo.setVagueArg(o.getDataType().vagueArg);

                    DasArgument.Direction dir = o.getArgumentDirection();

                    if (dir.isReturnOrResult()) {
                        spDbo.setType(sqlTypeDbo);
                    }
                    if (dir.isIn()) {
                        StoredProcedureArgDbo argDbo = new StoredProcedureArgDbo();
                        argDbo.setName(o.getName().isEmpty() ? "arg" + position.data : o.getName());
                        argDbo.setType(sqlTypeDbo);
                        spDbo.addArgument(argDbo);
                        ++position.data;
                    }
                    if (dir.isOut()) {
                        ColumnDataSourceDbo field = new ColumnDataSourceDbo();
                        field.setSelected(true);
                        field.setPrimary(false);
                        field.setDefaultValue(null);
                        field.setNotNull(false);
                        field.setName(o.getName());
                        field.setType(sqlTypeDbo);
                        returnType.addColumn(field);
                    }
                });
                if (returnType.getColumns().size() > 1) {
                    boolean isSetOf = sp.getReturnArgument().getDataType().typeName.startsWith("setof ");
                    SqlTypeDbo sqlTypeDbo = new SqlTypeDbo();
                    sqlTypeDbo.setName((isSetOf ? "setof " : "" ) + returnType.getName());
                    sqlTypeDbo.guessArrayAndSetOf();
                    sqlTypeDbo.setVagueArg(null);
                    spDbo.setType(sqlTypeDbo);
                    dataSourceDbo.addType(returnType);
                }
            }
        });
        return dataSourceDbo;
    }
}