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

package com.rthoni.intellij.codefromds.DataAccess;

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.setType(column.getDataType().typeName);
                type.setVagueArg(column.getDataType().vagueArg);
                columnDataSourceDbo.data.setSqlType(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());
//                System.out.println("Table: " + tableDataSourceDbo.data.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("fk_" + targetTable.getName());//TODO find if it already exists
                    fkDbo.setTargetTable(targetTable);
                    fkDbo.setTargetForeignKeyName(srcTable.getName() + "_fk");//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);

//                    System.out.println("FK: " + fkDbo.getName());
                });
            }
            else if (dasObject instanceof DbColumn) {
                DbColumn column = (DbColumn) dasObject;
                columnDataSourceDbo.data = tableDataSourceDbo.data.findColumn(column.getName());
//                System.out.println("Column: " + columnDataSourceDbo.data.getName());
            }
            else if (dasObject instanceof DbIndex) {
                DbIndex index = (DbIndex) dasObject;
//                System.out.println("Index: " + index.getName());
            }
            else if (dasObject instanceof DbConstraint) {
                DbConstraint constraint = (DbConstraint) dasObject;
//                System.out.println("Constraint: " + constraint.getName());
            }
            else if (dasObject instanceof DbRoutine) {
                DbRoutine sp = (DbRoutine) dasObject;
                StoredProcedureDbo spDbo = new StoredProcedureDbo();
                spDbo.setSelected(true);
                dataSourceDbo.addStoredProcedure(spDbo);
                spDbo.setName(sp.getName());
//                System.out.println("SP: " + spDbo.getName());
                DataHolder<Integer> position = new DataHolder<>();
                position.data = 0;
                sp.getArguments().forEach(o ->
                {
                    SqlTypeDbo sqlTypeDbo = new SqlTypeDbo();
                    sqlTypeDbo.setType(o.getDataType().typeName);
                    sqlTypeDbo.setVagueArg(o.getDataType().vagueArg);

                    if (o.getArgumentDirection().isReturnOrResult()) {
                        spDbo.setSqlType(sqlTypeDbo);
//                        System.out.println("SP return: " + sqlTypeDbo.getType());
                    }
                    else {
                        StoredProcedureArgDbo argDbo = new StoredProcedureArgDbo();
                        argDbo.setName(o.getName().isEmpty() ? "arg" + position.data : o.getName());
                        argDbo.setSqlType(sqlTypeDbo);
                        argDbo.setOut(o.getArgumentDirection().isOut());
                        spDbo.addArgument(argDbo);
                        ++position.data;
//                        System.out.println("SP arg: " + o.getName());
                    }
                });
            }
//            else {
//                System.out.println(dasObject.getClass().getName());
//            }
        });
        return dataSourceDbo;
    }
}