fix: fix #24

Shonminh · Oct 22, 2021 · ccb0ae5 · ccb0ae5
1 parent 9caa5fa
commit ccb0ae5
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 5 deletions.
diff --git a/helper/src/main/java/org/shonminh/helper/sql/Model.java b/helper/src/main/java/org/shonminh/helper/sql/Model.java
@@ -1,10 +1,7 @@
 package org.shonminh.helper.sql;
 
 import com.alibaba.druid.sql.ast.SQLDataTypeImpl;
-import com.alibaba.druid.sql.ast.statement.SQLColumnConstraint;
-import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
-import com.alibaba.druid.sql.ast.statement.SQLNotNullConstraint;
-import com.alibaba.druid.sql.ast.statement.SQLTableElement;
+import com.alibaba.druid.sql.ast.statement.*;
 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;
@@ -152,6 +149,18 @@ public void setAllColumnProperties(MySqlCreateTableStatement createTable) {
             column.setType(columnDefinition.getDataType().toString());
             column.setColumn(columnDefinition.getColumnName());
 
+
+            // if primary key is defined in the current column
+            List<SQLColumnConstraint> constraints = columnDefinition.getConstraints();
+            if (constraints != null && constraints.size() > 0) {
+                for (SQLColumnConstraint constraint : constraints) {
+                    if (constraint instanceof SQLColumnPrimaryKey && this.primaryKey == null) {
+                        this.setPrimaryKey(column.getName());
+                        break;
+                    }
+                }
+            }
+
             // set unsigned
             if (((SQLDataTypeImpl) columnDefinition.getDataType()).isUnsigned()) {
                 column.setUnsigned(true);

diff --git a/helper/src/test/java/org/shonminh/helper/sql/SqlParserTest.java b/helper/src/test/java/org/shonminh/helper/sql/SqlParserTest.java
@@ -41,7 +41,8 @@ public void TestExecute() {
     }
 
     @Test
-    public void TestIssue16_Can_not_work_with_goland_windows_version() {
+    // https://github.com/Shonminh/go-orm-code-helper/issues/16
+    public void TestIssue16_can_not_work_with_goland_windows_version() {
         String sql = "CREATE TABLE group_job (\n" +
                 "g_id int(10) unsigned NOT NULL AUTO_INCREMENT,\n" +
                 "g_type tinyint(4) NOT NULL DEFAULT '1' COMMENT '1:TYPEA; 2:TYPEB',\n" +
@@ -78,4 +79,37 @@ public void TestIssue16_Can_not_work_with_goland_windows_version() {
                 "}\n";
         assertEquals(expect, actual);
     }
+
+    @Test
+    // // https://github.com/Shonminh/go-orm-code-helper/issues/24
+    public void TestIssue24_support_primary_key_format() {
+        String sql = "CREATE TABLE `user`\n" +
+                "(\n" +
+                "    `id`                     int unsigned primary key auto_increment,\n" +
+                "    `name`                   varchar(24) not null unique,\n" +
+                "    `password`               char(60)    not null,\n" +
+                "    `randomSalt`             char(24)    not null,\n" +
+                "    `create_time`            timestamp NULL DEFAULT CURRENT_TIMESTAMP,\n" +
+                "    `update_time`            timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,\n" +
+                "    `api_permission_bitmask` int unsigned not null\n" +
+                ");";
+        SqlParser sqlParser = new SqlParser();
+        String actual = sqlParser.Execute(sql);
+        String expect = "package model\n" +
+                "\n" +
+                "import (\n" +
+                "\t\"time\"\n" +
+                ")\n" +
+                "\n" +
+                "type User struct {\n" +
+                "\tId                   uint32    `gorm:\"column:id;type:INT UNSIGNED;PRIMARY_KEY;AUTO_INCREMENT;\"`\n" +
+                "\tName                 string    `gorm:\"column:name;type:VARCHAR(24);NOT NULL\"`\n" +
+                "\tPassword             string    `gorm:\"column:password;type:CHAR(60);NOT NULL\"`\n" +
+                "\tRandomSalt           string    `gorm:\"column:randomSalt;type:CHAR(24);NOT NULL\"`\n" +
+                "\tCreateTime           time.Time `gorm:\"column:create_time;type:TIMESTAMP;\"`\n" +
+                "\tUpdateTime           time.Time `gorm:\"column:update_time;type:TIMESTAMP;\"`\n" +
+                "\tApiPermissionBitmask uint32    `gorm:\"column:api_permission_bitmask;type:INT UNSIGNED;NOT NULL\"`\n" +
+                "}\n";
+        assertEquals(expect, actual);
+    }
 }