From f612a01d3319877d1ec91ff2f400afda9c035a09 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=AD=A6=E6=96=8C?= <wubinwill@gmail.com>
Date: Wed, 16 Oct 2024 22:57:25 +0800
Subject: [PATCH] add --db_type_mapping flag

---
 cmd/crud.go                                 |  3 +++
 cmd/internal/svc/codegen/database/common.go |  1 +
 cmd/internal/svc/codegen/database/gorm.go   | 16 ++++++++++++++++
 cmd/internal/svc/svc.go                     | 12 +++++++-----
 version/version.go                          |  2 +-
 5 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/cmd/crud.go b/cmd/crud.go
index 97eda5cb..e0337faa 100644
--- a/cmd/crud.go
+++ b/cmd/crud.go
@@ -20,6 +20,7 @@ var dbTableGlob string
 var dbTableExcludeGlob string
 var dbGenGenGo bool
 var dbOmitempty bool
+var dbTypeMapping string
 
 var crudCmd = &cobra.Command{
 	Use:   "crud",
@@ -42,6 +43,7 @@ var crudCmd = &cobra.Command{
 			Soft:             dbSoft,
 			Service:          dbService,
 			Omitempty:        dbOmitempty,
+			TypeMapping:      dbTypeMapping,
 		}))
 		fn := strcase.ToLowerCamel
 		switch naming {
@@ -68,6 +70,7 @@ func init() {
 	crudCmd.Flags().StringVar(&dbTableGlob, "db_table_glob", "", `used to filter glob-matched tables`)
 	crudCmd.Flags().StringVar(&dbTableExcludeGlob, "db_table_exclude_glob", "", `used to filter glob-matched tables`)
 	crudCmd.Flags().StringVar(&naming, "case", "lowerCamel", `protobuf message field and json tag case, only support "lowerCamel" and "snake"`)
+	crudCmd.Flags().StringVar(&dbTypeMapping, "db_type_mapping", "", `Specify custom column type to go type mapping, multiple mappings can be joined by comma, eg. bigint:string,tinyint:bool`)
 	crudCmd.Flags().StringVar(&protocCmd, "grpc_gen_cmd", "protoc --proto_path=. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative --go-json_out=. --go-json_opt=paths=source_relative,allow_unknown=true", `command to generate grpc service and message code`)
 	crudCmd.Flags().BoolVar(&dbOmitempty, "db_omitempty", false, `whether add omitempty json tag to generated model field"`)
 }
diff --git a/cmd/internal/svc/codegen/database/common.go b/cmd/internal/svc/codegen/database/common.go
index 51480b8a..2e4ccad7 100644
--- a/cmd/internal/svc/codegen/database/common.go
+++ b/cmd/internal/svc/codegen/database/common.go
@@ -39,6 +39,7 @@ type OrmGeneratorConfig struct {
 	Soft             string
 	ProtoGenerator   v3.ProtoGenerator
 	Omitempty        bool
+	TypeMapping      string
 }
 
 type IOrmGenerator interface {
diff --git a/cmd/internal/svc/codegen/database/gorm.go b/cmd/internal/svc/codegen/database/gorm.go
index 0443f4ae..0cd1bc8d 100644
--- a/cmd/internal/svc/codegen/database/gorm.go
+++ b/cmd/internal/svc/codegen/database/gorm.go
@@ -1,6 +1,7 @@
 package database
 
 import (
+	"github.com/samber/lo"
 	"github.com/unionj-cloud/toolkit/astutils"
 	"path/filepath"
 	"strings"
@@ -155,6 +156,21 @@ func (gg *GormGenerator) Initialize(conf OrmGeneratorConfig) {
 		}
 		return tagContent
 	})
+	if stringutils.IsNotEmpty(conf.TypeMapping) {
+		typeMapping := make(map[string]func(gorm.ColumnType) (dataType string))
+		lo.ForEach(strings.Split(conf.TypeMapping, ","), func(item string, index int) {
+			kv := strings.Split(item, ":")
+			columnType := kv[0]
+			goType := kv[1]
+			typeMapping[columnType] = func(columnType gorm.ColumnType) (dataType string) {
+				if n, ok := columnType.Nullable(); ok && n {
+					return "*" + goType
+				}
+				return goType
+			}
+		})
+		g.WithDataTypeMap(typeMapping)
+	}
 	g.UseDB(db)
 	g.GenGenGo = gg.GenGenGo
 	var models []interface{}
diff --git a/cmd/internal/svc/svc.go b/cmd/internal/svc/svc.go
index 8597b617..87484478 100644
--- a/cmd/internal/svc/svc.go
+++ b/cmd/internal/svc/svc.go
@@ -107,11 +107,12 @@ type DbConfig struct {
 	TableGlob        string
 	TableExcludeGlob string
 	// whether generate gen.go file
-	GenGenGo  bool
-	Orm       string
-	Soft      string
-	Service   string
-	Omitempty bool
+	GenGenGo    bool
+	Orm         string
+	Soft        string
+	Service     string
+	Omitempty   bool
+	TypeMapping string
 }
 
 func (receiver *Svc) SetWatcher(w *watcher.Watcher) {
@@ -201,6 +202,7 @@ func (receiver *Svc) Crud() {
 		Soft:             receiver.DbConfig.Soft,
 		ProtoGenerator:   receiver.protoGenerator,
 		Omitempty:        receiver.DbConfig.Omitempty,
+		TypeMapping:      receiver.DbConfig.TypeMapping,
 	})
 	if stringutils.IsNotEmpty(receiver.DbConfig.Service) {
 		switch receiver.DbConfig.Service {
diff --git a/version/version.go b/version/version.go
index 0fee4987..3e1bdad1 100644
--- a/version/version.go
+++ b/version/version.go
@@ -1,3 +1,3 @@
 package version
 
-const Release = "v2.5.2"
+const Release = "v2.5.3"