Merge pull request #20 from Shonminh/shonminh/bugfix/issue-16

Shonminh/bugfix/issue 16
Shonminh · May 24, 2020 · c12beb1 · c12beb1
2 parents 908ba56 + eb258ab
commit c12beb1
Show file tree

Hide file tree

Showing 8 changed files with 156 additions and 11 deletions.
diff --git a/helper/build.gradle b/helper/build.gradle
@@ -4,7 +4,7 @@ plugins {
 }
 
 group 'org.shonminh.helper'
-version '1.4.1'
+version '1.4.2'
 
 sourceCompatibility = 1.8
 

diff --git a/helper/src/main/java/org/shonminh/helper/sql/Header.java b/helper/src/main/java/org/shonminh/helper/sql/Header.java
@@ -0,0 +1,50 @@
+package org.shonminh.helper.sql;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+public class Header {
+    private List<String> dependencyPackageList;
+    private Set<String> existDependencySet;
+
+    public Header() {
+        this.dependencyPackageList = new ArrayList<>();
+        this.existDependencySet = new HashSet<>();
+    }
+
+    public void appendDependencyPackage(String dependencyPackageName) {
+        if (this.existDependencySet.contains(dependencyPackageName)) {
+            return;
+        }
+        this.existDependencySet.add(dependencyPackageName);
+
+        List<String> arrayList = this.getDependencyPackageList();
+        arrayList.add(dependencyPackageName);
+        arrayList.sort(String::compareTo);
+        this.setDependencyPackageList(arrayList);
+
+    }
+
+    public void appendDependencyPackageList(List<String> dependencyPackageList) {
+        List<String> arrayList = new ArrayList<>(this.getDependencyPackageList());
+        for (String packageName : dependencyPackageList) {
+            if (this.existDependencySet.contains(packageName)) {
+                continue;
+            }
+            arrayList.add(packageName);
+        }
+        arrayList.sort(String::compareTo);
+        this.setDependencyPackageList(arrayList);
+    }
+
+
+    public List<String> getDependencyPackageList() {
+        return dependencyPackageList;
+    }
+
+    public void setDependencyPackageList(List<String> dependencyPackageList) {
+        this.dependencyPackageList = dependencyPackageList;
+    }
+}
diff --git a/helper/src/main/java/org/shonminh/helper/sql/Model.java b/helper/src/main/java/org/shonminh/helper/sql/Model.java
@@ -8,6 +8,7 @@
 import com.alibaba.druid.sql.dialect.mysql.ast.MySqlPrimaryKey;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement;
 import org.shonminh.helper.util.GoTypeUtil;
+import org.shonminh.helper.util.HeaderUtil;
 import org.shonminh.helper.util.StringUtil;
 
 import java.util.ArrayList;
@@ -18,6 +19,7 @@ public class Model {
     private String modelName;
     private String primaryKey;
     private List<Column> columns;
+    private Header header;
 
     public String getPrimaryKey() {
         return primaryKey;
@@ -29,6 +31,7 @@ public void setPrimaryKey(String primaryKey) {
 
     public Model(String modelName) {
         this.modelName = modelName;
+        this.header = new Header();
     }
 
     public String getModelName() {
@@ -47,6 +50,14 @@ public void setColumns(List<Column> columns) {
         this.columns = columns;
     }
 
+    public Header getHeader() {
+        return header;
+    }
+
+    public void setHeader(Header header) {
+        this.header = header;
+    }
+
     public void appendColumn(Column column) {
         if (null == this.columns) {
             this.columns = new ArrayList<>();
@@ -158,6 +169,15 @@ public void setAllColumnProperties(MySqlCreateTableStatement createTable) {
             if (columnDefinition.isAutoIncrement()) {
                 column.setAutoIncrement(true);
             }
+
+            // set model header
+            String dependencyPackageName = HeaderUtil.getDependencyPackageName(column.getType());
+            if (dependencyPackageName != null) {
+                Header h = this.getHeader();
+                h.appendDependencyPackage(dependencyPackageName);
+                this.setHeader(h);
+            }
+
             this.appendColumn(column);
         }
     }
@@ -182,6 +202,7 @@ public String toString() {
                 "modelName='" + modelName + '\'' +
                 ", primaryKey='" + primaryKey + '\'' +
                 ", columns=" + columns +
+                ", header=" + header +
                 '}';
     }
 }
diff --git a/helper/src/main/java/org/shonminh/helper/sql/SqlParser.java b/helper/src/main/java/org/shonminh/helper/sql/SqlParser.java
@@ -8,6 +8,7 @@
 import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement;
 import com.alibaba.druid.util.JdbcConstants;
+import org.shonminh.helper.util.HeaderUtil;
 import org.shonminh.helper.util.StringUtil;
 
 import java.util.ArrayList;
@@ -48,26 +49,24 @@ private void generateStatements(String str) {
 
     private String parseStatements() {
         StringBuilder resultStringBuilder = new StringBuilder();
-        boolean isFirst = true;
+        Header h = new Header();
         for (MySqlCreateTableStatement createTable : statements) {
             try {
                 String modelName = getModelName(createTable);
                 Model model = new Model(modelName);
                 // set all column properties
                 model.setAllColumnProperties(createTable);
                 this.setFileName(model.getModelName().replaceFirst("_tab$", "") + ".go");
-
-                // if is first append string then add golang package name
-                if (isFirst) {
-                    resultStringBuilder.append("package model\n\n");
-                    isFirst = false;
-                }
+                List<String> packageList = model.getHeader().getDependencyPackageList();
+                h.appendDependencyPackageList(packageList);
                 resultStringBuilder.append(model.generateGoStruct());
             } catch (Exception e) {
                 this.setErrMsg(e.getCause().getMessage());
             }
         }
-        return resultStringBuilder.toString();
+        List<String> totalPackageList = h.getDependencyPackageList();
+        String headerCodes = HeaderUtil.getHeaderCodes(totalPackageList);
+        return headerCodes + resultStringBuilder.toString();
     }
 
 

diff --git a/helper/src/main/java/org/shonminh/helper/util/HeaderUtil.java b/helper/src/main/java/org/shonminh/helper/util/HeaderUtil.java
@@ -0,0 +1,34 @@
+package org.shonminh.helper.util;
+
+import java.util.List;
+
+public class HeaderUtil {
+    public static String getDependencyPackageName(String goType) {
+        if ("timestamp".equals(goType)) {
+            return "time";
+        }
+        return null;
+    }
+
+    public static String getHeaderCodes(List<String> dependencyPackageList) {
+        StringBuilder sb = new StringBuilder();
+        sb.append("package model\n\n");
+        if (dependencyPackageList == null || dependencyPackageList.size() == 0) {
+            return sb.toString();
+        }
+        sb.append("import (\n");
+        for (String dependencyName :
+                dependencyPackageList) {
+            sb.append("\t");
+            sb.append("\"");
+            sb.append(dependencyName);
+            sb.append("\"");
+            sb.append("\n");
+        }
+        sb.append(")\n\n");
+        return sb.toString();
+    }
+}
+
+
+
diff --git a/helper/src/main/resources/META-INF/plugin.xml b/helper/src/main/resources/META-INF/plugin.xml
@@ -1,7 +1,7 @@
 <idea-plugin>
     <id>org.shonminh.helper.GenerateModelAction</id>
     <name>Go Orm Code Helper</name>
-    <version>1.4.1</version>
+    <version>1.4.2</version>
     <vendor email="shonminh@gmail.com" url="https://github.com/Shonminh">Shonminh Yang</vendor>
 
     <description><![CDATA[
@@ -26,7 +26,7 @@
     </description>
     <change-notes>
         <![CDATA[
-        <b>1.4.1</b>
+        <b>1.4.2</b>
         <ul>
           <li>optimize parse sql statements.</li>
         </ul>

diff --git a/helper/src/test/java/org/shonminh/helper/sql/HeaderTest.java b/helper/src/test/java/org/shonminh/helper/sql/HeaderTest.java
@@ -0,0 +1,37 @@
+package org.shonminh.helper.sql;
+
+import org.junit.Test;
+import org.shonminh.helper.util.HeaderUtil;
+
+import java.util.List;
+
+import static org.junit.Assert.*;
+
+public class HeaderTest {
+
+    @Test
+    public void appendDependencyPackage() {
+        Header header = new Header();
+        header.appendDependencyPackage("demo.b");
+        header.appendDependencyPackage("demo.a");
+        List<String> dependencyPackageList = header.getDependencyPackageList();
+        assertEquals(2, dependencyPackageList.size());
+        assertEquals("demo.a", dependencyPackageList.get(0));
+        assertEquals("demo.b", dependencyPackageList.get(1));
+    }
+
+    @Test
+    public void getHeaderCodes() {
+        Header header = new Header();
+        header.appendDependencyPackage("demo.a");
+        header.appendDependencyPackage("demo.b");
+        String headerCodes = HeaderUtil.getHeaderCodes(header.getDependencyPackageList());
+        assertEquals("package model\n" +
+                "\n" +
+                "import (\n" +
+                "\t\"demo.a\"\n" +
+                "\t\"demo.b\"\n" +
+                ")\n" +
+                "\n", headerCodes);
+    }
+}
diff --git a/helper/src/test/java/org/shonminh/helper/sql/SqlParserTest.java b/helper/src/test/java/org/shonminh/helper/sql/SqlParserTest.java
@@ -59,6 +59,10 @@ public void TestIssue16_Can_not_work_with_goland_windows_version() {
         SqlParser sqlParser = new SqlParser();
         String actual = sqlParser.Execute(sql);
         String expect = "package model\n" +
+                "\n" +
+                "import (\n" +
+                "\t\"time\"\n" +
+                ")\n" +
                 "\n" +
                 "type GroupJob struct {\n" +
                 "\tGId         uint32    `gorm:\"type:INT(10) UNSIGNED;PRIMARY_KEY;AUTO_INCREMENT;NOT NULL\"`\n" +