makemodel/create.go

package makemodel

import (
	"fmt"
	"gorm.io/gorm"
	"os"
	"path/filepath"
	"strings"
)

type Column struct {
	ColumnName    string `gorm:"column:COLUMN_NAME"`
	DataType      string `gorm:"column:DATA_TYPE"`
	ColumnKey     string `gorm:"column:COLUMN_KEY"`
	ColumnComment string `gorm:"column:COLUMN_COMMENT"`
}

// Create @Title 创建model
func Create(db *gorm.DB, table string, path ...string) error {
	var columns []Column
	sql := fmt.Sprintf("select * from information_schema.COLUMNS where table_name ='%s' order by ordinal_position", db.NamingStrategy.TableName(table))
	db.Raw(sql).Scan(&columns)
	tableName := camelString(table)

	strContent := "type " + tableName + " struct {\n"
	imports := map[string]bool{}
	for _, column := range columns {
		columnType := ""
		switch column.DataType {
		case "varchar", "text":
			columnType = "string"
		case "int", "tinyint":
			columnType = "uint"
		case "decimal":
			columnType = "decimal.Decimal"
			imports["github.com/shopspring/decimal"] = true
		case "datetime":
			if column.ColumnName == "deleted_at" {
				columnType = "gorm.DeletedAt"
				imports["gorm.io/gorm"] = true
			} else {
				columnType = "time.Time"
				imports["time"] = true
			}
		}
		pri := ""
		if column.ColumnKey == "PRI" {
			pri = "`gorm:\"primaryKey\"`"
		}
		comment := ""
		if column.ColumnComment != "" {
			comment = "// " + column.ColumnComment
		}
		strContent += fmt.Sprintf("\t%s %s %s%s\n", camelString(column.ColumnName), columnType, pri, comment)
	}
	strContent += "}"
	headerStr := "package model\n\nimport(\n"
	for key, _ := range imports {
		headerStr += fmt.Sprintf("\t\"%s\"\n", key)
	}
	headerStr += ")\n\n"
	strContent = headerStr + strContent

	fileName := strings.ToLower(tableName[0:1]) + tableName[1:]
	dir := ""
	if len(path) > 0 {
		dir = path[0]
		os.MkdirAll(dir, 0777)
	}
	f, err := os.OpenFile(filepath.Join(dir, fileName+".go"), os.O_CREATE|os.O_RDWR, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = f.WriteString(strContent)
	if err != nil {
		return err
	}
	return nil
}

// camel string, xx_yy to XxYy
func camelString(s string) string {
	data := make([]byte, 0, len(s))
	j := false
	k := false
	num := len(s) - 1
	for i := 0; i <= num; i++ {
		d := s[i]
		if k == false && d >= 'A' && d <= 'Z' {
			k = true
		}
		if d >= 'a' && d <= 'z' && (j || k == false) {
			d = d - 32
			j = false
			k = true
		}
		if k && d == '_' && num > i && s[i+1] >= 'a' && s[i+1] <= 'z' {
			j = true
			continue
		}
		data = append(data, d)
	}
	return string(data[:])
}