使用Goose迁移工具结合Ent框架实现Golang代码迁移的方案咨询

要实现用goose create生成带*sql.DB参数的迁移文件，适配Ent框架，需从自定义Goose二进制、替换迁移模板、集成Ent客户端三个方面入手，具体步骤如下：

1. 实现自定义Goose Provider（基于*sql.DB） #

Goose默认Provider基于*sql.Tx，我们需要实现一个返回*sql.DB的自定义Provider，让迁移函数能拿到DB实例初始化Ent客户端。

在项目根目录创建cmd/goose/main.go，内容如下：

package main

import (
	"context"
	"database/sql"
	"flag"
	"log"
	"os"

	_ "github.com/lib/pq"
	"github.com/pressly/goose/v3"
)

// DBProvider 实现goose.Provider接口，返回*sql.DB实例
type DBProvider struct {
	db *sql.DB
}

func NewDBProvider(ctx context.Context, dsn string) (*DBProvider, error) {
	db, err := sql.Open("postgres", dsn)
	if err != nil {
		return nil, err
	}
	if err := db.PingContext(ctx); err != nil {
		return nil, err
	}
	return &DBProvider{db: db}, nil
}

// Provide 返回*sql.DB，供迁移函数调用
func (p *DBProvider) Provide(ctx context.Context) (interface{}, error) {
	return p.db, nil
}

// Close 关闭数据库连接
func (p *DBProvider) Close() error {
	return p.db.Close()
}

func main() {
	flag.Parse()
	args := flag.Args()

	if len(args) < 1 {
		log.Fatalf("usage: goose <command> [args...]")
	}

	ctx := context.Background()
	// 从环境变量读取数据库DSN，也可改为命令行参数
	dsn := os.Getenv("DATABASE_URL")
	if dsn == "" {
		log.Fatal("DATABASE_URL environment variable not set")
	}

	provider, err := NewDBProvider(ctx, dsn)
	if err != nil {
		log.Fatalf("failed to create DB provider: %v", err)
	}
	defer provider.Close()

	// 将自定义Provider注册到Goose
	goose.SetProvider(provider)

	// 指定迁移文件存放目录（按你的项目结构调整）
	migrationsDir := "./migrations"

	// 执行Goose命令
	if err := goose.RunContext(ctx, args[0], migrationsDir, args[1:]...); err != nil {
		log.Fatalf("goose run failed: %v", err)
	}
}

2. 自定义迁移文件模板 #

复制Goose默认的Go迁移模板，修改函数签名为接收*sql.DB和context.Context：

1. 在项目根目录创建templates/goose/目录，新建go.tmpl文件：

package migrations

import (
	"context"
	"database/sql"

	"entgo.io/ent"
	entSQL "entgo.io/ent/dialect/sql"
)

func Up_{{.Version}}{{.UnderscoreName}}(ctx context.Context, db *sql.DB) error {
	// 初始化Ent客户端
	driver := entSQL.NewDriver("postgres", db.Conn(ctx))
	client := ent.NewClient(ent.Driver(driver))
	defer client.Close()

	// 在这里编写你的Ent操作逻辑
	// 示例：client.User.Create().SetName("demo").Save(ctx)
	return nil
}

func Down_{{.Version}}{{.UnderscoreName}}(ctx context.Context, db *sql.DB) error {
	// 初始化Ent客户端
	driver := entSQL.NewDriver("postgres", db.Conn(ctx))
	client := ent.NewClient(ent.Driver(driver))
	defer client.Close()

	// 在这里编写回滚逻辑
	return nil
}

2. 编译自定义Goose时，通过环境变量指定模板目录：

export GOOSE_TEMPLATES="./templates/goose"
go build -o ./bin/goose ./cmd/goose

3. 使用自定义Goose生成迁移 #

用编译好的二进制生成迁移文件：

./bin/goose create fetch_user_data go

生成的迁移文件会自动使用你定义的模板，Up/Down函数直接接收*sql.DB参数，可直接初始化Ent客户端进行操作。

关键注意事项 #

• 确保Ent驱动版本与Goose、数据库驱动版本兼容，避免连接适配问题。
• 迁移函数中使用Ent客户端时，务必通过defer client.Close()释放资源。
• 自定义Provider的DSN获取方式可根据项目需求调整，比如从配置文件读取。

内容的提问来源于stack exchange，提问作者boazbch