Merge pull request #19 from go-gorm/dev

Dev

do.go

@@ -176,7 +176,7 @@ func (d *DO) Order(columns ...field.Expr) Dao {
 
 func (d *DO) Distinct(columns ...field.Expr) Dao {
 	Emit(methodDistinct)
-	return NewDO(d.db.Distinct(toInterfaceSlice(toColNames(d.db.Statement, columns...))))
+	return NewDO(d.db.Distinct(toInterfaceSlice(toColNames(d.db.Statement, columns...))...))
 }
 
 func (d *DO) Omit(columns ...field.Expr) Dao {

generator.go

@@ -79,12 +79,12 @@ func (g *Generator) UseDB(db *gorm.DB) {
 }
 
 // GenerateModel catch table info from db, return a BaseStruct
-func (g *Generator) GenerateModel(name string) *check.BaseStruct {
-	structs, err := check.GenBaseStructs(g.db, g.Config.ModelPkgName, name)
+func (g *Generator) GenerateModel(tableName string, modelName string) *check.BaseStruct {
+	s, err := check.GenBaseStructs(g.db, g.Config.ModelPkgName, tableName, modelName)
 	if err != nil {
 		log.Fatalf("check struct error: %s", err)
 	}
-	return structs[0]
+	return s
 }
 
 // ApplyBasic specify models which will implement basic method

generator_test.go

@@ -15,4 +15,4 @@ func TestConfig(t *testing.T) {
 
 		queryPkgName: "query",
 	}
-}
+}

internal/check/export.go

@@ -3,7 +3,6 @@ package check
 import (
 	"fmt"
 	"reflect"
-	"strings"
 
 	"gorm.io/gorm"
 
@@ -31,7 +30,7 @@ func CheckStructs(db *gorm.DB, structs ...interface{}) (bases []*BaseStruct, err
 		base := &BaseStruct{
 			S:             GetSimpleName(name),
 			StructName:    name,
-			NewStructName: strings.ToLower(name),
+			NewStructName: uncaptialize(name),
 			StructInfo:    parser.Param{Type: name, Package: getPackageName(structType.String())},
 			Source:        Struct,
 			db:            db,

internal/check/gen_structs.go

@@ -57,43 +57,46 @@ var dataType = map[string]string{
 }
 
 // GenBaseStructs generate db model by table name
-func GenBaseStructs(db *gorm.DB, pkg string, tableName ...string) (bases []*BaseStruct, err error) {
+func GenBaseStructs(db *gorm.DB, pkg string, tableName, modelName string) (bases *BaseStruct, err error) {
 	if isDBUndefined(db) {
 		return nil, fmt.Errorf("gen config db is undefined")
 	}
+	if !isModelNameValid(modelName) {
+		return nil, fmt.Errorf("model name %q is invalid", modelName)
+	}
 	if pkg == "" {
 		pkg = ModelPkg
 	}
 	singular := singularModel(db.Config)
 	dbName := getSchemaName(db)
-	for _, tb := range tableName {
-		columns, err := getTbColumns(db, dbName, tb)
-		if err != nil {
-			return nil, err
-		}
-		var base BaseStruct
-		base.Source = TableName
-		base.GenBaseStruct = true
-		base.TableName = tb
-		base.StructName = convertToModelName(singular, tb)
-		base.StructInfo = parser.Param{Type: base.StructName, Package: pkg}
-		for _, field := range columns {
-			mt := dataType[field.DataType]
-			base.Members = append(base.Members, &Member{
-				Name:          nameToCamelCase(field.ColumnName),
-				Type:          mt,
-				ModelType:     mt,
-				ColumnName:    field.ColumnName,
-				ColumnComment: field.ColumnComment,
-			})
-		}
-
-		base.NewStructName = strings.ToLower(base.StructName)
-		base.S = string(base.NewStructName[0])
-		_ = base.check()
-		bases = append(bases, &base)
+	columns, err := getTbColumns(db, dbName, tableName)
+	if err != nil {
+		return nil, err
 	}
-	return
+	var base BaseStruct
+	base.Source = TableName
+	base.GenBaseStruct = true
+	base.TableName = tableName
+	base.StructName = convertToModelName(singular, tableName)
+	if modelName != "" {
+		base.StructName = captialize(modelName)
+	}
+	base.NewStructName = uncaptialize(base.StructName)
+	base.S = string(base.NewStructName[0])
+	base.StructInfo = parser.Param{Type: base.StructName, Package: pkg}
+	for _, field := range columns {
+		mt := dataType[field.DataType]
+		base.Members = append(base.Members, &Member{
+			Name:          nameToCamelCase(field.ColumnName),
+			Type:          mt,
+			ModelType:     mt,
+			ColumnName:    field.ColumnName,
+			ColumnComment: field.ColumnComment,
+		})
+	}
+
+	_ = base.check()
+	return &base, nil
 }
 
 //Mysql
@@ -145,3 +148,16 @@ func singularModel(conf *gorm.Config) bool {
 	}
 	return false
 }
+
+// get mysql db' name
+var modelNameReg = regexp.MustCompile(`^\w+$`)
+
+func isModelNameValid(name string) bool {
+	if name == "" {
+		return true
+	}
+	if !modelNameReg.MatchString(name) {
+		return false
+	}
+	return (name[0] > '9' || name[0] < '0') && name[0] != '_'
+}

internal/check/utils.go

@@ -89,3 +89,19 @@ func getStructName(t string) string {
 	list := strings.Split(t, ".")
 	return list[len(list)-1]
 }
+
+func uncaptialize(s string) string {
+	if s == "" {
+		return ""
+	}
+
+	return strings.ToLower(s[:1]) + s[1:]
+}
+
+func captialize(s string) string {
+	if s == "" {
+		return ""
+	}
+
+	return strings.ToUpper(s[:1]) + s[1:]
+}