db2code/main.go

package main

import (
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gen"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"strings"
)

type Config struct {
	DbAddr   string `json:"dbAddr,omitempty"`
	DbSchema string `json:"dbSchema,omitempty"`
	DbUser   string `json:"dbUser,omitempty"`
	DbPwd    string `json:"dbPwd,omitempty"`
}

func main() {
	fmt.Println("db2code started ....")
	defer fmt.Println("db2code stopped ....")

	cfg := &Config{
		DbAddr:   "158.8.12.102:3308",
		DbSchema: "ticket-cloud",
		DbUser:   "root",
		DbPwd:    "ticket@2025",
	}

	db := initDb(cfg)
	genStructs(db)
}

func initDb(cfg *Config) *gorm.DB {
	if cfg == nil {
		panic("can not init db wih nil config")
	}

	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
		cfg.DbUser,
		cfg.DbPwd,
		cfg.DbAddr,
		cfg.DbSchema)

	if db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Info),
	}); err != nil {
		panic(fmt.Sprintf("try connect to db fail, dsn: '%s', error info: %v",
			dsn, err))
	} else {
		return db
	}
}

func genStructs(db *gorm.DB) {
	g := gen.NewGenerator(gen.Config{
		OutPath:           "po",
		WithUnitTest:      false,
		FieldNullable:     true,
		FieldCoverable:    true,
		FieldSignable:     false,
		FieldWithIndexTag: true,
		FieldWithTypeTag:  true,
		Mode:              gen.WithDefaultQuery | gen.WithQueryInterface,
	})

	g.WithDataTypeMap(map[string]func(columnType gorm.ColumnType) (dataType string){
		"tinyint": func(filedType gorm.ColumnType) string {
			if ct, ok := filedType.ColumnType(); ok {
				if strings.HasPrefix(strings.ToLower(ct), "tinyint(1)") {
					return "int8"
				}
				return "int16"
			} else {
				return "int32"
			}
		},
	})

	g.UseDB(db)
	g.GenerateAllTable()
	g.Execute()
}