95 lines
2.4 KiB
Go
95 lines
2.4 KiB
Go
package cmd
|
|
|
|
import (
|
|
"ariga.io/atlas/sql/sqltool"
|
|
"context"
|
|
"entgo.io/ent/dialect"
|
|
"entgo.io/ent/dialect/sql/schema"
|
|
"fmt"
|
|
entmigrate "framework_v2/internal/ent/migrate"
|
|
_ "github.com/jackc/pgx/v5/stdlib"
|
|
"github.com/pressly/goose/v3"
|
|
"github.com/spf13/cobra"
|
|
"log"
|
|
"os"
|
|
)
|
|
|
|
var dsnCommand = &cobra.Command{
|
|
Use: "dsn",
|
|
Short: "生成 DSN",
|
|
Long: "生成 DSN",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
fmt.Print(config.DB.Driver + "://" + config.DB.DSN)
|
|
},
|
|
}
|
|
|
|
var migrateCommand = &cobra.Command{
|
|
Use: "migrate [command]",
|
|
Short: "goose 迁移，用法 <command>",
|
|
Long: "适用于生产环境的数据库迁移",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
if len(args) == 0 {
|
|
_ = cmd.Help()
|
|
return
|
|
}
|
|
RunMigrate(args)
|
|
},
|
|
}
|
|
|
|
var createMigrateCommand = &cobra.Command{
|
|
Use: "create-migrate",
|
|
Short: "新建迁移",
|
|
Long: "从 internal/ent 中新建迁移。在这之前，需要运行 go generate ./internal/ent",
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
generateMigration()
|
|
},
|
|
}
|
|
|
|
// RunMigrate 为数据库函数
|
|
func RunMigrate(args []string) {
|
|
db, err := goose.OpenDBWithDriver("postgres", config.DB.DSN2)
|
|
if err != nil {
|
|
log.Fatalf("goose: failed to open DB: %v\n", err)
|
|
}
|
|
|
|
defer func() {
|
|
if err := db.Close(); err != nil {
|
|
log.Fatalf("goose: failed to close DB: %v\n", err)
|
|
}
|
|
}()
|
|
|
|
command := args[0]
|
|
|
|
var arguments []string
|
|
if len(args) > 3 {
|
|
arguments = append(arguments, args[3:]...)
|
|
}
|
|
|
|
if err := goose.RunContext(context.Background(), command, db, "internal/migrations", arguments...); err != nil {
|
|
log.Fatalf("goose %v: %v", command, err)
|
|
}
|
|
}
|
|
|
|
func generateMigration() {
|
|
ctx := context.Background()
|
|
dir, err := sqltool.NewGooseDir("internal/migrations")
|
|
if err != nil {
|
|
log.Fatalf("failed creating atlas migration directory: %v", err)
|
|
}
|
|
// Migrate diff options.
|
|
opts := []schema.MigrateOption{
|
|
schema.WithDir(dir), // provide migration directory
|
|
schema.WithMigrationMode(schema.ModeInspect), // provide migration mode
|
|
schema.WithDialect(dialect.Postgres), // Ent dialect to use
|
|
|
|
}
|
|
if len(os.Args) != 3 {
|
|
log.Fatalln("migration name is required. Use: 'go run -mod=mod internal/ent/migrate/main.go <name>'")
|
|
}
|
|
|
|
err = entmigrate.NamedDiff(ctx, config.DB.Driver+"://"+config.DB.DSN, os.Args[2], opts...)
|
|
if err != nil {
|
|
log.Fatalf("failed generating migration file: %v", err)
|
|
}
|
|
}
|