使用GORM自动生成Go结构体文件的工具
在Go开发中,我们经常需要根据数据库表结构生成对应的Go结构体。虽然网上有很多现成的包,但大多过于复杂或不符合我们的需求。本文介绍一个简单实用的工具,可以根据MySQL数据库表自动生成Go结构体文件。
package main
import (
"flag"
"fmt"
"io"
"os"
"strings"
"gorm.io/gorm"
"gorm.io/driver/mysql"
)
var dbConnection *gorm.DB
func main() {
var dbHost string
var dbUser string
var dbPassword string
var dbName string
var targetTable string
var outputPackage string
var dbPort string
flag.StringVar(&dbHost, "host", "", "数据库主机地址")
flag.StringVar(&dbPassword, "pwd", "", "数据库密码")
flag.StringVar(&dbUser, "user", "", "数据库用户名")
flag.StringVar(&dbName, "db", "", "数据库名称")
flag.StringVar(&targetTable, "table", "", "目标表名")
flag.StringVar(&dbPort, "port", "3306", "数据库端口")
flag.StringVar(&outputPackage, "pkg", "", "输出包名")
flag.Parse()
dataSourceName := fmt.Sprintf(
"%s:%s@tcp(%s:%s)/%s?parseTime=True&loc=Local",
dbUser,
dbPassword,
dbHost,
dbPort,
dbName,
)
var err error
dbConnection, err = gorm.Open(mysql.Open(dataSourceName), &gorm.Config{})
if err != nil {
fmt.Printf("连接数据库失败: %v\n", err)
return
}
type ColumnInfo struct {
ColumnName string `gorm:"column:Column_Name"`
DataType string `gorm:"column:Data_Type"`
ColumnComment string `gorm:"column:Column_Comment"`
}
var columns []ColumnInfo
packageHeader := fmt.Sprintf("package %s\n\n", outputPackage)
structName := convertToPascalCase(targetTable)
structStart := fmt.Sprintf("type %s struct {", structName)
content := packageHeader + structStart
query := fmt.Sprintf("SELECT COLUMN_NAME, COLUMN_COMMENT as Column_Comment, DATA_TYPE as Data_Type FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA='%s' AND TABLE_NAME='%s';", dbName, targetTable)
result := dbConnection.Raw(query).Scan(&columns)
if result.Error != nil {
fmt.Printf("查询表结构失败: %v\n", result.Error)
return
}
for _, column := range columns {
field := fmt.Sprintf("\n\t%s %s `json:\"%s\"` // %s",
convertToPascalCase(column.ColumnName),
mapToGoType(column.DataType),
column.ColumnName,
column.ColumnComment)
content += field
}
content += "\n}"
tableNameMethod := fmt.Sprintf("\n\nfunc (%s) TableName() string {\n\treturn \"%s\"\n}", structName, targetTable)
content += tableNameMethod
outputFile := fmt.Sprintf("%s.go", structName)
// 检查文件是否已存在
if _, err := os.Stat(outputFile); err == nil {
// 文件存在且不为空,则不允许覆盖
if file, _ := os.Open(outputFile); file != nil {
defer file.Close()
if stat, _ := file.Stat(); stat.Size() > 0 {
fmt.Println("目标文件已存在且不为空,请先删除或重命名")
return
}
}
}
file, err := os.Create(outputFile)
if err != nil {
fmt.Printf("创建文件失败: %v\n", err)
return
}
defer file.Close()
_, err = io.WriteString(file, content)
if err != nil {
fmt.Printf("写入文件失败: %v\n", err)
return
}
fmt.Printf("成功生成文件: %s\n", outputFile)
}
// 将下划线命名转换为帕斯卡命名
func convertToPascalCase(input string) string {
var nextUpper bool
result := ""
for i, char := range input {
if i == 0 {
result += strings.ToUpper(string(char))
continue
}
if string(char) == "_" {
nextUpper = true
continue
}
if nextUpper {
result += strings.ToUpper(string(char))
nextUpper = false
} else {
result += string(char)
}
}
return result
}
// 将MySQL数据类型映射为Go类型
func mapToGoType(mysqlType string) string {
typeIndex := strings.Index(mysqlType, "(")
if typeIndex < 0 {
typeIndex = len(mysqlType) - 1
}
baseType := mysqlType[:typeIndex]
typeMapping := map[string]string{
"tinyint": "int64",
"smallint": "int64",
"mediumint": "int64",
"int": "int64",
"bigint": "int64",
"float": "float64",
"double": "float64",
"decimal": "string",
"bit": "string",
"year": "string",
"time": "string",
"date": "string",
"datetime": "string",
"timestamp": "string",
"char": "string",
"varchar": "string",
"tinytext": "string",
"text": "string",
"mediumtext": "string",
"longtext": "string",
"enum": "string",
}
if goType, exists := typeMapping[baseType]; exists {
return goType
}
return "string"
}
使用示例:
go run main.go -host 127.0.0.1 -port 3306 -user test -pwd password -db mydb -table users -pkg models
该工具会根据指定的数据库连接信息和表名,在当前目录下生成一个Go结构体文件,文件名与表名对应(首字母大写),并包含指定的包名和表名方法。
工具特点:
- 简单易用,无需复杂配置
- 支持命令行参数配置
- 自动将数据库字段类型转换为Go类型
- 自动添加JSON标签和注释
- 防止覆盖已存在的非空文件