db2code/services/gen_service.go

package services

import (
	"flag"
	"fmt"
	"github.com/pkg/errors"
	"gorm.io/driver/mysql"
	"gorm.io/gen"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"strings"
)

type GenService struct {
	dbInstance *gorm.DB
}

func (s *GenService) Init() error {
	if !flag.Parsed() {
		flag.Parse()
	}

	if dbConfigInstance.IsZero() {

	}

	return nil
}

func (s *GenService) Start() error {
	if s.dbInstance == nil {
		return errors.New("try start gen service with nil database instance")
	}

	g := gen.NewGenerator(gen.Config{
		OutPath:           "po",
		OutFile:           "po",
		ModelPkgPath:      "po",
		WithUnitTest:      false,
		FieldNullable:     true,
		FieldCoverable:    true,
		FieldSignable:     true,
		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(s.dbInstance)
	g.GenerateAllTable()
	g.Execute()

	return nil
}

func (s *GenService) Stop() {
	// do nothing
}

func (s *GenService) GetName() string {
	return "GenService"
}

func (s *GenService) GetInitOrder() int {
	return 0
}

func (s *GenService) initDatabase(dbCfg *DbConfig) error {
	if dbCfg == nil {
		return errors.New("can not init database with nil db config")
	}

	if err := dbCfg.FixDefaultValues(); err != nil {
		return errors.Wrap(err, "try fix default values of database config fail")
	}

	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local",
		dbCfg.User,
		dbCfg.Pwd,
		dbCfg.Host,
		dbCfg.Port,
		dbCfg.Schema)

	if db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Info),
	}); err != nil {
		return errors.Wrapf(err, "try open database with dsn '%s' fail", dsn)
	} else {
		s.dbInstance = db
	}

	return nil
}