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"
	"jollia.cn/jolib/jokode"
	"jollia.cn/jolib/jokode/log"
	"strings"
)

const (
	defCfgFilePath = "conf/db.yaml"
)

type GenService struct {
	dbInstance *gorm.DB
	tables     []string
}

func (s *GenService) Init() error {
	log.Warn("Enter: GenService.Init")
	defer log.Warn("Leave: GenService.Init")

	if !flag.Parsed() {
		flag.Parse()
	}

	if createCfgFileFlag {
		return nil
	}

	var (
		dbCfg       *DbConfig
		cfgFilePath string
		force       bool
	)

	if cfgFilePathFlag != "" {
		cfgFilePath = cfgFilePathFlag
		force = true
	} else {
		cfgFilePath = defCfgFilePath
		force = false
	}

	if cfg, err := loadConfigFromFile(cfgFilePath, force); err != nil {
		return errors.Wrapf(err, "try load db config from file '%s' fail", cfgFilePath)
	} else if cfg == nil {
		return errors.New("load db config returns nil")
	} else {
		dbCfg = cfg
		if len(cfg.Tables) > 0 {
			s.tables = cfg.Tables
		}
	}

	if !dbConfigFlag.IsZero() {
		if dbConfigFlag.Host != "" {
			dbCfg.Host = dbConfigFlag.Host
		}

		if dbConfigFlag.Port != 0 {
			dbCfg.Port = dbConfigFlag.Port
		}

		if dbConfigFlag.Schema != "" {
			dbCfg.Schema = dbConfigFlag.Schema
		}

		if dbConfigFlag.User != "" {
			dbCfg.User = dbConfigFlag.User
		}

		if dbConfigFlag.Pwd != "" {
			dbCfg.Pwd = dbConfigFlag.Pwd
		}
	}

	if err := s.initDatabase(dbCfg); err != nil {
		return errors.Wrap(err, "try init database fail")
	}

	return nil
}

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
}

func loadConfigFromFile(filePath string, force bool) (*DbConfig, error) {
	cfg := &DbConfig{}

	if err := jokode.LoadConfigFromFile(cfg, filePath); err != nil {
		if force {
			return nil, errors.Wrapf(err, "try load config from file '%s' fail", filePath)
		}
	}

	return cfg, nil
}

func (s *GenService) Start() error {
	log.Warn("Enter: GenService.Start")
	defer log.Warn("Leave: GenService.Start")

	if createCfgFileFlag {
		cfgFilePath := cfgFilePathFlag
		if cfgFilePath == "" {
			cfgFilePath = defCfgFilePath
		}

		if jokode.FileExists(cfgFilePath) {
			return errors.Errorf("try create default config file '%s', but the file exists already", cfgFilePath)
		}

		cfg := &DbConfig{
			Host:   "localhost",
			Port:   3306,
			User:   "root",
			Pwd:    "123456",
			Schema: "db",
			Tables: []string{"users", "roles"},
		}

		return errors.Wrapf(jokode.SaveConfigToFile(cfg, cfgFilePath), "try create default config file '%s' fail", cfgFilePath)
	}

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

	g := gen.NewGenerator(gen.Config{
		OutPath:           "model",
		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)

	if len(s.tables) > 0 {
		for _, tableName := range s.tables {
			g.GenerateModel(tableName)
		}
	} else {
		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
}