feat: support type mapping by colume (go-gorm#820)

* feat: support type mapping by colum * feat: test
miseyu · Apr 7, 2023 · 6d02fd9 · 6d02fd9
1 parent cdd627b
commit 6d02fd9
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 13 deletions.
diff --git a/config.go b/config.go
@@ -53,7 +53,7 @@ type Config struct {
 	modelNameNS func(tableName string) (modelName string)
 	fileNameNS  func(tableName string) (fileName string)
 
-	dataTypeMap    map[string]func(detailType string) (dataType string)
+	dataTypeMap    map[string]func(columnType gorm.ColumnType) (dataType string)
 	fieldJSONTagNS func(columnName string) (tagContent string)
 	fieldNewTagNS  func(columnName string) (tagContent string)
 
@@ -94,7 +94,7 @@ func (cfg *Config) WithFileNameStrategy(ns func(tableName string) (fileName stri
 }
 
 // WithDataTypeMap specify data type mapping relationship, only work when syncing table from db
-func (cfg *Config) WithDataTypeMap(newMap map[string]func(detailType string) (dataType string)) {
+func (cfg *Config) WithDataTypeMap(newMap map[string]func(columnType gorm.ColumnType) (dataType string)) {
 	cfg.dataTypeMap = newMap
 }
 

diff --git a/examples/cmd/sync_table/generate.go b/examples/cmd/sync_table/generate.go
@@ -6,6 +6,7 @@ import (
 	"gorm.io/gen"
 	"gorm.io/gen/examples/conf"
 	"gorm.io/gen/examples/dal"
+	"gorm.io/gorm"
 )
 
 func init() {
@@ -15,13 +16,14 @@ func init() {
 }
 
 // dataMap mapping relationship
-var dataMap = map[string]func(detailType string) (dataType string){
+var dataMap = map[string]func(gorm.ColumnType) (dataType string){
 	// int mapping
-	"int": func(detailType string) (dataType string) { return "int32" },
+	"int": func(columnType gorm.ColumnType) (dataType string) { return "int32" },
 
 	// bool mapping
-	"tinyint": func(detailType string) (dataType string) {
-		if strings.HasPrefix(detailType, "tinyint(1)") {
+	"tinyint": func(columnType gorm.ColumnType) (dataType string) {
+		ct, _ := columnType.ColumnType()
+		if strings.HasPrefix(ct, "tinyint(1)") {
 			return "bool"
 		}
 		return "byte"

diff --git a/examples/cmd/ultimate/generate.go b/examples/cmd/ultimate/generate.go
@@ -5,6 +5,7 @@ import (
 	"gorm.io/gen/examples/conf"
 	"gorm.io/gen/examples/dal"
 	"gorm.io/gen/examples/dal/model"
+	"gorm.io/gorm"
 )
 
 func init() {
@@ -13,9 +14,9 @@ func init() {
 	prepare(dal.DB) // prepare table for generate
 }
 
-var dataMap = map[string]func(detailType string) (dataType string){
-	"int":  func(detailType string) (dataType string) { return "int64" },
-	"json": func(string) string { return "json.RawMessage" },
+var dataMap = map[string]func(gorm.ColumnType) (dataType string){
+	"int":  func(columnType gorm.ColumnType) (dataType string) { return "int64" },
+	"json": func(columnType gorm.ColumnType) string { return "json.RawMessage" },
 }
 
 func main() {

diff --git a/internal/model/config.go b/internal/model/config.go
@@ -33,7 +33,7 @@ type NameStrategy struct {
 
 // FieldConfig field configuration
 type FieldConfig struct {
-	DataTypeMap map[string]func(detailType string) (dataType string)
+	DataTypeMap map[string]func(columnType gorm.ColumnType) (dataType string)
 
 	FieldNullable     bool // generate pointer when field is nullable
 	FieldCoverable    bool // generate pointer when field has default value

diff --git a/internal/model/tbl_column.go b/internal/model/tbl_column.go
@@ -15,20 +15,20 @@ type Column struct {
 	TableName   string                                               `gorm:"column:TABLE_NAME"`
 	Indexes     []*Index                                             `gorm:"-"`
 	UseScanType bool                                                 `gorm:"-"`
-	dataTypeMap map[string]func(detailType string) (dataType string) `gorm:"-"`
+	dataTypeMap map[string]func(columnType gorm.ColumnType) (dataType string) `gorm:"-"`
 	jsonTagNS   func(columnName string) string                       `gorm:"-"`
 	newTagNS    func(columnName string) string                       `gorm:"-"`
 }
 
 // SetDataTypeMap set data type map
-func (c *Column) SetDataTypeMap(m map[string]func(detailType string) (dataType string)) {
+func (c *Column) SetDataTypeMap(m map[string]func(columnType gorm.ColumnType) (dataType string)) {
 	c.dataTypeMap = m
 }
 
 // GetDataType get data type
 func (c *Column) GetDataType() (fieldtype string) {
 	if mapping, ok := c.dataTypeMap[c.DatabaseTypeName()]; ok {
-		return mapping(c.columnType())
+		return mapping(c.ColumnType)
 	}
 	if c.UseScanType && c.ScanType() != nil {
 		return c.ScanType().String()