feat: get table info

pkg/sql2code/parser/mongodb.go

@@ -261,10 +261,11 @@ func ConvertToSQLByMgoFields(tableName string, fields []*MgoField) (string, map[
 			protoObjectStrs = append(protoObjectStrs, field.ProtoObjectStr)
 		}
 	}
-	if isHaveID {
+
-		fieldStr = "    id bigint unsigned primary key,\n" + fieldStr
-	}
 	fieldStr = strings.TrimSuffix(fieldStr, ",\n")
+	if isHaveID {
+		fieldStr = "    `id` varchar(24),\n" + fieldStr + ",\n    PRIMARY KEY (id)"
+	}
 
 	if len(objectStrs) > 0 {
 		srcMongoTypeMap[SubStructKey] = strings.Join(objectStrs, "\n") + "\n"

pkg/sql2code/parser/option.go

@@ -30,6 +30,8 @@ type options struct {
 	IsEmbed        bool // is gorm.Model embedded
 	IsWebProto     bool // true: proto file include router path and swagger info, false: normal proto file without router and swagger
 	IsExtendedAPI  bool // true: extended api (9 api), false: basic api (5 api)
+
+	IsCustomTemplate bool // true: custom extend template, false: sponge template
 }
 
 var defaultOptions = options{
@@ -149,6 +151,13 @@ func WithExtendedAPI() Option {
 	}
 }
 
+// WithCustomTemplate set custom template
+func WithCustomTemplate() Option {
+	return func(o *options) {
+		o.IsCustomTemplate = true
+	}
+}
+
 func parseOption(options []Option) options {
 	o := defaultOptions
 	for _, f := range options {

pkg/sql2code/parser/parser.go

@@ -37,6 +37,8 @@ const (
 	CodeTypeService = "service"
 	// CodeTypeCrudInfo crud info json data
 	CodeTypeCrudInfo = "crud_info"
+	// CodeTypeTableInfo table info json data
+	CodeTypeTableInfo = "table_info"
 
 	// DBDriverMysql mysql driver
 	DBDriverMysql = "mysql"
@@ -87,6 +89,7 @@ func ParseSQL(sql string, options ...Option) (map[string]string, error) {
 	importPath := make(map[string]struct{})
 	tableNames := make([]string, 0, len(stmts))
 	primaryKeysCodes := make([]string, 0, len(stmts))
+	tableInfoCodes := make([]string, 0, len(stmts))
 	for _, stmt := range stmts {
 		if ct, ok := stmt.(*ast.CreateTableStmt); ok {
 			code, err2 := makeCode(ct, opt)
@@ -101,6 +104,7 @@ func ParseSQL(sql string, options ...Option) (map[string]string, error) {
 			modelJSONCodes = append(modelJSONCodes, code.modelJSON)
 			tableNames = append(tableNames, toCamel(ct.Table.Name.String()))
 			primaryKeysCodes = append(primaryKeysCodes, code.crudInfo)
+			tableInfoCodes = append(tableInfoCodes, string(code.tableInfo))
 			for _, s := range code.importPaths {
 				importPath[s] = struct{}{}
 			}
@@ -124,24 +128,27 @@ func ParseSQL(sql string, options ...Option) (map[string]string, error) {
 	}
 
 	var codesMap = map[string]string{
-		CodeTypeModel:    modelCode,
+		CodeTypeModel:     modelCode,
-		CodeTypeJSON:     strings.Join(modelJSONCodes, "\n\n"),
+		CodeTypeJSON:      strings.Join(modelJSONCodes, "\n\n"),
-		CodeTypeDAO:      strings.Join(updateFieldsCodes, "\n\n"),
+		CodeTypeDAO:       strings.Join(updateFieldsCodes, "\n\n"),
-		CodeTypeHandler:  strings.Join(handlerStructCodes, "\n\n"),
+		CodeTypeHandler:   strings.Join(handlerStructCodes, "\n\n"),
-		CodeTypeProto:    strings.Join(protoFileCodes, "\n\n"),
+		CodeTypeProto:     strings.Join(protoFileCodes, "\n\n"),
-		CodeTypeService:  strings.Join(serviceStructCodes, "\n\n"),
+		CodeTypeService:   strings.Join(serviceStructCodes, "\n\n"),
-		TableName:        strings.Join(tableNames, ", "),
+		TableName:         strings.Join(tableNames, ", "),
-		CodeTypeCrudInfo: strings.Join(primaryKeysCodes, "||||"),
+		CodeTypeCrudInfo:  strings.Join(primaryKeysCodes, "||||"),
+		CodeTypeTableInfo: strings.Join(tableInfoCodes, "||||"),
 	}
 
 	return codesMap, nil
 }
 
 type tmplData struct {
-	TableName       string
+	TableNamePrefix string
-	TName           string
+
+	RawTableName    string // raw table name, example: foo_bar
+	TableName       string // table name in camel case, example: FooBar
+	TName           string // table name first letter in lower case, example: fooBar
 	NameFunc        bool
-	RawTableName    string
 	Fields          []tmplField
 	Comment         string
 	SubStructs      string // sub structs for model
@@ -351,25 +358,31 @@ type codeText struct {
 	protoFile     string
 	serviceStruct string
 	crudInfo      string
+	tableInfo     []byte
 }
 
 // nolint
 func makeCode(stmt *ast.CreateTableStmt, opt options) (*codeText, error) {
 	importPath := make([]string, 0, 1)
 	data := tmplData{
-		TableName:    stmt.Table.Name.String(),
+		TableNamePrefix: opt.TablePrefix,
-		RawTableName: stmt.Table.Name.String(),
+		RawTableName:    stmt.Table.Name.String(),
-		//Fields:       make([]tmplField, 0, 1),
+		DBDriver:        opt.DBDriver,
-		DBDriver: opt.DBDriver,
 	}
-	tablePrefix := opt.TablePrefix
+
-	if tablePrefix != "" && strings.HasPrefix(data.TableName, tablePrefix) {
+	tablePrefix := data.TableNamePrefix
+	if tablePrefix != "" && strings.HasPrefix(data.RawTableName, tablePrefix) {
 		data.NameFunc = true
-		data.TableName = data.TableName[len(tablePrefix):]
+		data.TableName = toCamel(data.RawTableName[len(tablePrefix):])
+	} else {
+		data.TableName = toCamel(data.RawTableName)
 	}
+	data.TName = firstLetterToLower(data.TableName)
+
 	if opt.ForceTableName || data.RawTableName != inflection.Plural(data.RawTableName) {
 		data.NameFunc = true
 	}
+
 	switch opt.DBDriver {
 	case DBDriverMongodb:
 		if opt.JSONNamedType != 0 {
@@ -379,9 +392,6 @@ func makeCode(stmt *ast.CreateTableStmt, opt options) (*codeText, error) {
 		}
 	}
 
-	data.TableName = toCamel(data.TableName)
-	data.TName = firstLetterToLower(data.TableName)
-
 	// find table comment
 	for _, o := range stmt.Options {
 		if o.Tp == ast.TableOptionComment {
@@ -518,6 +528,13 @@ func makeCode(stmt *ast.CreateTableStmt, opt options) (*codeText, error) {
 		data.Fields = append(data.Fields, field)
 	}
 
+	if v, ok := opt.FieldTypes[SubStructKey]; ok {
+		data.SubStructs = v
+	}
+	if v, ok := opt.FieldTypes[ProtoSubStructKey]; ok {
+		data.ProtoSubStructs = v
+	}
+
 	if len(data.Fields) == 0 {
 		return nil, errors.New("no columns found in table " + data.TableName)
 	}
@@ -525,11 +542,9 @@ func makeCode(stmt *ast.CreateTableStmt, opt options) (*codeText, error) {
 	data.CrudInfo = newCrudInfo(data)
 	data.CrudInfo.IsCommonType = data.isCommonStyle(opt.IsEmbed)
 
-	if v, ok := opt.FieldTypes[SubStructKey]; ok {
+	if opt.IsCustomTemplate {
-		data.SubStructs = v
+		tableInfo := newTableInfo(data)
-	}
+		return &codeText{tableInfo: tableInfo.getCode()}, nil
-	if v, ok := opt.FieldTypes[ProtoSubStructKey]; ok {
-		data.ProtoSubStructs = v
 	}
 
 	updateFieldsCode, err := getUpdateFieldsCode(data, opt.IsEmbed)

pkg/sql2code/parser/parser_test.go

@@ -11,10 +11,8 @@ import (
 )
 
 func TestParseSQL(t *testing.T) {
-	sqls := []string{`create table user
+	sqls := []string{`create table user (
-(
+    id          bigint unsigned auto_increment,
-    id          bigint unsigned auto_increment
-        primary key,
     created_at  datetime        null,
     updated_at  datetime        null,
     deleted_at  datetime        null,
@@ -26,34 +24,32 @@ func TestParseSQL(t *testing.T) {
     gender      tinyint         not null comment '性别，1:男，2:女，3:未知',
     status      tinyint         not null comment '账号状态，1:未激活，2:已激活，3:封禁',
     login_state tinyint         not null comment '登录状态，1:未登录，2:已登录',
+    primary key (id),
     constraint user_email_uindex
         unique (email)
 );`,
 
-		`create table user_order
+		`create table user_order (
-(
+    id         varchar(36)     not null comment '订单id',
-    id         varchar(36)     not null comment '订单id'
-        primary key,
     product_id varchar(36)     not null comment '商品id',
     user_id    bigint unsigned not null comment '用户id',
     status     smallint        null comment '0:未支付, 1:已支付, 2:已取消',
     created_at timestamp       null comment '创建时间',
-    updated_at timestamp       null comment '更新时间'
+    updated_at timestamp       null comment '更新时间',
+    primary key (id)
 );`,
 
-		`create table user_str
+		`create table user_str (
-(
+    user_id    varchar(36)  not null comment '用户id',
-    user_id    varchar(36)  not null comment '用户id'
-        primary key,
     username   varchar(50)  not null comment '用户名',
     email      varchar(100) not null comment '邮箱',
     created_at datetime     null comment '创建时间',
+    primary key (user_id),
     constraint email
         unique (email)
 );`,
 
-		`create table user_no_primary
+		`create table user_no_primary (
-(
     username   varchar(50)  not null comment '用户名',
     email      varchar(100) not null comment '邮箱',
     user_id    varchar(36)  not null comment '用户id',
@@ -66,6 +62,9 @@ func TestParseSQL(t *testing.T) {
 		codes, err := ParseSQL(sql, WithJSONTag(0), WithEmbed())
 		assert.Nil(t, err)
 		for k, v := range codes {
+			if k == CodeTypeTableInfo {
+				continue
+			}
 			assert.NotEmpty(t, k)
 			assert.NotEmpty(t, v)
 		}
@@ -74,6 +73,9 @@ func TestParseSQL(t *testing.T) {
 		codes, err = ParseSQL(sql, WithJSONTag(1), WithWebProto(), WithDBDriver(DBDriverMysql))
 		assert.Nil(t, err)
 		for k, v := range codes {
+			if k == CodeTypeTableInfo {
+				continue
+			}
 			assert.NotEmpty(t, k)
 			assert.NotEmpty(t, v)
 		}
@@ -82,6 +84,9 @@ func TestParseSQL(t *testing.T) {
 		codes, err = ParseSQL(sql, WithJSONTag(0), WithDBDriver(DBDriverPostgresql))
 		assert.Nil(t, err)
 		for k, v := range codes {
+			if k == CodeTypeTableInfo {
+				continue
+			}
 			assert.NotEmpty(t, k)
 			assert.NotEmpty(t, v)
 		}
@@ -90,38 +95,70 @@ func TestParseSQL(t *testing.T) {
 		codes, err = ParseSQL(sql, WithJSONTag(0), WithDBDriver(DBDriverSqlite))
 		assert.Nil(t, err)
 		for k, v := range codes {
+			if k == CodeTypeTableInfo {
+				continue
+			}
 			assert.NotEmpty(t, k)
 			assert.NotEmpty(t, v)
 		}
 		//printCode(codes)
+
+		codes, err = ParseSQL(sql, WithDBDriver(DBDriverSqlite), WithCustomTemplate())
+		assert.Nil(t, err)
+		for k, v := range codes {
+			if k == CodeTypeTableInfo {
+				assert.NotEmpty(t, k)
+				assert.NotEmpty(t, v)
+				break
+			}
+		}
+		//printCode(codes)
 	}
 }
 
 func TestParseSqlWithTablePrefix(t *testing.T) {
 	sql := `CREATE TABLE t_person_info (
-  id BIGINT(11) PRIMARY KEY AUTO_INCREMENT NOT NULL COMMENT 'id',
+  id BIGINT(11) AUTO_INCREMENT NOT NULL COMMENT 'id',
   age INT(11) unsigned NULL,
   name VARCHAR(30) NOT NULL DEFAULT 'default_name' COMMENT 'name',
   created_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
   login_at datetime NOT NULL DEFAULT CURRENT_TIMESTAMP,
   gender INT(8) NULL,
   num INT(11) DEFAULT 3 NULL,
-  comment TEXT
+  comment TEXT,
+  PRIMARY KEY (id)
   ) COMMENT="person info";`
 
 	codes, err := ParseSQL(sql, WithTablePrefix("t_"), WithJSONTag(0), WithNullStyle(NullDisable))
 	assert.Nil(t, err)
 	for k, v := range codes {
+		if k == CodeTypeTableInfo {
+			continue
+		}
 		assert.NotEmpty(t, k)
 		assert.NotEmpty(t, v)
 	}
-	t.Log(codes[CodeTypeJSON])
-
 	//printCode(codes)
 
+	codes, err = ParseSQL(sql, WithTablePrefix("t_"), WithJSONTag(0), WithCustomTemplate())
+	assert.Nil(t, err)
+	for k, v := range codes {
+		if k != CodeTypeTableInfo {
+			continue
+		}
+		assert.NotEmpty(t, k)
+		assert.NotEmpty(t, v)
+	}
+	jsonData := codes[CodeTypeTableInfo]
+	t.Log(jsonData)
+	t.Log(UnMarshalTableInfo(jsonData))
+
 	codes, err = ParseSQL(sql, WithTablePrefix("t_"), WithJSONTag(0), WithEmbed())
 	assert.Nil(t, err)
 	for k, v := range codes {
+		if k == CodeTypeTableInfo {
+			continue
+		}
 		assert.NotEmpty(t, k)
 		assert.NotEmpty(t, v)
 	}
@@ -130,6 +167,9 @@ func TestParseSqlWithTablePrefix(t *testing.T) {
 	codes, err = ParseSQL(sql, WithTablePrefix("t_"), WithJSONTag(0), WithWebProto())
 	assert.Nil(t, err)
 	for k, v := range codes {
+		if k == CodeTypeTableInfo {
+			continue
+		}
 		assert.NotEmpty(t, k)
 		assert.NotEmpty(t, v)
 	}
@@ -138,6 +178,9 @@ func TestParseSqlWithTablePrefix(t *testing.T) {
 	codes, err = ParseSQL(sql, WithTablePrefix("t_"), WithJSONTag(0), WithDBDriver(DBDriverPostgresql))
 	assert.Nil(t, err)
 	for k, v := range codes {
+		if k == CodeTypeTableInfo {
+			continue
+		}
 		assert.NotEmpty(t, k)
 		assert.NotEmpty(t, v)
 	}

pkg/sql2code/parser/tableInfo.go

@@ -0,0 +1,129 @@
+package parser
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/huandu/xstrings"
+	"github.com/jinzhu/inflection"
+)
+
+// TableInfo is the struct for extend template
+type TableInfo struct {
+	TableNamePrefix string // table name prefix, example: t_
+
+	TableName               string // original table name, example: foo_bar
+	TableNameCamel          string // camel case, example: FooBar
+	TableNameCamelFCL       string // camel case and first character lower, example: fooBar
+	TableNamePluralCamel    string // plural, camel case, example: FooBars
+	TableNamePluralCamelFCL string // plural, camel case and first character lower, example: fooBars
+	TableNameSnake          string // snake case, example: foo_bar
+
+	TableComment string // table comment
+
+	Columns    []Field     // columns of the table
+	PrimaryKey *PrimaryKey // primary key information
+
+	DBDriver string // database driver, example: mysql, postgresql, sqlite3, mongodb
+
+	ColumnSubStructure string // column sub structure for model
+	ColumnSubMessage   string // sub message for protobuf
+}
+
+// Field is the struct for column information
+type Field struct {
+	ColumnName         string // original column name, example: foo_bar
+	ColumnNameCamel    string // first character lower, example: FooBar
+	ColumnNameCamelFCL string // first character lower, example: fooBar
+
+	ColumnComment string // column comment
+	IsPrimaryKey  bool   // is primary key
+
+	GoType string // convert to go type
+	Tag    string // tag for model struct field, default gorm tag
+}
+
+// PrimaryKey is the struct for primary key information, it used for generate CRUD code
+type PrimaryKey struct {
+	Name               string // primary key name, example: foo_bar
+	NameCamel          string // primary key name, camel case, example: FooBar
+	NameCamelFCL       string // primary key name, camel case and first character lower, example: fooBar
+	NamePluralCamel    string // primary key name, plural, camel case, example: FooBars
+	NamePluralCamelFCL string // primary key name, plural, camel case and first character lower, example: fooBars
+
+	GoType    string // go type, example:  int, string
+	GoTypeFCU string // go type, first character upper, example: Int64, String
+
+	IsStringType bool // go type is string or not
+}
+
+func newTableInfo(data tmplData) TableInfo {
+	pluralName := inflection.Plural(data.TableName)
+	return TableInfo{
+		TableNamePrefix:         data.TableNamePrefix,
+		TableName:               data.RawTableName,
+		TableNameCamel:          data.TableName,
+		TableNameCamelFCL:       data.TName,
+		TableNamePluralCamel:    customEndOfLetterToLower(data.TableName, pluralName),
+		TableNamePluralCamelFCL: customFirstLetterToLower(customEndOfLetterToLower(data.TableName, pluralName)),
+		TableNameSnake:          xstrings.ToSnakeCase(data.TName),
+		TableComment:            data.Comment,
+		Columns:                 getColumns(data.Fields),
+		PrimaryKey:              getPrimaryKeyInfo(data.CrudInfo),
+		DBDriver:                data.DBDriver,
+		ColumnSubStructure:      data.SubStructs,
+		ColumnSubMessage:        data.ProtoSubStructs,
+	}
+}
+
+func (table TableInfo) getCode() []byte {
+	code, err := json.Marshal(&table)
+	if err != nil {
+		fmt.Printf("table: %v, json.Marshal error: %v\n", table.TableName, err)
+	}
+	return code
+}
+
+func getColumns(fields []tmplField) []Field {
+	var columns []Field
+
+	for _, field := range fields {
+		columns = append(columns, Field{
+			ColumnName:         field.ColName,
+			ColumnNameCamel:    field.Name,
+			ColumnNameCamelFCL: customFirstLetterToLower(field.Name),
+			ColumnComment:      field.Comment,
+			IsPrimaryKey:       field.IsPrimaryKey,
+			GoType:             field.GoType,
+			Tag:                field.Tag,
+		})
+	}
+
+	return columns
+}
+
+func getPrimaryKeyInfo(info *CrudInfo) *PrimaryKey {
+	if info == nil {
+		return nil
+	}
+	return &PrimaryKey{
+		Name:               info.ColumnName,
+		NameCamel:          info.ColumnNameCamel,
+		NameCamelFCL:       info.ColumnNameCamelFCL,
+		NamePluralCamel:    info.ColumnNamePluralCamel,
+		NamePluralCamelFCL: info.ColumnNamePluralCamelFCL,
+		GoType:             info.GoType,
+		GoTypeFCU:          info.GoTypeFCU,
+		IsStringType:       info.IsStringType,
+	}
+}
+
+// UnMarshalTableInfo unmarshal the json data to TableInfo struct
+func UnMarshalTableInfo(data string) (map[string]interface{}, error) {
+	info := map[string]interface{}{}
+	err := json.Unmarshal([]byte(data), &info)
+	if err != nil {
+		return info, err
+	}
+	return info, nil
+}

pkg/sql2code/sql2code.go

@@ -40,6 +40,8 @@ type Args struct {
 	NoNullType     bool
 	NullStyle      string
 	IsExtendedAPI  bool // true: generate extended api (9 api), false: generate basic api (5 api)
+
+	IsCustomTemplate bool // whether to use custom template, default is false
 }
 
 func (a *Args) checkValid() error {
@@ -181,6 +183,9 @@ func setOptions(args *Args) []parser.Option {
 	if args.IsExtendedAPI {
 		opts = append(opts, parser.WithExtendedAPI())
 	}
+	if args.IsCustomTemplate {
+		opts = append(opts, parser.WithCustomTemplate())
+	}
 
 	return opts
 }

pkg/sql2code/sql2code_test.go

@@ -85,14 +85,43 @@ func TestGenerate(t *testing.T) {
 			}},
 			wantErr: false,
 		},
-		//{
+		{
-		//	name: "sql from db",
+			name: "sql from mysql",
-		//	args: args{args: &Args{
+			args: args{args: &Args{
-		//		DBDsn:   "root:123456@(127.0.0.1:3306)/test",
+				DBDsn:    "root:123456@(192.168.3.37:3306)/account",
-		//		DBTable: "user",
+				DBTable:  "user",
-		//	}},
+				DBDriver: "mysql",
-		//	wantErr: false,
+			}},
-		//},
+			wantErr: false,
+		},
+		{
+			name: "sql from postgresql",
+			args: args{args: &Args{
+				DBDsn:    "root:123456@(192.168.3.37:5432)/account",
+				DBTable:  "user",
+				DBDriver: "postgresql",
+			}},
+			wantErr: false,
+		},
+		{
+			name: "sql from sqlite",
+			args: args{args: &Args{
+				DBDsn:    "C:\\Users\\zhuyasen\\Desktop\\genTest\\sql\\sqlite\\sponge.db",
+				DBTable:  "user",
+				DBDriver: "sqlite",
+			}},
+			wantErr: false,
+		},
+		{
+			name: "sql from mongodb",
+			args: args{args: &Args{
+				DBDsn:            "root:123456@(192.168.3.37:27017)/account",
+				DBTable:          "people",
+				DBDriver:         "mongodb",
+				IsCustomTemplate: true,
+			}},
+			wantErr: false,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {