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() } 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 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 }