added sp support

TODO

@@ -1,3 +1,5 @@
 proc stock
 many to many
-table inheritance
+table inheritance
+generate custom types (sp return)
+UI for sp selection

src/com/rthoni/intellij/codefromds/DataAccess/DataSourcesDataAccess.java

@@ -1,5 +1,6 @@
 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;
@@ -64,7 +65,6 @@ public class DataSourcesDataAccess {
             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) {
@@ -77,9 +77,9 @@ public class DataSourcesDataAccess {
                     ForeignKeyDbo fkDbo = new ForeignKeyDbo();
                     fkDbo.setName(fk.getName());
                     fkDbo.setSourceTable(srcTable);
-                    fkDbo.setSourceForeignKeyName("fk_" + targetTable.getName());//TODO find if it already exists
+                    fkDbo.setSourceForeignKeyName(fkDbo.getName());//TODO find if it already exists
                     fkDbo.setTargetTable(targetTable);
-                    fkDbo.setTargetForeignKeyName(srcTable.getName() + "_fk");//TODO find if it already exists
+                    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)));
@@ -87,29 +87,25 @@ public class DataSourcesDataAccess {
                     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());
+                spDbo.setFullName(sp.getText());
+                dataSourceDbo.addStoredProcedure(spDbo);
                 DataHolder<Integer> position = new DataHolder<>();
                 position.data = 0;
                 sp.getArguments().forEach(o ->
@@ -118,24 +114,26 @@ public class DataSourcesDataAccess {
                     sqlTypeDbo.setType(o.getDataType().typeName);
                     sqlTypeDbo.setVagueArg(o.getDataType().vagueArg);
 
-                    if (o.getArgumentDirection().isReturnOrResult()) {
+                    DasArgument.Direction dir = o.getArgumentDirection();
+
+
+
+                    if (dir.isReturnOrResult()) {
                         spDbo.setSqlType(sqlTypeDbo);
-//                        System.out.println("SP return: " + sqlTypeDbo.getType());
                     }
-                    else {
+                    if (dir.isIn()) {
                         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());
                     }
                 });
+//                DasArgument arg = sp.getReturnArgument();
+//                if (arg != null)
+//                    arg.toString();
             }
-//            else {
-//                System.out.println(dasObject.getClass().getName());
-//            }
         });
         return dataSourceDbo;
     }

src/com/rthoni/intellij/codefromds/business/Generator.java

@@ -2,10 +2,7 @@ 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.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.options.*;
 import com.rthoni.intellij.codefromds.dbo.template.*;
 import groovy.json.StringEscapeUtils;
 import org.json.JSONArray;
@@ -70,13 +67,18 @@ public abstract class Generator {
         return StringEscapeUtils.escapeJava(str);
     }
 
-    public static void convertSqlType(SqlTypeDbo type, TypesCastOptions options)
+    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();
 
@@ -90,11 +92,20 @@ public abstract class Generator {
             }
         }
         if (typeName == null) {
-            typeName = types.get("*").get("*");
+            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));
     }
@@ -130,16 +141,20 @@ public abstract class Generator {
                 ColumnSelection columnSelection = tableSelection.getColumns().stream()
                         .filter(c -> c.getColumn().getName().equals(column.getName()))
                         .findFirst().orElse(null);
-                convertSqlType(column.getSqlType(), types);
+                convertSqlType(column.getSqlType(), types, dataSourceDbo);
                 column.setSelected(columnSelection.isSelected());
             }
         }
 
         for (StoredProcedureDbo spDbo : dataSourceDbo.getStoredProcedures())
         {
-            convertSqlType(spDbo.getSqlType(), types);
+            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);
+                convertSqlType(arg.getSqlType(), types, dataSourceDbo);
             }
         }
 
@@ -186,6 +201,7 @@ public abstract class Generator {
         dbo.setReservedWords(reservedWords);
 
         dbo.setArrayTemplate(obj.getString("arrayTemplate"));
+        dbo.setSetOfTemplate(obj.getString("setOfTemplate"));
 
         return dbo;
     }
@@ -205,7 +221,12 @@ public abstract class Generator {
 
             @Override
             public Object execute(FunctionRequest functionRequest) {
-                String arg = (String) functionRequest.get(0);
+                Object obj = functionRequest.get(0);
+                if (obj == null)
+                {
+                    return null;
+                }
+                String arg = (String) obj;
                 String v = escapeReservedWords(arg, options);
                 return v;
             }

src/com/rthoni/intellij/codefromds/business/Helper.java

@@ -1,12 +1,15 @@
 package com.rthoni.intellij.codefromds.business;
 
+import com.intellij.database.model.DasRoutine;
 import com.intellij.database.psi.DbDataSource;
 import com.intellij.database.psi.DbPsiFacade;
+import com.intellij.database.psi.DbRoutine;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.project.ProjectManager;
 import org.json.JSONArray;
 import org.json.JSONObject;
 
+import javax.sql.DataSource;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
@@ -15,6 +18,7 @@ import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
+import java.util.Vector;
 import java.util.stream.Collectors;
 
 /**
@@ -22,6 +26,18 @@ import java.util.stream.Collectors;
  */
 public abstract class Helper {
 
+    public static List<DbRoutine> getRoutines(DbDataSource dataSource)
+    {
+        List<DbRoutine> routines = new Vector<>();
+        dataSource.getModel().traverser().forEach(dasObject ->
+        {
+            if (dasObject instanceof DbRoutine) {
+                routines.add((DbRoutine) dasObject);
+            }
+        });
+        return routines;
+    }
+
     public static List<DbDataSource> getDataSources()
     {
         ProjectManager pm = ProjectManager.getInstance();
@@ -35,25 +51,25 @@ public abstract class Helper {
         return getDataSources().stream().filter(d -> d.getName().equals(name)).findFirst().orElse(null);
     }
 
-    public static JSONObject findTableInJson(JSONArray tables, String name)
+    public static JSONObject findInJson(JSONArray tables, String key, String name)
     {
         for (int i = 0; i < tables.length(); ++i) {
-            if (tables.getJSONObject(i).getString("table").equals(name)) {
+            if (tables.getJSONObject(i).getString(key).equals(name)) {
                 return tables.getJSONObject(i);
             }
         }
         return null;
     }
 
-    public static JSONObject findColumnInJson(JSONArray tables, String name)
-    {
-        for (int i = 0; i < tables.length(); ++i) {
-            if (tables.getJSONObject(i).getString("column").equals(name)) {
-                return tables.getJSONObject(i);
-            }
-        }
-        return null;
-    }
+//    public static JSONObject findColumnInJson(JSONArray tables, String name)
+//    {
+//        for (int i = 0; i < tables.length(); ++i) {
+//            if (tables.getJSONObject(i).getString("column").equals(name)) {
+//                return tables.getJSONObject(i);
+//            }
+//        }
+//        return null;
+//    }
 
     public static String readFile(String path) throws IOException {
         String data = Files.readAllLines(Paths.get(path), StandardCharsets.UTF_8)

src/com/rthoni/intellij/codefromds/dbo/options/DataSourceSelection.java

@@ -20,9 +20,12 @@ public class DataSourceSelection {
 
     private List<TableSelection> _tables;
 
+    private List<StoredProcedureSelection> _storedProcedures;
+
     public DataSourceSelection(DbDataSource source) {
         _source = source;
         _tables = DasUtil.getTables(source).toList().stream().map(TableSelection::new).collect(Collectors.toList());
+        _storedProcedures = Helper.getRoutines(source).stream().map(StoredProcedureSelection::new).collect(Collectors.toList());
     }
 
     public HashMap<String, Object> toMap()
@@ -30,12 +33,19 @@ public class DataSourceSelection {
         HashMap<String, Object> map = new HashMap<>();
 
         map.put("source", _source == null ? null : _source.getName());
+
         List<Object> tables = new Vector<>();
         for (TableSelection table : _tables) {
             tables.add(table.toMap());
         }
         map.put("tables", tables);
 
+        List<Object> sps = new Vector<>();
+        for (StoredProcedureSelection sp : _storedProcedures) {
+            sps.add(sp.toMap());
+        }
+        map.put("storedProcedures", sps);
+
         return map;
     }
 
@@ -43,7 +53,7 @@ public class DataSourceSelection {
     {
         JSONArray array = json.getJSONArray("tables");
         for (TableSelection table : _tables) {
-            JSONObject obj = Helper.findTableInJson(array, table.getTable().getName());
+            JSONObject obj = Helper.findInJson(array, "table", table.getTable().getName());
             if (obj != null) {
                 table.fromJson(obj);
             } else {
@@ -52,6 +62,14 @@ public class DataSourceSelection {
                 }
             }
         }
+
+        array = json.getJSONArray("storedProcedures");
+        for (StoredProcedureSelection sp : _storedProcedures) {
+            JSONObject obj = Helper.findInJson(array, "name", sp.getStoredProcedure().getText());
+            if (obj != null) {
+                sp.fromJson(obj);
+            }
+        }
     }
 
     public DbDataSource getSource() {
@@ -69,4 +87,12 @@ public class DataSourceSelection {
     public void setTables(List<TableSelection> tables) {
         _tables = tables;
     }
+
+    public List<StoredProcedureSelection> getStoredProcedures() {
+        return _storedProcedures;
+    }
+
+    public void setStoredProcedures(List<StoredProcedureSelection> storedProcedures) {
+        _storedProcedures = storedProcedures;
+    }
 }

src/com/rthoni/intellij/codefromds/dbo/options/StoredProcedureSelection.java

@@ -0,0 +1,52 @@
+package com.rthoni.intellij.codefromds.dbo.options;
+
+import com.intellij.database.psi.DbRoutine;
+import org.json.JSONObject;
+
+import java.util.HashMap;
+
+/**
+ * Created by robin on 4/14/17.
+ */
+public class StoredProcedureSelection {
+
+    private DbRoutine _storedProcedure;
+
+    private boolean _isSelected;
+
+    public StoredProcedureSelection(DbRoutine storedProcedure) {
+        _storedProcedure = storedProcedure;
+        _isSelected = true;
+    }
+
+    public DbRoutine getStoredProcedure() {
+        return _storedProcedure;
+    }
+
+    public void setStoredProcedure(DbRoutine storedProcedure) {
+        _storedProcedure = storedProcedure;
+    }
+
+    public boolean isSelected() {
+        return _isSelected;
+    }
+
+    public void setSelected(boolean selected) {
+        _isSelected = selected;
+    }
+
+    public HashMap<String, Object> toMap()
+    {
+        HashMap<String, Object> map = new HashMap<>();
+
+        map.put("name", _storedProcedure == null ? null : _storedProcedure.getText());
+        map.put("selected", _isSelected);
+
+        return map;
+    }
+
+    public void fromJson(JSONObject json)
+    {
+        _isSelected = json.getBoolean("selected");
+    }
+}

src/com/rthoni/intellij/codefromds/dbo/options/TableSelection.java

@@ -44,7 +44,7 @@ public class TableSelection {
     {
         JSONArray array = json.getJSONArray("columns");
         for (ColumnSelection column : _columns) {
-            JSONObject obj = Helper.findColumnInJson(array, column.getColumn().getName());
+            JSONObject obj = Helper.findInJson(array, "column", column.getColumn().getName());
             if (obj != null) {
                 column.fromJson(obj);
             } else {

src/com/rthoni/intellij/codefromds/dbo/options/TypesCastOptions.java

@@ -10,6 +10,8 @@ public class TypesCastOptions {
 
     private String _arrayTemplate;
 
+    private String _setOfTemplate;
+
     private HashMap<String, HashMap<String, String>> _types;
 
     private List<String> _nonNullableTypes;
@@ -47,4 +49,12 @@ public class TypesCastOptions {
     public void setReservedWords(List<String> reservedWords) {
         _reservedWords = reservedWords;
     }
+
+    public String getSetOfTemplate() {
+        return _setOfTemplate;
+    }
+
+    public void setSetOfTemplate(String setOfTemplate) {
+        _setOfTemplate = setOfTemplate;
+    }
 }

src/com/rthoni/intellij/codefromds/dbo/template/StoredProcedureDbo.java

@@ -16,6 +16,8 @@ public class StoredProcedureDbo {
 
     private boolean _isSelected;
 
+    private String _fullName;
+
     public String getName() {
         return _name;
     }
@@ -51,4 +53,12 @@ public class StoredProcedureDbo {
     public void setSelected(boolean selected) {
         _isSelected = selected;
     }
+
+    public String getFullName() {
+        return _fullName;
+    }
+
+    public void setFullName(String fullName) {
+        _fullName = fullName;
+    }
 }

src/com/rthoni/intellij/codefromds/ui/dialogs/GenerateDialog.java

@@ -11,10 +11,7 @@ import com.intellij.ui.ColoredListCellRenderer;
 import com.intellij.ui.JBColor;
 import com.intellij.ui.components.JBList;
 import com.rthoni.intellij.codefromds.business.Helper;
-import com.rthoni.intellij.codefromds.dbo.options.ColumnSelection;
-import com.rthoni.intellij.codefromds.dbo.options.DataSourceSelection;
-import com.rthoni.intellij.codefromds.dbo.options.GenerateOptions;
-import com.rthoni.intellij.codefromds.dbo.options.TableSelection;
+import com.rthoni.intellij.codefromds.dbo.options.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -50,6 +47,7 @@ public class GenerateDialog extends DialogWrapper {
     private JLabel _lblModelsTemplatePath;
     private TextFieldWithBrowseButton _textCastFile;
     private JLabel _lblCastFile;
+    private JBList _listStoredProcedure;
 
     private GenerateOptions _options;
 
@@ -165,6 +163,12 @@ public class GenerateDialog extends DialogWrapper {
             }
         });
 
+        _listStoredProcedure.addListSelectionListener(e -> {
+            if (!e.getValueIsAdjusting()) {
+                updateSpSelection();
+            }
+        });
+
         setupTextField(_textModels, null, true, "Models");
         setupTextField(_textDataSource, null, false, "Data Source");
         setupTextField(_textDataSourceTemplate, null, false, "Data Source Template");
@@ -213,6 +217,20 @@ public class GenerateDialog extends DialogWrapper {
             }
         });
 
+        _listStoredProcedure.setCellRenderer(new ColoredListCellRenderer() {
+            @Override
+            protected void customizeCellRenderer(@NotNull JList jList, Object o, int i, boolean b, boolean b1) {
+                StoredProcedureSelection spSelection = _options.getSelection().getStoredProcedures().get(i);
+                if (spSelection.isSelected()) {
+                    setBackground(JBColor.GREEN);
+                }
+                else {
+                    setBackground(JBColor.RED);
+                }
+                append(spSelection.getStoredProcedure().getText());
+            }
+        });
+
         if (dataSources.size() > 0) {
             _listDatasources.setSelectedIndices(new int[]{0});
         }
@@ -264,9 +282,12 @@ public class GenerateDialog extends DialogWrapper {
             else {
                 showTable(null);
             }
+            List<String> sps = source.getStoredProcedures().stream().map(t -> t.getStoredProcedure().getText()).collect(Collectors.toList());
+            _listStoredProcedure.setListData(sps.toArray(new String[sps.size()]));
         }
         else {
             _listTables.setListData(new String[]{});
+            _listStoredProcedure.setListData(new String[]{});
             showTable(null);
         }
     }
@@ -283,7 +304,7 @@ public class GenerateDialog extends DialogWrapper {
         return indices.stream().mapToInt(Integer::intValue).toArray();
     }
 
-    private void updateSelection(final TableSelection table)
+    private void updateColumnSelection(final TableSelection table)
     {
         List<ColumnSelection> columns = table.getColumns();
         for (int i = 0; i < columns.size(); ++i) {
@@ -291,6 +312,14 @@ public class GenerateDialog extends DialogWrapper {
         }
     }
 
+    private void updateSpSelection()
+    {
+        List<StoredProcedureSelection> sps = _options.getSelection().getStoredProcedures();
+        for (int i = 0; i < sps.size(); ++i) {
+            sps.get(i).setSelected(_listStoredProcedure.isSelectedIndex(i));
+        }
+    }
+
     private void showTable(TableSelection table)
     {
         for (ListSelectionListener e : _listColumns.getListSelectionListeners()) {
@@ -303,7 +332,7 @@ public class GenerateDialog extends DialogWrapper {
             _listColumns.addListSelectionListener(e -> {
                 if (!e.getValueIsAdjusting()) {
                     _listTables.updateUI();
-                    updateSelection(table);
+                    updateColumnSelection(table);
                 }
             });
         }

src/com/rthoni/intellij/codefromds/ui/forms/GenerateForm.form

@@ -1,22 +1,22 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <form xmlns="http://www.intellij.com/uidesigner/form/" version="1" bind-to-class="com.rthoni.intellij.codefromds.ui.dialogs.GenerateDialog">
-  <grid id="27dc6" binding="_panel" layout-manager="GridLayoutManager" row-count="4" column-count="2" same-size-horizontally="false" same-size-vertically="false" hgap="-1" vgap="-1">
+  <grid id="27dc6" binding="_panel" layout-manager="GridLayoutManager" row-count="5" column-count="2" same-size-horizontally="false" same-size-vertically="false" hgap="-1" vgap="-1">
     <margin top="0" left="0" bottom="0" right="0"/>
     <constraints>
-      <xy x="20" y="20" width="934" height="526"/>
+      <xy x="20" y="20" width="934" height="652"/>
     </constraints>
     <properties/>
     <border type="none"/>
     <children>
       <vspacer id="cffbd">
         <constraints>
-          <grid row="3" column="0" row-span="1" col-span="2" vsize-policy="6" hsize-policy="1" anchor="0" fill="2" indent="0" use-parent-layout="false"/>
+          <grid row="4" column="0" row-span="1" col-span="2" vsize-policy="6" hsize-policy="1" anchor="0" fill="2" indent="0" use-parent-layout="false"/>
         </constraints>
       </vspacer>
       <grid id="6946b" layout-manager="GridLayoutManager" row-count="12" column-count="2" same-size-horizontally="false" same-size-vertically="false" hgap="-1" vgap="-1">
         <margin top="0" left="0" bottom="0" right="0"/>
         <constraints>
-          <grid row="2" column="0" row-span="1" col-span="2" vsize-policy="3" hsize-policy="3" anchor="0" fill="3" indent="0" use-parent-layout="false"/>
+          <grid row="3" column="0" row-span="1" col-span="2" vsize-policy="3" hsize-policy="3" anchor="0" fill="3" indent="0" use-parent-layout="false"/>
         </constraints>
         <properties/>
         <border type="none"/>
@@ -237,6 +237,29 @@
           </scrollpane>
         </children>
       </grid>
+      <grid id="484d8" layout-manager="GridLayoutManager" row-count="1" column-count="1" same-size-horizontally="false" same-size-vertically="false" hgap="-1" vgap="-1">
+        <margin top="0" left="0" bottom="0" right="0"/>
+        <constraints>
+          <grid row="2" column="0" row-span="1" col-span="2" vsize-policy="3" hsize-policy="3" anchor="0" fill="3" indent="0" use-parent-layout="false"/>
+        </constraints>
+        <properties/>
+        <border type="none" title="Select Stored Procedures To Map"/>
+        <children>
+          <scrollpane id="705ec">
+            <constraints>
+              <grid row="0" column="0" row-span="1" col-span="1" vsize-policy="7" hsize-policy="7" anchor="0" fill="3" indent="0" use-parent-layout="false"/>
+            </constraints>
+            <properties/>
+            <border type="none"/>
+            <children>
+              <component id="77e75" class="com.intellij.ui.components.JBList" binding="_listStoredProcedure">
+                <constraints/>
+                <properties/>
+              </component>
+            </children>
+          </scrollpane>
+        </children>
+      </grid>
     </children>
   </grid>
 </form>